From: Matt Joiner <anacrolix@gmail.com>
Date: Sun, 18 Feb 2024 01:18:54 +0000 (+1100)
Subject: Include smart ban block recording in BenchmarkConnectionMainReadLoop
X-Git-Tag: v1.54.1^2~4
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=a21ae78802252549d5d7328140608931ca14b8bf;p=btrtrc.git

Include smart ban block recording in BenchmarkConnectionMainReadLoop
---

diff --git a/peerconn_test.go b/peerconn_test.go
index 3fdcbff1..8a640544 100644
--- a/peerconn_test.go
+++ b/peerconn_test.go
@@ -4,8 +4,10 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"golang.org/x/sync/errgroup"
 	"io"
 	"net"
+	"net/netip"
 	"sync"
 	"testing"
 
@@ -107,29 +109,32 @@ func BenchmarkConnectionMainReadLoop(b *testing.B) {
 	t.onSetInfo()
 	t._pendingPieces.Add(0)
 	r, w := net.Pipe()
+	c.Logf("pipe reader remote addr: %v", r.RemoteAddr())
 	cn := cl.newConnection(r, newConnectionOpts{
-		outgoing:   true,
-		remoteAddr: r.RemoteAddr(),
+		outgoing: true,
+		// TODO: This is a hack to give the pipe a bannable remote address.
+		remoteAddr: netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 2, 3, 4}), 1234),
 		network:    r.RemoteAddr().Network(),
 		connString: regularNetConnPeerConnConnString(r),
 	})
+	c.Assert(cn.bannableAddr.Ok, qt.IsTrue)
 	cn.setTorrent(t)
-	mrlErrChan := make(chan error)
 	msg := pp.Message{
 		Type:  pp.Piece,
 		Piece: make([]byte, defaultChunkSize),
 	}
-	go func() {
+	var errGroup errgroup.Group
+	errGroup.Go(func() error {
 		cl.lock()
 		err := cn.mainReadLoop()
-		if err != nil {
-			mrlErrChan <- err
+		if errors.Is(err, io.EOF) {
+			err = nil
 		}
-		close(mrlErrChan)
-	}()
+		return err
+	})
 	wb := msg.MustMarshalBinary()
 	b.SetBytes(int64(len(msg.Piece)))
-	go func() {
+	errGroup.Go(func() error {
 		ts.writeSem.Lock()
 		for i := 0; i < b.N; i += 1 {
 			cl.lock()
@@ -143,17 +148,19 @@ func BenchmarkConnectionMainReadLoop(b *testing.B) {
 			n, err := w.Write(wb)
 			require.NoError(b, err)
 			require.EqualValues(b, len(wb), n)
+			// This is unlocked by a successful write to storage. So this unblocks when that is
+			// done.
 			ts.writeSem.Lock()
 		}
 		if err := w.Close(); err != nil {
 			panic(err)
 		}
-	}()
-	mrlErr := <-mrlErrChan
-	if mrlErr != nil && !errors.Is(mrlErr, io.EOF) {
-		c.Fatal(mrlErr)
-	}
+		return nil
+	})
+	err := errGroup.Wait()
+	c.Assert(err, qt.IsNil)
 	c.Assert(cn._stats.ChunksReadUseful.Int64(), quicktest.Equals, int64(b.N))
+	c.Assert(t.smartBanCache.HasBlocks(), qt.IsTrue)
 }
 
 func TestConnPexPeerFlags(t *testing.T) {
diff --git a/smartban/smartban.go b/smartban/smartban.go
index 5a9050eb..ba568c98 100644
--- a/smartban/smartban.go
+++ b/smartban/smartban.go
@@ -52,3 +52,9 @@ func (me *Cache[Peer, BlockKey, Hash]) ForgetBlock(key BlockKey) {
 	defer me.lock.Unlock()
 	delete(me.blocks, key)
 }
+
+func (me *Cache[Peer, BlockKey, Hash]) HasBlocks() bool {
+	me.lock.RLock()
+	defer me.lock.RUnlock()
+	return len(me.blocks) != 0
+}