Skip to content

Commit

Permalink
perf(encoding): speed up EncodeVarint with io.ByteWriter+hand rolled …
Browse files Browse the repository at this point in the history
…varintEncode

This change speeds up EncodeVarint by testing if the input writer
implements io.ByteWriter and if so, goes to use our hand-rolled
varint encoder, instead of using the awkward standard libary
encoding/binary.PutVarint that requires a byteslice, which we
also retrofitted using a bytearray pool.
While here, added parity tests to ensure that we get the exact
same results as with the Go standard library's encoding/binary
package with caution from https://cyber.orijtech.com/advisory/varint-decode-limitless
and also added benchmarks whose results reflect the change in just
the benchmark initially

```shell
$ benchstat before.txt after.txt
name            old time/op    new time/op    delta
EncodeVarint-8     360ns ± 3%     245ns ± 3%  -31.80%  (p=0.000 n=10+10)

name            old alloc/op   new alloc/op   delta
EncodeVarint-8     0.00B          0.00B          ~     (all equal)

name            old allocs/op  new allocs/op  delta
EncodeVarint-8      0.00           0.00          ~     (all equal)
```

but when ran over all the benchmarks, produce even more stark results.
```shell
```

Fixes #891
  • Loading branch information
odeke-em committed Mar 15, 2024
1 parent 17cbe3f commit 9679c03
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 0 deletions.
59 changes: 59 additions & 0 deletions internal/encoding/bench_test.go
Original file line number Diff line number Diff line change
@@ -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()
}
}
}
24 changes: 24 additions & 0 deletions internal/encoding/encoding.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down

0 comments on commit 9679c03

Please sign in to comment.