Support sending HaveAll

This commit is contained in:
Matt Joiner 2018-02-04 19:10:25 +11:00
parent ec7b90db45
commit e0c2265cb0
4 changed files with 47 additions and 30 deletions

View File

@ -893,13 +893,20 @@ func (cl *Client) sendInitialMessages(conn *connection, torrent *Torrent) {
}(), }(),
}) })
} }
if torrent.haveAnyPieces() { func() {
conn.Bitfield(torrent.bitfield()) if conn.fastEnabled() {
} else if cl.extensionBytes.SupportsFast() && conn.PeerExtensionBytes.SupportsFast() { if torrent.haveAllPieces() {
conn.Post(pp.Message{ conn.Post(pp.Message{Type: pp.HaveAll})
Type: pp.HaveNone, conn.sentHaves.AddRange(0, conn.t.NumPieces())
}) return
} else if !torrent.haveAnyPieces() {
conn.Post(pp.Message{Type: pp.HaveNone})
conn.sentHaves.Clear()
return
} }
}
conn.PostBitfield()
}()
if conn.PeerExtensionBytes.SupportsDHT() && cl.extensionBytes.SupportsDHT() && cl.dHT != nil { if conn.PeerExtensionBytes.SupportsDHT() && cl.extensionBytes.SupportsDHT() && cl.dHT != nil {
conn.Post(pp.Message{ conn.Post(pp.Message{
Type: pp.Port, Type: pp.Port,
@ -1101,7 +1108,7 @@ func (cl *Client) allTorrentsCompleted() bool {
if !t.haveInfo() { if !t.haveInfo() {
return false return false
} }
if t.numPiecesCompleted() != t.numPieces() { if !t.haveAllPieces() {
return false return false
} }
} }

View File

@ -66,7 +66,7 @@ type connection struct {
// Indexed by metadata piece, set to true if posted and pending a // Indexed by metadata piece, set to true if posted and pending a
// response. // response.
metadataRequests []bool metadataRequests []bool
sentHaves []bool sentHaves bitmap.Bitmap
// Stuff controlled by the remote peer. // Stuff controlled by the remote peer.
PeerID PeerID PeerID PeerID
@ -141,7 +141,7 @@ func (cn *connection) completedString() string {
// invalid, such as by receiving badly sized BITFIELD, or invalid HAVE // invalid, such as by receiving badly sized BITFIELD, or invalid HAVE
// messages. // messages.
func (cn *connection) setNumPieces(num int) error { func (cn *connection) setNumPieces(num int) error {
cn.peerPieces.RemoveRange(num, -1) cn.peerPieces.RemoveRange(num, bitmap.ToEnd)
cn.peerPiecesChanged() cn.peerPiecesChanged()
return nil return nil
} }
@ -468,31 +468,28 @@ func (cn *connection) writer(keepAliveTimeout time.Duration) {
} }
func (cn *connection) Have(piece int) { func (cn *connection) Have(piece int) {
for piece >= len(cn.sentHaves) { if cn.sentHaves.Get(piece) {
cn.sentHaves = append(cn.sentHaves, false)
}
if cn.sentHaves[piece] {
return return
} }
cn.Post(pp.Message{ cn.Post(pp.Message{
Type: pp.Have, Type: pp.Have,
Index: pp.Integer(piece), Index: pp.Integer(piece),
}) })
cn.sentHaves[piece] = true cn.sentHaves.Add(piece)
} }
func (cn *connection) Bitfield(haves []bool) { func (cn *connection) PostBitfield() {
if cn.sentHaves != nil { if cn.sentHaves.Len() != 0 {
panic("bitfield must be first have-related message sent") panic("bitfield must be first have-related message sent")
} }
if !cn.t.haveAnyPieces() {
return
}
cn.Post(pp.Message{ cn.Post(pp.Message{
Type: pp.Bitfield, Type: pp.Bitfield,
Bitfield: haves, Bitfield: cn.t.bitfield(),
}) })
// Make a copy of haves, as that's read when the message is marshalled cn.sentHaves = cn.t.completedPieces.Copy()
// without the lock. Also it obviously shouldn't change in the Msg due to
// changes in .sentHaves.
cn.sentHaves = append([]bool(nil), haves...)
} }
// Determines interest and requests to send to a connected peer. // Determines interest and requests to send to a connected peer.
@ -553,8 +550,11 @@ func iterBitmapsDistinct(skip bitmap.Bitmap, bms ...bitmap.Bitmap) iter.Func {
func (cn *connection) unbiasedPieceRequestOrder() iter.Func { func (cn *connection) unbiasedPieceRequestOrder() iter.Func {
now, readahead := cn.t.readerPiecePriorities() now, readahead := cn.t.readerPiecePriorities()
var skip bitmap.Bitmap
if !cn.peerSentHaveAll {
// Pieces to skip include pieces the peer doesn't have // Pieces to skip include pieces the peer doesn't have
skip := bitmap.Flip(cn.peerPieces, 0, cn.t.numPieces()) skip = bitmap.Flip(cn.peerPieces, 0, cn.t.numPieces())
}
// And pieces that we already have. // And pieces that we already have.
skip.Union(cn.t.completedPieces) skip.Union(cn.t.completedPieces)
// Return an iterator over the different priority classes, minus the skip // Return an iterator over the different priority classes, minus the skip

View File

@ -24,11 +24,15 @@ func TestSendBitfieldThenHave(t *testing.T) {
cl.initLogger() cl.initLogger()
c := cl.newConnection(nil) c := cl.newConnection(nil)
c.setTorrent(cl.newTorrent(metainfo.Hash{}, nil)) c.setTorrent(cl.newTorrent(metainfo.Hash{}, nil))
c.t.setInfo(&metainfo.Info{
Pieces: make([]byte, metainfo.HashSize*3),
})
c.r = r c.r = r
c.w = w c.w = w
go c.writer(time.Minute) go c.writer(time.Minute)
c.mu().Lock() c.mu().Lock()
c.Bitfield([]bool{false, true, false}) c.t.completedPieces.Add(1)
c.PostBitfield( /*[]bool{false, true, false}*/ )
c.mu().Unlock() c.mu().Unlock()
c.mu().Lock() c.mu().Lock()
c.Have(2) c.Have(2)

View File

@ -735,6 +735,10 @@ type Peer struct {
} }
func (t *Torrent) pieceLength(piece int) pp.Integer { func (t *Torrent) pieceLength(piece int) pp.Integer {
if t.info.PieceLength == 0 {
// There will be no variance amongst pieces. Only pain.
return 0
}
if piece == t.numPieces()-1 { if piece == t.numPieces()-1 {
ret := pp.Integer(*t.length % t.info.PieceLength) ret := pp.Integer(*t.length % t.info.PieceLength)
if ret != 0 { if ret != 0 {
@ -762,13 +766,15 @@ func (t *Torrent) hashPiece(piece int) (ret metainfo.Hash) {
} }
func (t *Torrent) haveAnyPieces() bool { func (t *Torrent) haveAnyPieces() bool {
for i := range t.pieces { return t.completedPieces.Len() != 0
if t.pieceComplete(i) {
return true
}
} }
func (t *Torrent) haveAllPieces() bool {
if !t.haveInfo() {
return false return false
} }
return t.completedPieces.Len() == t.numPieces()
}
func (t *Torrent) havePiece(index int) bool { func (t *Torrent) havePiece(index int) bool {
return t.haveInfo() && t.pieceComplete(index) return t.haveInfo() && t.pieceComplete(index)