perf(encoding): speed up EncodeVarint with io.ByteWriter+hand rolled varintEncode

internal/encoding/bench_test.go

@@ -0,0 +1,59 @@
+package encoding
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"math"
+	"testing"
+)
+
+var encValues = []int64{
+	-1, -100, -1 << 32,
+	0, 1, 100, 1 << 32,
+	-1 << 52, 1 << 52, 17,
+	19, 28, 37, 388888888,
+	-99999999999, 99999999999,
+	math.MaxInt64, math.MinInt64,
+}
+
+// This tests that the results from directly invoking binary.PutVarint match
+// exactly those that we get from invoking EncodeVarint and its internals.
+func TestEncodeVarintParity(t *testing.T) {
+	buf := new(bytes.Buffer)
+	var board [binary.MaxVarintLen64]byte
+
+	for _, val := range encValues {
+		val := val
+		name := fmt.Sprintf("%d", val)
+
+		buf.Reset()
+		t.Run(name, func(t *testing.T) {
+			if err := EncodeVarint(buf, val); err != nil {
+				t.Fatal(err)
+			}
+
+			n := binary.PutVarint(board[:], val)
+			got := buf.Bytes()
+			want := board[:n]
+			if !bytes.Equal(got, want) {
+				t.Fatalf("Result mismatch\n\tGot:  %d\n\tWant: %d", got, want)
+			}
+		})
+	}
+}
+
+func BenchmarkEncodeVarint(b *testing.B) {
+	buf := new(bytes.Buffer)
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		for _, val := range encValues {
+			if err := EncodeVarint(buf, val); err != nil {
+				b.Fatal(err)
+			}
+			buf.Reset()
+		}
+	}
+}

internal/encoding/encoding.go

@@ -139,6 +139,10 @@ func EncodeUvarintSize(u uint64) int {
 
 // EncodeVarint writes a varint-encoded integer to an io.Writer.
 func EncodeVarint(w io.Writer, i int64) error {
+	if bw, ok := w.(io.ByteWriter); ok {
+		return fVarintEncode(bw, i)
+	}
+
 	// Use a pool here to reduce allocations.
 	//
 	// Though this allocates just 10 bytes on the stack, doing allocation for every calls
@@ -157,6 +161,26 @@ func EncodeVarint(w io.Writer, i int64) error {
 	return err
 }
 
+func fVarintEncode(bw io.ByteWriter, x int64) error {
+	// Firstly convert it into a uvarint
+	ux := uint64(x) << 1
+	if x < 0 {
+		ux = ^ux
+	}
+	for ux >= 0x80 { // While there are 7 or more bits in the value, keep going
+		// Convert it into a byte then toggle the
+		// 7th bit to indicate that more bytes coming.
+		// byte(x & 0x7f) is redundant but useful for illustrative
+		// purposes when translating to other languages
+		if err := bw.WriteByte(byte(ux&0x7f) | 0x80); err != nil {
+			return err
+		}
+		ux >>= 7
+	}
+
+	return bw.WriteByte(byte(ux & 0x7f))
+}
+
 // EncodeVarintSize returns the byte size of the given integer as a varint.
 func EncodeVarintSize(i int64) int {
 	ux := uint64(i) << 1