fix: Improve jwriter throughput with append-based buffer internals

jwriter/streamable_buffer.go

@@ -1,46 +1,94 @@
 package jwriter
 
 import (
-	"bytes"
 	"io"
+	"unicode/utf8"
 )
 
+// streamableBuffer is a byte buffer that can optionally flush its contents to an io.Writer
+// whenever they reach a chunk size. The zero value is a ready-to-use in-memory buffer.
+//
+// Output accumulates in a plain byte slice via append operations, which the compiler can
+// inline at call sites, rather than through bytes.Buffer method calls. For the same reason,
+// tokenWriter appends directly to the buf field in its hottest code paths; any code that
+// does so must call maybeFlush afterward so that streaming mode keeps flushing incrementally.
+//
+// Token-level writes — strings, numbers, and the multi-byte Write used for keyword and
+// raw tokens — reserve capacity before appending: bare append grows large slices by only
+// ~1.25x, and the at-least-doubling policy in reserve keeps the number of reallocations
+// (and the total bytes copied) logarithmic for large outputs. WriteByte and WriteRune
+// deliberately do not reserve: a capacity check on every single-byte delimiter write
+// measurably slows encoding, and a growth triggered by one is made rare by the reserves
+// on the token writes around it.
+//
+// The chunk-size check happens after a write, never during one, so a single write larger
+// than the chunk size — a raw JSON value, or a long escape-free string segment — is
+// buffered whole before it is flushed.
 type streamableBuffer struct {
-	buf       bytes.Buffer
+	buf       []byte
 	dest      io.Writer
 	destErr   error
 	chunkSize int
 }
 
 func (b *streamableBuffer) Bytes() []byte {
-	return b.buf.Bytes()
+	return b.buf
 }
 
+// Grow ensures that the buffer has room for at least n more bytes, reallocating it if
+// necessary. It panics if n is negative or if the buffer cannot grow that large.
 func (b *streamableBuffer) Grow(n int) {
-	b.buf.Grow(n)
+	if n < 0 {
+		panic("jwriter: cannot grow buffer by a negative count")
+	}
+	b.reserve(n)
+}
+
+// reserve ensures that the buffer has room for at least n more bytes (n must not be
+// negative), at least doubling the capacity when it must reallocate so that repeated
+// appends remain amortized.
+func (b *streamableBuffer) reserve(n int) {
+	if cap(b.buf)-len(b.buf) < n {
+		if len(b.buf)+n < 0 {
+			panic("jwriter: buffer too large")
+		}
+		newCap := 2 * cap(b.buf)
+		if newCap < len(b.buf)+n {
+			newCap = len(b.buf) + n
+		}
+		newBuf := make([]byte, len(b.buf), newCap)
+		copy(newBuf, b.buf)
+		b.buf = newBuf
+	}
 }
 
 func (b *streamableBuffer) SetStreamingWriter(w io.Writer, chunkSize int) {
 	b.dest = w
 	b.chunkSize = chunkSize
 }
 
