diff --git a/peer.go b/peer.go index 23bf5468f0..d3ea15161a 100644 --- a/peer.go +++ b/peer.go @@ -6,6 +6,7 @@ import ( "io" "net" "strings" + "sync" "time" "github.com/RoaringBitmap/roaring" @@ -600,9 +601,13 @@ func (c *Peer) receiveChunk(msg *pp.Message) error { } req := c.t.requestIndexFromRequest(ppReq) - if c.bannableAddr.Ok { - t.smartBanCache.RecordBlock(c.bannableAddr.Value, req, msg.Piece) - } + recordBlockForSmartBan := sync.OnceFunc(func() { + c.recordBlockForSmartBan(req, msg.Piece) + }) + // This needs to occur before we return, but we try to do it when the client is unlocked. It + // can't be done before checking if chunks are valid because they won't be deallocated by piece + // hashing if they're out of bounds. + defer recordBlockForSmartBan() if c.peerChoking { chunksReceived.Add("while choked", 1) @@ -683,6 +688,8 @@ func (c *Peer) receiveChunk(msg *pp.Message) error { err = func() error { cl.unlock() defer cl.lock() + // Opportunistically do this here while we aren't holding the client lock. + recordBlockForSmartBan() concurrentChunkWrites.Add(1) defer concurrentChunkWrites.Add(-1) // Write the chunk out. Note that the upper bound on chunk writing concurrency will be the @@ -875,3 +882,9 @@ func (p *Peer) decPeakRequests() { // } p.peakRequests-- } + +func (p *Peer) recordBlockForSmartBan(req RequestIndex, blockData []byte) { + if p.bannableAddr.Ok { + p.t.smartBanCache.RecordBlock(p.bannableAddr.Value, req, blockData) + } +}