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
```

bson/bsoncodec/codec_cache.go

@@ -0,0 +1,157 @@
+// 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"
+)
+
+// 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
+}

bson/bsoncodec/codec_cache_test.go

@@ -0,0 +1,174 @@
+// 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"
+ "strconv"
+ "strings"
+ "testing"
+)
+
+var codecCacheTestTypes = [16]reflect.Type{
+ reflect.TypeOf(uint8(0)),
+ reflect.TypeOf(uint16(0)),
+ reflect.TypeOf(uint32(0)),
+ reflect.TypeOf(uint64(0)),
+ reflect.TypeOf(uint(0)),
+ reflect.TypeOf(uintptr(0)),
+ reflect.TypeOf(int8(0)),
+ reflect.TypeOf(int16(0)),
+ reflect.TypeOf(int32(0)),
+ reflect.TypeOf(int64(0)),
+ reflect.TypeOf(int(0)),
+ reflect.TypeOf(float32(0)),
+ reflect.TypeOf(float64(0)),
+ reflect.TypeOf(true),
+ reflect.TypeOf(struct{ A int }{}),
+ reflect.TypeOf(map[int]int{}),
+}
+
+func TestTypeCache(t *testing.T) {
+ rt := reflect.TypeOf(int(0))
+ ec := new(typeEncoderCache)
+ dc := new(typeDecoderCache)
+
+ codec := new(fakeCodec)
+ ec.Store(rt, codec)
+ dc.Store(rt, codec)
+ if v, ok := ec.Load(rt); !ok || !reflect.DeepEqual(v, codec) {
+ t.Errorf("Load(%s) = %v, %t; want: %v, %t", rt, v, ok, codec, true)
+ }
+ if v, ok := dc.Load(rt); !ok || !reflect.DeepEqual(v, codec) {
+ t.Errorf("Load(%s) = %v, %t; want: %v, %t", rt, v, ok, codec, true)
+ }
+
+ // Make sure we overwrite the stored value with nil
+ ec.Store(rt, nil)
+ dc.Store(rt, nil)
+ if v, ok := ec.Load(rt); ok || v != nil {
+ t.Errorf("Load(%s) = %v, %t; want: %v, %t", rt, v, ok, nil, false)
+ }
+ if v, ok := dc.Load(rt); ok || v != nil {
+ t.Errorf("Load(%s) = %v, %t; want: %v, %t", rt, v, ok, nil, false)
+ }
+}
+
+func TestTypeCacheClone(t *testing.T) {
+ codec := new(fakeCodec)
+ ec1 := new(typeEncoderCache)
+ dc1 := new(typeDecoderCache)
+ for _, rt := range codecCacheTestTypes {
+ ec1.Store(rt, codec)
+ dc1.Store(rt, codec)
+ }
+ ec2 := ec1.Clone()
+ dc2 := dc1.Clone()
+ for _, rt := range codecCacheTestTypes {
+ if v, _ := ec2.Load(rt); !reflect.DeepEqual(v, codec) {
+ t.Errorf("Load(%s) = %#v; want: %#v", rt, v, codec)
+ }
+ if v, _ := dc2.Load(rt); !reflect.DeepEqual(v, codec) {
+ t.Errorf("Load(%s) = %#v; want: %#v", rt, v, codec)
+ }
+ }
+}
+
+func TestKindCacheArray(t *testing.T) {
+ // Check array bounds
+ var c kindEncoderCache
+ codec := new(fakeCodec)
+ c.Store(reflect.UnsafePointer, codec) // valid
+ c.Store(reflect.UnsafePointer+1, codec) // ignored
+ if v, ok := c.Load(reflect.UnsafePointer); !ok || v != codec {
+ t.Errorf("Load(reflect.UnsafePointer) = %v, %t; want: %v, %t", v, ok, codec, true)
+ }
+ if v, ok := c.Load(reflect.UnsafePointer + 1); ok || v != nil {
+ t.Errorf("Load(reflect.UnsafePointer + 1) = %v, %t; want: %v, %t", v, ok, nil, false)
+ }
+
+ // Make sure that reflect.UnsafePointer is the last/largest reflect.Type.
+ //
+ // The String() method of invalid reflect.Type types are of the format
+ // "kind{NUMBER}".
+ for rt := reflect.UnsafePointer + 1; rt < reflect.UnsafePointer+16; rt++ {
+ s := rt.String()
+ if !strings.Contains(s, strconv.Itoa(int(rt))) {
+ t.Errorf("reflect.Type(%d) appears to be valid: %q", rt, s)
+ }
+ }
+}
+
+func TestKindCacheClone(t *testing.T) {
+ e1 := new(kindEncoderCache)
+ d1 := new(kindDecoderCache)
+ codec := new(fakeCodec)
+ for k := reflect.Invalid; k <= reflect.UnsafePointer; k++ {
+ e1.Store(k, codec)
+ d1.Store(k, codec)
+ }
+ e2 := e1.Clone()
+ for k := reflect.Invalid; k <= reflect.UnsafePointer; k++ {
+ v1, ok1 := e1.Load(k)
+ v2, ok2 := e2.Load(k)
+ if ok1 != ok2 || !reflect.DeepEqual(v1, v2) || v1 == nil || v2 == nil {
+ t.Errorf("Encoder(%s): %#v, %t != %#v, %t", k, v1, ok1, v2, ok2)
+ }
+ }
+ d2 := d1.Clone()
+ for k := reflect.Invalid; k <= reflect.UnsafePointer; k++ {
+ v1, ok1 := d1.Load(k)
+ v2, ok2 := d2.Load(k)
+ if ok1 != ok2 || !reflect.DeepEqual(v1, v2) || v1 == nil || v2 == nil {
+ t.Errorf("Decoder(%s): %#v, %t != %#v, %t", k, v1, ok1, v2, ok2)
+ }
+ }
+}
+
+func TestKindCacheEncoderNilEncoder(t *testing.T) {
+ t.Run("Encoder", func(t *testing.T) {
+ c := new(kindEncoderCache)
+ c.Store(reflect.Invalid, ValueEncoder(nil))
+ v, ok := c.Load(reflect.Invalid)
+ if v != nil || ok {
+ t.Errorf("Load of nil ValueEncoder should return: nil, false; got: %v, %t", v, ok)
+ }
+ })
+ t.Run("Decoder", func(t *testing.T) {
+ c := new(kindDecoderCache)
+ c.Store(reflect.Invalid, ValueDecoder(nil))
+ v, ok := c.Load(reflect.Invalid)
+ if v != nil || ok {
+ t.Errorf("Load of nil ValueDecoder should return: nil, false; got: %v, %t", v, ok)
+ }
+ })
+}
+
+func BenchmarkEncoderCacheLoad(b *testing.B) {
+ c := new(typeEncoderCache)
+ codec := new(fakeCodec)
+ typs := codecCacheTestTypes
+ for _, t := range typs {
+ c.Store(t, codec)
+ }
+ b.RunParallel(func(pb *testing.PB) {
+ for i := 0; pb.Next(); i++ {
+ c.Load(typs[i%len(typs)])
+ }
+ })
+}
+
+func BenchmarkEncoderCacheStore(b *testing.B) {
+ c := new(typeEncoderCache)
+ codec := new(fakeCodec)
+ b.RunParallel(func(pb *testing.PB) {
+ typs := codecCacheTestTypes
+ for i := 0; pb.Next(); i++ {
+ c.Store(typs[i%len(typs)], codec)
+ }
+ })
+}