diff --git a/protocol/bt/client.go b/protocol/bt/client.go index d4336ebd..c7b82bb4 100644 --- a/protocol/bt/client.go +++ b/protocol/bt/client.go @@ -46,12 +46,12 @@ func (client *Client) AddTorrent(rawurl string) error { } torrent := torrent.Torrent{ - client: client, + // client: client, MetaInfo: metaInfo, } - client.torrents = append(client.torrents, torrent) + // client.torrents = append(client.torrents, torrent) - torrent.Download() + torrent.Download("") return nil } diff --git a/protocol/bt/torrent/peer_conn.go b/protocol/bt/torrent/peer_conn.go index 62b1f1b2..d43e04ab 100644 --- a/protocol/bt/torrent/peer_conn.go +++ b/protocol/bt/torrent/peer_conn.go @@ -10,14 +10,12 @@ import ( "github.com/monkeyWie/gopeed/protocol/bt/metainfo" "github.com/monkeyWie/gopeed/protocol/bt/peer" "github.com/monkeyWie/gopeed/protocol/bt/peer/message" - "path/filepath" - "sync/atomic" - "io" "math" "net" "os" "path" + "path/filepath" "time" ) @@ -37,14 +35,13 @@ type peerConn struct { // peer is interested in this client peerInterested bool - bitfield *message.Bitfield - readyComplete bool + bitfield *message.Bitfield + readyComplete bool + downloadComplete bool - downloadedCh chan bool + downloadedCh chan error // block个数 blockCount int - // 已经下载的block个数 - blockDownloadedCount int32 // block下载队列，官方推荐为5 blockQueue chan interface{} } @@ -114,12 +111,13 @@ func (pc *peerConn) handshake() (*peer.Handshake, error) { // 准备下载 func (pc *peerConn) ready() error { if err := pc.dialMse(); err != nil { - return err + return fmt.Errorf("tcp dial error %w", err) } if _, err := pc.handshake(); err != nil { - return err + return fmt.Errorf("handshake error %w", err) } pc.readyComplete = false + pc.downloadComplete = false readyCh := make(chan bool) defer close(readyCh) go func() { @@ -157,6 +155,14 @@ func (pc *peerConn) ready() error { } } } + // 还未下载完成时连接断开 + if pc.downloadedCh != nil && !pc.downloadComplete { + err := scanner.Err() + if err == nil { + err = io.EOF + } + pc.downloadedCh <- err + } }() err := func() error { select { @@ -183,18 +189,17 @@ func (pc *peerConn) ready() error { // 下载指定piece func (pc *peerConn) downloadPiece(index int) error { pieceLength := pc.torrent.MetaInfo.GetPieceLength(index) - lastOffset := 0 - pc.blockCount = int(math.Ceil(float64(pieceLength) / BlockSize)) - pc.downloadedCh = make(chan bool) + pc.downloadedCh = make(chan error) defer close(pc.downloadedCh) - - waitDuration := time.Second * 30 - timer := time.NewTimer(waitDuration) - defer timer.Stop() // 按块下载分片 for i := 0; i < pc.blockCount; i++ { - offset := pieceLength * i + offset := BlockSize * i + // 如果已经下载过就跳过 + if pc.torrent.pieces.isBlockDownloaded(index, offset) { + continue + } + var blockLength uint32 if i == pc.blockCount-1 { // 最后一个block大小需要计算出来 @@ -202,24 +207,20 @@ func (pc *peerConn) downloadPiece(index int) error { } else { blockLength = BlockSize } - // 标记block开始下载 - offset32 := uint32(lastOffset) - // block下载排队，并做超时处理 - timer.Reset(waitDuration) + // block下载排队，如果连接出现问题直接返回异常 select { case pc.blockQueue <- nil: - case <-timer.C: - return errors.New("download timeout") + break + case err := <-pc.downloadedCh: + return err } // 发起request，对方会响应piece - _, err := pc.conn.Write(message.NewRequest(uint32(index), offset32, blockLength).Encode()) + _, err := pc.conn.Write(message.NewRequest(uint32(index), uint32(offset), blockLength).Encode()) if err != nil { fmt.Println(err) } } - // 监听是否下载完成 - <-pc.downloadedCh - return nil + return <-pc.downloadedCh } func (pc *peerConn) handleUnchoke(readyCh chan<- bool) { @@ -317,11 +318,13 @@ func (pc *peerConn) handlePiece(buf []byte) { } }() } + pc.torrent.pieces.setBlockDownloaded(int(piece.Index), int(piece.Begin)) + fmt.Printf("piece %d downloaded block %d,count %d\n", piece.Index, piece.Begin, pc.torrent.pieces.blockSize(int(piece.Index))) // 出队 <-pc.blockQueue // piece全部下载完 - if atomic.AddInt32(&pc.blockDownloadedCount, 1) == int32(pc.blockCount) { + if pc.torrent.pieces.blockSize(int(piece.Index)) == pc.blockCount { // 计算piece对应的文件偏移 sha1 := sha1.New() pieceLength := pc.torrent.MetaInfo.GetPieceLength(int(piece.Index)) @@ -360,9 +363,15 @@ func (pc *peerConn) handlePiece(buf []byte) { fmt.Printf("piece %d 校验通过\n", piece.Index) } else { fmt.Printf("piece %d 校验失败\n", piece.Index) + // 清空该piece上所有的block下载记录 + pc.torrent.pieces.clearBlocks(int(piece.Index)) } + // 标记下载完成 + pc.downloadComplete = true + // 断开连接 TODO 用连接池进行复用 + pc.conn.Close() // piece下载完成 - pc.downloadedCh <- true + pc.downloadedCh <- nil } } @@ -385,9 +394,9 @@ type fileBlock struct { // 获取peer能提供需要下载的文件分片 func (pc *peerConn) getHavePieces(bitfield *message.Bitfield) []int { - states := make([]bool, pc.torrent.PiecesState.Size()) + states := make([]bool, pc.torrent.pieces.size()) for i := range states { - states[i] = pc.torrent.PiecesState.getState(i) == stateFinished + states[i] = pc.torrent.pieces.getState(i) == stateFinished } return bitfield.Have(states) } diff --git a/protocol/bt/torrent/peer_pool.go b/protocol/bt/torrent/peer_pool.go index 41e444a1..ae03637f 100644 --- a/protocol/bt/torrent/peer_pool.go +++ b/protocol/bt/torrent/peer_pool.go @@ -6,77 +6,63 @@ import ( ) type peerPool struct { - peers []peerState - lock *sync.Mutex + states map[peer.Peer]*peerState + lock *sync.Mutex } type peerState struct { - peer *peer.Peer + // 是否在被使用 using bool - ready bool + // 下载失败次数 + errors int } func newPeerPool() *peerPool { return &peerPool{ - lock: &sync.Mutex{}, + lock: &sync.Mutex{}, + states: map[peer.Peer]*peerState{}, } } func (pp *peerPool) put(peers []peer.Peer) { pp.lock.Lock() defer pp.lock.Unlock() - for i := range peers { - pp.peers = append(pp.peers, peerState{ - peer: &peers[i], - using: false, - }) + for _, peer := range peers { + if _, ok := pp.states[peer]; !ok { + pp.states[peer] = &peerState{ + using: false, + errors: 0, + } + } + } } func (pp *peerPool) get() *peer.Peer { pp.lock.Lock() defer pp.lock.Unlock() - for i, p := range pp.peers { - if !p.using { - pp.peers[i].using = true - return pp.peers[i].peer + for peer, state := range pp.states { + if !state.using { + state.using = true + return &peer } } return nil } -func (pp *peerPool) ready(peer *peer.Peer) { - pp.lock.Lock() - defer pp.lock.Unlock() - i := findPeer(pp, peer) - if i != -1 { - pp.peers[i].ready = true - } -} - func (pp *peerPool) release(peer *peer.Peer) { pp.lock.Lock() defer pp.lock.Unlock() - i := findPeer(pp, peer) - if i != -1 { - pp.peers[i].using = false - } + pp.states[*peer].using = false + pp.states[*peer].errors = 0 } -func (pp *peerPool) remove(peer *peer.Peer) { +func (pp *peerPool) unavailable(peer *peer.Peer) { pp.lock.Lock() defer pp.lock.Unlock() - i := findPeer(pp, peer) - if i != -1 { - pp.peers = append(pp.peers[0:i], pp.peers[i+1:]...) + pp.states[*peer].errors++ + if pp.states[*peer].errors > 3 { + delete(pp.states, *peer) } -} -func findPeer(pp *peerPool, peer *peer.Peer) int { - for i := range pp.peers { - if pp.peers[i].peer == peer { - return i - } - } - return -1 } diff --git a/protocol/bt/torrent/peer_pool_test.go b/protocol/bt/torrent/peer_pool_test.go index 765b4080..64926789 100644 --- a/protocol/bt/torrent/peer_pool_test.go +++ b/protocol/bt/torrent/peer_pool_test.go @@ -32,7 +32,7 @@ func Test_PeerPool(t *testing.T) { }) get := pool.get() pool.release(get) - pool.remove(get) + pool.unavailable(get) for i := 0; i < 5; i++ { fmt.Println(pool.get()) } diff --git a/protocol/bt/torrent/piece_state.go b/protocol/bt/torrent/piece_state.go index 34d245c8..7d507912 100644 --- a/protocol/bt/torrent/piece_state.go +++ b/protocol/bt/torrent/piece_state.go @@ -1,6 +1,9 @@ package torrent -import "sync" +import ( + "github.com/RoaringBitmap/roaring" + "sync" +) type state int @@ -11,46 +14,73 @@ const ( ) // 所有piece的下载状态 -type piecesState struct { - rw *sync.RWMutex - states []state +type pieces struct { + rw *sync.RWMutex + arr []piece } -func NewPiecesState(size int) *piecesState { - ps := &piecesState{ - rw: &sync.RWMutex{}, - states: make([]state, size), +type piece struct { + // piece的下载状态 + state state + // piece中block的下载状态 + blocks *roaring.Bitmap +} + +func newPiecesState(size int) *pieces { + ps := &pieces{ + rw: &sync.RWMutex{}, + arr: make([]piece, size), } for i := 0; i < size; i++ { - ps.states[i] = stateReady + ps.arr[i].state = stateReady + ps.arr[i].blocks = roaring.NewBitmap() } return ps } -func (ps *piecesState) getState(index int) state { +func (ps *pieces) getState(index int) state { ps.rw.RLock() defer ps.rw.RUnlock() - return ps.states[index] + return ps.arr[index].state } -func (ps *piecesState) setState(index int, state state) { +func (ps *pieces) setState(index int, state state) { ps.rw.Lock() defer ps.rw.Unlock() - ps.states[index] = state + ps.arr[index].state = state } -func (ps *piecesState) getReadyAndDownload() int { +func (ps *pieces) getReady() int { ps.rw.Lock() defer ps.rw.Unlock() - for i, s := range ps.states { - if s == stateReady { - ps.states[i] = stateDownloading + for i, s := range ps.arr { + if s.state == stateReady { + ps.arr[i].state = stateDownloading return i } } return -1 } -func (ps *piecesState) Size() int { - return len(ps.states) +func (ps *pieces) size() int { + return len(ps.arr) +} + +// 指定piece下的某个block是否下载完成 +func (ps *pieces) isBlockDownloaded(pieceIndex int, blockIndex int) bool { + return ps.arr[pieceIndex].blocks.ContainsInt(blockIndex) +} + +// 指定piece下的某个block下载完成 +func (ps *pieces) setBlockDownloaded(pieceIndex int, blockIndex int) { + ps.arr[pieceIndex].blocks.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()) } diff --git a/protocol/bt/torrent/torrent.go b/protocol/bt/torrent/torrent.go index 6e669348..a18713cf 100644 --- a/protocol/bt/torrent/torrent.go +++ b/protocol/bt/torrent/torrent.go @@ -4,7 +4,6 @@ import ( "fmt" "github.com/monkeyWie/gopeed/protocol/bt/metainfo" "github.com/monkeyWie/gopeed/protocol/bt/tracker" - "sync" "time" ) @@ -12,17 +11,17 @@ type Torrent struct { PeerID [20]byte MetaInfo *metainfo.MetaInfo - PiecesState *piecesState - peerPool *peerPool + pieces *pieces + peerPool *peerPool Path string } func NewTorrent(peerID [20]byte, metaInfo *metainfo.MetaInfo) *Torrent { torrent := &Torrent{ - PeerID: peerID, - MetaInfo: metaInfo, - PiecesState: NewPiecesState(len(metaInfo.Info.Pieces)), + PeerID: peerID, + MetaInfo: metaInfo, + pieces: newPiecesState(len(metaInfo.Info.Pieces)), } return torrent } @@ -32,10 +31,11 @@ func (t *Torrent) Download(path string) { t.peerPool = newPeerPool() t.fetchPeers() - taskCh := make(chan interface{}, 16) + taskCh := make(chan interface{}, 128) defer close(taskCh) for { - index := t.PiecesState.getReadyAndDownload() + // 获取一个待下载的piece + index := t.pieces.getReady() if index == -1 { return } @@ -43,28 +43,30 @@ func (t *Torrent) Download(path string) { go func(index int) { peer := t.peerPool.get() if peer == nil { + fmt.Println("no peer") time.Sleep(time.Second * 10) - t.PiecesState.setState(index, stateReady) + t.pieces.setState(index, stateReady) <-taskCh return } pc := NewPeerConn(t, peer) err := pc.ready() if err != nil { - t.PiecesState.setState(index, stateReady) - t.peerPool.remove(peer) + t.pieces.setState(index, stateReady) + t.peerPool.unavailable(peer) } else { fmt.Printf("peer ready:%s,download %d\n", peer.Address(), index) err = pc.downloadPiece(index) if err != nil { - fmt.Println(err) + fmt.Println("down piece error:" + err.Error()) // 下载失败 - t.PiecesState.setState(index, stateReady) + t.pieces.setState(index, stateReady) + t.peerPool.unavailable(peer) } else { // 下载完成 - t.PiecesState.setState(index, stateFinished) + t.pieces.setState(index, stateFinished) + t.peerPool.release(peer) } - t.peerPool.release(peer) } <-taskCh }(index) @@ -78,23 +80,12 @@ func (t *Torrent) fetchPeers() { MetaInfo: t.MetaInfo, } urls := tracker.MetaInfo.AnnounceList[0] - var wg sync.WaitGroup - ch := make(chan interface{}) - wg.Add(len(urls)) for _, url := range urls { go func(url string) { - defer wg.Done() peers, err := tracker.DoTracker(url) if err == nil { - fmt.Printf("Tracker end,url:%s,count:%d\n", url, len(peers)) - ch <- nil t.peerPool.put(peers) } }(url) } - go func() { - wg.Wait() - close(ch) - }() - <-ch } diff --git a/protocol/bt/tracker/tracker.go b/protocol/bt/tracker/tracker.go index 5dd6f114..df9417c7 100644 --- a/protocol/bt/tracker/tracker.go +++ b/protocol/bt/tracker/tracker.go @@ -28,7 +28,7 @@ const ( ) /* - Offset Size Name Value + Offset size Name Value 0 64-bit integer protocol_id 0x41727101980 // magic constant 8 32-bit integer action 0 // connect 12 32-bit integer transaction_id @@ -58,7 +58,7 @@ func (req *udpConnectRequest) encode() []byte { } /* - Offset Size Name Value + Offset size Name Value 0 32-bit integer action 0 // connect 4 32-bit integer transaction_id 8 64-bit integer connection_id @@ -79,7 +79,7 @@ func newUdpConnectResponse(buf []byte) *udpConnectResponse { } /* - Offset Size Name Value + Offset size Name Value 0 64-bit integer connection_id 8 32-bit integer action 1 // announce 12 32-bit integer transaction_id @@ -139,7 +139,7 @@ func (req *udpAnnounceRequest) encode() []byte { } /* - Offset Size Name Value + Offset size Name Value 0 32-bit integer action 1 // announce 4 32-bit integer transaction_id 8 32-bit integer interval