nprasad2077/gopeed
1
0 Fork 0
mirror of https://github.com/GopeedLab/gopeed.git synced 2026-11-09 18:10:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

block下载记录支持

Browse Source
This commit is contained in:
liwei committed 2020-02-26 18:53:39 +08:00
1 parent 8298051c81
commit 7bb3e38b21
7 files changed
+140 -124
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

protocol/bt/client.go
+3 -3
View File
@@ -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
}
protocol/bt/torrent/peer_conn.go
+41 -32
View File
@@ -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)
}
protocol/bt/torrent/peer_pool.go
+25 -39
View File
@@ -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
}
protocol/bt/torrent/peer_pool_test.go
+1 -1
View File
@@ -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())
}
protocol/bt/torrent/piece_state.go
+49 -19
View File
@@ -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())
}
protocol/bt/torrent/torrent.go
+17 -26
View File
@@ -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
}
protocol/bt/tracker/tracker.go
+4 -4
View File
@@ -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