+// Flush writes any buffered output to the destination, if there is one. Once a destination
+// write has failed, the failure is remembered and returned by every subsequent Flush call,
+// even after the buffered data that could not be delivered has been discarded.
 func (b *streamableBuffer) Flush() error {
-	if b.dest != nil {
-		if b.buf.Len() > 0 {
-			if b.destErr == nil {
-				data := b.buf.Bytes()
-				_, b.destErr = b.dest.Write(data)
+	if b.dest == nil {
+		return nil
+	}
+	if len(b.buf) > 0 {
+		if b.destErr == nil {
+			n, err := b.dest.Write(b.buf)
+			if err == nil && n < len(b.buf) {
+				err = io.ErrShortWrite
 			}
-			b.buf.Reset()
-			return b.destErr
+			b.destErr = err
 		}
+		b.buf = b.buf[:0]
 	}
-	return nil
+	return b.destErr
 }
 
 func (b *streamableBuffer) maybeFlush() {
-	if b.dest != nil && b.buf.Len() >= b.chunkSize {
+	if b.dest != nil && len(b.buf) >= b.chunkSize {
 		_ = b.Flush()
 	}
 }
@@ -50,21 +98,23 @@ func (b *streamableBuffer) GetWriterError() error {
 }
 
 func (b *streamableBuffer) Write(data []byte) {
-	_, _ = b.buf.Write(data)
+	b.reserve(len(data))
+	b.buf = append(b.buf, data...)
 	b.maybeFlush()
 }
 
 func (b *streamableBuffer) WriteByte(data byte) { //nolint:govet
-	_ = b.buf.WriteByte(data)
+	b.buf = append(b.buf, data)
 	b.maybeFlush()
 }
 
 func (b *streamableBuffer) WriteRune(ch rune) {
-	_, _ = b.buf.WriteRune(ch)
+	b.buf = utf8.AppendRune(b.buf, ch)
 	b.maybeFlush()
 }
 
 func (b *streamableBuffer) WriteString(s string) {
-	_, _ = b.buf.WriteString(s)
+	b.reserve(len(s))
+	b.buf = append(b.buf, s...)
 	b.maybeFlush()
 }

jwriter/streamable_buffer_test.go

@@ -2,6 +2,7 @@ package jwriter
 
 import (
 	"bytes"
+	"math"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -72,3 +73,17 @@ func writeTestDataToBuffer(b *streamableBuffer) string {
 
 	return expected
 }
+
+func TestStreamableBufferGrow(t *testing.T) {
+	var b streamableBuffer
+	b.WriteString("abc")
+	b.Grow(100)
+	require.GreaterOrEqual(t, cap(b.Bytes())-len(b.Bytes()), 100)
+	require.Equal(t, "abc", string(b.Bytes()))
+	b.WriteString("def")
+	require.Equal(t, "abcdef", string(b.Bytes()))
+	require.PanicsWithValue(t, "jwriter: cannot grow buffer by a negative count",
+		func() { b.Grow(-1) })
+	require.PanicsWithValue(t, "jwriter: buffer too large",
+		func() { b.Grow(math.MaxInt) })
+}

jwriter/token_writer_default.go

@@ -7,7 +7,6 @@ import (
 	"encoding/json"
 	"io"
 	"strconv"
-	"unicode/utf8"
 )
 
 // This file defines the default implementation of the low-level JSON token writer. If the launchdarkly_easyjson
@@ -22,13 +21,46 @@ var (
 	tokenFalse = []byte("false") //nolint:gochecknoglobals
 )
 
+const hexDigits = "0123456789abcdef"
+
+// plainStringChars marks the bytes that writeQuotedString copies into a JSON string
+// verbatim. Unlike the reader's equivalent table, bytes outside the ASCII range are
+// plain here: this writer never escapes them, so multi-byte characters pass through
+// without inspection. A single table lookup per byte is measurably faster than the
+// equivalent range and equality comparisons in this loop.
+var plainStringChars = makePlainStringChars() //nolint:gochecknoglobals
+
+func makePlainStringChars() (t [256]bool) {
+	for c := 0x20; c < 256; c++ {
+		if c != '"' && c != '\\' {
+			t[c] = true
+		}
+	}
+	return
+}
+
+// initialBufferCapacity is the buffer capacity preallocated by newTokenWriter. Paying for one
+// small allocation up front keeps the append growth ladder short for typical outputs; it is
+// the same minimum that bytes.Buffer uses.
+const initialBufferCapacity = 64
+
+// maxNumberLength is the longest text that Int or Float64 can produce for one value: a
+// float64 in Go's shortest 'g' representation needs at most 24 characters (for example
+// "-1.7976931348623157e+308"), and an int64 needs at most 20.
+const maxNumberLength = 24
+
 type tokenWriter struct {
-	buf       streamableBuffer
-	tempBytes [50]byte
+	buf streamableBuffer
 }
 
 func newTokenWriter() tokenWriter {
-	return tokenWriter{}
+	return newTokenWriterWithCapacity(initialBufferCapacity)
+}
+
+func newTokenWriterWithCapacity(capacity int) tokenWriter {
+	tw := tokenWriter{}
+	tw.buf.buf = make([]byte, 0, capacity)
+	return tw
 }
 
 func newStreamingTokenWriter(dest io.Writer, bufferSize int) tokenWriter {
@@ -46,12 +78,6 @@ func (tw *tokenWriter) Bytes() []byte {
 	return tw.buf.Bytes()
 }
 
-// Grow expands the internal buffer by the specified number of bytes. It is the same as calling Grow
-// on a bytes.Buffer.
-func (tw *tokenWriter) Grow(n int) {
-	tw.buf.Grow(n)
-}
-
 // Flush writes any remaining in-memory output to the underlying Writer, if this is a streaming buffer
 // created with newStreamingTokenWriter. It has no effect otherwise.
 func (tw *tokenWriter) Flush() error {
@@ -78,13 +104,9 @@ func (tw *tokenWriter) Bool(value bool) error {
 
 // Int writes an integer JSON number.
 func (tw *tokenWriter) Int(value int) error {
-	if value == 0 {
-		tw.buf.WriteByte('0')
-	} else {
-		out := tw.tempBytes[0:0]
-		out = strconv.AppendInt(out, int64(value), 10)
-		tw.buf.Write(out)
-	}
+	tw.buf.reserve(maxNumberLength)
+	tw.buf.buf = strconv.AppendInt(tw.buf.buf, int64(value), 10)
+	tw.buf.maybeFlush()
 	return tw.buf.GetWriterError()
 }
 
@@ -97,9 +119,9 @@ func (tw *tokenWriter) Float64(value float64) error {
 		if float64(i) == value {
 			return tw.Int(i)
 		}
-		out := tw.tempBytes[0:0]
-		out = strconv.AppendFloat(out, value, 'g', -1, 64)
-		tw.buf.Write(out)
+		tw.buf.reserve(maxNumberLength)
+		tw.buf.buf = strconv.AppendFloat(tw.buf.buf, value, 'g', -1, 64)
+		tw.buf.maybeFlush()
 	}
 	return tw.buf.GetWriterError()
 }
@@ -131,62 +153,64 @@ func (tw *tokenWriter) Delimiter(delimiter byte) error {
 }
 
 func (tw *tokenWriter) writeQuotedString(s string) error {
-	// This is basically the same logic used internally by json.Marshal
-	tw.buf.WriteByte('"')
+	// This is basically the same logic used internally by json.Marshal: scan for the next byte
+	// that requires escaping, and copy the whole clean segment before it in one append. Bytes
+	// outside the ASCII range are copied through verbatim without being decoded, since this
+	// writer only ever escapes control characters, quotes, and backslashes.
+	//
+	// In-memory mode reserves len(s)+2 up front — the exact encoded length when nothing needs
+	// escaping; escape expansion beyond that grows through append. Streaming mode must not
+	// reserve by input length: it reserves nothing and instead
+	// flushes at chunk boundaries as it scans, so that the buffer stays near the chunk size
+	// no matter how long or escape-heavy the input string is. A clean segment is still
+	// appended in one piece, so a segment longer than the chunk size overshoots it by that
+	// segment's length, which the buffer permits.
+	if tw.buf.dest == nil {
+		tw.buf.reserve(len(s) + 2)
+	}
+	// dst is a local copy of the buffer's slice header; it must be stored back before any
+	// buffer method runs (and reloaded after), or the flushed bytes would reappear.
+	dst := tw.buf.buf
+	dst = append(dst, '"')
 	start := 0
-	for i := 0; i < len(s); {
+	for i := 0; i < len(s); i++ {
 		aByte := s[i]
-		if aByte < ' ' || aByte == '"' || aByte == '\\' {
-			if i > start {
-				tw.buf.WriteString(s[start:i])
-			}
-			tw.writeEscapedChar(aByte)
-			i++
-			start = i
-		} else {
-			if aByte < utf8.RuneSelf { // single-byte character
-				i++
-			} else {
-				_, size := utf8.DecodeRuneInString(s[i:])
-				i += size
-			}
+		if plainStringChars[aByte] {
+			continue
+		}
+		dst = append(dst, s[start:i]...)
+		dst = appendEscapedChar(dst, aByte)
+		start = i + 1
+		if tw.buf.dest != nil && len(dst) >= tw.buf.chunkSize {
+			tw.buf.buf = dst
+			tw.buf.maybeFlush()
+			dst = tw.buf.buf
 		}
 	}
-	if start < len(s) {
-		tw.buf.WriteString(s[start:])
-	}
-	tw.buf.WriteByte('"')
+	dst = append(dst, s[start:]...)
+	dst = append(dst, '"')
+	tw.buf.buf = dst
+	tw.buf.maybeFlush()
 	return tw.buf.GetWriterError()
 }
 
-func (tw *tokenWriter) writeEscapedChar(ch byte) {
-	out := tw.tempBytes[0:2]
-	out[0] = '\\'
+func appendEscapedChar(dst []byte, ch byte) []byte {
 	switch ch {
 	case '\b':
-		out[1] = 'b'
+		return append(dst, '\\', 'b')
 	case '\t':
-		out[1] = 't'
+		return append(dst, '\\', 't')
 	case '\n':
-		out[1] = 'n'
+		return append(dst, '\\', 'n')
 	case '\f':
-		out[1] = 'f'
+		return append(dst, '\\', 'f')
 	case '\r':
-		out[1] = 'r'
+		return append(dst, '\\', 'r')
 	case '"':
-		out[1] = '"'
+		return append(dst, '\\', '"')
 	case '\\':
-		out[1] = '\\'
+		return append(dst, '\\', '\\')
 	default:
-		out[1] = 'u'
-		out = append(out, '0')
-		out = append(out, '0')
-		hexChars := make([]byte, 0, 4)
-		hexChars = strconv.AppendInt(hexChars, int64(ch), 16)
-		if len(hexChars) < 2 {
-			out = append(out, '0')
-		}
-		out = append(out, hexChars...)
+		return append(dst, '\\', 'u', '0', '0', hexDigits[ch>>4], hexDigits[ch&0xf])
 	}
-	tw.buf.Write(out)
 }

jwriter/token_writer_easyjson.go

@@ -32,6 +32,14 @@ func newTokenWriter() tokenWriter {
 	return tokenWriter{}
 }
 
+// newTokenWriterWithCapacity is equivalent to newTokenWriter, but pre-sizes the output
+// buffer for callers that know the approximate size of the document in advance.
+func newTokenWriterWithCapacity(capacity int) tokenWriter {
+	tw := tokenWriter{}
+	tw.inlineWriter.Buffer.EnsureSpace(capacity)
+	return tw
+}
+
 func newTokenWriterFromEasyjsonWriter(writer *ejwriter.Writer) tokenWriter {
 	return tokenWriter{pWriter: writer}
 }