diff --git a/protocol/bt/metainfo/metainfo_test.go b/protocol/bt/metainfo/metainfo_test.go index 36266723..2d8b99f0 100644 --- a/protocol/bt/metainfo/metainfo_test.go +++ b/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()) +} diff --git a/protocol/bt/peer/message/msg_bitfield.go b/protocol/bt/peer/message/msg_bitfield.go index 9bfc8d63..ee4b1199 100644 --- a/protocol/bt/peer/message/msg_bitfield.go +++ b/protocol/bt/peer/message/msg_bitfield.go @@ -1,15 +1,18 @@ package message +import "github.com/RoaringBitmap/roaring" + // 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 - if index >= len(b.payload) { - return false - } - return b.payload[index]&(1<<(7-i%8)) > 0 + return b.payload.ContainsInt(i) } -// 给定一组分片下载状态,计算出当前peer能提供下载的分片下标 -func (b *Bitfield) Have(pieces []bool) []int { - arr := make([]int, 0) - length := len(pieces) / 8 - if len(pieces)%8 != 0 { - length++ - } - - 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 +// 给定一组分片下载状态,计算出当前peer能提供下载的分片下标,如果已经拥有则不用提供 +// had = 0b10010000,has = 0b01111111 +// had & has = 0b10010000 & 0b01111111 = 0b00010000 +// had & has ^ has = 0b00010000 ^ 0b01111111 = 0b01101111 +func (b *Bitfield) Provide(had *roaring.Bitmap) []uint32 { + xor := roaring.Xor(roaring.And(had, b.payload), b.payload) + return xor.ToArray() } diff --git a/protocol/bt/peer/message/msg_bitfield_test.go b/protocol/bt/peer/message/msg_bitfield_test.go index 0208dc01..b9ff7479 100644 --- a/protocol/bt/peer/message/msg_bitfield_test.go +++ b/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, - }, - 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, + 5, }, 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}, - }, - { - "case-4", - fields{ - payload: []byte{0b11010101}, - }, - args{ - []bool{false, true, false, false, true, false, false, true, false}, - }, - []int{0, 3, 5}, + []uint32{}, }, } 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) { - t.Errorf("Bitfield.Have() = %v, want %v", got, tt.want) + if got := b.Provide(tt.args.had); !reflect.DeepEqual(got, tt.want) { + t.Errorf("Bitfield.Provide() = %v, want %v", got, tt.want) } }) } diff --git a/protocol/bt/torrent/peer_conn.go b/protocol/bt/torrent/peer_conn.go index 2a262d7f..11015016 100644 --- a/protocol/bt/torrent/peer_conn.go +++ b/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 { - states := make([]bool, pc.torrent.pieceStates.size()) - for i := range states { - states[i] = pc.torrent.pieceStates.getState(i) == stateFinish +func (pc *peerConn) getHavePieces(bitfield *message.Bitfield) []uint32 { + had := roaring.New() + for i := 0; i < pc.torrent.pieceStates.size(); i++ { + if pc.torrent.pieceStates.getState(i) == stateFinish { + had.AddInt(i) + } } - return bitfield.Have(states) + return bitfield.Provide(had) } // 获取要写入到的文件 diff --git a/protocol/bt/torrent/peer_conn_test.go b/protocol/bt/torrent/peer_conn_test.go index 288e4549..795313df 100644 --- a/protocol/bt/torrent/peer_conn_test.go +++ b/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 }