From 923d7983faf7da86a910529d19ca0e6443811841 Mon Sep 17 00:00:00 2001 From: Raveesh Agarwal Date: Mon, 26 Oct 2026 03:08:45 +0530 Subject: [PATCH] internal/json: decode in place instead of through a streaming Decoder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Unmarshal and UnmarshalUseNumber wrapped data in a bytes.Reader and decoded it with a json.Decoder, which copies its input into a read buffer of at least 32 KiB on every call. Unmarshal runs several times per request, so these buffers made up most of the memory allocated to serve a tool call. Decode in place instead: find the extent of the first JSON value with the same parse the Decoder runs on its buffer, then decode just that value with the same flags. Malformed input, the rare case, still goes through a Decoder, so results are unchanged: the decoded value, trailing data ignored, io.EOF on empty input, and each error by type and text. One function, newDecoder, configures the Decoder; parseFlags gives json.Parse the same rules, and a differential test and fuzz target compare the two paths, so the rules cannot drift apart. Tests also check that decoded values never share memory with the input, and that UnmarshalUseNumber, which decodes every JSON-RPC envelope, matches field names case-sensitively. On an Apple M1 Max: name old time/op new time/op delta StreamableServing 368.7µs ± 2% 219.3µs ± 1% -40.51% Unmarshal 4.015µs ± 1% 1.337µs ± 1% -66.71% name old alloc/op new alloc/op delta StreamableServing 528.3KiB ± 0% 60.1KiB ± 1% -88.62% Unmarshal 37.41KiB ± 0% 5.39KiB ± 0% -85.59% Fixes #1453 --- internal/json/json.go | 44 ++++++++++-- internal/json/json_test.go | 135 +++++++++++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+), 7 deletions(-) diff --git a/internal/json/json.go b/internal/json/json.go index 3a0e52b2a..131243ec4 100644 --- a/internal/json/json.go +++ b/internal/json/json.go @@ -28,29 +28,59 @@ type Decoder struct { } func NewDecoder(r io.Reader) *Decoder { + return newDecoder(r, false) +} + +// newDecoder returns a Decoder that decodes as every function of this +// package does: field names match case-sensitively, and with useNumber, +// numbers decode to json.Number. [parseFlags] must stay in step with it. +func newDecoder(r io.Reader, useNumber bool) *Decoder { dec := json.NewDecoder(r) dec.DontMatchCaseInsensitiveStructFields() + if useNumber { + dec.UseNumber() + } return &Decoder{dec: dec} } +// parseFlags returns the [json.Parse] flags that decode as a Decoder from +// [newDecoder] does. +func parseFlags(useNumber bool) json.ParseFlags { + flags := json.DontMatchCaseInsensitiveStructFields + if useNumber { + flags |= json.UseNumber + } + return flags +} + func (d *Decoder) Decode(v any) error { return d.dec.Decode(v) } func Unmarshal(data []byte, v any) error { - if err := checkMaxDepth(data, defaultMaxDepth); err != nil { - return err - } - return NewDecoder(bytes.NewReader(data)).Decode(v) + return unmarshal(data, v, false) } func UnmarshalUseNumber(data []byte, v any) error { + return unmarshal(data, v, true) +} + +// unmarshal decodes the first JSON value in data into v exactly as a +// [Decoder] reading data does, but without the Decoder's read buffer of at +// least 32 KiB. Like the Decoder, it first finds the extent of the value, +// here in place, then decodes just that value. Malformed input still goes +// through a Decoder, so its error is exactly the Decoder's. +func unmarshal(data []byte, v any, useNumber bool) error { if err := checkMaxDepth(data, defaultMaxDepth); err != nil { return err } - dec := NewDecoder(bytes.NewReader(data)) - dec.dec.UseNumber() - return dec.Decode(v) + flags := parseFlags(useNumber) + var value json.RawMessage + if _, err := json.Parse(data, &value, flags|json.DontCopyRawMessage); err != nil { + return newDecoder(bytes.NewReader(data), useNumber).Decode(v) + } + _, err := json.Parse(value, v, flags) + return err } // checkMaxDepth scans data once and reports [errMaxDepthExceeded] if the diff --git a/internal/json/json_test.go b/internal/json/json_test.go index b1d110ce6..2c3dcbfaf 100644 --- a/internal/json/json_test.go +++ b/internal/json/json_test.go @@ -9,6 +9,7 @@ import ( "encoding/json" "errors" "fmt" + "reflect" "strings" "testing" "time" @@ -257,3 +258,137 @@ func FuzzUnmarshalFieldIsolation(f *testing.F) { } }) } + +// decodeStream is how Unmarshal decoded before it parsed data in place: +// through the package's Decoder, which allocates a read buffer of at least +// 32 KiB for every call. Comparing with it also keeps [parseFlags] in step +// with [newDecoder]. +func decodeStream(data []byte, v any, useNumber bool) error { + if err := checkMaxDepth(data, defaultMaxDepth); err != nil { + return err + } + return newDecoder(bytes.NewReader(data), useNumber).Decode(v) +} + +type unmarshalTarget struct { + ID any `json:"id"` + Method string `json:"method"` + Params map[string]any `json:"params"` + Raw json.RawMessage `json:"raw"` + Nested *unmarshalTarget `json:"nested"` + List []int `json:"list"` + Map map[string]string `json:"map"` +} + +// sameAsStream checks that Unmarshal and UnmarshalUseNumber decode data +// exactly as the streaming Decoder did: the same value, and the same error, +// by type and by text. +func sameAsStream(t *testing.T, data []byte) { + t.Helper() + for _, useNumber := range []bool{false, true} { + for _, newTarget := range []func() any{ + func() any { return new(any) }, + func() any { return new(unmarshalTarget) }, + func() any { return new(map[string]any) }, + func() any { return new([]any) }, + func() any { return new(string) }, + } { + want, got := newTarget(), newTarget() + wantErr := decodeStream(data, want, useNumber) + var gotErr error + if useNumber { + gotErr = UnmarshalUseNumber(data, got) + } else { + gotErr = Unmarshal(data, got) + } + if fmt.Sprintf("%T %v", gotErr, gotErr) != fmt.Sprintf("%T %v", wantErr, wantErr) { + t.Fatalf("%q into %T (useNumber %v): error %T %v, want %T %v", data, got, useNumber, gotErr, gotErr, wantErr, wantErr) + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("%q into %T (useNumber %v): got %#v, want %#v", data, got, useNumber, got, want) + } + } + } +} + +var unmarshalCases = []string{ + ``, ` `, "\n\t ", `{}`, `[]`, `null`, `true`, `1`, `-0.5e3`, `"s"`, + `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"get","arguments":{"a":[1,2,{"b":null}]}}}`, + ` {"id":"x"} `, `{"id":1} {"id":2}`, `{"id":1}garbage`, `{"id":1}]`, `[1,2]x`, + `{"id":`, `{"id":1`, `[1,2`, `"unterminated`, `{"id":1,}`, `{"id" 1}`, `{,}`, `]`, `}`, `x`, `nul`, `tru`, `1.`, + `{"list":[1,"two"]}`, `{"method":1}`, `{"nested":{"nested":{"method":[]}}}`, `{"raw":{"keep":"as is"}}`, + `{"Method":"case"}`, `{"METHOD":"case"}`, `{"map":{"a":"b","a":"c"}}`, `{"id":12345678901234567890}`, + `{"id":1.5}`, `{"params":{"n":1e400}}`, "{\"method\":\"nul\x00l\"}", `{"method":"é😀"}`, + `{"method":"\ud800"}`, "\xff", `{"method":"` + "\xff" + `"}`, `[[[[[]]]]]`, "\ufeff{}", +} + +func TestUnmarshalSameAsStream(t *testing.T) { + for _, c := range unmarshalCases { + sameAsStream(t, []byte(c)) + } +} + +func FuzzUnmarshalSameAsStream(f *testing.F) { + for _, c := range unmarshalCases { + f.Add([]byte(c)) + } + f.Fuzz(sameAsStream) +} + +// UnmarshalUseNumber matches field names case-sensitively, like Unmarshal: +// jsonrpc2.DecodeMessage decodes every message envelope with it, and a +// case-insensitive match would let "Method" override "method". +func TestUnmarshalUseNumberCaseSensitivity(t *testing.T) { + var got struct { + Method string `json:"method"` + ID any `json:"id"` + } + data := []byte(`{"method":"tools/list","Method":"tools/call","ID":2,"id":9007199254740993}`) + if err := UnmarshalUseNumber(data, &got); err != nil { + t.Fatal(err) + } + if got.Method != "tools/list" || got.ID != json.Number("9007199254740993") { + t.Errorf("got method %q, id %#v; want \"tools/list\" and json.Number(\"9007199254740993\")", got.Method, got.ID) + } +} + +// Decoded values never share memory with the input: Unmarshal finds the +// value in place, and the caller may reuse the input buffer afterwards. +func TestUnmarshalDoesNotAliasInput(t *testing.T) { + const input = `{"s":"hello","raw":{"a":"b"},"n":12345,"any":{"k":"v"}}` + data := []byte(input) + var v struct { + S string `json:"s"` + Raw json.RawMessage `json:"raw"` + N json.Number `json:"n"` + Any any `json:"any"` + } + if err := UnmarshalUseNumber(data, &v); err != nil { + t.Fatal(err) + } + var whole json.RawMessage + if err := Unmarshal(data, &whole); err != nil { + t.Fatal(err) + } + for i := range data { + data[i] = 'X' + } + if v.S != "hello" || string(v.Raw) != `{"a":"b"}` || v.N != "12345" || string(whole) != input { + t.Errorf("decoded values changed with the input: %+v, %s", v, whole) + } + if m, ok := v.Any.(map[string]any); !ok || m["k"] != "v" { + t.Errorf("decoded map changed with the input: %v", v.Any) + } +} + +// Unmarshal allocates in proportion to its input, not a stream buffer. +func BenchmarkUnmarshal(b *testing.B) { + data := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"get_todo","arguments":{"name":"todos/1"}}}`) + b.ReportAllocs() + for b.Loop() { + var v unmarshalTarget + if err := Unmarshal(data, &v); err != nil { + b.Fatal(err) + } + } +}