encoding/json: give it a chance to put encodeState back in pool when error occurs

name                       old time/op    new time/op    delta
CodeEncoderError-10           688µs ± 8%     496µs ±15%   -27.92%  (p=0.000 n=10+9)
CodeMarshalError-10           747µs ± 6%     546µs ± 4%   -26.86%  (p=0.000 n=10+10)
MarshalBytesError/32-10       284µs ± 2%     273µs ± 1%    -3.84%  (p=0.000 n=10+10)
MarshalBytesError/256-10      281µs ± 2%     278µs ± 4%      ~     (p=0.053 n=9+10)
MarshalBytesError/4096-10     290µs ± 1%     279µs ± 3%    -3.52%  (p=0.000 n=10+10)

name                       old speed      new speed      delta
CodeEncoderError-10        2.83GB/s ± 8%  3.84GB/s ±20%   +36.03%  (p=0.000 n=10+10)
CodeMarshalError-10        2.60GB/s ± 5%  3.56GB/s ± 4%   +36.61%  (p=0.000 n=10+10)

name                       old alloc/op   new alloc/op   delta
CodeEncoderError-10          4.05MB ± 1%    0.00MB ± 1%  -100.00%  (p=0.000 n=10+9)
CodeMarshalError-10          6.05MB ± 0%    1.99MB ± 1%   -67.13%  (p=0.000 n=10+10)
MarshalBytesError/32-10      66.0kB ± 0%     0.2kB ± 0%   -99.67%  (p=0.000 n=9+8)
MarshalBytesError/256-10     50.1kB ± 0%     0.9kB ± 0%   -98.23%  (p=0.000 n=9+9)
MarshalBytesError/4096-10    87.4kB ± 0%     7.5kB ± 0%   -91.47%  (p=0.000 n=8+10)

name                       old allocs/op  new allocs/op  delta
CodeEncoderError-10            25.0 ± 0%       4.0 ± 0%   -84.00%  (p=0.000 n=9+10)
CodeMarshalError-10            27.0 ± 0%       6.0 ± 0%   -77.78%  (p=0.000 n=10+10)
MarshalBytesError/32-10        18.0 ± 0%       5.0 ± 0%   -72.22%  (p=0.000 n=10+10)
MarshalBytesError/256-10       17.0 ± 0%       6.0 ± 0%   -64.71%  (p=0.000 n=10+10)
MarshalBytesError/4096-10      16.0 ± 0%       6.0 ± 0%   -62.50%  (p=0.000 n=10+10)

Change-Id: I48070bb05f55707251c694e40d2570403bbf61f8
Reviewed-on: https://go-review.googlesource.com/c/go/+/423694
Auto-Submit: Ian Lance Taylor <iant@google.com>
Reviewed-by: Daniel Martí <mvdan@mvdan.cc>
Run-TryBot: Ian Lance Taylor <iant@google.com>
Reviewed-by: Ian Lance Taylor <iant@google.com>
Reviewed-by: David Chase <drchase@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

5 files changed, 178 insertions, 3 deletions

diff --git a/src/encoding/json/bench_test.go b/src/encoding/json/bench_test.go
index 95609140b0..133084976b 100644
--- a/src/encoding/json/bench_test.go
+++ b/src/encoding/json/bench_test.go
@@ -99,6 +99,36 @@ func BenchmarkCodeEncoder(b *testing.B) {
 	b.SetBytes(int64(len(codeJSON)))
 }
 
