From f23d0975024dff2d27fce389ede380021d3a93b6 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Sat, 3 Feb 2018 12:08:16 +1100
Subject: [PATCH] Include closed connections in dirtiers

---
 connection.go | 17 +++++++++++++----
 piece.go      |  4 ++++
 torrent.go    | 11 +++++------
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/connection.go b/connection.go
index da18c5bc..70afb2a8 100644
--- a/connection.go
+++ b/connection.go
@@ -1138,15 +1138,24 @@ func (c *connection) receiveChunk(msg *pp.Message) {
 		t.pendAllChunkSpecs(index)
 	}
 
-	if c.peerTouchedPieces == nil {
-		c.peerTouchedPieces = make(map[int]struct{})
-	}
-	c.peerTouchedPieces[index] = struct{}{}
+	c.onDirtiedPiece(index)
 
 	cl.event.Broadcast()
 	t.publishPieceChange(int(req.Index))
 }
 
+func (c *connection) onDirtiedPiece(piece int) {
+	if c.peerTouchedPieces == nil {
+		c.peerTouchedPieces = make(map[int]struct{})
+	}
+	c.peerTouchedPieces[piece] = struct{}{}
+	ds := &c.t.pieces[piece].dirtiers
+	if *ds == nil {
+		*ds = make(map[*connection]struct{})
+	}
+	(*ds)[c] = struct{}{}
+}
+
 func (c *connection) uploadAllowed() bool {
 	if c.t.cl.config.NoUpload {
 		return false
diff --git a/piece.go b/piece.go
index 8e89ba62..1f094813 100644
--- a/piece.go
+++ b/piece.go
@@ -59,6 +59,10 @@ type Piece struct {
 	pendingWritesMutex sync.Mutex
 	pendingWrites      int
 	noPendingWrites    sync.Cond
+
+	// Connections that have written data to this piece since its last check.
+	// This can include connections that have closed.
+	dirtiers map[*connection]struct{}
 }
 
 func (p *Piece) String() string {
diff --git a/torrent.go b/torrent.go
index a58d6b8f..2846d275 100644
--- a/torrent.go
+++ b/torrent.go
@@ -1623,15 +1623,14 @@ func (t *Torrent) verifyPiece(piece int) {
 	t.publishPieceChange(piece)
 }
 
-// Return the connections that touched a piece, and clear the entry while
+// Return the connections that touched a piece, and clear the entries while
 // doing it.
 func (t *Torrent) reapPieceTouchers(piece int) (ret []*connection) {
-	for c := range t.conns {
-		if _, ok := c.peerTouchedPieces[piece]; ok {
-			ret = append(ret, c)
-			delete(c.peerTouchedPieces, piece)
-		}
+	for c := range t.pieces[piece].dirtiers {
+		delete(c.peerTouchedPieces, piece)
+		ret = append(ret, c)
 	}
+	t.pieces[piece].dirtiers = nil
 	return
 }
 
-- 
2.50.0