mirror of
https://github.com/GopeedLab/gopeed.git
synced 2026-11-07 17:10:29 +00:00
refactor: impl piece state with bitmap
This commit is contained in:
liwei committed
1 parent
b823dfbb36
commit
1d2f38d7ae
5 files changed
+79
-146
No files matched your search
@@ -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
|
||||
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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user