Optimize memory usage by avoiding intermediate buffer in message serialization

peer-conn-msg-writer.go

@@ -117,10 +117,21 @@ func (cn *peerConnMsgWriter) run(keepAliveTimeout time.Duration) {
 	}
 }
 
+func (cn *peerConnMsgWriter) writeToBuffer(msg pp.Message) (err error) {
+	originalLen := cn.writeBuffer.Len()
+	defer func() {
+		if err != nil {
+			// Since an error occurred during buffer write, revert buffer to its original state before the write.
+			cn.writeBuffer.Truncate(originalLen)
+		}
+	}()
+	return msg.WriteTo(cn.writeBuffer)
+}
+
 func (cn *peerConnMsgWriter) write(msg pp.Message) bool {
 	cn.mu.Lock()
 	defer cn.mu.Unlock()
-	cn.writeBuffer.Write(msg.MustMarshalBinary())
+	cn.writeToBuffer(msg)
 	cn.writeCond.Broadcast()
 	return !cn.writeBufferFull()
 }

peer-conn-msg-writer_test.go

@@ -0,0 +1,68 @@
+package torrent
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/dustin/go-humanize"
+
+	pp "github.com/anacrolix/torrent/peer_protocol"
+)
+
+func PieceMsg(length int64) pp.Message {
+	return pp.Message{
+		Type:  pp.Piece,
+		Index: pp.Integer(0),
+		Begin: pp.Integer(0),
+		Piece: make([]byte, length),
+	}
+}
+
+var benchmarkPieceLengths = []int{defaultChunkSize, 1 << 20, 4 << 20, 8 << 20}
+
+func runBenchmarkWriteToBuffer(b *testing.B, length int64) {
+	writer := &peerConnMsgWriter{
+		writeBuffer: &bytes.Buffer{},
+	}
+	msg := PieceMsg(length)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		//b.StopTimer()
+		writer.writeBuffer.Reset()
+		//b.StartTimer()
+		writer.writeToBuffer(msg)
+	}
+}
+
+func BenchmarkWritePieceMsg(b *testing.B) {
+	for _, length := range benchmarkPieceLengths {
+		b.Run(humanize.IBytes(uint64(length)), func(b *testing.B) {
+			b.Run("ToBuffer", func(b *testing.B) {
+				b.SetBytes(int64(length))
+				runBenchmarkWriteToBuffer(b, int64(length))
+			})
+			b.Run("MarshalBinary", func(b *testing.B) {
+				b.SetBytes(int64(length))
+				runBenchmarkMarshalBinaryWrite(b, int64(length))
+			})
+		})
+	}
+}
+
+func runBenchmarkMarshalBinaryWrite(b *testing.B, length int64) {
+	writer := &peerConnMsgWriter{
+		writeBuffer: &bytes.Buffer{},
+	}
+	msg := PieceMsg(length)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		//b.StopTimer()
+		writer.writeBuffer.Reset()
+		//b.StartTimer()
+		writer.writeBuffer.Write(msg.MustMarshalBinary())
+	}
+}

peer_protocol/msg.go

@@ -6,6 +6,7 @@ import (
 	"encoding"
 	"encoding/binary"
 	"fmt"
+	"io"
 )
 
 // This is a lazy union representing all the possible fields for messages. Go doesn't have ADTs, and
@@ -61,13 +62,14 @@ func (msg Message) MustMarshalBinary() []byte {
 	return b
 }
 
