refactor: impl piece state with bitmap

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)
 }
 
 // 获取要写入到的文件

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
 	}