+func BenchmarkCodeEncoderError(b *testing.B) {
+	b.ReportAllocs()
+	if codeJSON == nil {
+		b.StopTimer()
+		codeInit()
+		b.StartTimer()
+	}
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+
+	b.RunParallel(func(pb *testing.PB) {
+		enc := NewEncoder(io.Discard)
+		for pb.Next() {
+			if err := enc.Encode(&codeStruct); err != nil {
+				b.Fatal("Encode:", err)
+			}
+			if _, err := Marshal(dummy); err == nil {
+				b.Fatal("expect an error here")
+			}
+		}
+	})
+	b.SetBytes(int64(len(codeJSON)))
+}
+
 func BenchmarkCodeMarshal(b *testing.B) {
 	b.ReportAllocs()
 	if codeJSON == nil {
@@ -116,6 +146,35 @@ func BenchmarkCodeMarshal(b *testing.B) {
 	b.SetBytes(int64(len(codeJSON)))
 }
 
+func BenchmarkCodeMarshalError(b *testing.B) {
+	b.ReportAllocs()
+	if codeJSON == nil {
+		b.StopTimer()
+		codeInit()
+		b.StartTimer()
+	}
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			if _, err := Marshal(&codeStruct); err != nil {
+				b.Fatal("Marshal:", err)
+			}
+			if _, err := Marshal(dummy); err == nil {
+				b.Fatal("expect an error here")
+			}
+		}
+	})
+	b.SetBytes(int64(len(codeJSON)))
+}
+
 func benchMarshalBytes(n int) func(*testing.B) {
 	sample := []byte("hello world")
 	// Use a struct pointer, to avoid an allocation when passing it as an
@@ -134,6 +193,36 @@ func benchMarshalBytes(n int) func(*testing.B) {
 	}
 }
 