-func (msg Message) MarshalBinary() (data []byte, err error) {
-	// It might look like you could have a pool of buffers and preallocate the message length
-	// prefix, but because we have to return []byte, it becomes non-trivial to make this fast. You
-	// will need a benchmark.
-	var buf bytes.Buffer
+type MessageWriter interface {
+	io.ByteWriter
+	io.Writer
+}
+
+func (msg *Message) writePayloadTo(buf MessageWriter) (err error) {
 	mustWrite := func(data any) {
-		err := binary.Write(&buf, binary.BigEndian, data)
+		err := binary.Write(buf, binary.BigEndian, data)
 		if err != nil {
 			panic(err)
 		}
@@ -85,10 +87,10 @@ func (msg Message) MarshalBinary() (data []byte, err error) {
 		switch msg.Type {
 		case Choke, Unchoke, Interested, NotInterested, HaveAll, HaveNone:
 		case Have, AllowedFast, Suggest:
-			err = binary.Write(&buf, binary.BigEndian, msg.Index)
+			err = binary.Write(buf, binary.BigEndian, msg.Index)
 		case Request, Cancel, Reject:
 			for _, i := range []Integer{msg.Index, msg.Begin, msg.Length} {
-				err = binary.Write(&buf, binary.BigEndian, i)
+				err = binary.Write(buf, binary.BigEndian, i)
 				if err != nil {
 					break
 				}
@@ -97,7 +99,7 @@ func (msg Message) MarshalBinary() (data []byte, err error) {
 			_, err = buf.Write(marshalBitfield(msg.Bitfield))
 		case Piece:
 			for _, i := range []Integer{msg.Index, msg.Begin} {
-				err = binary.Write(&buf, binary.BigEndian, i)
+				err = binary.Write(buf, binary.BigEndian, i)
 				if err != nil {
 					return
 				}
@@ -116,19 +118,43 @@ func (msg Message) MarshalBinary() (data []byte, err error) {
 			}
 			_, err = buf.Write(msg.ExtendedPayload)
 		case Port:
-			err = binary.Write(&buf, binary.BigEndian, msg.Port)
+			err = binary.Write(buf, binary.BigEndian, msg.Port)
 		case HashRequest:
 			buf.Write(msg.PiecesRoot[:])
 			writeConsecutive(msg.BaseLayer, msg.Index, msg.Length, msg.ProofLayers)
 		default:
 			err = fmt.Errorf("unknown message type: %v", msg.Type)
 		}
 	}
-	data = make([]byte, 4+buf.Len())
-	binary.BigEndian.PutUint32(data, uint32(buf.Len()))
-	if buf.Len() != copy(data[4:], buf.Bytes()) {
-		panic("bad copy")
+	return
+}
+
+func (msg *Message) WriteTo(w MessageWriter) (err error) {
+	length, err := msg.getPayloadLength()
+	if err != nil {
+		return
+	}
+	err = binary.Write(w, binary.BigEndian, length)
+	if err != nil {
+		return
 	}
+	return msg.writePayloadTo(w)
+}
+
+func (msg *Message) getPayloadLength() (length Integer, err error) {
+	var lw lengthWriter
+	err = msg.writePayloadTo(&lw)
+	length = lw.n
+	return
+}
+
+func (msg Message) MarshalBinary() (data []byte, err error) {
+	// It might look like you could have a pool of buffers and preallocate the message length
+	// prefix, but because we have to return []byte, it becomes non-trivial to make this fast. You
+	// will need a benchmark.
+	var buf bytes.Buffer
+	err = msg.WriteTo(&buf)
+	data = buf.Bytes()
 	return
 }
 
@@ -158,3 +184,18 @@ func (me *Message) UnmarshalBinary(b []byte) error {
 	}
 	return nil
 }
+
+type lengthWriter struct {
+	n Integer
+}
+
+func (l *lengthWriter) WriteByte(c byte) error {
+	l.n++
+	return nil
+}
+
+func (l *lengthWriter) Write(p []byte) (n int, err error) {
+	n = len(p)
+	l.n += Integer(n)
+	return
+}

requesting.go

@@ -9,9 +9,8 @@ import (
 	"time"
 	"unsafe"
 
-	g "github.com/anacrolix/generics"
-
 	"github.com/RoaringBitmap/roaring"
+	g "github.com/anacrolix/generics"
 	"github.com/anacrolix/generics/heap"
 	"github.com/anacrolix/log"
 	"github.com/anacrolix/multiless"