Optimize getting of cached decoder

codec.go

@@ -0,0 +1,112 @@
+package json
+
+import (
+	"fmt"
+	"reflect"
+	"sync/atomic"
+	"unsafe"
+)
+
+const (
+	maxAcceptableTypeAddrRange = 1024 * 1024 * 2 // 2 Mib
+)
+
+var (
+	cachedOpcodeSets       []*opcodeSet
+	cachedOpcodeMap        unsafe.Pointer // map[uintptr]*opcodeSet
+	existsCachedOpcodeSets bool
+	cachedDecoder          []decoder
+	cachedDecoderMap       unsafe.Pointer // map[uintptr]decoder
+	existsCachedDecoder    bool
+	baseTypeAddr           uintptr
+)
+
+//go:linkname typelinks reflect.typelinks
+func typelinks() ([]unsafe.Pointer, [][]int32)
+
+//go:linkname rtypeOff reflect.rtypeOff
+func rtypeOff(unsafe.Pointer, int32) unsafe.Pointer
+
+func setupCodec() error {
+	sections, offsets := typelinks()
+	if len(sections) != 1 {
+		return fmt.Errorf("failed to get sections")
+	}
+	if len(offsets) != 1 {
+		return fmt.Errorf("failed to get offsets")
+	}
+	section := sections[0]
+	offset := offsets[0]
+	var (
+		min uintptr = uintptr(^uint(0))
+		max uintptr = 0
+	)
+	for i := 0; i < len(offset); i++ {
+		typ := (*rtype)(rtypeOff(section, offset[i]))
+		addr := uintptr(unsafe.Pointer(typ))
+		if min > addr {
+			min = addr
+		}
+		if max < addr {
+			max = addr
+		}
+		if typ.Kind() == reflect.Ptr {
+			addr = uintptr(unsafe.Pointer(typ.Elem()))
+			if min > addr {
+				min = addr
+			}
+			if max < addr {
+				max = addr
+			}
+		}
+	}
+	addrRange := max - min
+	if addrRange == 0 {
+		return fmt.Errorf("failed to get address range of types")
+	}
+	if addrRange > maxAcceptableTypeAddrRange {
+		return fmt.Errorf("too big address range %d", addrRange)
+	}
+	cachedOpcodeSets = make([]*opcodeSet, addrRange)
+	existsCachedOpcodeSets = true
+	cachedDecoder = make([]decoder, addrRange)
+	existsCachedDecoder = true
+	baseTypeAddr = min
+	return nil
+}
+
+func init() {
+	_ = setupCodec()
+}
+
+func loadOpcodeMap() map[uintptr]*opcodeSet {
+	p := atomic.LoadPointer(&cachedOpcodeMap)
+	return *(*map[uintptr]*opcodeSet)(unsafe.Pointer(&p))
+}
+
+func storeOpcodeSet(typ uintptr, set *opcodeSet, m map[uintptr]*opcodeSet) {
+	newOpcodeMap := make(map[uintptr]*opcodeSet, len(m)+1)
+	newOpcodeMap[typ] = set
+
+	for k, v := range m {
+		newOpcodeMap[k] = v
+	}
+
+	atomic.StorePointer(&cachedOpcodeMap, *(*unsafe.Pointer)(unsafe.Pointer(&newOpcodeMap)))
+}
+
+func loadDecoderMap() map[uintptr]decoder {
+	p := atomic.LoadPointer(&cachedDecoderMap)
+	return *(*map[uintptr]decoder)(unsafe.Pointer(&p))
+}
+
+func storeDecoder(typ uintptr, dec decoder, m map[uintptr]decoder) {
+	newDecoderMap := make(map[uintptr]decoder, len(m)+1)
+	newDecoderMap[typ] = dec
+
+	for k, v := range m {
+		newDecoderMap[k] = v
+	}
+
+	atomic.StorePointer(&cachedDecoderMap, *(*unsafe.Pointer)(unsafe.Pointer(&newDecoderMap)))
+}

decode.go

@@ -5,7 +5,6 @@ import (
 	"io"
 	"reflect"
 	"strconv"
-	"sync"
 	"unsafe"
 )
 
