fix: improve json marshal performance, reduce necessary allocations

pkg/json/json.go

@@ -1,7 +1,6 @@
 package json
 
 import (
-	"reflect"
 	"sync"
 
 	"github.com/bytedance/sonic"
@@ -21,9 +20,7 @@ var (
 //
 // It's like json.Marshal, but with some differences:
 //
-//   - It's ~4-5x faster in most cases.
+//   - It supports custom Marshaler interface (MarshalJSONTo(buf []byte) []byte)
-//
-//   - It also supports custom Marshaler interface (MarshalJSONTo(buf []byte) []byte)
 //     to allow further optimizations.
 //
 //   - It leverages the strutils library.
@@ -36,20 +33,20 @@ var (
 //
 //     `use_marshaler` to force using the custom marshaler for primitive types declaration (e.g. `type Status int`).
 //
-//   - It correct the behavior of *url.URL and time.Duration.
+//   - It corrects the behavior of *url.URL and time.Duration.
 //
 //   - It does not support maps other than string-keyed maps.
 func Marshal(v any) ([]byte, error) {
 	buf := newBytes()
 	defer putBytes(buf)
-	return cloneBytes(appendMarshal(reflect.ValueOf(v), buf)), nil
+	return cloneBytes(appendMarshalAny(v, buf)), nil
 }
 
 func MarshalTo(v any, buf []byte) []byte {
-	return appendMarshal(reflect.ValueOf(v), buf)
+	return appendMarshalAny(v, buf)
 }
 
-const bufSize = 1024
+const bufSize = 8192
 
 var bytesPool = sync.Pool{
 	New: func() any {

pkg/json/marshal.go

@@ -20,7 +20,6 @@ var (
 	marshalFuncByKind map[reflect.Kind]marshalFunc
 
 	marshalFuncsByType = newCacheMap[reflect.Type, marshalFunc]()
-	flattenFieldsCache = newCacheMap[reflect.Type, []*field]()
 
 	nilValue = reflect.ValueOf(nil)
 )
@@ -44,6 +43,7 @@ func init() {
 		reflect.Map:       appendMap,
 		reflect.Slice:     appendArray,
 		reflect.Array:     appendArray,
+		reflect.Struct:    appendStruct,
 		reflect.Pointer:   appendPtrInterface,
 		reflect.Interface: appendPtrInterface,
 	}
@@ -69,17 +69,15 @@ func must(buf []byte, err error) []byte {
 	return buf
 }
 
+func appendMarshalAny(v any, buf []byte) []byte {
+	return appendMarshal(reflect.ValueOf(v), buf)
+}
+
 func appendMarshal(v reflect.Value, buf []byte) []byte {
 	if v == nilValue {
 		return append(buf, "null"...)
 	}
 	kind := v.Kind()
-	if kind == reflect.Struct {
-		if res, ok := appendWithCachedFunc(v, buf); ok {
-			return res
-		}
-		return appendStruct(v, buf)
-	}
 	marshalFunc, ok := marshalFuncByKind[kind]
 	if !ok {
 		panic(fmt.Errorf("unsupported type: %s", v.Type()))
@@ -87,6 +85,10 @@ func appendMarshal(v reflect.Value, buf []byte) []byte {
 	return marshalFunc(v, buf)
 }
 
+func cacheMarshalFunc(t reflect.Type, marshalFunc marshalFunc) {
+	marshalFuncsByType.Store(t, marshalFunc)
+}
+
 func appendWithCachedFunc(v reflect.Value, buf []byte) (res []byte, ok bool) {
 	marshalFunc, ok := marshalFuncsByType.Load(v.Type())
 	if ok {
@@ -154,40 +156,6 @@ func appendKV(k reflect.Value, v reflect.Value, buf []byte) []byte {
 	return appendMarshal(v, buf)
 }
 
-func appendStruct(v reflect.Value, buf []byte) []byte {
-	if res, ok := appendWithCustomMarshaler(v, buf); ok {
-		return res
-	}
-	buf = append(buf, '{')
-	oldN := len(buf)
-	fields := flattenFields(v.Type())
-
-	for _, f := range fields {
-		cur := v.Field(f.index)
-		if f.omitEmpty && f.checkEmpty(cur) {
-			continue
-		}
-		if !f.hasInner {
-			buf = f.appendKV(cur, buf)
-			buf = append(buf, ',')
-		} else {
-			if f.isPtr {
-				cur = cur.Elem()
-			}
-			for _, inner := range f.inner {
-				buf = inner.appendKV(cur.Field(inner.index), buf)
-				buf = append(buf, ',')
-			}
-		}
-	}
-
-	n := len(buf)
-	if oldN != n {
-		buf = buf[:n-1]
-	}
-	return append(buf, '}')
-}
-
 func appendMap(v reflect.Value, buf []byte) []byte {
 	if v.Type().Key().Kind() != reflect.String {
 		panic(fmt.Errorf("map key must be string: %s", v.Type()))
@@ -230,10 +198,6 @@ func appendArray(v reflect.Value, buf []byte) []byte {
 	return append(buf, ']')
 }
 
-func cacheMarshalFunc(t reflect.Type, marshalFunc marshalFunc) {
-	marshalFuncsByType.Store(t, marshalFunc)
-}
-
 func appendPtrInterface(v reflect.Value, buf []byte) []byte {
 	return appendMarshal(v.Elem(), buf)
 }

pkg/json/struct.go

@@ -20,6 +20,10 @@ type field struct {
 	marshal    marshalFunc
 }
 
+func (f *field) appendKV(v reflect.Value, buf []byte) []byte {
+	return f.marshal(v, append(buf, f.quotedNameWithCol...))
+}
+
 const (
 	tagOmitEmpty    = "omitempty"
 	tagString       = "string" // https://pkg.go.dev/github.com/yusing/go-proxy/pkg/json#Marshal
@@ -28,13 +32,56 @@ const (
 	tagUseMarshaler = "use_marshaler"
 )
 
-func flattenFields(t reflect.Type) []*field {
+func appendStruct(v reflect.Value, buf []byte) []byte {
-	fields, ok := flattenFieldsCache.Load(t)
+	if res, ok := appendWithCachedFunc(v, buf); ok {
-	if ok {
+		return res
-		return fields
 	}
 
-	fields = make([]*field, 0, t.NumField())
+	if res, ok := appendWithCustomMarshaler(v, buf); ok {
+		return res
+	}
+
+	t := v.Type()
+	fields := flattenFields(t)
+	marshalFn := func(v reflect.Value, buf []byte) []byte {
+		return appendFields(v, fields, buf)
+	}
+	cacheMarshalFunc(t, marshalFn)
+	return marshalFn(v, buf)
+}
+
+func appendFields(v reflect.Value, fields []*field, buf []byte) []byte {
+	buf = append(buf, '{')
+	oldN := len(buf)
+
+	for _, f := range fields {
+		cur := v.Field(f.index)
+		if f.omitEmpty && f.checkEmpty(cur) {
+			continue
+		}
+		if !f.hasInner {
+			buf = f.appendKV(cur, buf)
+			buf = append(buf, ',')
+		} else {
+			if f.isPtr {
+				cur = cur.Elem()
+			}
+			for _, inner := range f.inner {
+				buf = inner.appendKV(cur.Field(inner.index), buf)
+				buf = append(buf, ',')
+			}
+		}
+	}
+
+	n := len(buf)
+	if oldN != n {
+		buf = buf[:n-1]
+	}
+	return append(buf, '}')
+}
+
+func flattenFields(t reflect.Type) []*field {
+	fields := make([]*field, 0, t.NumField())
 	for i := range t.NumField() {
 		structField := t.Field(i)
 		if !structField.IsExported() {
@@ -95,11 +142,5 @@ func flattenFields(t reflect.Type) []*field {
 		}
 		f.quotedNameWithCol = strconv.Quote(f.quotedNameWithCol) + ":"
 	}
-
-	flattenFieldsCache.Store(t, fields)
 	return fields
 }
-
-func (f *field) appendKV(v reflect.Value, buf []byte) []byte {
-	return f.marshal(v, append(buf, f.quotedNameWithCol...))
-}