Try to do smart ban block hashing without client lock

[btrtrc.git] / peer.go

diff --git a/peer.go b/peer.go
index 23bf5468f08bd1ede676c55a69ea1db75346a7ab..d3ea15161ab8261d5c0df48a6eddd4a05f9a2df0 100644 (file)
--- 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)
+       }
+}