Vai encoding benchmark (#742)

Also includes some misc changes.

pkg/sqlcache/db/encoding.go

@@ -0,0 +1,169 @@
+package db
+
+import (
+	"bytes"
+	"compress/gzip"
+	"encoding/gob"
+	"encoding/json"
+	"fmt"
+	"io"
+	"reflect"
+	"strings"
+	"sync"
+)
+
+func init() {
+	// necessary in order to gob/ungob unstructured.Unstructured objects
+	gob.Register(map[string]any{})
+	gob.Register([]any{})
+}
+
+type encoding interface {
+	// Encode serializes an object into the provided writer
+	Encode(io.Writer, any) error
+	// Decode reads from a provided reader and deserializes into an object
+	Decode(io.Reader, any) error
+}
+
+type gobEncoding struct {
+	writeLock, readLock sync.Mutex
+	writeBuf, readBuf   bytes.Buffer
+	encoder             *gob.Encoder
+	decoder             *gob.Decoder
+	seenTypes           map[reflect.Type]struct{}
+}
+
+func (g *gobEncoding) Encode(w io.Writer, obj any) error {
+	g.writeLock.Lock()
+	defer g.writeLock.Unlock()
+
+	if g.encoder == nil {
+		g.encoder = gob.NewEncoder(&g.writeBuf)
+	}
+
+	g.writeBuf.Reset()
+	if err := g.encoder.Encode(obj); err != nil {
+		return err
+	}
+
+	if err := g.registerTypeIfNeeded(obj); err != nil {
+		return err
+	}
+
+	_, err := g.writeBuf.WriteTo(w)
+	return err
+}
+
+// registerTypeIfNeeded prevents future decoding errors by running Decode right after the first Encode for an object type
+// This is needed when reusing a gob.Encoder, as it assumes the receiving end of those messages is always the same gob.Decoder.
+// Due to this assumption, it applies some optimizations, like avoiding sending complete type's information (needed for decoding) if it already did it before.
+// This means the first object for each type encoded by a gob.Encoder will have a bigger payload, but more importantly,
+// the decoding will fail if this is not the first object being decoded later. This function forces that to prevent this from happening.
+func (g *gobEncoding) registerTypeIfNeeded(obj any) error {
+	if g.seenTypes == nil {
+		g.seenTypes = make(map[reflect.Type]struct{})
+	}
+
+	typ := reflect.TypeOf(obj)
+	if _, ok := g.seenTypes[typ]; !ok {
+		g.seenTypes[typ] = struct{}{}
+		// Consume the current write buffer and re-generate it (will produce a smaller version)
+		newObj := reflect.New(typ).Interface()
+		if err := g.Decode(bytes.NewReader(g.writeBuf.Bytes()), newObj); err != nil {
+			return fmt.Errorf("could not decode %T: %w", obj, err)
+		}
+
+		g.writeBuf.Reset()
+		return g.encoder.Encode(obj)
+	}
+	return nil
+}
+
+func (g *gobEncoding) Decode(r io.Reader, into any) error {
+	g.readLock.Lock()
+	defer g.readLock.Unlock()
+
+	if g.decoder == nil {
+		g.decoder = gob.NewDecoder(&g.readBuf)
+	}
+
+	g.readBuf.Reset()
+	if _, err := g.readBuf.ReadFrom(r); err != nil {
+		return err
+	}
+	return g.decoder.Decode(into)
+}
+
+type jsonEncoding struct {
+	indentLevel int
+}
+
+func (j jsonEncoding) Encode(w io.Writer, obj any) error {
+	enc := json.NewEncoder(w)
+	if j.indentLevel > 0 {
+		enc.SetIndent("", strings.Repeat(" ", j.indentLevel))
+	}
+
+	if err := enc.Encode(obj); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (j jsonEncoding) Decode(r io.Reader, into any) error {
+	return json.NewDecoder(r).Decode(into)
+}
+
+type gzipEncoding struct {
+	encoding
+	writers sync.Pool
+	readers sync.Pool
+}
+
+func gzipped(wrapped encoding) *gzipEncoding {
+	gz := gzipEncoding{
+		encoding: wrapped,
+	}
+	return &gz
+}
+
+func (gz *gzipEncoding) Encode(w io.Writer, obj any) error {
+	gzw, ok := gz.writers.Get().(*gzip.Writer)
+	if !ok {
+		gzw = gzip.NewWriter(io.Discard)
+	}
+	gzw.Reset(w)
+	defer func() {
+		gzw.Reset(nil)
+		gz.writers.Put(gzw)
+	}()
+
+	if err := gz.encoding.Encode(gzw, obj); err != nil {
+		return err
+	}
+	return gzw.Close()
+}
+
+func (gz *gzipEncoding) Decode(r io.Reader, into any) error {
+	gzr, ok := gz.readers.Get().(*gzip.Reader)
+	if ok {
+		if err := gzr.Reset(r); err != nil {
+			return err
+		}
+	} else {
+		var err error
+		gzr, err = gzip.NewReader(r)
+		if err != nil {
+			return err
+		}
+	}
+	defer func() {
+		gzr.Close()
+		gz.readers.Put(gzr)
+	}()
+
+	if err := gz.encoding.Decode(gzr, into); err != nil {
+		return err
+	}
+	return gzr.Close()
+}