refactor: piece download logic

protocol/bt/metainfo/metainfo_test.go

@@ -3,7 +3,6 @@ package metainfo
 import (
 	"encoding/hex"
 	"fmt"
-	"github.com/RoaringBitmap/roaring"
 	"testing"
 )
 
@@ -18,12 +17,3 @@ 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) {
-	bitmap := roaring.NewBitmap()
-	bitmap.AddMany([]uint32{1, 2, 3, 4, 5, 6, 7, 10, 20, 9, 9, 1})
-	fmt.Println(bitmap.Contains(1))
-	fmt.Println(bitmap.Contains(2))
-	fmt.Println(bitmap.Contains(8))
-	fmt.Println(bitmap.GetCardinality())
-}

protocol/bt/torrent/peer_conn.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math"
 	"net"
 	"os"
 	"path"
@@ -49,9 +48,6 @@ type peerConn struct {
 
 	downloadedCh chan error
 	disconnectCh chan error
-	keepaliveCh  chan error
-	// block个数
-	blockCount int
 	// block下载队列，官方推荐为5
 	blockQueueCh chan interface{}
 }
@@ -134,8 +130,6 @@ func (pc *peerConn) handleKeepalive() {
 			break
 		}
 	}
-	pc.keepaliveCh <- fmt.Errorf("keepalive fail")
-	close(pc.keepaliveCh)
 }
 
 // 准备下载
