Cache version specific encoders/decoders. (segmentio#652)

Fixes segmentio#651.
deeplon · May 3, 2021 · 8462374 · 8462374
1 parent d15a635
commit 8462374
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/protocol/decode.go b/protocol/decode.go
@@ -498,8 +498,9 @@ func readInt64(b []byte) int64 {
 
 func Unmarshal(data []byte, version int16, value interface{}) error {
 	typ := elemTypeOf(value)
-	cache, _ := unmarshalers.Load().(map[_type]decodeFunc)
-	decode := cache[typ]
+	cache, _ := unmarshalers.Load().(map[versionedType]decodeFunc)
+	key := versionedType{typ: typ, version: version}
+	decode := cache[key]
 
 	if decode == nil {
 		decode = decodeFuncOf(reflect.TypeOf(value).Elem(), version, false, structTag{
@@ -510,8 +511,8 @@ func Unmarshal(data []byte, version int16, value interface{}) error {
 			Nullable:   true,
 		})
 
-		newCache := make(map[_type]decodeFunc, len(cache)+1)
-		newCache[typ] = decode
+		newCache := make(map[versionedType]decodeFunc, len(cache)+1)
+		newCache[key] = decode
 
 		for typ, fun := range cache {
 			newCache[typ] = fun
@@ -541,5 +542,5 @@ func Unmarshal(data []byte, version int16, value interface{}) error {
 
 var (
 	decoders     sync.Pool    // *decoder
-	unmarshalers atomic.Value // map[_type]decodeFunc
+	unmarshalers atomic.Value // map[versionedType]decodeFunc
 )
diff --git a/protocol/encode.go b/protocol/encode.go
@@ -587,8 +587,9 @@ func writeInt64(b []byte, i int64) {
 
 func Marshal(version int16, value interface{}) ([]byte, error) {
 	typ := typeOf(value)
-	cache, _ := marshalers.Load().(map[_type]encodeFunc)
-	encode := cache[typ]
+	cache, _ := marshalers.Load().(map[versionedType]encodeFunc)
+	key := versionedType{typ: typ, version: version}
+	encode := cache[key]
 
 	if encode == nil {
 		encode = encodeFuncOf(reflect.TypeOf(value), version, false, structTag{
@@ -599,8 +600,8 @@ func Marshal(version int16, value interface{}) ([]byte, error) {
 			Nullable:   true,
 		})
 
-		newCache := make(map[_type]encodeFunc, len(cache)+1)
-		newCache[typ] = encode
+		newCache := make(map[versionedType]encodeFunc, len(cache)+1)
+		newCache[key] = encode
 
 		for typ, fun := range cache {
 			newCache[typ] = fun
@@ -633,7 +634,12 @@ func Marshal(version int16, value interface{}) ([]byte, error) {
 	return out, nil
 }
 
+type versionedType struct {
+	typ     _type
+	version int16
+}
+
 var (
 	encoders   sync.Pool    // *encoder
-	marshalers atomic.Value // map[_type]encodeFunc
+	marshalers atomic.Value // map[versionedType]encodeFunc
 )