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

This commit improves BSON marshaling performance by ~58% and
unmarshaling performance by ~29% by replacing the mutex based
decoder/encoder caches with sync.Map, which can often avoid locking and
is ideally suited for caches that only grow.

The commit also adds the BenchmarkCodeMarshal and BenchmarkCodeUnmarshal
benchmarks from the Go standard library's encoding/json package since
they do an excellent job of stress testing parallel encoding/decoding (a
common use case in a database driver) and are how the lock contention
that led to this commit were discovered.

```
goos: darwin
goarch: arm64
pkg: go.mongodb.org/mongo-driver/bson
                      │ base.20.txt │             new.20.txt              │
                      │   sec/op    │   sec/op     vs base                │
CodeUnmarshal/BSON-10   3.192m ± 1%   2.246m ± 1%  -29.64% (p=0.000 n=20)
CodeUnmarshal/JSON-10   2.735m ± 1%   2.737m ± 0%        ~ (p=0.640 n=20)
CodeMarshal/BSON-10     2.972m ± 0%   1.221m ± 3%  -58.93% (p=0.000 n=20)
CodeMarshal/JSON-10     471.0µ ± 1%   464.6µ ± 0%   -1.36% (p=0.000 n=20)
geomean                 1.870m        1.366m       -26.92%

                      │ base.20.txt  │               new.20.txt               │
                      │     B/s      │      B/s       vs base                 │
CodeUnmarshal/BSON-10   579.7Mi ± 1%    823.9Mi ± 1%   +42.13% (p=0.000 n=20)
CodeUnmarshal/JSON-10   676.6Mi ± 1%    676.2Mi ± 0%         ~ (p=0.640 n=20)
CodeMarshal/BSON-10     622.7Mi ± 0%   1516.2Mi ± 3%  +143.46% (p=0.000 n=20)
CodeMarshal/JSON-10     3.837Gi ± 1%    3.890Gi ± 0%    +1.38% (p=0.000 n=20)
geomean                 989.8Mi         1.323Gi        +36.84%

                      │ base.20.txt  │             new.20.txt              │
                      │     B/op     │     B/op      vs base               │
CodeUnmarshal/BSON-10   4.219Mi ± 0%   4.219Mi ± 0%       ~ (p=0.077 n=20)
CodeUnmarshal/JSON-10   2.904Mi ± 0%   2.904Mi ± 0%       ~ (p=0.672 n=20)
CodeMarshal/BSON-10     2.821Mi ± 1%   2.776Mi ± 2%  -1.59% (p=0.023 n=20)
CodeMarshal/JSON-10     1.857Mi ± 0%   1.859Mi ± 0%       ~ (p=0.331 n=20)
geomean                 2.830Mi        2.820Mi       -0.37%

                      │ base.20.txt │              new.20.txt              │
                      │  allocs/op  │  allocs/op   vs base                 │
CodeUnmarshal/BSON-10   230.4k ± 0%   230.4k ± 0%       ~ (p=1.000 n=20)
CodeUnmarshal/JSON-10   92.67k ± 0%   92.67k ± 0%       ~ (p=1.000 n=20) ¹
CodeMarshal/BSON-10     94.07k ± 0%   94.07k ± 0%       ~ (p=0.112 n=20)
CodeMarshal/JSON-10      1.000 ± 0%    1.000 ± 0%       ~ (p=1.000 n=20) ¹
geomean                 6.694k        6.694k       +0.00%
¹ all samples are equal
```

Author: charlievieth

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(codeJSON)))
+	})
+	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(codeJSON)))
+	})
+	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,151 @@
+package bsoncodec
+
+import (
+	"reflect"
+	"sync"
+	"sync/atomic"
+)
+
+// 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
+}