+func benchMarshalBytesError(n int) func(*testing.B) {
+	sample := []byte("hello world")
+	// Use a struct pointer, to avoid an allocation when passing it as an
+	// interface parameter to Marshal.
+	v := &struct {
+		Bytes []byte
+	}{
+		bytes.Repeat(sample, (n/len(sample))+1)[:n],
+	}
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+
+	return func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			if _, err := Marshal(v); err != nil {
+				b.Fatal("Marshal:", err)
+			}
+			if _, err := Marshal(dummy); err == nil {
+				b.Fatal("expect an error here")
+			}
+		}
+	}
+}
+
 func BenchmarkMarshalBytes(b *testing.B) {
 	b.ReportAllocs()
 	// 32 fits within encodeState.scratch.
@@ -145,6 +234,17 @@ func BenchmarkMarshalBytes(b *testing.B) {
 	b.Run("4096", benchMarshalBytes(4096))
 }
 
+func BenchmarkMarshalBytesError(b *testing.B) {
+	b.ReportAllocs()
+	// 32 fits within encodeState.scratch.
+	b.Run("32", benchMarshalBytesError(32))
+	// 256 doesn't fit in encodeState.scratch, but is small enough to
+	// allocate and avoid the slower base64.NewEncoder.
+	b.Run("256", benchMarshalBytesError(256))
+	// 4096 is large enough that we want to avoid allocating for it.
+	b.Run("4096", benchMarshalBytesError(4096))
+}
+
 func BenchmarkCodeDecoder(b *testing.B) {
 	b.ReportAllocs()
 	if codeJSON == nil {
diff --git a/src/encoding/json/encode.go b/src/encoding/json/encode.go
index 5b67251fbb..9d59b0ff2b 100644
--- a/src/encoding/json/encode.go
+++ b/src/encoding/json/encode.go
@@ -156,6 +156,7 @@ import (
 // an error.
 func Marshal(v any) ([]byte, error) {
 	e := newEncodeState()
+	defer encodeStatePool.Put(e)
 
 	err := e.marshal(v, encOpts{escapeHTML: true})
 	if err != nil {
@@ -163,8 +164,6 @@ func Marshal(v any) ([]byte, error) {
 	}
 	buf := append([]byte(nil), e.Bytes()...)
 
-	encodeStatePool.Put(e)
-
 	return buf, nil
 }
 
diff --git a/src/encoding/json/encode_test.go b/src/encoding/json/encode_test.go
index 0b021f0074..c1b9ed2676 100644
--- a/src/encoding/json/encode_test.go
+++ b/src/encoding/json/encode_test.go
@@ -12,6 +12,7 @@ import (
 	"math"
 	"reflect"
 	"regexp"
+	"runtime/debug"
 	"strconv"
 	"testing"
 	"unicode"
@@ -760,6 +761,41 @@ func TestIssue10281(t *testing.T) {
 	}
 }
 
+func TestMarshalErrorAndReuseEncodeState(t *testing.T) {
+	// Disable the GC temporarily to prevent encodeState's in Pool being cleaned away during the test.
+	percent := debug.SetGCPercent(-1)
+	defer debug.SetGCPercent(percent)
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+	if b, err := Marshal(dummy); err == nil {
+		t.Errorf("Marshal(dummy) = %#q; want error", b)
+	}
+
+	type Data struct {
+		A string
+		I int
+	}
+	data := Data{A: "a", I: 1}
+	b, err := Marshal(data)
+	if err != nil {
+		t.Errorf("Marshal(%v) = %v", data, err)
+	}
+
+	var data2 Data
+	if err := Unmarshal(b, &data2); err != nil {
+		t.Errorf("Unmarshal(%v) = %v", data2, err)
+	}
+	if data2 != data {
+		t.Errorf("expect: %v, but get: %v", data, data2)
+	}
+}
+
 func TestHTMLEscape(t *testing.T) {
 	var b, want bytes.Buffer
 	m := `{"M":"<html>foo &` + "\xe2\x80\xa8 \xe2\x80\xa9" + `</html>"}`
diff --git a/src/encoding/json/stream.go b/src/encoding/json/stream.go
index b278ee4013..1442ef29ef 100644
--- a/src/encoding/json/stream.go
+++ b/src/encoding/json/stream.go
@@ -202,7 +202,10 @@ func (enc *Encoder) Encode(v any) error {
 	if enc.err != nil {
 		return enc.err
 	}
+
 	e := newEncodeState()
+	defer encodeStatePool.Put(e)
+
 	err := e.marshal(v, encOpts{escapeHTML: enc.escapeHTML})
 	if err != nil {
 		return err
@@ -231,7 +234,6 @@ func (enc *Encoder) Encode(v any) error {
 	if _, err = enc.w.Write(b); err != nil {
 		enc.err = err
 	}
-	encodeStatePool.Put(e)
 	return err
 }
 
diff --git a/src/encoding/json/stream_test.go b/src/encoding/json/stream_test.go
index 0e156d98e9..1f40c79670 100644
--- a/src/encoding/json/stream_test.go
+++ b/src/encoding/json/stream_test.go
@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"reflect"
+	"runtime/debug"
 	"strings"
 	"testing"
 )
@@ -59,6 +60,43 @@ func TestEncoder(t *testing.T) {
 	}
 }
 
+func TestEncoderErrorAndReuseEncodeState(t *testing.T) {
+	// Disable the GC temporarily to prevent encodeState's in Pool being cleaned away during the test.
+	percent := debug.SetGCPercent(-1)
+	defer debug.SetGCPercent(percent)
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+
+	var buf bytes.Buffer
+	enc := NewEncoder(&buf)
+	if err := enc.Encode(dummy); err == nil {
+		t.Errorf("Encode(dummy) == nil; want error")
+	}
+
+	type Data struct {
+		A string
+		I int
+	}
+	data := Data{A: "a", I: 1}
+	if err := enc.Encode(data); err != nil {
+		t.Errorf("Marshal(%v) = %v", data, err)
+	}
+
+	var data2 Data
+	if err := Unmarshal(buf.Bytes(), &data2); err != nil {
+		t.Errorf("Unmarshal(%v) = %v", data2, err)
+	}
+	if data2 != data {
+		t.Errorf("expect: %v, but get: %v", data, data2)
+	}
+}
+
 var streamEncodedIndent = `0.1
 "hello"
 null