diff --git a/jwriter/streamable_buffer.go b/jwriter/streamable_buffer.go index 209f56e..a67cc37 100644 --- a/jwriter/streamable_buffer.go +++ b/jwriter/streamable_buffer.go @@ -1,23 +1,65 @@ 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) { @@ -25,22 +67,28 @@ func (b *streamableBuffer) SetStreamingWriter(w io.Writer, chunkSize int) { 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() } diff --git a/jwriter/streamable_buffer_test.go b/jwriter/streamable_buffer_test.go index bacc626..7c27532 100644 --- a/jwriter/streamable_buffer_test.go +++ b/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) }) +} diff --git a/jwriter/token_writer_default.go b/jwriter/token_writer_default.go index b551f3d..00d9a52 100644 --- a/jwriter/token_writer_default.go +++ b/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) } diff --git a/jwriter/token_writer_easyjson.go b/jwriter/token_writer_easyjson.go index aecf22c..e349a1d 100644 --- a/jwriter/token_writer_easyjson.go +++ b/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} } diff --git a/jwriter/writer.go b/jwriter/writer.go index 57c6ba3..9d1f912 100644 --- a/jwriter/writer.go +++ b/jwriter/writer.go @@ -18,6 +18,10 @@ import ( // when using NewStreamingWriter), or if an error is explicitly raised with AddError, the Writer // permanently enters a failed state and remembers that error; all subsequent method calls for // producing output will be ignored. +// +// Like a bytes.Buffer, a Writer must not be copied: copies share the same output buffer and +// will corrupt each other's output. Use a pointer, or write through the states returned by +// Array and Object. type Writer struct { tw tokenWriter err error @@ -54,9 +58,15 @@ func (w *Writer) AddError(err error) { } // Flush writes any remaining in-memory output to the underlying io.Writer, if this is a streaming -// writer created with NewStreamingWriter. It has no effect otherwise. +// writer created with NewStreamingWriter. It has no effect otherwise. If the underlying io.Writer +// has failed, either now or during an earlier automatic flush, Flush returns that error and the +// Writer also records it, so Error() will report it as well. func (w *Writer) Flush() error { - return w.tw.Flush() + if err := w.tw.Flush(); err != nil { + w.AddError(err) + return err + } + return nil } // Null writes a JSON null value to the output. diff --git a/jwriter/writer_default_test.go b/jwriter/writer_default_test.go new file mode 100644 index 0000000..9773b96 --- /dev/null +++ b/jwriter/writer_default_test.go @@ -0,0 +1,75 @@ +//go:build !launchdarkly_easyjson +// +build !launchdarkly_easyjson + +package jwriter + +// These tests pin buffer-management behaviors of the default token writer (growth +// amortization and the pre-sized marshal buffer), so they apply only to the default +// implementation. + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestWriterReallocationsAreAmortized(t *testing.T) { + // Growing the output buffer must at least double its capacity when tokens are + // written, so the number of reallocations stays logarithmic in the output size. + // 5000 tokens of each kind produce far under 64 KiB of output, so a doubling ladder + // starting at 64 bytes can take at most ~11 steps; the bound below fails if a token + // path falls back to append's ~1.25x growth. (The bound has slack because delimiter + // writes do not reserve; see the streamableBuffer comment.) + const maxSteps = 14 + writeToken := map[string]func(arr *ArrayState, i int){ + "string": func(arr *ArrayState, i int) { arr.String("value\twith\n\"escaped chars\"") }, + "int": func(arr *ArrayState, i int) { arr.Int(123456789 + i) }, + "bool": func(arr *ArrayState, i int) { arr.Bool(i%2 == 0) }, + "raw": func(arr *ArrayState, i int) { arr.Raw(json.RawMessage(`{"k":[1,2,3]}`)) }, + } + for kind, write := range writeToken { + t.Run(kind, func(t *testing.T) { + w := NewWriter() + arr := w.Array() + steps := 0 + lastCap := -1 + for i := 0; i < 5000; i++ { + write(&arr, i) + if c := cap(w.tw.buf.Bytes()); c != lastCap { + steps++ + lastCap = c + } + } + arr.End() + require.NoError(t, w.Error()) + require.LessOrEqual(t, steps, maxSteps) + }) + } +} + +// marshalAllocationTestWritable emits a document of several hundred bytes: larger than the +// 64-byte buffer a plain NewWriter starts with, but within MarshalJSONWithWriter's initial +// capacity, so the allocation assertion below fails if that capacity is ever lost. +type marshalAllocationTestWritable struct{} + +func (marshalAllocationTestWritable) WriteToJSONWriter(w *Writer) { + arr := w.Array() + for i := 0; i < 20; i++ { + arr.String("string value for the allocation test") + } + arr.End() +} + +func TestMarshalJSONWithWriterAllocations(t *testing.T) { + var writable Writable = marshalAllocationTestWritable{} + allocs := testing.AllocsPerRun(500, func() { + data, err := MarshalJSONWithWriter(writable) + if err != nil || len(data) == 0 { + t.Fail() + } + }) + // One allocation for the Writer (it escapes through the Writable interface) and one + // for the output buffer, allocated once at full size. + require.LessOrEqual(t, allocs, 2.0) +} diff --git a/jwriter/writer_init_default.go b/jwriter/writer_init_default.go index 76bb0af..433b0a4 100644 --- a/jwriter/writer_init_default.go +++ b/jwriter/writer_init_default.go @@ -7,6 +7,9 @@ import "io" // This function returns the struct by value (Writer, not *Writer). This avoids the overhead of a // heap allocation since, in typical usage, the Writer will not escape the scope in which it was // declared and can remain on the stack. +// +// Like a bytes.Buffer, the returned Writer must not be copied: it contains a preallocated +// output buffer, so copies of it would write into the same memory. func NewWriter() Writer { return Writer{tw: newTokenWriter()} } @@ -14,6 +17,7 @@ func NewWriter() Writer { // NewStreamingWriter creates a Writer that will buffer a limited amount of its output in memory // and dump the output to the specified io.Writer whenever the buffer is full. You should also // call Flush at the end of your output to ensure that any remaining buffered output is flushed. +// It panics if bufferSize is negative. // // If the Writer returns an error at any point, it enters a failed state and will not try to // write any more data to the target. @@ -21,6 +25,9 @@ func NewWriter() Writer { // This function returns the struct by value (Writer, not *Writer). This avoids the overhead of a // heap allocation since, in typical usage, the Writer will not escape the scope in which it was // declared and can remain on the stack. +// +// Like a bytes.Buffer, the returned Writer must not be copied: it contains a preallocated +// output buffer, so copies of it would write into the same memory. func NewStreamingWriter(target io.Writer, bufferSize int) Writer { return Writer{tw: newStreamingTokenWriter(target, bufferSize)} } diff --git a/jwriter/writer_marshal.go b/jwriter/writer_marshal.go index f3d569c..3a4f8d4 100644 --- a/jwriter/writer_marshal.go +++ b/jwriter/writer_marshal.go @@ -1,10 +1,13 @@ package jwriter +// marshalInitialCapacity is the output buffer capacity that MarshalJSONWithWriter starts +// with, sized so that most documents marshal without reallocating. +const marshalInitialCapacity = 1000 + // MarshalJSONWithWriter is a convenience method for implementing json.Marshaler to marshal to a // byte slice with the default TokenWriter implementation. func MarshalJSONWithWriter(writable Writable) ([]byte, error) { - w := NewWriter() - w.tw.Grow(1000) + w := Writer{tw: newTokenWriterWithCapacity(marshalInitialCapacity)} writable.WriteToJSONWriter(&w) if err := w.Error(); err != nil { return nil, err diff --git a/jwriter/writer_streaming_default_test.go b/jwriter/writer_streaming_default_test.go index 8e15455..a2ea417 100644 --- a/jwriter/writer_streaming_default_test.go +++ b/jwriter/writer_streaming_default_test.go @@ -5,6 +5,11 @@ package jwriter import ( "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "strings" "testing" "github.com/stretchr/testify/require" @@ -23,7 +28,7 @@ func TestStreamingWriterWritesToTargetInChunks(t *testing.T) { require.Equal(t, expected, buf.String()) arr.String("abc") - expected += `[true,"abc` + expected += `[true,"abc"` require.Equal(t, expected, buf.String()) arr.Int(33) @@ -33,13 +38,141 @@ func TestStreamingWriterWritesToTargetInChunks(t *testing.T) { require.Equal(t, expected, buf.String()) arr.Float64(2.5) - expected += `",33,null,` + expected += `,33,null,2.5` require.Equal(t, expected, buf.String()) arr.End() require.Equal(t, expected, buf.String()) require.NoError(t, w.Flush()) - expected += `2.5]` + expected += `]` require.Equal(t, expected, buf.String()) } + +var errFailingDestination = errors.New("destination failed") + +type failingDestination struct { + failAfter int // number of successful writes before failures begin + writes int +} + +func (f *failingDestination) Write(p []byte) (int, error) { + f.writes++ + if f.writes > f.failAfter { + return 0, errFailingDestination + } + return len(p), nil +} + +// truncatingDestination reports success but consumes one byte less than it was given, +// violating the io.Writer contract. +type truncatingDestination struct{} + +func (truncatingDestination) Write(p []byte) (int, error) { + n := len(p) - 1 + if n < 0 { + n = 0 + } + return n, nil +} + +func writeStreamingTestDocument(w *Writer) { + arr := w.Array() + arr.String("hello world") + arr.Int(1) + arr.String("value\twith\n\"escaped chars\"") + arr.Float64(2.5) + arr.Null() + arr.Bool(true) + obj := arr.Object() + obj.Name("prop").Int(-1234567890) + obj.End() + arr.Raw(json.RawMessage(`{"raw":[1,2,3]}`)) + arr.End() +} + +func TestStreamingWriterSurfacesDestinationError(t *testing.T) { + // Whenever the destination has failed, both Error() and the final Flush() must report it, + // including when a failed mid-stream flush left the buffer empty. + cases := []struct{ chunkSize, failAfter int }{ + {1, 0}, {1, 1}, {1, 2}, {1, 5}, + {2, 0}, {2, 1}, + {10, 0}, {10, 1}, + {50, 0}, + {1000, 0}, // nothing flushes until the final Flush + } + for _, c := range cases { + t.Run(fmt.Sprintf("chunkSize=%d failAfter=%d", c.chunkSize, c.failAfter), func(t *testing.T) { + f := &failingDestination{failAfter: c.failAfter} + w := NewStreamingWriter(f, c.chunkSize) + writeStreamingTestDocument(&w) + require.True(t, errors.Is(w.Flush(), errFailingDestination)) + require.True(t, errors.Is(w.Error(), errFailingDestination)) + // The destination must see exactly one failing write and nothing afterward, and + // the undeliverable data must have been discarded rather than retained. + require.Equal(t, c.failAfter+1, f.writes) + require.Empty(t, w.tw.buf.Bytes()) + }) + } +} + +func TestStreamingWriterOutputMatchesInMemoryWriter(t *testing.T) { + expectedWriter := NewWriter() + writeStreamingTestDocument(&expectedWriter) + require.NoError(t, expectedWriter.Error()) + expected := string(expectedWriter.Bytes()) + + for _, chunkSize := range []int{0, 1, 2, 3, 5, 8, 13, 64, 100, 1000} { + t.Run(fmt.Sprintf("chunkSize=%d", chunkSize), func(t *testing.T) { + var target bytes.Buffer + w := NewStreamingWriter(&target, chunkSize) + writeStreamingTestDocument(&w) + require.NoError(t, w.Flush()) + require.Equal(t, expected, target.String()) + }) + } +} + +func TestStreamingWriterFlushesAfterEachToken(t *testing.T) { + // With a chunk size of 1, every completed token must reach the destination immediately; + // this pins the invariant that each token-writing code path checks for a flush. + var target bytes.Buffer + w := NewStreamingWriter(&target, 1) + + arr := w.Array() + arr.Int(12345) + require.Equal(t, `[12345`, target.String()) + arr.Float64(2.5) + require.Equal(t, `[12345,2.5`, target.String()) + arr.String("ab\tc") + require.Equal(t, `[12345,2.5,"ab\tc"`, target.String()) + arr.Bool(true) + require.Equal(t, `[12345,2.5,"ab\tc",true`, target.String()) + arr.Raw(json.RawMessage(`{}`)) + require.Equal(t, `[12345,2.5,"ab\tc",true,{}`, target.String()) + arr.End() + require.Equal(t, `[12345,2.5,"ab\tc",true,{}]`, target.String()) + + require.NoError(t, w.Flush()) + require.Equal(t, `[12345,2.5,"ab\tc",true,{}]`, target.String()) +} + +func TestStreamingWriterBoundsBufferForEscapeHeavyStrings(t *testing.T) { + // A string full of escaped characters must not make the internal buffer grow in + // proportion to the string: the writer flushes at chunk boundaries while scanning. + const chunkSize = 1024 + const inputLen = 200000 + var target bytes.Buffer + w := NewStreamingWriter(&target, chunkSize) + w.String(strings.Repeat("\n", inputLen)) + require.NoError(t, w.Flush()) + require.Equal(t, `"`+strings.Repeat(`\n`, inputLen)+`"`, target.String()) + require.LessOrEqual(t, cap(w.tw.buf.Bytes()), 4*chunkSize) +} + +func TestStreamingWriterSurfacesShortWrite(t *testing.T) { + w := NewStreamingWriter(truncatingDestination{}, 10) + writeStreamingTestDocument(&w) + require.True(t, errors.Is(w.Flush(), io.ErrShortWrite)) + require.True(t, errors.Is(w.Error(), io.ErrShortWrite)) +}