Add smart ban hash benchmarks and reduce allocations

diff --git a/go.mod b/go.mod
index 68120007ebcff1b0ef08b38faef530fbdd64fd30..1b0f5816fee11a28622fffb7e94d20305efc2bf7 100644 (file)
--- a/go.mod
+++ b/go.mod
@@ -25,6 +25,7 @@ require (
        github.com/anacrolix/utp v0.1.0
        github.com/bahlo/generic-list-go v0.2.0
        github.com/bradfitz/iter v0.0.0-20191230175014-e8f45d346db8
+       github.com/cespare/xxhash v1.1.0
        github.com/davecgh/go-spew v1.1.1
        github.com/dustin/go-humanize v1.0.0
        github.com/edsrzf/mmap-go v1.1.0

diff --git a/go.sum b/go.sum
index ecf851997883f090e7d8e17ce9bfe8d4672e235a..c8ce30f99d6495b690a4055859456b85ea883c53 100644 (file)
--- a/go.sum
+++ b/go.sum
@@ -142,6 +142,7 @@ github.com/bradfitz/iter v0.0.0-20191230175014-e8f45d346db8/go.mod h1:spo1JLcs67
 github.com/cenkalti/backoff/v4 v4.1.3 h1:cFAlzYUlVYDysBEH2T5hyJZMh3+5+WCBvSnK6Q8UtC4=
 github.com/cenkalti/backoff/v4 v4.1.3/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
+github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cespare/xxhash/v2 v2.1.2 h1:YRXhKfTDauu4ajMg1TPgFO5jnlC2HCbmLXMcTG5cbYE=

diff --git a/smartban/smartban.go b/smartban/smartban.go
index 96e9b759a5d17aa1eba9f95530e65862237f2570..5a9050ebcb6cbb44d1cc1d5141a89a1ff2b55d62 100644 (file)
--- a/smartban/smartban.go
+++ b/smartban/smartban.go
@@ -1,6 +1,7 @@
 package smartban
 
 import (
+       g "github.com/anacrolix/generics"
        "sync"
 )
 
@@ -8,7 +9,7 @@ type Cache[Peer, BlockKey, Hash comparable] struct {
        Hash func([]byte) Hash
 
        lock   sync.RWMutex
-       blocks map[BlockKey]map[Peer]Hash
+       blocks map[BlockKey][]peerAndHash[Peer, Hash]
 }
 
 type Block[Key any] struct {
@@ -16,8 +17,13 @@ type Block[Key any] struct {
        Data []byte
 }
 
+type peerAndHash[Peer, Hash any] struct {
+       Peer Peer
+       Hash Hash
+}
+
 func (me *Cache[Peer, BlockKey, Hash]) Init() {
-       me.blocks = make(map[BlockKey]map[Peer]Hash)
+       g.MakeMap(&me.blocks)
 }
 
 func (me *Cache[Peer, BlockKey, Hash]) RecordBlock(peer Peer, key BlockKey, data []byte) {
@@ -25,20 +31,17 @@ func (me *Cache[Peer, BlockKey, Hash]) RecordBlock(peer Peer, key BlockKey, data
        me.lock.Lock()
        defer me.lock.Unlock()
        peers := me.blocks[key]
-       if peers == nil {
-               peers = make(map[Peer]Hash)
-               me.blocks[key] = peers
-       }
-       peers[peer] = hash
+       peers = append(peers, peerAndHash[Peer, Hash]{peer, hash})
+       me.blocks[key] = peers
 }
 
 func (me *Cache[Peer, BlockKey, Hash]) CheckBlock(key BlockKey, data []byte) (bad []Peer) {
        correct := me.Hash(data)
        me.lock.RLock()
        defer me.lock.RUnlock()
-       for peer, hash := range me.blocks[key] {
-               if hash != correct {
-                       bad = append(bad, peer)
+       for _, item := range me.blocks[key] {
+               if item.Hash != correct {
+                       bad = append(bad, item.Peer)
                }
        }
        return

diff --git a/smartban_test.go b/smartban_test.go
new file mode 100644 (file)
index 0000000..2947f52
--- /dev/null
+++ b/smartban_test.go
@@ -0,0 +1,39 @@
+package torrent
+
+import (
+       "crypto/sha1"
+       "github.com/anacrolix/missinggo/v2/iter"
+       "github.com/anacrolix/torrent/smartban"
+       "github.com/cespare/xxhash"
+       "net/netip"
+       "testing"
+)
+
+func benchmarkSmartBanRecordBlock[Sum comparable](b *testing.B, hash func([]byte) Sum) {
+       var cache smartban.Cache[bannableAddr, RequestIndex, Sum]
+       cache.Hash = hash
+       cache.Init()
+       var data [defaultChunkSize]byte
+       var addr netip.Addr
+       b.SetBytes(int64(len(data)))
+       for i := range iter.N(b.N) {
+               cache.RecordBlock(addr, RequestIndex(i), data[:])
+       }
+}
+
+func BenchmarkSmartBanRecordBlock(b *testing.B) {
+       b.Run("xxHash", func(b *testing.B) {
+               var salt [8]byte
+               benchmarkSmartBanRecordBlock(b, func(block []byte) uint64 {
+                       h := xxhash.New()
+                       // xxHash is not cryptographic, and so we're salting it so attackers can't know a priori
+                       // where block data collisions are.
+                       h.Write(salt[:])
+                       h.Write(block)
+                       return h.Sum64()
+               })
+       })
+       b.Run("Sha1", func(b *testing.B) {
+               benchmarkSmartBanRecordBlock(b, sha1.Sum)
+       })
+}