fix: improve json marshal performance, reduce necessary allocations

pkg/json/json.go

@@ -21,9 +21,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,7 +34,7 @@ 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) {
@@ -49,7 +47,7 @@ func MarshalTo(v any, buf []byte) []byte {
 	return appendMarshal(reflect.ValueOf(v), buf)
 }
 
-const bufSize = 1024
+const bufSize = 8192
 
 var bytesPool = sync.Pool{
 	New: func() any {

pkg/json/map.go

@@ -7,18 +7,17 @@ import (
 type Map[V any] map[string]V
 
 func (m Map[V]) MarshalJSONTo(buf []byte) []byte {
+	oldN := len(buf)
 	buf = append(buf, '{')
-	i := 0
-	n := len(m)
 	for k, v := range m {
 		buf = AppendString(buf, k)
 		buf = append(buf, ':')
 		buf = appendMarshal(reflect.ValueOf(v), buf)
-		if i != n-1 {
 		buf = append(buf, ',')
 	}
-		i++
+	n := len(buf)
+	if oldN != n {
+		buf = buf[:n-1]
 	}
-	buf = append(buf, '}')
+	return append(buf, '}')
-	return buf
 }

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,
 	}
@@ -73,20 +73,17 @@ func appendMarshal(v reflect.Value, buf []byte) []byte {
 	if v == nilValue {
 		return append(buf, "null"...)
 	}
-	kind := v.Kind()
+	marshalFunc, ok := marshalFuncByKind[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()))
 	}
 	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,46 +151,11 @@ 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()))
 	}
 	buf = append(buf, '{')
-	i := 0
 	oldN := len(buf)
 	iter := v.MapRange()
 	for iter.Next() {
@@ -201,7 +163,6 @@ func appendMap(v reflect.Value, buf []byte) []byte {
 		v := iter.Value()
 		buf = appendKV(k, v, buf)
 		buf = append(buf, ',')
-		i++
 	}
 	n := len(buf)
 	if oldN != n {
@@ -230,10 +191,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/marshal_test.go

@@ -147,12 +147,12 @@ func TestMarshal(t *testing.T) {
 		{
 			name:     "slice_of_struct",
 			input:    []testStruct{{Name: "John", Age: 30, Score: 8.5}, {Name: "Jane", Age: 25, Score: 9.5}},
-			expected: `[{"name":"John","age":30,"score":8.50},{"name":"Jane","age":25,"score":9.50}]`,
+			expected: `[{"name":"John","age":30,"score":8.5},{"name":"Jane","age":25,"score":9.5}]`,
 		},
 		{
 			name:     "slice_of_struct_pointer",
 			input:    []*testStruct{{Name: "John", Age: 30, Score: 8.5}, {Name: "Jane", Age: 25, Score: 9.5}},
-			expected: `[{"name":"John","age":30,"score":8.50},{"name":"Jane","age":25,"score":9.50}]`,
+			expected: `[{"name":"John","age":30,"score":8.5},{"name":"Jane","age":25,"score":9.5}]`,
 		},
 		{
 			name:     "slice_of_map",
@@ -162,12 +162,12 @@ func TestMarshal(t *testing.T) {
 		{
 			name:     "struct",
 			input:    testStruct{Name: "John", Age: 30, Score: 8.5},
-			expected: `{"name":"John","age":30,"score":8.50}`,
+			expected: `{"name":"John","age":30,"score":8.5}`,
 		},
 		{
 			name:     "struct_pointer",
 			input:    &testStruct{Name: "Jane", Age: 25, Score: 9.5},
-			expected: `{"name":"Jane","age":25,"score":9.50}`,
+			expected: `{"name":"Jane","age":25,"score":9.5}`,
 		},
 		{
 			name:     "byte_slice",

pkg/json/struct.go

@@ -3,7 +3,6 @@ package json
 import (
 	"fmt"
 	"reflect"
-	"strconv"
 
 	"github.com/yusing/go-proxy/internal/utils/strutils"
 )
@@ -20,6 +19,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,19 +31,64 @@ 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 {
-	for i := range t.NumField() {
+		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, ',')
+			continue
+		}
+		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 {
+	n := t.NumField()
+	fields := make([]*field, 0, n)
+	for i := range n {
 		structField := t.Field(i)
 		if !structField.IsExported() {
 			continue
 		}
-		kind := structField.Type.Kind()
+		t := structField.Type
+		kind := t.Kind()
 		f := &field{
 			index: i,
 			isPtr: kind == reflect.Pointer,
@@ -55,6 +103,7 @@ func flattenFields(t reflect.Type) []*field {
 				switch parts[1] {
 				case tagOmitEmpty:
 					f.omitEmpty = true
+					f.checkEmpty = checkEmptyFuncs[kind]
 				case tagString:
 					f.marshal = appendStringRepr
 				case tagByteSize:
@@ -79,27 +128,21 @@ func flattenFields(t reflect.Type) []*field {
 			f.marshal = appendMarshal
 		}
 		if structField.Anonymous {
-			t := structField.Type
+			if kind == reflect.Pointer {
-			if t.Kind() == reflect.Pointer {
 				t = t.Elem()
+				kind = t.Kind()
+				f.isPtr = true
 				f.omitEmpty = true
+				f.checkEmpty = checkEmptyFuncs[kind]
 			}
-			if t.Kind() == reflect.Struct {
+			if kind == reflect.Struct {
 				f.inner = flattenFields(t)
 				f.hasInner = len(f.inner) > 0
 			}
 		}
 		fields = append(fields, f)
-		if f.omitEmpty {
+		quotedNameWithCol := AppendString(nil, f.quotedNameWithCol)
-			f.checkEmpty = checkEmptyFuncs[kind]
+		f.quotedNameWithCol = string(quotedNameWithCol) + ":"
 	}
-		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...))
-}