wire: Optimize writes to bytes.Buffer

This special cases writes to bytes.Buffer, which is always the writer type
written to by WriteMessageN.  There are several optimizations that can be
implemented by special casing this type:

First, pulling temporary short buffers from binary freelist can be skipped
entirely, and instead the binary encoding of integers can be appended directly
to its existing capacity.  This avoids the synchronization cost to add and
remove buffers from the free list, and for applications which only ever write
wire messages with WriteMessageN, the allocation and ongoing garbage
collection scanning cost to for these buffers can be completely skipped.

Second, special casing the buffer type in WriteVarString saves us from
creating a temporary heap copy of a string, as the buffer's WriteString method
can be used instead.

Third, special casing the buffer allows WriteMessageN to calculate the
serialize size and grow its buffer so all remaining appends for writing block
and transactions will not have to reallocate the buffer's backing allocation.
This same optimization can be applied to other messages in the future.

Author: jrick

wire/common.go

@@ -6,6 +6,7 @@
 package wire
 
 import (
+	"bytes"
 	"crypto/rand"
 	"encoding/binary"
 	"fmt"
@@ -131,49 +132,6 @@ func (l binaryFreeList) Uint64(r io.Reader, byteOrder binary.ByteOrder) (uint64,
 	return rv, nil
 }
 
-// PutUint8 copies the provided uint8 into a buffer from the free list and
-// writes the resulting byte to the given writer.
-func (l binaryFreeList) PutUint8(w io.Writer, val uint8) error {
-	buf := l.Borrow()[:1]
-	buf[0] = val
-	_, err := w.Write(buf)
-	l.Return(buf)
-	return err
-}
-
-// PutUint16 serializes the provided uint16 using the given byte order into a
-// buffer from the free list and writes the resulting two bytes to the given
-// writer.
-func (l binaryFreeList) PutUint16(w io.Writer, byteOrder binary.ByteOrder, val uint16) error {
-	buf := l.Borrow()[:2]
-	byteOrder.PutUint16(buf, val)
-	_, err := w.Write(buf)
-	l.Return(buf)
-	return err
-}
-
-// PutUint32 serializes the provided uint32 using the given byte order into a
-// buffer from the free list and writes the resulting four bytes to the given
-// writer.
-func (l binaryFreeList) PutUint32(w io.Writer, byteOrder binary.ByteOrder, val uint32) error {
-	buf := l.Borrow()[:4]
-	byteOrder.PutUint32(buf, val)
-	_, err := w.Write(buf)
-	l.Return(buf)
-	return err
-}
-
-// PutUint64 serializes the provided uint64 using the given byte order into a
-// buffer from the free list and writes the resulting eight bytes to the given
-// writer.
-func (l binaryFreeList) PutUint64(w io.Writer, byteOrder binary.ByteOrder, val uint64) error {
-	buf := l.Borrow()[:8]
-	byteOrder.PutUint64(buf, val)
-	_, err := w.Write(buf)
-	l.Return(buf)
-	return err
-}
-
 // binarySerializer provides a free list of buffers to use for serializing and
 // deserializing primitive integer values to and from io.Readers and io.Writers.
 var binarySerializer binaryFreeList = make(chan []byte, binaryFreeListMaxItems)
@@ -412,48 +370,108 @@ func readElements(r io.Reader, elements ...interface{}) error {
 	return nil
 }
 
+// shortWrite optimizes short (<= 8 byte) writes to w by special casing
+// buffer allocations for specific writer types.
+//
+// The callback returns a short buffer to 8 bytes in length and a size
+// specifying how much of the buffer to write.
+//
+// For longer writes and writes of byte arrays, dynamic dispatch to w.Write
+// should be used instead.
+func shortWrite(w io.Writer, cb func() (data [8]byte, size int)) error {
+	data, size := cb()
+
+	switch w := w.(type) {
+	// The most common case (called through WriteMessageN) is that the writer is a
+	// *bytes.Buffer.  Optimize for that case by appending binary serializations
+	// to its existing capacity instead of paying the synchronization cost to
+	// serialize to temporary buffers pulled from the binary freelist.
+	case *bytes.Buffer:
+		w.Write(data[:size])
+		return nil
+
+	default:
+		p := binarySerializer.Borrow()[:size]
+		copy(p, data[:size])
+		_, err := w.Write(p)
+		return err
+	}
+}
+
+// writeUint8 writes the byte value to the writer.
+func writeUint8(w io.Writer, value uint8) error {
+	return shortWrite(w, func() (buf [8]byte, size int) {
+		buf[0] = value
+		return buf, 1
+	})
+}
+
+// writeUint16 writes the little endian encoding of value to the writer.
+func writeUint16(w io.Writer, value uint16) error {
+	return shortWrite(w, func() (buf [8]byte, size int) {
+		littleEndian.PutUint16(buf[:], value)
+		return buf, 2
+	})
+}
+
+// writeUint32 writes the little endian encoding of value to the writer.
+func writeUint32(w io.Writer, value uint32) error {
+	return shortWrite(w, func() (buf [8]byte, size int) {
+		littleEndian.PutUint32(buf[:], value)
+		return buf, 4
+	})
+}
+
+// writeUint64 writes the little endian encoding of value to the writer.
+func writeUint64(w io.Writer, value uint64) error {
+	return shortWrite(w, func() (buf [8]byte, size int) {
+		littleEndian.PutUint64(buf[:], value)
+		return buf, 8
+	})
+}
+
 // writeElement writes the little endian representation of element to w.
 func writeElement(w io.Writer, element interface{}) error {
 	// Attempt to write the element based on the concrete type via fast
 	// type assertions first.
 	switch e := element.(type) {
 	case *uint8:
-		err := binarySerializer.PutUint8(w, *e)
+		err := writeUint8(w, *e)
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *uint16:
-		err := binarySerializer.PutUint16(w, littleEndian, *e)
+		err := writeUint16(w, *e)
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *int32:
-		err := binarySerializer.PutUint32(w, littleEndian, uint32(*e))
+		err := writeUint32(w, uint32(*e))
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *uint32:
-		err := binarySerializer.PutUint32(w, littleEndian, *e)
+		err := writeUint32(w, *e)
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *int64:
-		err := binarySerializer.PutUint64(w, littleEndian, uint64(*e))
+		err := writeUint64(w, uint64(*e))
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *uint64:
-		err := binarySerializer.PutUint64(w, littleEndian, *e)
+		err := writeUint64(w, *e)
 		if err != nil {
 			return err
 		}
@@ -462,9 +480,9 @@ func writeElement(w io.Writer, element interface{}) error {
 	case *bool:
 		var err error
 		if *e {
-			err = binarySerializer.PutUint8(w, 0x01)
+			err = writeUint8(w, 0x01)
 		} else {
-			err = binarySerializer.PutUint8(w, 0x00)
+			err = writeUint8(w, 0x00)
 		}
 		if err != nil {
 			return err
@@ -558,28 +576,28 @@ func writeElement(w io.Writer, element interface{}) error {
 		return nil
 
 	case *ServiceFlag:
-		err := binarySerializer.PutUint64(w, littleEndian, uint64(*e))
+		err := writeUint64(w, uint64(*e))
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *InvType:
-		err := binarySerializer.PutUint32(w, littleEndian, uint32(*e))
+		err := writeUint32(w, uint32(*e))
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *CurrencyNet:
-		err := binarySerializer.PutUint32(w, littleEndian, uint32(*e))
+		err := writeUint32(w, uint32(*e))
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *RejectCode:
-		err := binarySerializer.PutUint8(w, uint8(*e))
+		err := writeUint8(w, uint8(*e))
 		if err != nil {
 			return err
 		}
@@ -669,30 +687,34 @@ func ReadVarInt(r io.Reader, pver uint32) (uint64, error) {
 // on its value.
 func WriteVarInt(w io.Writer, pver uint32, val uint64) error {
 	if val < 0xfd {
-		return binarySerializer.PutUint8(w, uint8(val))
+		return shortWrite(w, func() (p [8]byte, size int) {
+			p[0] = uint8(val)
+			return p, 1
+		})
 	}
 
 	if val <= math.MaxUint16 {
-		err := binarySerializer.PutUint8(w, 0xfd)
-		if err != nil {
-			return err
-		}
-		return binarySerializer.PutUint16(w, littleEndian, uint16(val))
+		return shortWrite(w, func() (p [8]byte, size int) {
+			p[0] = 0xfd
+			littleEndian.PutUint16(p[1:], uint16(val))
+			return p, 3
+		})
 	}
 
 	if val <= math.MaxUint32 {
-		err := binarySerializer.PutUint8(w, 0xfe)
-		if err != nil {
-			return err
-		}
-		return binarySerializer.PutUint32(w, littleEndian, uint32(val))
+		return shortWrite(w, func() (p [8]byte, size int) {
+			p[0] = 0xfe
+			littleEndian.PutUint32(p[1:], uint32(val))
+			return p, 5
+		})
 	}
 
-	err := binarySerializer.PutUint8(w, 0xff)
+	// shortWrite is not designed for writes > 8 bytes.
+	err := writeUint8(w, 0xff)
 	if err != nil {
 		return err
 	}
-	return binarySerializer.PutUint64(w, littleEndian, val)
+	return writeUint64(w, val)
 }
 
 // VarIntSerializeSize returns the number of bytes it would take to serialize
@@ -798,7 +820,13 @@ func WriteVarString(w io.Writer, pver uint32, str string) error {
 	if err != nil {
 		return err
 	}
-	_, err = w.Write([]byte(str))
+
+	switch w := w.(type) {
+	case *bytes.Buffer:
+		_, err = w.WriteString(str)
+	default:
+		_, err = w.Write([]byte(str))
+	}
 	return err
 }
 

wire/message.go

@@ -306,14 +306,23 @@ func WriteMessageN(w io.Writer, msg Message, pver uint32, dcrnet CurrencyNet) (i
 	}
 	copy(elems.command[:], []byte(cmd))
 
+	// Allocate enough buffer space for the entire message size if it is
+	// known.  When it is not known, use an extra size hint of 64 bytes,
+	// which matches the default small allocation size of a bytes.Buffer
+	// as of Go 1.25.
+	extraCap := 64
+	switch msg := msg.(type) {
+	case interface{ SerializeSize() int }:
+		extraCap = msg.SerializeSize()
+	}
+
 	// Initialize buffer with zeroed bytes for the message header (to be
 	// filled in, with checksum, after appending the payload
-	// serialization).
-	var buf bytes.Buffer
-	buf.Write(make([]byte, MessageHeaderSize))
+	// serialization), plus additional capacity for writing the payload.
+	buf := bytes.NewBuffer(make([]byte, MessageHeaderSize, MessageHeaderSize+extraCap))
 
 	// Encode the message payload.
-	err := msg.BtcEncode(&buf, pver)
+	err := msg.BtcEncode(buf, pver)
 	if err != nil {
 		return 0, err
 	}
@@ -342,7 +351,7 @@ func WriteMessageN(w io.Writer, msg Message, pver uint32, dcrnet CurrencyNet) (i
 	cksumHash := chainhash.HashH(payload)
 	copy(elems.checksum[:], cksumHash[0:4])
 	buf.Reset()
-	writeElements(&buf, &elems.dcrnet, &elems.command, &elems.lenp, &elems.checksum)
+	writeElements(buf, &elems.dcrnet, &elems.command, &elems.lenp, &elems.checksum)
 	if buf.Len() != MessageHeaderSize {
 		// The length of data written for the header is always
 		// constant, is not dependent on the message being serialized,

wire/msgcfheaders.go

@@ -116,7 +116,7 @@ func (msg *MsgCFHeaders) BtcEncode(w io.Writer, pver uint32) error {
 	}
 
 	// Write filter type
-	err = binarySerializer.PutUint8(w, uint8(msg.FilterType))
+	err = writeUint8(w, uint8(msg.FilterType))
 	if err != nil {
 		return err
 	}

wire/msgcfilter.go

@@ -80,7 +80,7 @@ func (msg *MsgCFilter) BtcEncode(w io.Writer, pver uint32) error {
 		return err
 	}
 
-	err = binarySerializer.PutUint8(w, uint8(msg.FilterType))
+	err = writeUint8(w, uint8(msg.FilterType))
 	if err != nil {
 		return err
 	}

wire/msgcftypes.go

@@ -92,7 +92,7 @@ func (msg *MsgCFTypes) BtcEncode(w io.Writer, pver uint32) error {
 	}
 
 	for i := range msg.SupportedFilters {
-		err = binarySerializer.PutUint8(w, uint8(msg.SupportedFilters[i]))
+		err = writeUint8(w, uint8(msg.SupportedFilters[i]))
 		if err != nil {
 			return err
 		}

wire/msggetcfheaders.go

@@ -115,7 +115,7 @@ func (msg *MsgGetCFHeaders) BtcEncode(w io.Writer, pver uint32) error {
 		return err
 	}
 
-	return binarySerializer.PutUint8(w, uint8(msg.FilterType))
+	return writeUint8(w, uint8(msg.FilterType))
 }
 
 // Command returns the protocol command string for the message.  This is part

wire/msggetcfilter.go

@@ -54,7 +54,7 @@ func (msg *MsgGetCFilter) BtcEncode(w io.Writer, pver uint32) error {
 	if err != nil {
 		return err
 	}
-	return binarySerializer.PutUint8(w, uint8(msg.FilterType))
+	return writeUint8(w, uint8(msg.FilterType))
 }
 
 // Command returns the protocol command string for the message.  This is part