From 1296c07c821a59a0d22dae9df5e1982b8e188673 Mon Sep 17 00:00:00 2001 From: CMGS Date: Mon, 5 Oct 2026 15:47:52 +0800 Subject: [PATCH] wire, sdk/go: streaming consumers decode bulk frames into one reused buffer Every stdout, stderr and data frame was base64-decoded into a fresh 256 KiB buffer, so a Pull or ReadFileTo allocated its whole size again and spent the GC on it. wire.DecodeResponseInto decodes a bulk frame into a caller-owned scratch buffer, and silkd.Conn.RecvTransient uses one per connection; Recv and DecodeResponse keep their fresh-buffer contract. drainData, pumpStdio and the pty drain hand each payload to a writer or sink that does not keep it, so they read transiently; ReadFile, which kept the frames, now appends them into its own buffer. --- protocol/wire/frame.go | 99 ++++++++++++++++++++++--------- protocol/wire/frame_bench_test.go | 10 ++++ protocol/wire/frame_test.go | 50 ++++++++++++++++ sdk/go/files.go | 13 ++-- sdk/go/pty.go | 2 +- sdk/go/silkd/conn.go | 24 ++++++-- sdk/go/silkd/conn_test.go | 30 ++++++++++ sdk/go/utils.go | 19 ++++-- 8 files changed, 198 insertions(+), 49 deletions(-) diff --git a/protocol/wire/frame.go b/protocol/wire/frame.go index 2779dc10..a8539f1d 100644 --- a/protocol/wire/frame.go +++ b/protocol/wire/frame.go @@ -105,15 +105,15 @@ var ( // responseDecoders dispatches on the type tag first: one key differs in shape across variants. responseDecoders = map[string]respDecoder{ "started": decodeResp[Started], - "stdout": fastBulk("stdout", decodeResp[Stdout], func(d []byte) Response { return &Stdout{Data: d} }), - "stderr": fastBulk("stderr", decodeResp[Stderr], func(d []byte) Response { return &Stderr{Data: d} }), + "stdout": decodeResp[Stdout], + "stderr": decodeResp[Stderr], "exit": decodeResp[Exit], "done": decodeResp[Done], "ready": decodeResp[Ready], "error": decodeResp[ErrorResp], "info": decodeResp[InfoResp], "procs": decodeResp[Procs], - "data": fastBulk("data", decodeResp[DataResp], func(d []byte) Response { return &DataResp{Data: d} }), + "data": decodeResp[DataResp], "entries": decodeResp[Entries], "stat": decodeResp[Stat], "session_created": decodeResp[SessionCreated], @@ -126,6 +126,13 @@ var ( "git_commit_result": decodeResp[GitCommitResult], "git_branches": decodeResp[GitBranches], } + + // bulkResponses decode their canonical shape without json.Unmarshal; any other shape takes responseDecoders. + bulkResponses = bulkDecoders( + func(d []byte) Response { return &Stdout{Data: d} }, + func(d []byte) Response { return &Stderr{Data: d} }, + func(d []byte) Response { return &DataResp{Data: d} }, + ) ) // Request is a client→server frame; Op is its wire tag. @@ -136,6 +143,11 @@ type Response interface{ RespType() string } type respDecoder func([]byte) (Response, error) +type bulkResponse struct { + head []byte + mk func([]byte) Response +} + // reqPtr and respPtr turn a mistyped decoder table entry into a compile error. type reqPtr[T any] interface { *T @@ -667,15 +679,12 @@ func DecodeRequest(line []byte) (Request, error) { // DecodeResponse parses one frame into its type's concrete Go type. Byte // fields are freshly allocated per frame, so callers may retain them. func DecodeResponse(line []byte) (Response, error) { - typ, err := frameTag(line, respTagHead, "type") - if err != nil { - return nil, fmt.Errorf("parse response frame: %w", err) - } - dec, ok := responseDecoders[string(typ)] - if !ok { - return nil, fmt.Errorf("unknown response type %q", typ) - } - return dec(line) + return decodeResponse(line, nil) +} + +// DecodeResponseInto is DecodeResponse for a streaming consumer: a bulk frame's bytes land in *scratch, grown as needed, and stay valid only until the next call with it. +func DecodeResponseInto(line []byte, scratch *[]byte) (Response, error) { + return decodeResponse(line, scratch) } // NewFrameScanner wraps r for newline-delimited frames capped at MaxFrame. @@ -707,26 +716,58 @@ func IsContinuation(line []byte) bool { return false } -// fastBulk decodes a canonical bulk frame without json.Unmarshal; any other shape falls back to slow. -func fastBulk(tag string, slow respDecoder, mk func([]byte) Response) respDecoder { - head := []byte(`{"type":"` + tag + `","data":"`) - return func(line []byte) (Response, error) { - after, ok := bytes.CutPrefix(line, head) - if !ok { - return slow(line) - } - b64, rest, ok := bytes.Cut(after, []byte(`"`)) - if !ok || !bytes.Equal(rest, []byte(`}`)) { - return slow(line) +func bulkDecoders(mks ...func([]byte) Response) map[string]bulkResponse { + out := make(map[string]bulkResponse, len(mks)) + for _, mk := range mks { + tag := mk(nil).RespType() + out[tag] = bulkResponse{head: []byte(`{"type":"` + tag + `","data":"`), mk: mk} + } + return out +} + +func decodeResponse(line []byte, scratch *[]byte) (Response, error) { + typ, err := frameTag(line, respTagHead, "type") + if err != nil { + return nil, fmt.Errorf("parse response frame: %w", err) + } + if bulk, ok := bulkResponses[string(typ)]; ok { + if data, ok := decodeBulk(line, bulk.head, scratch); ok { + return bulk.mk(data), nil } - out := make([]byte, base64.StdEncoding.DecodedLen(len(b64))) - n, err := base64.StdEncoding.Decode(out, b64) - // Decode skips CR/LF, so a length mismatch means non-canonical input. - if err != nil || base64.StdEncoding.EncodedLen(n) != len(b64) { - return slow(line) + } + dec, ok := responseDecoders[string(typ)] + if !ok { + return nil, fmt.Errorf("unknown response type %q", typ) + } + return dec(line) +} + +// decodeBulk decodes a canonical bulk frame's payload into scratch, or a fresh buffer when scratch is nil; ok is false for any other shape. +func decodeBulk(line, head []byte, scratch *[]byte) ([]byte, bool) { + after, ok := bytes.CutPrefix(line, head) + if !ok { + return nil, false + } + b64, rest, ok := bytes.Cut(after, []byte(`"`)) + if !ok || !bytes.Equal(rest, []byte(`}`)) { + return nil, false + } + size := base64.StdEncoding.DecodedLen(len(b64)) + var out []byte + if scratch == nil { + out = make([]byte, size) + } else { + if cap(*scratch) < size { + *scratch = make([]byte, size) } - return mk(out[:n]), nil + out = (*scratch)[:size] + } + n, err := base64.StdEncoding.Decode(out, b64) + // Decode skips CR/LF, so a length mismatch means non-canonical input. + if err != nil || base64.StdEncoding.EncodedLen(n) != len(b64) { + return nil, false } + return out[:n], true } // encodeTagged splices the tag head in front of v's fields, so wire tags never live on the structs; the spare capacity byte lets Conn.Send append the newline without a copy. diff --git a/protocol/wire/frame_bench_test.go b/protocol/wire/frame_bench_test.go index d3017289..65153df6 100644 --- a/protocol/wire/frame_bench_test.go +++ b/protocol/wire/frame_bench_test.go @@ -19,6 +19,16 @@ func BenchmarkDecodeBulk(b *testing.B) { } } }) + b.Run(fmt.Sprintf("%dKiB-into", size>>10), func(b *testing.B) { + b.ReportAllocs() + b.SetBytes(int64(size)) + var scratch []byte + for b.Loop() { + if _, err := DecodeResponseInto(frame, &scratch); err != nil { + b.Fatalf("decode: %v", err) + } + } + }) b.Run(fmt.Sprintf("%dKiB-json", size>>10), func(b *testing.B) { b.SetBytes(int64(size)) for b.Loop() { diff --git a/protocol/wire/frame_test.go b/protocol/wire/frame_test.go index 77c69262..95b39b44 100644 --- a/protocol/wire/frame_test.go +++ b/protocol/wire/frame_test.go @@ -247,6 +247,56 @@ func TestBulkDecodeStrictShape(t *testing.T) { } } +func TestDecodeResponseIntoReusesScratch(t *testing.T) { + var scratch []byte + first, err := DecodeResponseInto([]byte(`{"type":"data","data":"aGVsbG8="}`), &scratch) + if err != nil || string(first.(*DataResp).Data) != "hello" { + t.Fatalf("first frame: %#v, %v", first, err) + } + second, err := DecodeResponseInto([]byte(`{"type":"stdout","data":"d29ybGQ="}`), &scratch) + if err != nil || string(second.(*Stdout).Data) != "world" { + t.Fatalf("second frame: %#v, %v", second, err) + } + if &first.(*DataResp).Data[0] != &second.(*Stdout).Data[0] { + t.Error("second frame did not reuse the scratch buffer") + } + shadowed, err := DecodeResponseInto([]byte(`{"type":"stdout","meta":{"data":"WFhY"},"data":"aGk="}`), &scratch) + if err != nil || string(shadowed.(*Stdout).Data) != "hi" { + t.Errorf("non-canonical frame: %#v, %v", shadowed, err) + } + for _, frame := range []string{ + `{"type":"stdout","data":"aGk="}garbage`, + "{\"type\":\"stdout\",\"data\":\"Q\nQ==\"}", + } { + if _, err := DecodeResponseInto([]byte(frame), &scratch); err == nil { + t.Errorf("malformed frame accepted: %s", frame) + } + } +} + +func TestDecodeResponseIntoMatchesDecodeResponse(t *testing.T) { + entries, err := os.ReadDir(fixtureDir) + if err != nil { + t.Fatal(err) + } + var scratch []byte + for _, entry := range entries { + if !strings.HasPrefix(entry.Name(), "resp_") { + continue + } + raw, err := os.ReadFile(filepath.Join(fixtureDir, entry.Name())) + if err != nil { + t.Fatal(err) + } + raw = bytes.TrimSpace(raw) + want, wantErr := DecodeResponse(raw) + got, gotErr := DecodeResponseInto(raw, &scratch) + if (wantErr == nil) != (gotErr == nil) || !reflect.DeepEqual(got, want) { + t.Errorf("%s: Into = %#v, %v; DecodeResponse = %#v, %v", entry.Name(), got, gotErr, want, wantErr) + } + } +} + func TestTagAfterOtherKeys(t *testing.T) { resp, err := DecodeResponse([]byte(`{"data":"aGk=","type":"stdout"}`)) if err != nil { diff --git a/sdk/go/files.go b/sdk/go/files.go index 92b881ef..b1ed4236 100644 --- a/sdk/go/files.go +++ b/sdk/go/files.go @@ -4,7 +4,6 @@ import ( "bytes" "context" "io" - "slices" "github.com/cocoonstack/sandbox/protocol/wire" ) @@ -17,19 +16,15 @@ func (s *Sandbox) WriteFile(ctx context.Context, path string, data []byte, mode // ReadFile returns the contents of path. func (s *Sandbox) ReadFile(ctx context.Context, path string) ([]byte, error) { - // Retaining b is safe: fastBulk decodes each frame into a fresh buffer. - var chunks [][]byte + var out []byte err := s.downloadRPC(ctx, &wire.FsRead{Path: path}, func(b []byte) error { - chunks = append(chunks, b) + out = append(out, b...) return nil }) - if err != nil || len(chunks) == 0 { + if err != nil { return nil, err } - if len(chunks) == 1 { - return chunks[0], nil - } - return slices.Concat(chunks...), nil + return out, nil } // ReadFileTo streams path into w; an error from w ends the read there. diff --git a/sdk/go/pty.go b/sdk/go/pty.go index 117ab49c..a46a05e5 100644 --- a/sdk/go/pty.go +++ b/sdk/go/pty.go @@ -68,7 +68,7 @@ func (p *Pty) ExitCode() (code int, ok bool) { func (p *Pty) drain(ctx context.Context, pw *io.PipeWriter, stop func()) { defer stop() for { - resp, err := recv(ctx, p.conn) + resp, err := recvTransient(ctx, p.conn) if err != nil { _ = pw.CloseWithError(err) return diff --git a/sdk/go/silkd/conn.go b/sdk/go/silkd/conn.go index 64fdc9bc..6d693168 100644 --- a/sdk/go/silkd/conn.go +++ b/sdk/go/silkd/conn.go @@ -3,6 +3,7 @@ package silkd import ( "bufio" + "cmp" "fmt" "io" "sync" @@ -16,6 +17,7 @@ import ( type Conn struct { wmu sync.Mutex wbuf []byte + rbuf []byte rwc io.ReadWriteCloser sc *bufio.Scanner } @@ -46,19 +48,31 @@ func (c *Conn) Send(r wire.Request) error { // Recv reads the next response frame; io.EOF signals a clean close. func (c *Conn) Recv() (wire.Response, error) { - if !c.sc.Scan() { - if err := c.sc.Err(); err != nil { - return nil, err - } - return nil, io.EOF + if err := c.scan(); err != nil { + return nil, err } return wire.DecodeResponse(c.sc.Bytes()) } +// RecvTransient is Recv for a streaming consumer: a bulk frame's Data is valid only until the next RecvTransient. +func (c *Conn) RecvTransient() (wire.Response, error) { + if err := c.scan(); err != nil { + return nil, err + } + return wire.DecodeResponseInto(c.sc.Bytes(), &c.rbuf) +} + func (c *Conn) Close() error { return c.rwc.Close() } +func (c *Conn) scan() error { + if c.sc.Scan() { + return nil + } + return cmp.Or(c.sc.Err(), io.EOF) +} + // sendBulk renders a data/stdin frame into the reused buffer. func (c *Conn) sendBulk(op string, payload []byte) error { c.wmu.Lock() diff --git a/sdk/go/silkd/conn_test.go b/sdk/go/silkd/conn_test.go index 1a5b7757..9e703445 100644 --- a/sdk/go/silkd/conn_test.go +++ b/sdk/go/silkd/conn_test.go @@ -84,6 +84,36 @@ func TestConnExecStdinEcho(t *testing.T) { } } +func TestRecvKeepsDataWhileRecvTransientReusesIt(t *testing.T) { + client, server := net.Pipe() + t.Cleanup(func() { client.Close(); server.Close() }) + go func() { + for _, payload := range []string{"aGVsbG8=", "d29ybGQ=", "Zmlyc3Q=", "c2Vjb25k"} { + _, _ = io.WriteString(server, `{"type":"data","data":"`+payload+`"}`+"\n") + } + }() + conn := silkd.NewConn(client) + recvData := func(recv func() (wire.Response, error)) []byte { + t.Helper() + resp, err := recv() + if err != nil { + t.Fatalf("recv: %v", err) + } + return resp.(*wire.DataResp).Data + } + kept := recvData(conn.Recv) + if got := recvData(conn.Recv); string(kept) != "hello" || string(got) != "world" { + t.Errorf("Recv frames = %q, %q; want hello, world", kept, got) + } + first := recvData(conn.RecvTransient) + if string(first) != "first" { + t.Errorf("first transient frame = %q", first) + } + if second := recvData(conn.RecvTransient); string(second) != "second" || &first[0] != &second[0] { + t.Errorf("second transient frame = %q at a fresh buffer, want second in the reused one", second) + } +} + func TestConnRejectsOversizedFrame(t *testing.T) { client, server := net.Pipe() t.Cleanup(func() { client.Close(); server.Close() }) diff --git a/sdk/go/utils.go b/sdk/go/utils.go index be0ed9f1..be77a812 100644 --- a/sdk/go/utils.go +++ b/sdk/go/utils.go @@ -133,7 +133,7 @@ func pumpStdio(ctx context.Context, conn *silkd.Conn, stdout, stderr io.Writer) stdout = cmp.Or(stdout, io.Discard) stderr = cmp.Or(stderr, io.Discard) for { - resp, err := recv(ctx, conn) + resp, err := recvTransient(ctx, conn) if err != nil { return pid, 0, false, err } @@ -160,10 +160,10 @@ func pumpStdio(ctx context.Context, conn *silkd.Conn, stdout, stderr io.Writer) } } -// drainData consumes Data frames into sink until Done; any other frame is an error. +// drainData consumes Data frames into sink until Done; sink must not retain its argument. func drainData(ctx context.Context, conn *silkd.Conn, sink func([]byte) error) error { for { - resp, err := recv(ctx, conn) + resp, err := recvTransient(ctx, conn) if err != nil { return err } @@ -202,9 +202,18 @@ func expect[T any, PT respPtr[T]](ctx context.Context, conn *silkd.Conn) (*T, er return nil, unexpected(resp) } -// recv reads one frame, translating a canceled ctx and an early EOF. func recv(ctx context.Context, conn *silkd.Conn) (wire.Response, error) { - resp, err := conn.Recv() + return recvWith(ctx, conn.Recv) +} + +// recvTransient is recv for a consumer done with a bulk frame's Data before its next read. +func recvTransient(ctx context.Context, conn *silkd.Conn) (wire.Response, error) { + return recvWith(ctx, conn.RecvTransient) +} + +// recvWith reads one frame, translating a canceled ctx and an early EOF. +func recvWith(ctx context.Context, next func() (wire.Response, error)) (wire.Response, error) { + resp, err := next() if err != nil { if ctxErr := ctx.Err(); ctxErr != nil { return nil, ctxErr