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 two 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.

Author: jrick

wire/common.go

@@ -6,6 +6,7 @@
 package wire
 
 import (
+	"bytes"
 	"crypto/rand"
 	"encoding/binary"
 	"fmt"
@@ -412,60 +413,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 +599,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 +718,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 +853,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
 }