GODRIVER-2914 bson: improve marshal/unmarshal performance by ~58% and ~29%

bson/benchmark_test.go

@@ -7,12 +7,15 @@
 package bson
 
 import (
+	"bytes"
 	"compress/gzip"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"path"
+	"sync"
 	"testing"
 )
 
@@ -129,10 +132,20 @@ var nestedInstance = nestedtest1{
 
 const extendedBSONDir = "../testdata/extended_bson"
 
+var (
+	extJSONFiles   map[string]map[string]interface{}
+	extJSONFilesMu sync.Mutex
+)
+
 // readExtJSONFile reads the GZIP-compressed extended JSON document from the given filename in the
 // "extended BSON" test data directory (../testdata/extended_bson) and returns it as a
 // map[string]interface{}. It panics on any errors.
 func readExtJSONFile(filename string) map[string]interface{} {
+	extJSONFilesMu.Lock()
+	defer extJSONFilesMu.Unlock()
+	if v, ok := extJSONFiles[filename]; ok {
+		return v
+	}
 	filePath := path.Join(extendedBSONDir, filename)
 	file, err := os.Open(filePath)
 	if err != nil {
@@ -161,6 +174,10 @@ func readExtJSONFile(filename string) map[string]interface{} {
 		panic(fmt.Sprintf("error unmarshalling extended JSON: %s", err))
 	}
 
+	if extJSONFiles == nil {
+		extJSONFiles = make(map[string]map[string]interface{})
+	}
+	extJSONFiles[filename] = v
 	return v
 }
 
@@ -305,3 +322,128 @@ func BenchmarkUnmarshal(b *testing.B) {
 		})
 	}
 }
+
+// The following benchmarks are copied from the Go standard library's
+// encoding/json package.
+
+type codeResponse struct {
+	Tree     *codeNode `json:"tree"`
+	Username string    `json:"username"`
+}
+
+type codeNode struct {
+	Name     string      `json:"name"`
+	Kids     []*codeNode `json:"kids"`
+	CLWeight float64     `json:"cl_weight"`
+	Touches  int         `json:"touches"`
+	MinT     int64       `json:"min_t"`
+	MaxT     int64       `json:"max_t"`
+	MeanT    int64       `json:"mean_t"`
+}
+
+var codeJSON []byte
+var codeBSON []byte
+var codeStruct codeResponse
+
+func codeInit() {
+	f, err := os.Open("testdata/code.json.gz")
+	if err != nil {
+		panic(err)
+	}
+	defer f.Close()
+	gz, err := gzip.NewReader(f)
+	if err != nil {
+		panic(err)
+	}
+	data, err := io.ReadAll(gz)
+	if err != nil {
+		panic(err)
+	}
+
+	codeJSON = data
+
+	if err := json.Unmarshal(codeJSON, &codeStruct); err != nil {
+		panic("json.Unmarshal code.json: " + err.Error())
+	}
+
+	if data, err = json.Marshal(&codeStruct); err != nil {
+		panic("json.Marshal code.json: " + err.Error())
+	}
+
+	if codeBSON, err = Marshal(&codeStruct); err != nil {
+		panic("Marshal code.json: " + err.Error())
+	}
+
+	if !bytes.Equal(data, codeJSON) {
+		println("different lengths", len(data), len(codeJSON))
+		for i := 0; i < len(data) && i < len(codeJSON); i++ {
+			if data[i] != codeJSON[i] {
+				println("re-marshal: changed at byte", i)
+				println("orig: ", string(codeJSON[i-10:i+10]))
+				println("new: ", string(data[i-10:i+10]))
+				break
+			}
+		}
+		panic("re-marshal code.json: different result")
+	}
+}
+
+func BenchmarkCodeUnmarshal(b *testing.B) {
+	b.ReportAllocs()
+	if codeJSON == nil {
+		b.StopTimer()
+		codeInit()
+		b.StartTimer()
+	}
+	b.Run("BSON", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				var r codeResponse
+				if err := Unmarshal(codeBSON, &r); err != nil {
+					b.Fatal("Unmarshal:", err)
+				}
+			}
+		})
+		b.SetBytes(int64(len(codeBSON)))
+	})
+	b.Run("JSON", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				var r codeResponse
+				if err := json.Unmarshal(codeJSON, &r); err != nil {
+					b.Fatal("json.Unmarshal:", err)
+				}
+			}
+		})
+		b.SetBytes(int64(len(codeJSON)))
+	})
+}
+
+func BenchmarkCodeMarshal(b *testing.B) {
+	b.ReportAllocs()
+	if codeJSON == nil {
+		b.StopTimer()
+		codeInit()
+		b.StartTimer()
+	}
+	b.Run("BSON", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				if _, err := Marshal(&codeStruct); err != nil {
+					b.Fatal("Marshal:", err)
+				}
+			}
+		})
+		b.SetBytes(int64(len(codeBSON)))
+	})
+	b.Run("JSON", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				if _, err := json.Marshal(&codeStruct); err != nil {
+					b.Fatal("json.Marshal:", err)
+				}
+			}
+		})
+		b.SetBytes(int64(len(codeJSON)))
+	})
+}

bson/bsoncodec/codec_cache.go

