encoding/json: periodically flush decodeState’s buffer to Encoder’s io.Writer (addresses golang#7872)

Author: ydnar

src/encoding/json/encode.go

@@ -14,6 +14,7 @@ import (
 	"bytes"
 	"encoding"
 	"encoding/base64"
+	"io"
 	"math"
 	"reflect"
 	"runtime"
@@ -141,7 +142,7 @@ func Marshal(v interface{}) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	return e.Bytes(), nil
+	return e.buf.Bytes(), nil
 }
 
 // MarshalIndent is like Marshal but applies Indent to format the output.
@@ -243,23 +244,68 @@ func (e *MarshalerError) Error() string {
 
 var hex = "0123456789abcdef"
 
-// An encodeState encodes JSON into a bytes.Buffer.
+// An encodeState encodes JSON into a bytes.Buffer,
+// optionally flushing to an io.Writer.
 type encodeState struct {
-	bytes.Buffer // accumulated output
-	scratch      [64]byte
+	buf     bytes.Buffer // accumulated output
+	w       io.Writer    // optional
+	scratch [64]byte
 }
 
 var encodeStatePool sync.Pool
 
 func newEncodeState() *encodeState {
 	if v := encodeStatePool.Get(); v != nil {
 		e := v.(*encodeState)
-		e.Reset()
+		e.buf.Reset()
 		return e
 	}
 	return new(encodeState)
 }
 
+const maxBuffer = 256 // bytes
+
+func (e *encodeState) flush(force bool) error {
+	// @ydnar: Non-streaming encoders don’t have a writer.
+	if e.w == nil {
+		return nil
+	}
+	if force || e.buf.Len() > maxBuffer {
+		b := e.buf.Bytes()
+		n, err := e.w.Write(b)
+		e.buf.Reset()
+		if n != len(b) {
+			e.buf.Write(b[n:]) // Keep remaining bytes in the buffer
+		}
+		return err
+	}
+	return nil
+}
+
+func (e *encodeState) Write(p []byte) (int, error) {
+	n, err := e.buf.Write(p)
+	if err == nil {
+		err = e.flush(false)
+	}
+	return n, err
+}
+
+func (e *encodeState) WriteString(s string) (int, error) {
+	n, err := e.buf.WriteString(s)
+	if err == nil {
+		err = e.flush(false)
+	}
+	return n, err
+}
+
+func (e *encodeState) WriteByte(c byte) error {
+	err := e.buf.WriteByte(c)
+	if err == nil {
+		err = e.flush(false)
+	}
+	return err
+}
+
 func (e *encodeState) marshal(v interface{}) (err error) {
 	defer func() {
 		if r := recover(); r != nil {
@@ -423,7 +469,7 @@ func marshalerEncoder(e *encodeState, v reflect.Value, quoted bool) {
 	b, err := m.MarshalJSON()
 	if err == nil {
 		// copy JSON into buffer, checking validity.
-		err = compact(&e.Buffer, b, true)
+		err = compact(&e.buf, b, true)
 	}
 	if err != nil {
 		e.error(&MarshalerError{v.Type(), err})
@@ -440,7 +486,7 @@ func addrMarshalerEncoder(e *encodeState, v reflect.Value, quoted bool) {
 	b, err := m.MarshalJSON()
 	if err == nil {
 		// copy JSON into buffer, checking validity.
-		err = compact(&e.Buffer, b, true)
+		err = compact(&e.buf, b, true)
 	}
 	if err != nil {
 		e.error(&MarshalerError{v.Type(), err})
@@ -818,7 +864,7 @@ func (sv stringValues) get(i int) string   { return sv[i].String() }
 
 // NOTE: keep in sync with stringBytes below.
 func (e *encodeState) string(s string) (int, error) {
-	len0 := e.Len()
+	len0 := e.buf.Len()
 	e.WriteByte('"')
 	start := 0
 	for i := 0; i < len(s); {
@@ -889,12 +935,12 @@ func (e *encodeState) string(s string) (int, error) {
 		e.WriteString(s[start:])
 	}
 	e.WriteByte('"')
-	return e.Len() - len0, nil
+	return e.buf.Len() - len0, nil
 }
 
 // NOTE: keep in sync with string above.
 func (e *encodeState) stringBytes(s []byte) (int, error) {
-	len0 := e.Len()
+	len0 := e.buf.Len()
 	e.WriteByte('"')
 	start := 0
 	for i := 0; i < len(s); {
@@ -965,7 +1011,7 @@ func (e *encodeState) stringBytes(s []byte) (int, error) {
 		e.Write(s[start:])
 	}
 	e.WriteByte('"')
-	return e.Len() - len0, nil
+	return e.buf.Len() - len0, nil
 }
 
 // A field represents a single field found in a struct.

src/encoding/json/encode_test.go

@@ -392,8 +392,8 @@ func TestStringBytes(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	enc := es.Buffer.String()
-	encBytes := esBytes.Buffer.String()
+	enc := es.buf.String()
+	encBytes := esBytes.buf.String()
 	if enc != encBytes {
 		i := 0
 		for i < len(enc) && i < len(encBytes) && enc[i] == encBytes[i] {

src/encoding/json/stream.go

@@ -157,6 +157,7 @@ func (enc *Encoder) Encode(v interface{}) error {
 		return enc.err
 	}
 	e := newEncodeState()
+	e.w = enc.w
 	err := e.marshal(v)
 	if err != nil {
 		return err
@@ -170,7 +171,8 @@ func (enc *Encoder) Encode(v interface{}) error {
 	// digits coming.
 	e.WriteByte('\n')
 
-	if _, err = enc.w.Write(e.Bytes()); err != nil {
+	// @ydnar: Flush any bytes remaining.
+	if err = e.flush(true); err != nil {
 		enc.err = err
 	}
 	encodeStatePool.Put(e)