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 MsgBlock and MsgTx BtcEncode method to
calculate the serialize size and grow the buffer being written to, so all
remaining appends 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,60 +370,100 @@ func readElements(r io.Reader, elements ...interface{}) error {
 	return nil
 }
 
+// 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.  This is
+// wrapped by this helper function with callback as only binary encoding short
+// writes (<= 8 bytes) requires the append (in all other cases, Write can be
+// called directly).
+func appendWriter(w io.Writer, do func(p []byte) []byte) (rerr error) {
+	var p []byte
+	switch w := w.(type) {
+	case *bytes.Buffer:
+		p = w.Bytes()[w.Len():]
+		defer func() {
+			w.Write(p)
+		}()
+	default:
+		p = binarySerializer.Borrow()[:0]
+		defer func() {
+			_, rerr = w.Write(p)
+			binarySerializer.Return(p)
+		}()
+	}
+
+	p = do(p)
+	return nil
+}
+
 // 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 := appendWriter(w, func(p []byte) []byte {
+			return append(p, *e)
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *uint16:
-		err := binarySerializer.PutUint16(w, littleEndian, *e)
+		err := appendWriter(w, func(p []byte) []byte {
+			return littleEndian.AppendUint16(p, *e)
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *int32:
-		err := binarySerializer.PutUint32(w, littleEndian, uint32(*e))
+		err := appendWriter(w, func(p []byte) []byte {
+			return littleEndian.AppendUint32(p, uint32(*e))
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *uint32:
-		err := binarySerializer.PutUint32(w, littleEndian, *e)
+		err := appendWriter(w, func(p []byte) []byte {
+			return littleEndian.AppendUint32(p, *e)
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *int64:
-		err := binarySerializer.PutUint64(w, littleEndian, uint64(*e))
+		err := appendWriter(w, func(p []byte) []byte {
+			return littleEndian.AppendUint64(p, uint64(*e))
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *uint64:
-		err := binarySerializer.PutUint64(w, littleEndian, *e)
+		err := appendWriter(w, func(p []byte) []byte {
+			return littleEndian.AppendUint64(p, *e)
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *bool:
-		var err error
-		if *e {
-			err = binarySerializer.PutUint8(w, 0x01)
-		} else {
-			err = binarySerializer.PutUint8(w, 0x00)
-		}
+		err := appendWriter(w, func(p []byte) []byte {
+			var value byte
+			if *e {
+				value = 1
+			}
+			return append(p, value)
+		})
 		if err != nil {
 			return err
 		}
@@ -558,28 +556,36 @@ func writeElement(w io.Writer, element interface{}) error {
 		return nil
 
 	case *ServiceFlag:
-		err := binarySerializer.PutUint64(w, littleEndian, uint64(*e))
+		err := appendWriter(w, func(p []byte) []byte {
+			return littleEndian.AppendUint64(p, uint64(*e))
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *InvType:
-		err := binarySerializer.PutUint32(w, littleEndian, uint32(*e))
+		err := appendWriter(w, func(p []byte) []byte {
+			return littleEndian.AppendUint32(p, uint32(*e))
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *CurrencyNet:
-		err := binarySerializer.PutUint32(w, littleEndian, uint32(*e))
+		err := appendWriter(w, func(p []byte) []byte {
+			return littleEndian.AppendUint32(p, uint32(*e))
+		})
 		if err != nil {
 			return err
 		}
 		return nil
 
 	case *RejectCode:
-		err := binarySerializer.PutUint8(w, uint8(*e))
+		err := appendWriter(w, func(p []byte) []byte {
+			return append(p, uint8(*e))
+		})
 		if err != nil {
 			return err
 		}
@@ -669,30 +675,36 @@ 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 appendWriter(w, func(p []byte) []byte {
+			return append(p, uint8(val))
+		})
 	}
 
 	if val <= math.MaxUint16 {
-		err := binarySerializer.PutUint8(w, 0xfd)
-		if err != nil {
-			return err
-		}
-		return binarySerializer.PutUint16(w, littleEndian, uint16(val))
+		return appendWriter(w, func(p []byte) []byte {
+			p = append(p, 0xfd)
+			p = littleEndian.AppendUint16(p, uint16(val))
+			return p
+		})
 	}
 
 	if val <= math.MaxUint32 {
-		err := binarySerializer.PutUint8(w, 0xfe)
-		if err != nil {
-			return err
-		}
-		return binarySerializer.PutUint32(w, littleEndian, uint32(val))
+		return appendWriter(w, func(p []byte) []byte {
+			p = append(p, 0xfe)
+			p = littleEndian.AppendUint32(p, uint32(val))
+			return p
+		})
 	}
 
-	err := binarySerializer.PutUint8(w, 0xff)
+	err := appendWriter(w, func(p []byte) []byte {
+		return append(p, 0xff)
+	})
 	if err != nil {
 		return err
 	}
-	return binarySerializer.PutUint64(w, littleEndian, val)
+	return appendWriter(w, func(p []byte) []byte {
+		return littleEndian.AppendUint64(p, val)
+	})
 }
 
 // VarIntSerializeSize returns the number of bytes it would take to serialize
@@ -798,7 +810,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/msgblock.go

@@ -252,6 +252,11 @@ func (msg *MsgBlock) DeserializeTxLoc(r *bytes.Buffer) ([]TxLoc, []TxLoc, error)
 // See Serialize for encoding blocks to be stored to disk, such as in a
 // database, as opposed to encoding blocks for the wire.
 func (msg *MsgBlock) BtcEncode(w io.Writer, pver uint32) error {
+	switch w := w.(type) {
+	case *bytes.Buffer:
+		w.Grow(msg.SerializeSize())
+	}
+
 	err := writeBlockHeader(w, pver, &msg.Header)
 	if err != nil {
 		return err
@@ -263,7 +268,7 @@ func (msg *MsgBlock) BtcEncode(w io.Writer, pver uint32) error {
 	}
 
 	for _, tx := range msg.Transactions {
-		err = tx.BtcEncode(w, pver)
+		err = tx.btcEncode(w, pver)
 		if err != nil {
 			return err
 		}
@@ -275,7 +280,7 @@ func (msg *MsgBlock) BtcEncode(w io.Writer, pver uint32) error {
 	}
 
 	for _, tx := range msg.STransactions {
-		err = tx.BtcEncode(w, pver)
+		err = tx.btcEncode(w, pver)
 		if err != nil {
 			return err
 		}

wire/msgcfheaders.go

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

wire/msgcfilter.go

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

wire/msgcftypes.go

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

wire/msggetcfheaders.go

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

wire/msggetcfilter.go

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

wire/msgmixconfirm.go

@@ -145,6 +145,7 @@ func (msg *MsgMixConfirm) writeMessageNoSignature(op string, w io.Writer, pver u
 		return err
 	}
 
+	// XXX: use btcEncode after providing our own initial size hint.
 	err = msg.Mix.BtcEncode(w, pver)
 	if err != nil {
 		return err