refactor: impl piece state with bitmap

protocol/bt/metainfo/metainfo_test.go

@@ -3,6 +3,7 @@ package metainfo
 import (
 	"encoding/hex"
 	"fmt"
+	"github.com/RoaringBitmap/roaring"
 	"testing"
 )
 
@@ -17,3 +18,15 @@ func TestParseFromFile(t *testing.T) {
 	fmt.Printf("%+v\n", metaInfo.Info.PieceLength)
 	fmt.Printf("%+v\n", hex.EncodeToString(metaInfo.infoHash[:]))
 }
+
+func TestParseFromFile2(t *testing.T) {
+	a := 0b10010000
+	b := 0b01111111
+	fmt.Printf("%08b\n", (a&b)^b)
+
+	have := roaring.BitmapOf(0, 3)
+	peer := roaring.BitmapOf(1, 2, 3, 4, 5, 6, 7)
+
+	need := roaring.Xor(roaring.And(have, peer), peer)
+	fmt.Println(need.ToArray())
+}

protocol/bt/peer/message/msg_bitfield.go

@@ -1,15 +1,18 @@
 package message
 
+import "github.com/RoaringBitmap/roaring"
+
 // bitfield: <len=0001+X><id=5><bitfield>
 type Bitfield struct {
 	*Message
-	payload []byte
+	payload *roaring.Bitmap
 }
 
