From b823dfbb36d00b65745835e2c9fb8d0e00f6fb17 Mon Sep 17 00:00:00 2001 From: liwei Date: Tue, 3 Mar 2020 10:08:39 +0800 Subject: [PATCH] refactor: piece download logic --- protocol/bt/metainfo/metainfo_test.go | 10 --- protocol/bt/torrent/peer_conn.go | 41 ++++------ protocol/bt/torrent/piece_state.go | 113 +++++++++++++------------- protocol/bt/torrent/torrent.go | 94 ++++++++++----------- protocol/bt/torrent/torrent_test.go | 3 +- 5 files changed, 119 insertions(+), 142 deletions(-) diff --git a/protocol/bt/metainfo/metainfo_test.go b/protocol/bt/metainfo/metainfo_test.go index 10a03465..36266723 100644 --- a/protocol/bt/metainfo/metainfo_test.go +++ b/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()) -} diff --git a/protocol/bt/torrent/peer_conn.go b/protocol/bt/torrent/peer_conn.go index 975f13b0..2a262d7f 100644 --- a/protocol/bt/torrent/peer_conn.go +++ b/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) } diff --git a/protocol/bt/torrent/piece_state.go b/protocol/bt/torrent/piece_state.go index 91f3e5da..8a49a523 100644 --- a/protocol/bt/torrent/piece_state.go +++ b/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() } diff --git a/protocol/bt/torrent/torrent.go b/protocol/bt/torrent/torrent.go index 72847fdb..bf95165c 100644 --- a/protocol/bt/torrent/torrent.go +++ b/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") } diff --git a/protocol/bt/torrent/torrent_test.go b/protocol/bt/torrent/torrent_test.go index 66360f33..b9ce4e13 100644 --- a/protocol/bt/torrent/torrent_test.go +++ b/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() {