@@ -149,7 +143,6 @@ func (pc *peerConn) ready() error {
 	pc.readyEnd = false
 	readyCh := make(chan bool)
 	pc.disconnectCh = make(chan error)
-	pc.keepaliveCh = make(chan error)
 	go pc.handleKeepalive()
 	go func() {
 		scanner := bufio.NewScanner(pc.conn)
@@ -217,20 +210,20 @@ func (pc *peerConn) ready() error {
 // 下载指定piece
 func (pc *peerConn) downloadPiece(index int) (err error) {
 	pieceLength := pc.torrent.MetaInfo.GetPieceLength(index)
-	pc.blockCount = int(math.Ceil(float64(pieceLength) / blockSize))
 	pc.downloadedCh = make(chan error)
 	pc.blockQueueCh = make(chan interface{}, 5)
 	defer close(pc.blockQueueCh)
 	// 按块下载分片
-	for i := 0; i < pc.blockCount; i++ {
+	blockCount := pc.torrent.pieceStates.states[index].blockCount
+	for i := 0; i < blockCount; i++ {
 		offset := blockSize * i
 		// 如果已经下载过就跳过
-		if pc.torrent.pieces.isBlockDownloaded(index, offset) {
+		if pc.torrent.pieceStates.isBlockDone(index, offset) {
 			continue
 		}
 
 		var blockLength uint32
-		if i == pc.blockCount-1 {
+		if i == blockCount-1 {
 			// 最后一个block大小需要计算出来
 			blockLength = uint32(pieceLength - offset)
 		} else {
@@ -244,8 +237,6 @@ func (pc *peerConn) downloadPiece(index int) (err error) {
 			break
 		case err = <-pc.disconnectCh:
 			break
-		case err = <-pc.keepaliveCh:
-			break
 		}
 		if err != nil {
 			pc.conn.Close()
@@ -262,8 +253,6 @@ func (pc *peerConn) downloadPiece(index int) (err error) {
 		break
 	case err = <-pc.disconnectCh:
 		break
-	case err = <-pc.keepaliveCh:
-		break
 	}
 	if err != nil {
 		pc.conn.Close()
@@ -355,27 +344,33 @@ func (pc *peerConn) handlePiece(buf []byte) {
 	}
 	// 开始写入文件
 	for _, f := range fileBlocks {
-		func() {
+		err := func() error {
 			name := filepath.Join(pc.torrent.Path, f.filepath)
 			file, err := os.OpenFile(name, os.O_RDWR|os.O_CREATE, 0644)
 			if err != nil {
-				panic(err)
+				return err
 			}
 			defer file.Close()
 			_, err = file.WriteAt(piece.Block[f.blockBegin:f.blockEnd], f.fileSeek)
 			if err != nil {
-				panic(err)
+				return err
 			}
+			return nil
 		}()
+		if err != nil {
+			pc.downloadedCh <- errPieceCheckFailed
+			pc.conn.Close()
+			return
+		}
 	}
-	pc.torrent.pieces.setBlockDownloaded(int(piece.Index), int(piece.Begin))
+	pc.torrent.pieceStates.setBlockDone(int(piece.Index), int(piece.Begin))
 	// 出队
 	<-pc.blockQueueCh
 
-	log.Debugf("piece:%d block:%d already:%d", piece.Index, piece.Begin, pc.torrent.pieces.blockSize(int(piece.Index)))
+	log.Debugf("piece:%d block:%d", piece.Index, piece.Begin)
 
 	// piece全部下载完
-	if pc.torrent.pieces.blockSize(int(piece.Index)) == pc.blockCount {
+	if pc.torrent.pieceStates.isPieceDone(int(piece.Index)) {
 		// 计算piece对应的文件偏移
 		sha1 := sha1.New()
 		pieceLength := pc.torrent.MetaInfo.GetPieceLength(int(piece.Index))
@@ -442,9 +437,9 @@ type fileBlock struct {
 
 // 获取peer能提供需要下载的文件分片
 func (pc *peerConn) getHavePieces(bitfield *message.Bitfield) []int {
-	states := make([]bool, pc.torrent.pieces.size())
+	states := make([]bool, pc.torrent.pieceStates.size())
 	for i := range states {
-		states[i] = pc.torrent.pieces.getState(i) == stateFinished
+		states[i] = pc.torrent.pieceStates.getState(i) == stateFinish
 	}
 	return bitfield.Have(states)
 }

protocol/bt/torrent/piece_state.go

@@ -1,6 +1,8 @@
 package torrent
 
 import (
+	"github.com/monkeyWie/gopeed/protocol/bt/metainfo"
+	"math"
 	"sync"
 
 	"github.com/RoaringBitmap/roaring"
@@ -11,100 +13,97 @@ type state int
 const (
 	stateReady = iota
 	stateDownloading
-	stateFinished
+	stateFinish
 )
 
 // 所有piece的下载状态
-type pieces struct {
-	rw  *sync.RWMutex
-	arr []piece
+type pieceStates struct {
+	lock *sync.RWMutex
+	// 所有piece的block下载进度
+	states map[int]*pieceState
 }
 
-type piece struct {
-	// piece的下载状态
+type pieceState struct {
 	state state
-	// piece中block的下载状态
-	blocks *roaring.Bitmap
+	// piece中block的个数
+	blockCount int
+	// piece中block的下载记录
+	blockRecord *roaring.Bitmap
 }
 
-func newPiecesState(size int) *pieces {
-	ps := &pieces{
-		rw:  &sync.RWMutex{},
-		arr: make([]piece, size),
+func newPiecesState(meta *metainfo.MetaInfo) *pieceStates {
+	pieceCount := len(meta.Info.Pieces)
+	ps := &pieceStates{
+		lock:   &sync.RWMutex{},
+		states: make(map[int]*pieceState, pieceCount),
 	}
-	for i := 0; i < size; i++ {
-		ps.arr[i].state = stateReady
-		ps.arr[i].blocks = roaring.NewBitmap()
+	for i := 0; i < pieceCount; i++ {
+		ps.states[i] = &pieceState{
+			state:       stateReady,
+			blockCount:  int(math.Ceil(float64(meta.GetPieceLength(i)) / blockSize)),
+			blockRecord: roaring.New(),
+		}
 	}
 	return ps
 }
 
-func (ps *pieces) getState(index int) state {
-	ps.rw.RLock()
-	defer ps.rw.RUnlock()
-	return ps.arr[index].state
+func (ps *pieceStates) getState(index int) state {
+	ps.lock.RLock()
+	defer ps.lock.RUnlock()
+	return ps.states[index].state
 }
 
-func (ps *pieces) setState(index int, state state) {
-	ps.rw.Lock()
-	defer ps.rw.Unlock()
-	ps.arr[index].state = state
+func (ps *pieceStates) setState(index int, state state) {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+	ps.states[index].state = state
 }
 
-func (ps *pieces) getReady() int {
-	ps.rw.Lock()
-	defer ps.rw.Unlock()
-	for i, s := range ps.arr {
-		if s.state == stateReady {
-			ps.arr[i].state = stateDownloading
-			return i
-		}
-	}
-	return -1
+func (ps *pieceStates) size() int {
+	return len(ps.states)
 }
 
-func (ps *pieces) getLeft() int {
-	ps.rw.RLock()
-	defer ps.rw.RUnlock()
+func (ps *pieceStates) getLeft() int {
+	ps.lock.RLock()
+	defer ps.lock.RUnlock()
 	count := 0
-	for _, s := range ps.arr {
-		if s.state != stateFinished {
+	for _, v := range ps.states {
+		if v.state == stateReady || v.state == stateDownloading {
 			count++
 		}
 	}
 	return count
 }
 
-func (ps *pieces) size() int {
-	return len(ps.arr)
-}
-
-func (ps *pieces) isDone() bool {
-	ps.rw.RLock()
-	defer ps.rw.RUnlock()
-	for _, s := range ps.arr {
-		if s.state != stateFinished {
+func (ps *pieceStates) isDone() bool {
+	ps.lock.RLock()
+	defer ps.lock.RUnlock()
+	for _, v := range ps.states {
+		if v.state == stateReady || v.state == stateDownloading {
 			return false
 		}
 	}
 	return true
 }
 
+func (ps *pieceStates) isPieceDone(index int) bool {
+	ps.lock.RLock()
+	defer ps.lock.RUnlock()
+	state := ps.states[index]
+	return uint64(state.blockCount) == state.blockRecord.GetCardinality()
+}
+
 // 指定piece下的某个block是否下载完成
-func (ps *pieces) isBlockDownloaded(pieceIndex int, blockIndex int) bool {
-	return ps.arr[pieceIndex].blocks.ContainsInt(blockIndex)
+func (ps *pieceStates) isBlockDone(pieceIndex int, blockIndex int) bool {
+	return ps.states[pieceIndex].blockRecord.ContainsInt(blockIndex)
 }
 
 // 指定piece下的某个block下载完成
-func (ps *pieces) setBlockDownloaded(pieceIndex int, blockIndex int) {
-	ps.arr[pieceIndex].blocks.AddInt(blockIndex)
+func (ps *pieceStates) setBlockDone(pieceIndex int, blockIndex int) {
+	ps.states[pieceIndex].blockRecord.AddInt(blockIndex)
 }
 
 // 清空指定piece下所有的block下载记录
-func (ps *pieces) clearBlocks(index int) {
-	ps.arr[index].blocks.Clear()
-}
-
-func (ps *pieces) blockSize(index int) int {
-	return int(ps.arr[index].blocks.GetCardinality())
+func (ps *pieceStates) clearBlocks(index int) {
+	ps.states[index].blockRecord.Clear()
 }

protocol/bt/torrent/torrent.go

@@ -12,18 +12,16 @@ type Torrent struct {
 	PeerID   [20]byte
 	MetaInfo *metainfo.MetaInfo
 
-	pieces     *pieces
-	peerPool   *peerPool
-	completeCh chan interface{}
-
-	Path string
+	pieceStates *pieceStates
+	peerPool    *peerPool
+	Path        string
 }
 
 func NewTorrent(peerID [20]byte, metaInfo *metainfo.MetaInfo) *Torrent {
 	torrent := &Torrent{
-		PeerID:   peerID,
-		MetaInfo: metaInfo,
-		pieces:   newPiecesState(len(metaInfo.Info.Pieces)),
+		PeerID:      peerID,
+		MetaInfo:    metaInfo,
+		pieceStates: newPiecesState(metaInfo),
 	}
 	return torrent
 }
@@ -33,57 +31,53 @@ func (t *Torrent) Download(path string) {
 	t.peerPool = newPeerPool(t)
 	t.peerPool.fetch()
 
-	pieceQueueCh := make(chan interface{}, 256)
+	pieceQueueCh := make(chan interface{}, 128)
 	defer close(pieceQueueCh)
-	t.completeCh = make(chan interface{})
-	defer close(t.completeCh)
-	for {
+	downloadedCh := make(chan interface{})
+	defer close(downloadedCh)
+
+	for i := 0; i < len(t.MetaInfo.Info.Pieces); i++ {
+		// 加入下载队列
 		pieceQueueCh <- nil
-		// 获取一个待下载的piece
-		index := t.pieces.getReady()
-		// 没有待下载的piece了
-		if index == -1 {
-			log.Debugf("no piece to download")
-			<-t.completeCh
-			// 下载完成
-			break
-		}
 		go func(index int) {
-			peer := t.peerPool.get()
-			if peer == nil {
-				time.Sleep(time.Second * 10)
-				t.pieces.setState(index, stateReady)
-				<-pieceQueueCh
-				return
-			}
-			pc := newPeerConn(t, peer)
-			err := pc.ready()
-			if err != nil {
-				t.pieces.setState(index, stateReady)
-				t.peerPool.unavailable(peer)
-			} else {
-				err = pc.downloadPiece(index)
+			t.pieceStates.setState(index, stateDownloading)
+			// 一直下载到成功
+			for {
+				peer := t.peerPool.get()
+				if peer == nil {
+					time.Sleep(time.Second * 10)
+					continue
+				}
+				pc := newPeerConn(t, peer)
+				err := pc.ready()
 				if err != nil {
-					log.Debugf("piece %d download fail:%s", index, err.Error())
-					// 下载失败
-					t.pieces.setState(index, stateReady)
 					t.peerPool.unavailable(peer)
-					if errors.Is(err, errPieceCheckFailed) {
-						// piece校验失败，需要重新下载过一遍，清空block记录
-						t.pieces.clearBlocks(index)
-					}
 				} else {
-					// 下载完成
-					t.pieces.setState(index, stateFinished)
-					t.peerPool.release(peer)
-					log.Debugf("piece %d download success,total:%d,left:%d", index, t.pieces.size(), t.pieces.getLeft())
-					// 检查是否所有piece下载完成
-					if t.pieces.isDone() {
-						t.completeCh <- nil
+					err = pc.downloadPiece(index)
+					if err != nil {
+						log.Debugf("piece %d download fail:%s", index, err.Error())
+						// 下载失败
+						t.peerPool.unavailable(peer)
+						if errors.Is(err, errPieceCheckFailed) {
+							// piece校验失败，需要重新下载过一遍，清空block记录
+							t.pieceStates.clearBlocks(index)
+						}
+						continue
+					} else {
+						// 下载完成
+						t.pieceStates.setState(index, stateFinish)
+						t.peerPool.release(peer)
+						log.Debugf("piece %d download success,total:%d,left:%d", index, t.pieceStates.size(), t.pieceStates.getLeft())
+						if t.pieceStates.isDone() {
+							downloadedCh <- nil
+						}
+						break
 					}
 				}
 			}
 			<-pieceQueueCh
-		}(index)
+		}(i)
 	}
+	<-downloadedCh
+	log.Debug("download end")
 }

protocol/bt/torrent/torrent_test.go

@@ -1,11 +1,10 @@
 package torrent
 
 import (
-	"testing"
-
 	"github.com/monkeyWie/gopeed/protocol/bt/metainfo"
 	"github.com/monkeyWie/gopeed/protocol/bt/peer"
 	log "github.com/sirupsen/logrus"
+	"testing"
 )
 
 func init() {