@@ -25,31 +24,11 @@ type Decoder struct {
 	structTypeToDecoder map[uintptr]decoder
 }
 
-type decoderMap struct {
-	sync.Map
-}
-
-func (m *decoderMap) get(k uintptr) decoder {
-	if v, ok := m.Load(k); ok {
-		return v.(decoder)
-	}
-	return nil
-}
-
-func (m *decoderMap) set(k uintptr, dec decoder) {
-	m.Store(k, dec)
-}
-
 var (
-	cachedDecoder     decoderMap
 	unmarshalJSONType = reflect.TypeOf((*Unmarshaler)(nil)).Elem()
 	unmarshalTextType = reflect.TypeOf((*encoding.TextUnmarshaler)(nil)).Elem()
 )
 
-func init() {
-	cachedDecoder = decoderMap{}
-}
-
 const (
 	nul = '\000'
 )
@@ -89,16 +68,10 @@ func (d *Decoder) decode(src []byte, header *interfaceHeader) error {
 	if err := d.validateType(copiedType, ptr); err != nil {
 		return err
 	}
-	dec := cachedDecoder.get(typeptr)
+	dec, err := d.compileToGetDecoder(typeptr, typ)
-	if dec == nil {
-		d.structTypeToDecoder = map[uintptr]decoder{}
-		compiledDec, err := d.compileHead(copiedType)
 	if err != nil {
 		return err
 	}
-		cachedDecoder.set(typeptr, compiledDec)
-		dec = compiledDec
-	}
 	if _, err := dec.decode(src, 0, header.ptr); err != nil {
 		return err
 	}
@@ -154,16 +127,10 @@ func (d *Decoder) Decode(v interface{}) error {
 		return err
 	}
 
-	dec := cachedDecoder.get(typeptr)
+	dec, err := d.compileToGetDecoder(typeptr, typ)
-	if dec == nil {
-		d.structTypeToDecoder = map[uintptr]decoder{}
-		compiledDec, err := d.compileHead(typ)
 	if err != nil {
 		return err
 	}
-		cachedDecoder.set(typeptr, compiledDec)
-		dec = compiledDec
-	}
 	if err := d.prepareForDecode(); err != nil {
 		return err
 	}

decode_compile.go

@@ -6,6 +6,21 @@ import (
 	"unsafe"
 )
 
+func (d *Decoder) compileToGetDecoderSlowPath(typeptr uintptr, typ *rtype) (decoder, error) {
+	decoderMap := loadDecoderMap()
+	if dec, exists := decoderMap[typeptr]; exists {
+		return dec, nil
+	}
+
+	d.structTypeToDecoder = map[uintptr]decoder{}
+	dec, err := d.compileHead(typ)
+	if err != nil {
+		return nil, err
+	}
+	storeDecoder(typeptr, dec, decoderMap)
+	return dec, nil
+}
+
 func (d *Decoder) compileHead(typ *rtype) (decoder, error) {
 	switch {
 	case rtype_ptrTo(typ).Implements(unmarshalJSONType):

decode_compile_norace.go

@@ -0,0 +1,22 @@
+// +build !race
+
+package json
+
+func (d *Decoder) compileToGetDecoder(typeptr uintptr, typ *rtype) (decoder, error) {
+	if !existsCachedDecoder {
+		return d.compileToGetDecoderSlowPath(typeptr, typ)
+	}
+
+	index := typeptr - baseTypeAddr
+	if dec := cachedDecoder[index]; dec != nil {
+		return dec, nil
+	}
+
+	d.structTypeToDecoder = map[uintptr]decoder{}
+	dec, err := d.compileHead(typ)
+	if err != nil {
+		return nil, err
+	}
+	cachedDecoder[index] = dec
+	return dec, nil
+}

decode_compile_race.go

@@ -0,0 +1,31 @@
+// +build race
+
+package json
+
+import "sync"
+
+var decMu sync.RWMutex
+
+func (d *Decoder) compileToGetDecoder(typeptr uintptr, typ *rtype) (decoder, error) {
+	if !existsCachedDecoder {
+		return d.compileToGetDecoderSlowPath(typeptr, typ)
+	}
+
+	index := typeptr - baseTypeAddr
+	decMu.RLock()
+	if dec := cachedDecoder[index]; dec != nil {
+		decMu.RUnlock()
+		return dec, nil
+	}
+	decMu.RUnlock()
+
+	d.structTypeToDecoder = map[uintptr]decoder{}
+	dec, err := d.compileHead(typ)
+	if err != nil {
+		return nil, err
+	}
+	decMu.Lock()
+	cachedDecoder[index] = dec
+	decMu.Unlock()
+	return dec, nil
+}

encode_compile.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
-	"sync/atomic"
 	"unsafe"
 )
 
@@ -24,72 +23,8 @@ type opcodeSet struct {
 var (
 	marshalJSONType = reflect.TypeOf((*Marshaler)(nil)).Elem()
 	marshalTextType = reflect.TypeOf((*encoding.TextMarshaler)(nil)).Elem()
-	cachedOpcode           unsafe.Pointer // map[uintptr]*opcodeSet
-	baseTypeAddr           uintptr
-	cachedOpcodeSets       []*opcodeSet
-	existsCachedOpcodeSets bool
 )
 
-const (
-	maxAcceptableTypeAddrRange = 1024 * 1024 * 2 // 2 Mib
-)
-
-//go:linkname typelinks reflect.typelinks
-func typelinks() ([]unsafe.Pointer, [][]int32)
-
-//go:linkname rtypeOff reflect.rtypeOff
-func rtypeOff(unsafe.Pointer, int32) unsafe.Pointer
-
-func setupOpcodeSets() error {
-	sections, offsets := typelinks()
-	if len(sections) != 1 {
-		return fmt.Errorf("failed to get sections")
-	}
-	if len(offsets) != 1 {
-		return fmt.Errorf("failed to get offsets")
-	}
-	section := sections[0]
-	offset := offsets[0]
-	var (
-		min uintptr = uintptr(^uint(0))
-		max uintptr = 0
-	)
-	for i := 0; i < len(offset); i++ {
-		typ := (*rtype)(rtypeOff(section, offset[i]))
-		addr := uintptr(unsafe.Pointer(typ))
-		if min > addr {
-			min = addr
-		}
-		if max < addr {
-			max = addr
-		}
-		if typ.Kind() == reflect.Ptr {
-			addr = uintptr(unsafe.Pointer(typ.Elem()))
-			if min > addr {
-				min = addr
-			}
-			if max < addr {
-				max = addr
-			}
-		}
-	}
-	addrRange := max - min
-	if addrRange == 0 {
-		return fmt.Errorf("failed to get address range of types")
-	}
-	if addrRange > maxAcceptableTypeAddrRange {
-		return fmt.Errorf("too big address range %d", addrRange)
-	}
-	cachedOpcodeSets = make([]*opcodeSet, addrRange)
-	existsCachedOpcodeSets = true
-	baseTypeAddr = min
-	return nil
-}
-
-func init() {
-	_ = setupOpcodeSets()
-}
-
 func encodeCompileToGetCodeSetSlowPath(typeptr uintptr) (*opcodeSet, error) {
 	opcodeMap := loadOpcodeMap()
 	if codeSet, exists := opcodeMap[typeptr]; exists {
@@ -117,22 +52,6 @@ func encodeCompileToGetCodeSetSlowPath(typeptr uintptr) (*opcodeSet, error) {
 	return codeSet, nil
 }
 
-func loadOpcodeMap() map[uintptr]*opcodeSet {
-	p := atomic.LoadPointer(&cachedOpcode)
-	return *(*map[uintptr]*opcodeSet)(unsafe.Pointer(&p))
-}
-
-func storeOpcodeSet(typ uintptr, set *opcodeSet, m map[uintptr]*opcodeSet) {
-	newOpcodeMap := make(map[uintptr]*opcodeSet, len(m)+1)
-	newOpcodeMap[typ] = set
-
-	for k, v := range m {
-		newOpcodeMap[k] = v
-	}
-
-	atomic.StorePointer(&cachedOpcode, *(*unsafe.Pointer)(unsafe.Pointer(&newOpcodeMap)))
-}
-
 func encodeCompileHead(ctx *encodeCompileContext) (*opcode, error) {
 	typ := ctx.typ
 	switch {