-func NewBitfield(payload []byte) *Bitfield {
+func NewBitfield(payload *roaring.Bitmap) *Bitfield {
 	return &Bitfield{
 		Message: &Message{
-			Length: uint32(len(payload) + 1),
+			// TODO 长度处理
+			Length: 1,
 			ID:     IdBitfield,
 		},
 		payload: payload,
@@ -17,44 +20,46 @@ func NewBitfield(payload []byte) *Bitfield {
 }
 
 func (b *Bitfield) Encode() []byte {
-	return append(b.Message.Encode(), b.payload...)
+	buf := b.Message.Encode()
+	for i := 0; i < int(b.Length)-len(buf); i++ {
+		temp := byte(0)
+		for j := 0; j < 8; j++ {
+			// 如果piece存在,位数标识为1
+			if b.payload.ContainsInt(i*8 + j) {
+				temp = temp | (1 << (7 - j))
+			}
+		}
+		buf = append(buf, temp)
+	}
+	return buf
 }
 
 func (b *Bitfield) Decode(buf []byte) Serialize {
 	b.Message = &Message{}
 	b.Message.Decode(buf)
-	b.payload = buf[5:]
+	temp := buf[5:]
+	b.payload = roaring.New()
+	// 初始化bitmap
+	for i := 0; i < len(temp); i++ {
+		for j := 0; j < 8; j++ {
+			if temp[i]&(1<<(7-j)) > 0 {
+				b.payload.AddInt(i*8 + j)
+			}
+		}
+	}
 	return b
 }
 
 // 某个分片是否下载完成
 func (b *Bitfield) IsComplete(i int) bool {
-	index := i / 8
+	return b.payload.ContainsInt(i)
-	if index >= len(b.payload) {
-		return false
-	}
-	return b.payload[index]&(1<<(7-i%8)) > 0
 }
 
-// 给定一组分片下载状态，计算出当前peer能提供下载的分片下标
+// 给定一组分片下载状态，计算出当前peer能提供下载的分片下标，如果已经拥有则不用提供
-func (b *Bitfield) Have(pieces []bool) []int {
+// had = 0b10010000,has = 0b01111111
-	arr := make([]int, 0)
+// had & has = 0b10010000 & 0b01111111 = 0b00010000
-	length := len(pieces) / 8
+// had & has ^ has = 0b00010000 ^ 0b01111111 = 0b01101111
-	if len(pieces)%8 != 0 {
+func (b *Bitfield) Provide(had *roaring.Bitmap) []uint32 {
-		length++
+	xor := roaring.Xor(roaring.And(had, b.payload), b.payload)
-	}
+	return xor.ToArray()
-
-	for i := 0; i < length; i++ {
-		for j := 0; j < 8; j++ {
-			index := i*8 + j
-			// 如果此分片在本地还未下载，检查peer是否能提供该分片下载
-			if index < len(pieces) &&
-				!pieces[index] &&
-				i < len(b.payload) &&
-				b.payload[i]&(1<<(7-j)) > 0 {
-				arr = append(arr, index)
-			}
-		}
-	}
-	return arr
 }

protocol/bt/peer/message/msg_bitfield_test.go

@@ -1,6 +1,7 @@
 package message
 
 import (
+	"github.com/RoaringBitmap/roaring"
 	"reflect"
 	"testing"
 )
@@ -8,7 +9,7 @@ import (
 func TestBitfield_IsComplete(t *testing.T) {
 	type fields struct {
 		Message Message
-		payload []byte
+		payload *roaring.Bitmap
 	}
 	type args struct {
 		i int
@@ -19,20 +20,10 @@ func TestBitfield_IsComplete(t *testing.T) {
 		args   args
 		want   bool
 	}{
-		{
-			"case-0",
-			fields{
-				payload: []byte{0b10000000},
-			},
-			args{
-				0,
-			},
-			true,
-		},
 		{
 			"case-1",
 			fields{
-				payload: []byte{0b01000000},
+				payload: roaring.BitmapOf(1, 2, 3, 4),
 			},
 			args{
 				1,
@@ -42,40 +33,10 @@ func TestBitfield_IsComplete(t *testing.T) {
 		{
 			"case-2",
 			fields{
-				payload: []byte{0b00000000, 0b00100000},
+				payload: roaring.BitmapOf(1, 2, 3, 4),
 			},
 			args{
-				2,
+				5,
-			},
-			false,
-		},
-		{
-			"case-10",
-			fields{
-				payload: []byte{0b00000000, 0b00100000},
-			},
-			args{
-				10,
-			},
-			true,
-		},
-		{
-			"case-11",
-			fields{
-				payload: []byte{0b00000000, 0b00100000},
-			},
-			args{
-				11,
-			},
-			false,
-		},
-		{
-			"case-20",
-			fields{
-				payload: []byte{0b00000000, 0b00100000},
-			},
-			args{
-				20,
 			},
 			false,
 		},
@@ -96,56 +57,46 @@ func TestBitfield_IsComplete(t *testing.T) {
 func TestBitfield_Have(t *testing.T) {
 	type fields struct {
 		Message Message
-		payload []byte
+		payload *roaring.Bitmap
 	}
 	type args struct {
-		pieces []bool
+		had *roaring.Bitmap
 	}
 	tests := []struct {
 		name   string
 		fields fields
 		args   args
-		want   []int
+		want   []uint32
 	}{
 		{
 			"case-1",
 			fields{
-				payload: []byte{0b11111111, 0b00000000},
+				payload: roaring.BitmapOf(0, 1, 2, 3, 4),
 			},
 			args{
-				[]bool{false, false, false, false, false, false, false, false},
+				roaring.BitmapOf(0),
 			},
-			[]int{0, 1, 2, 3, 4, 5, 6, 7},
+			[]uint32{1, 2, 3, 4},
 		},
 		{
 			"case-2",
 			fields{
-				payload: []byte{0b01010101, 0b01000000},
+				payload: roaring.BitmapOf(0, 1, 2, 3, 4, 5, 6, 7, 8),
 			},
 			args{
-				[]bool{false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false},
+				roaring.BitmapOf(0, 1, 3, 5, 6, 8),
 			},
-			[]int{1, 3, 5, 7, 9},
+			[]uint32{2, 4, 7},
 		},
 		{
 			"case-3",
 			fields{
-				payload: []byte{0b11010101, 0b01000000},
+				payload: roaring.BitmapOf(),
 			},
 			args{
-				[]bool{false, true, false},
+				roaring.BitmapOf(0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
 			},
-			[]int{0},
+			[]uint32{},
-		},
-		{
-			"case-4",
-			fields{
-				payload: []byte{0b11010101},
-			},
-			args{
-				[]bool{false, true, false, false, true, false, false, true, false},
-			},
-			[]int{0, 3, 5},
 		},
 	}
 	for _, tt := range tests {
@@ -154,8 +105,8 @@ func TestBitfield_Have(t *testing.T) {
 				Message: &tt.fields.Message,
 				payload: tt.fields.payload,
 			}
-			if got := b.Have(tt.args.pieces); !reflect.DeepEqual(got, tt.want) {
+			if got := b.Provide(tt.args.had); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("Bitfield.Have() = %v, want %v", got, tt.want)
+				t.Errorf("Bitfield.Provide() = %v, want %v", got, tt.want)
 			}
 		})
 	}

protocol/bt/torrent/peer_conn.go

@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"github.com/RoaringBitmap/roaring"
 	"io"
 	"net"
 	"os"
@@ -229,7 +230,7 @@ func (pc *peerConn) downloadPiece(index int) (err error) {
 		} else {
 			blockLength = blockSize
 		}
-		// block下载排队，如果连接出现问题直接返回异常
+		// block下载排队
 		select {
 		case pc.blockQueueCh <- nil:
 			break
@@ -238,6 +239,7 @@ func (pc *peerConn) downloadPiece(index int) (err error) {
 		case err = <-pc.disconnectCh:
 			break
 		}
+		// 如果连接出现问题或下载失败直接返回异常
 		if err != nil {
 			pc.conn.Close()
 			return
@@ -436,12 +438,14 @@ type fileBlock struct {
 }
 
 // 获取peer能提供需要下载的文件分片
-func (pc *peerConn) getHavePieces(bitfield *message.Bitfield) []int {
+func (pc *peerConn) getHavePieces(bitfield *message.Bitfield) []uint32 {
-	states := make([]bool, pc.torrent.pieceStates.size())
+	had := roaring.New()
-	for i := range states {
+	for i := 0; i < pc.torrent.pieceStates.size(); i++ {
-		states[i] = pc.torrent.pieceStates.getState(i) == stateFinish
+		if pc.torrent.pieceStates.getState(i) == stateFinish {
+			had.AddInt(i)
+		}
 	}
-	return bitfield.Have(states)
+	return bitfield.Provide(had)
 }
 
 // 获取要写入到的文件

protocol/bt/torrent/peer_conn_test.go

@@ -25,46 +25,6 @@ func Test_peerState_ready(t *testing.T) {
 
 }
 
-func Test_peerState_download(t *testing.T) {
-
-	psCh := getUsablePeerMore()
-
-	lock := sync.Mutex{}
-	state := map[int]bool{}
-
-	for {
-		select {
-		case ps := <-psCh:
-			go func() {
-				// 下载前准备
-				err := ps.ready()
-				if err != nil {
-					fmt.Println("ready error")
-					return
-				}
-				have := ps.getHavePieces(ps.bitfield)
-				fmt.Printf("download is ready,have %d\n", len(have))
-				if len(have) > 0 {
-					// 获取分片的长度
-					for _, index := range have {
-						lock.Lock()
-						if state[index] {
-							lock.Unlock()
-							continue
-						} else {
-							state[index] = true
-							lock.Unlock()
-						}
-						fmt.Printf("download index %d\n", index)
-						ps.downloadPiece(index)
-					}
-					return
-				}
-			}()
-		}
-	}
-}
-
 func Test_peerState_downloadPiece(t *testing.T) {
 
 	ps := getUsablePeer()
@@ -81,7 +41,7 @@ func Test_peerState_downloadPiece(t *testing.T) {
 		// 获取分片的长度
 		for _, index := range have {
 			fmt.Printf("download index %d\n", index)
-			ps.downloadPiece(index)
+			ps.downloadPiece(int(index))
 		}
 		return
 	}