@@ -0,0 +1,166 @@
+// Copyright (C) MongoDB, Inc. 2017-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package bsoncodec
+
+import (
+	"reflect"
+	"sync"
+	"sync/atomic"
+)
+
+// Runtime check that the kind encoder and decoder caches can store any valid
+// reflect.Kind constant.
+func init() {
+	if s := reflect.Kind(len(kindEncoderCache{}.entries)).String(); s != "kind27" {
+		panic("The capacity of kindEncoderCache is too small.\n" +
+			"This is due to a new type being added to reflect.Kind.")
+	}
+}
+
+// statically assert array size
+var _ = (kindEncoderCache{}).entries[reflect.UnsafePointer]
+var _ = (kindDecoderCache{}).entries[reflect.UnsafePointer]
+
+type typeEncoderCache struct {
+	cache sync.Map // map[reflect.Type]ValueEncoder
+}
+
+func (c *typeEncoderCache) Store(rt reflect.Type, enc ValueEncoder) {
+	c.cache.Store(rt, enc)
+}
+
+func (c *typeEncoderCache) Load(rt reflect.Type) (ValueEncoder, bool) {
+	if v, _ := c.cache.Load(rt); v != nil {
+		return v.(ValueEncoder), true
+	}
+	return nil, false
+}
+
+func (c *typeEncoderCache) LoadOrStore(rt reflect.Type, enc ValueEncoder) ValueEncoder {
+	if v, loaded := c.cache.LoadOrStore(rt, enc); loaded {
+		enc = v.(ValueEncoder)
+	}
+	return enc
+}
+
+func (c *typeEncoderCache) Clone() *typeEncoderCache {
+	cc := new(typeEncoderCache)
+	c.cache.Range(func(k, v interface{}) bool {
+		if k != nil && v != nil {
+			cc.cache.Store(k, v)
+		}
+		return true
+	})
+	return cc
+}
+
+type typeDecoderCache struct {
+	cache sync.Map // map[reflect.Type]ValueDecoder
+}
+
+func (c *typeDecoderCache) Store(rt reflect.Type, dec ValueDecoder) {
+	c.cache.Store(rt, dec)
+}
+
+func (c *typeDecoderCache) Load(rt reflect.Type) (ValueDecoder, bool) {
+	if v, _ := c.cache.Load(rt); v != nil {
+		return v.(ValueDecoder), true
+	}
+	return nil, false
+}
+
+func (c *typeDecoderCache) LoadOrStore(rt reflect.Type, dec ValueDecoder) ValueDecoder {
+	if v, loaded := c.cache.LoadOrStore(rt, dec); loaded {
+		dec = v.(ValueDecoder)
+	}
+	return dec
+}
+
+func (c *typeDecoderCache) Clone() *typeDecoderCache {
+	cc := new(typeDecoderCache)
+	c.cache.Range(func(k, v interface{}) bool {
+		if k != nil && v != nil {
+			cc.cache.Store(k, v)
+		}
+		return true
+	})
+	return cc
+}
+
+// atomic.Value requires that all calls to Store() have the same concrete type
+// so we wrap the ValueEncoder with a kindEncoderCacheEntry to ensure the type
+// is always the same (since different concrete types may implement the
+// ValueEncoder interface).
+type kindEncoderCacheEntry struct {
+	enc ValueEncoder
+}
+
+type kindEncoderCache struct {
+	entries [reflect.UnsafePointer + 1]atomic.Value // *kindEncoderCacheEntry
+}
+
+func (c *kindEncoderCache) Store(rt reflect.Kind, enc ValueEncoder) {
+	if enc != nil && rt < reflect.Kind(len(c.entries)) {
+		c.entries[rt].Store(&kindEncoderCacheEntry{enc: enc})
+	}
+}
+
+func (c *kindEncoderCache) Load(rt reflect.Kind) (ValueEncoder, bool) {
+	if rt < reflect.Kind(len(c.entries)) {
+		if ent, ok := c.entries[rt].Load().(*kindEncoderCacheEntry); ok {
+			return ent.enc, ent.enc != nil
+		}
+	}
+	return nil, false
+}
+
+func (c *kindEncoderCache) Clone() *kindEncoderCache {
+	cc := new(kindEncoderCache)
+	for i, v := range c.entries {
+		if val := v.Load(); val != nil {
+			cc.entries[i].Store(val)
+		}
+	}
+	return cc
+}
+
+// atomic.Value requires that all calls to Store() have the same concrete type
+// so we wrap the ValueDecoder with a kindDecoderCacheEntry to ensure the type
+// is always the same (since different concrete types may implement the
+// ValueDecoder interface).
+type kindDecoderCacheEntry struct {
+	dec ValueDecoder
+}
+
+type kindDecoderCache struct {
+	entries [reflect.UnsafePointer + 1]atomic.Value // *kindDecoderCacheEntry
+}
+
+func (c *kindDecoderCache) Store(rt reflect.Kind, dec ValueDecoder) {
+	if rt < reflect.Kind(len(c.entries)) {
+		c.entries[rt].Store(&kindDecoderCacheEntry{dec: dec})
+	}
+}
+
+func (c *kindDecoderCache) Load(rt reflect.Kind) (ValueDecoder, bool) {
+	if rt < reflect.Kind(len(c.entries)) {
+		if ent, ok := c.entries[rt].Load().(*kindDecoderCacheEntry); ok {
+			return ent.dec, ent.dec != nil
+		}
+	}
+	return nil, false
+}
+
+func (c *kindDecoderCache) Clone() *kindDecoderCache {
+	cc := new(kindDecoderCache)
+	for i, v := range c.entries {
+		if val := v.Load(); val != nil {
+			cc.entries[i].Store(val)
+		}
+	}
+	return cc
+}