client.go | 1 + connection.go | 2 +- torrent.go | 81 +++++++++++++++++++++++++---------------------------- diff --git a/client.go b/client.go index 8c356eea324b2dc1981eb442f30d4d079d739511..c457bef6c85d4db12afef1f0cb5df88a4296a769 100644 --- a/client.go +++ b/client.go @@ -1147,6 +1147,7 @@ t = &Torrent{ cl: cl, infoHash: ih, peers: make(map[peersKey]Peer), + conns: make(map[*connection]struct{}, 2*defaultEstablishedConnsPerTorrent), halfOpen: make(map[string]struct{}), pieceStateChanges: pubsub.NewPubSub(), diff --git a/connection.go b/connection.go index 06cb7ca8fb3ba312fc1ddc828ab0603daead393a..3ba41076fbe3ba895adc0a754e08dd6b84c04a4b 100644 --- a/connection.go +++ b/connection.go @@ -964,7 +964,7 @@ // Record that we have the chunk. piece.unpendChunkIndex(chunkIndex(req.chunkSpec, t.chunkSize)) // Cancel pending requests for this chunk. - for _, c := range t.conns { + for c := range t.conns { if cl.connCancel(t, c, req) { c.updateRequests() } diff --git a/torrent.go b/torrent.go index c2baf9c10dd5a0eeaf2421d8488776f53122592c..f6f723bf42b74a5eb553f442ec4161a77830ff49 100644 --- a/torrent.go +++ b/torrent.go @@ -66,7 +66,7 @@ // The info dict. nil if we don't have it (yet). info *metainfo.Info // Active peer connections, running message stream loops. - conns []*connection + conns map[*connection]struct{} maxEstablishedConns int // Set of addrs to which we're attempting to connect. Connections are // half-open until all handshakes are completed. @@ -140,7 +140,7 @@ return t.pieces[piece].Storage().GetIsComplete() } func (t *Torrent) numConnsUnchoked() (num int) { - for _, c := range t.conns { + for c := range t.conns { if !c.PeerChoked { num++ } @@ -153,7 +153,7 @@ func (t *Torrent) addrActive(addr string) bool { if _, ok := t.halfOpen[addr]; ok { return true } - for _, c := range t.conns { + for c := range t.conns { if c.remoteAddr().String() == addr { return true } @@ -163,7 +163,7 @@ } func (t *Torrent) worstUnclosedConns() (ret []*connection) { ret = make([]*connection, 0, len(t.conns)) - for _, c := range t.conns { + for c := range t.conns { if !c.closed.IsSet() { ret = append(ret, c) } @@ -278,7 +278,7 @@ } t.metadataBytes = b t.metadataCompletedChunks = nil t.makePieces() - for _, conn := range t.conns { + for conn := range t.conns { if err := conn.setNumPieces(t.numPieces()); err != nil { log.Printf("closing connection: %s", err) conn.Close() @@ -325,7 +325,7 @@ return } t.metadataBytes = make([]byte, bytes) t.metadataCompletedChunks = make([]bool, (bytes+(1<<14)-1)/(1<<14)) - for _, c := range t.conns { + for c := range t.conns { c.requestPendingMetadata() } return @@ -477,8 +477,9 @@ fmt.Fprintf(w, "Pending peers: %d\n", len(t.peers)) fmt.Fprintf(w, "Half open: %d\n", len(t.halfOpen)) fmt.Fprintf(w, "Active peers: %d\n", len(t.conns)) - slices.Sort(t.conns, worseConn) - for i, c := range t.conns { + conns := t.connsAsSlice() + slices.Sort(conns, worseConn) + for i, c := range conns { fmt.Fprintf(w, "%2d. ", i+1) c.WriteStatus(w, t) } @@ -558,7 +559,7 @@ t.closed.Set() if t.storage != nil { t.storage.Close() } - for _, conn := range t.conns { + for conn := range t.conns { conn.Close() } t.pieceStateChanges.Close() @@ -873,7 +874,7 @@ t.openNewConns() } func (t *Torrent) piecePriorityChanged(piece int) { - for _, c := range t.conns { + for c := range t.conns { c.updatePiecePriority(piece) } t.maybeNewConns() @@ -1004,13 +1005,12 @@ ci := chunkIndex(req.chunkSpec, t.chunkSize) t.pieces[req.Index].pendChunkIndex(ci) } -func (t *Torrent) pieceChanged(piece int) { - correct := t.pieceComplete(piece) - defer t.cl.event.Broadcast() - if correct { - t.onCompletedPiece(piece) +func (t *Torrent) pieceCompletionChanged(piece int) { + t.cl.event.Broadcast() + if t.pieceComplete(piece) { + t.onPieceCompleted(piece) } else { - t.onFailedPiece(piece) + t.onIncompletePiece(piece) } t.updatePiecePriority(piece) } @@ -1039,7 +1039,7 @@ pcu := t.pieceCompleteUncached(piece) changed := t.completedPieces.Get(piece) != pcu t.completedPieces.Set(piece, pcu) if changed { - t.pieceChanged(piece) + t.pieceCompletionChanged(piece) } } @@ -1160,19 +1160,10 @@ return t.setInfoBytes(b) } // Returns true if connection is removed from torrent.Conns. -func (t *Torrent) deleteConnection(c *connection) bool { - for i0, _c := range t.conns { - if _c != c { - continue - } - i1 := len(t.conns) - 1 - if i0 != i1 { - t.conns[i0] = t.conns[i1] - } - t.conns = t.conns[:i1] - return true - } - return false +func (t *Torrent) deleteConnection(c *connection) (ret bool) { + _, ret = t.conns[c] + delete(t.conns, c) + return } func (t *Torrent) dropConnection(c *connection) { @@ -1356,7 +1347,7 @@ } if !t.wantConns() { return false } - for _, c0 := range t.conns { + for c0 := range t.conns { if c.PeerID == c0.PeerID { // Already connected to a client with that ID. duplicateClientConns.Add(1) @@ -1377,7 +1368,6 @@ } if len(t.conns) >= t.maxEstablishedConns { panic(len(t.conns)) } - t.conns = append(t.conns, c) if c.t != nil { panic("connection already associated with a torrent") } @@ -1386,6 +1376,7 @@ // torrent. t.stats.wroteBytes(c.stats.BytesWritten) t.stats.readBytes(c.stats.BytesRead) c.t = t + t.conns[c] = struct{}{} return true } @@ -1407,7 +1398,7 @@ t.cl.mu.Lock() defer t.cl.mu.Unlock() oldMax = t.maxEstablishedConns t.maxEstablishedConns = max - wcs := slices.HeapInterface(append([]*connection(nil), t.conns...), worseConn) + wcs := slices.HeapInterface(slices.FromMapKeys(t.conns), worseConn) for len(t.conns) > t.maxEstablishedConns && wcs.Len() > 0 { t.dropConnection(wcs.Pop().(*connection)) } @@ -1435,7 +1426,7 @@ pieceHashedNotCorrect.Add(1) } } p.EverHashed = true - touchers := t.reapPieceTouches(piece) + touchers := t.reapPieceTouchers(piece) if correct { for _, c := range touchers { c.goodPiecesDirtied++ @@ -1453,13 +1444,12 @@ t.cl.banPeerIP(missinggo.AddrIP(c.remoteAddr())) t.dropConnection(c) } } - t.pieceChanged(piece) } -func (t *Torrent) onCompletedPiece(piece int) { +func (t *Torrent) onPieceCompleted(piece int) { t.pendingPieces.Remove(piece) t.pendAllChunkSpecs(piece) - for _, conn := range t.conns { + for conn := range t.conns { conn.Have(piece) for r := range conn.Requests { if int(r.Index) == piece { @@ -1472,16 +1462,14 @@ conn.upload() } } -func (t *Torrent) onFailedPiece(piece int) { - cl := t.cl +func (t *Torrent) onIncompletePiece(piece int) { if t.pieceAllDirty(piece) { t.pendAllChunkSpecs(piece) } if !t.wantPieceIndex(piece) { return } - cl.openNewConns(t) - for _, conn := range t.conns { + for conn := range t.conns { if conn.PeerHasPiece(piece) { conn.updateRequests() } @@ -1512,8 +1500,8 @@ } // Return the connections that touched a piece, and clear the entry while // doing it. -func (t *Torrent) reapPieceTouches(piece int) (ret []*connection) { - for _, c := range t.conns { +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) @@ -1521,3 +1509,10 @@ } } return } + +func (t *Torrent) connsAsSlice() (ret []*connection) { + for c := range t.conns { + ret = append(ret, c) + } + return +}