Missed piece availability code

diff --git a/torrent.go b/torrent.go
index c56ddbd181ecaf24d7969bca7db3d6d4b1a7d337..ae78625fcffae9f2bf0a1bd7bdbbc291ca5c2d19 100644 (file)
--- a/torrent.go
+++ b/torrent.go
@@ -147,7 +147,7 @@ type Torrent struct {
        pex pexState
 }
 
-func (t *Torrent) pieceAvailability(i pieceIndex) (count int) {
+func (t *Torrent) pieceAvailabilityFromPeers(i pieceIndex) (count int) {
        t.iterPeers(func(peer *Peer) {
                if peer.peerHasPiece(i) {
                        count++
@@ -156,6 +156,19 @@ func (t *Torrent) pieceAvailability(i pieceIndex) (count int) {
        return
 }
 
+func (t *Torrent) decPieceAvailability(i pieceIndex) {
+       p := t.piece(i)
+       if p.availability <= 0 {
+               panic(p.availability)
+       }
+       p.availability--
+}
+
+func (t *Torrent) incPieceAvailability(i pieceIndex) {
+       p := t.piece(i)
+       p.availability++
+}
+
 func (t *Torrent) numConns() int {
        return len(t.conns)
 }
@@ -422,6 +435,10 @@ func (t *Torrent) onSetInfo() {
                        // t.logger.Printf("piece %s completion unknown, queueing check", p)
                        t.queuePieceCheck(pieceIndex(i))
                }
+               if p.availability != 0 {
+                       panic(p.availability)
+               }
+               p.availability = int64(t.pieceAvailabilityFromPeers(i))
        }
        t.cl.event.Broadcast()
        t.gotMetainfo.Set()
@@ -535,6 +552,22 @@ func (t *Torrent) newMetadataExtensionMessage(c *PeerConn, msgType int, piece in
        }
 }
 
+type pieceAvailabilityRun struct {
+       availability int64
+       count        pieceIndex
+}
+
+func (t *Torrent) pieceAvailabilityRuns() (ret []pieceAvailabilityRun) {
+       rle := missinggo.NewRunLengthEncoder(func(el interface{}, count uint64) {
+               ret = append(ret, pieceAvailabilityRun{el.(int64), int(count)})
+       })
+       for _, p := range t.pieces {
+               rle.Append(p.availability, 1)
+       }
+       rle.Flush()
+       return
+}
+
 func (t *Torrent) pieceStateRuns() (ret PieceStateRuns) {
        rle := missinggo.NewRunLengthEncoder(func(el interface{}, count uint64) {
                ret = append(ret, PieceStateRun{
@@ -605,17 +638,21 @@ func (t *Torrent) writeStatus(w io.Writer) {
                }
                fmt.Fprintln(w)
        }
-       fmt.Fprintf(w, "Piece length: %s\n", func() string {
-               if t.haveInfo() {
-                       return fmt.Sprint(t.usualPieceSize())
-               } else {
-                       return "?"
-               }
-       }())
+       fmt.Fprintf(w, "Piece length: %s\n",
+               func() string {
+                       if t.haveInfo() {
+                               return fmt.Sprintf("%v (%v chunks)",
+                                       t.usualPieceSize(),
+                                       float64(t.usualPieceSize())/float64(t.chunkSize))
+                       } else {
+                               return "no info"
+                       }
+               }(),
+       )
        if t.info != nil {
                fmt.Fprintf(w, "Num Pieces: %d (%d completed)\n", t.numPieces(), t.numPiecesCompleted())
-               fmt.Fprintf(w, "Piece States: %s", t.pieceStateRuns())
-               fmt.Fprintln(w)
+               fmt.Fprintf(w, "Piece States: %s\n", t.pieceStateRuns())
+               fmt.Fprintf(w, "Piece availability: %v\n", t.pieceAvailabilityRuns())
        }
        fmt.Fprintf(w, "Reader Pieces:")
        t.forReaderOffsetPieces(func(begin, end pieceIndex) (again bool) {
@@ -888,11 +925,14 @@ func chunkIndex(cs ChunkSpec, chunkSize pp.Integer) int {
 }
 
 func (t *Torrent) wantPieceIndex(index pieceIndex) bool {
-       if !t.haveInfo() {
-               return false
-       }
-       if index < 0 || index >= t.numPieces() {
-               return false
+       // TODO: Are these overly conservative, should we be guarding this here?
+       {
+               if !t.haveInfo() {
+                       return false
+               }
+               if index < 0 || index >= t.numPieces() {
+                       return false
+               }
        }
        p := &t.pieces[index]
        if p.queuedForHash() {
@@ -1317,6 +1357,13 @@ func (t *Torrent) deletePeerConn(c *PeerConn) (ret bool) {
        return
 }
 
+func (t *Torrent) decPeerPieceAvailability(p *Peer) {
+       p.newPeerPieces().IterTyped(func(i int) bool {
+               p.t.decPieceAvailability(i)
+               return true
+       })
+}
+
 func (t *Torrent) numActivePeers() (num int) {
        t.iterPeers(func(*Peer) {
                num++
@@ -2154,6 +2201,7 @@ func (t *Torrent) addWebSeed(url string) {
                ws.onGotInfo(t.info)
        }
        t.webSeeds[url] = &ws.peer
+       ws.peer.onPeerHasAllPieces()
 }
 
 func (t *Torrent) peerIsActive(p *Peer) (active bool) {