Support sending HaveAll

diff --git a/client.go b/client.go
index d2dac0fa45949807a2cc93863e297649a9a4fd75..81a7936b3b48dda103450140f0dde68b378390af 100644 (file)
--- a/client.go
+++ b/client.go
@@ -893,13 +893,20 @@ func (cl *Client) sendInitialMessages(conn *connection, torrent *Torrent) {
                        }(),
                })
        }
-       if torrent.haveAnyPieces() {
-               conn.Bitfield(torrent.bitfield())
-       } else if cl.extensionBytes.SupportsFast() && conn.PeerExtensionBytes.SupportsFast() {
-               conn.Post(pp.Message{
-                       Type: pp.HaveNone,
-               })
-       }
+       func() {
+               if conn.fastEnabled() {
+                       if torrent.haveAllPieces() {
+                               conn.Post(pp.Message{Type: pp.HaveAll})
+                               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 {
                conn.Post(pp.Message{
                        Type: pp.Port,
@@ -1101,7 +1108,7 @@ func (cl *Client) allTorrentsCompleted() bool {
                if !t.haveInfo() {
                        return false
                }
-               if t.numPiecesCompleted() != t.numPieces() {
+               if !t.haveAllPieces() {
                        return false
                }
        }

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

diff --git a/connection_test.go b/connection_test.go
index 4d4c7465ca4ba739fd3d1be2c30a09d94d880581..6b4ccd6b349df5e968d00821f3a7dd190ca58a2d 100644 (file)
--- a/connection_test.go
+++ b/connection_test.go
@@ -24,11 +24,15 @@ func TestSendBitfieldThenHave(t *testing.T) {
        cl.initLogger()
        c := cl.newConnection(nil)
        c.setTorrent(cl.newTorrent(metainfo.Hash{}, nil))
+       c.t.setInfo(&metainfo.Info{
+               Pieces: make([]byte, metainfo.HashSize*3),
+       })
        c.r = r
        c.w = w
        go c.writer(time.Minute)
        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().Lock()
        c.Have(2)

diff --git a/torrent.go b/torrent.go
index 12bc48f559d6d4890bda12f6244c77587b2c8eb2..fca9179899f53fbf73b47bd4e04ee8550ee49e79 100644 (file)
--- a/torrent.go
+++ b/torrent.go
@@ -735,6 +735,10 @@ type Peer struct {
 }
 
 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 {
                ret := pp.Integer(*t.length % t.info.PieceLength)
                if ret != 0 {
@@ -762,12 +766,14 @@ func (t *Torrent) hashPiece(piece int) (ret metainfo.Hash) {
 }
 
 func (t *Torrent) haveAnyPieces() bool {
-       for i := range t.pieces {
-               if t.pieceComplete(i) {
-                       return true
-               }
+       return t.completedPieces.Len() != 0
+}
+
+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 {