Add a benchmark for readMainLoop piece decoding

connection_test.go

@@ -5,12 +5,19 @@ import (
 	"io"
 	"io/ioutil"
 	"net"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/anacrolix/missinggo/bitmap"
+	"github.com/anacrolix/missinggo/pubsub"
+	"github.com/bradfitz/iter"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/anacrolix/torrent/metainfo"
+	"github.com/anacrolix/torrent/peer_protocol"
+	"github.com/anacrolix/torrent/storage"
 )
 
 func TestCancelRequestOptimized(t *testing.T) {
@@ -92,3 +99,92 @@ func TestSendBitfieldThenHave(t *testing.T) {
 	// arrive in the following Have message.
 	require.EqualValues(t, "\x00\x00\x00\x02\x05@\x00\x00\x00\x05\x04\x00\x00\x00\x02", string(b))
 }
+
+type torrentStorage struct {
+	writeSem sync.Mutex
+}
+
+func (me *torrentStorage) Close() error { return nil }
+
+func (me *torrentStorage) Piece(mp metainfo.Piece) storage.PieceImpl {
+	return me
+}
+
+func (me *torrentStorage) GetIsComplete() bool {
+	return false
+}
+
+func (me *torrentStorage) MarkComplete() error {
+	return nil
+}
+
+func (me *torrentStorage) MarkNotComplete() error {
+	return nil
+}
+
+func (me *torrentStorage) ReadAt([]byte, int64) (int, error) {
+	panic("shouldn't be called")
+}
+
+func (me *torrentStorage) WriteAt(b []byte, _ int64) (int, error) {
+	if len(b) != defaultChunkSize {
+		panic(len(b))
+	}
+	me.writeSem.Unlock()
+	return len(b), nil
+}
+
+func BenchmarkConnectionMainReadLoop(b *testing.B) {
+	cl := &Client{}
+	ts := &torrentStorage{}
+	t := &Torrent{
+		cl: cl,
+		info: &metainfo.Info{
+			Pieces:      make([]byte, 20),
+			Length:      1 << 20,
+			PieceLength: 1 << 20,
+		},
+		chunkSize:         defaultChunkSize,
+		storage:           &storage.Torrent{ts},
+		pieceStateChanges: pubsub.NewPubSub(),
+	}
+	t.makePieces()
+	t.pendingPieces.Add(0)
+	r, w := io.Pipe()
+	cn := &connection{
+		t: t,
+		rw: struct {
+			io.Reader
+			io.Writer
+		}{r, nil},
+	}
+	mrlErr := make(chan error)
+	cl.mu.Lock()
+	go func() {
+		err := cn.mainReadLoop()
+		if err != nil {
+			mrlErr <- err
+		}
+		close(mrlErr)
+	}()
+	msg := peer_protocol.Message{
+		Type:  peer_protocol.Piece,
+		Piece: make([]byte, defaultChunkSize),
+	}
+	wb, err := msg.MarshalBinary()
+	require.NoError(b, err)
+	b.SetBytes(int64(len(msg.Piece)))
+	ts.writeSem.Lock()
+	for range iter.N(b.N) {
+		cl.mu.Lock()
+		t.pieces[0].DirtyChunks.Clear()
+		cl.mu.Unlock()
+		n, err := w.Write(wb)
+		require.NoError(b, err)
+		require.EqualValues(b, len(wb), n)
+		ts.writeSem.Lock()
+	}
+	w.Close()
+	require.NoError(b, <-mrlErr)
+	require.EqualValues(b, b.N, cn.UsefulChunksReceived)
+}