From f68c00949a04552b318fbb0639fd9bccddf42969 Mon Sep 17 00:00:00 2001 From: Sung Yoon Whang Date: Tue, 7 Jul 2026 13:37:44 -0700 Subject: [PATCH 1/3] core/storage: stream graph writes instead of buffering in full WriteGraphStream and WriteChangedTargetsStream marshaled the entire response set into a bytes.Buffer before Put, keeping a second full copy of the serialized graph resident in memory. Replace this with a shared generic writeStream helper that encodes length-delimited protobuf into an io.Pipe while Put consumes the reader, so the payload is never fully buffered a second time. The writer goroutine checks ctx before each message so a cancellation unwinds the encode loop promptly and propagates the context error to the reader; if Put returns first, the reader is closed to unblock the goroutine. Both writers now delegate to the helper, dropping the redundant bytes.NewReader(buf.Bytes()) wrapper. There's another writer we can dedup this, but that will be a follow up. --- core/storage/changedtargetsreader.go | 12 +------- core/storage/graphwriter.go | 43 ++++++++++++++++++++++------ 2 files changed, 35 insertions(+), 20 deletions(-) diff --git a/core/storage/changedtargetsreader.go b/core/storage/changedtargetsreader.go index ca38347..d490a13 100644 --- a/core/storage/changedtargetsreader.go +++ b/core/storage/changedtargetsreader.go @@ -15,11 +15,8 @@ package storage import ( - "bytes" "context" - "fmt" - gogio "github.com/gogo/protobuf/io" pb "github.com/uber/tango/tangopb" ) @@ -43,12 +40,5 @@ func NewChangedTargetsReader(ctx context.Context, st Storage, key string) (Chang // WriteChangedTargetsStream writes a list of GetChangedTargetsResponse messages to storage. // The messages are written as length-delimited protobuf, allowing streaming reads. func WriteChangedTargetsStream(ctx context.Context, st Storage, key string, responses []*pb.GetChangedTargetsResponse) error { - buf := &bytes.Buffer{} - w := gogio.NewDelimitedWriter(buf) - for _, r := range responses { - if err := w.WriteMsg(r); err != nil { - return fmt.Errorf("write delimited: %w", err) - } - } - return st.Put(ctx, UploadRequest{Key: key, Reader: bytes.NewReader(buf.Bytes())}) + return writeStream[pb.GetChangedTargetsResponse](ctx, st, key, responses) } diff --git a/core/storage/graphwriter.go b/core/storage/graphwriter.go index 9bde41d..42c943f 100644 --- a/core/storage/graphwriter.go +++ b/core/storage/graphwriter.go @@ -15,24 +15,49 @@ package storage import ( - "bytes" "context" "fmt" + "io" gogio "github.com/gogo/protobuf/io" pb "github.com/uber/tango/tangopb" ) +// writeStream marshals msgs as length-delimited protobuf and streams them to +// storage under key. It uses an io.Pipe so the serialized payload is never +// buffered in full a second time: a writer goroutine encodes into the pipe +// while Put consumes from it. +// +// The writer goroutine checks ctx before each message so a cancellation +// unwinds the encode loop promptly instead of waiting for Put to notice and +// stop reading; the context error is propagated to the reader. If Put returns +// before draining the pipe, the reader is closed to unblock (and terminate) +// the writer goroutine either way. +func writeStream[T any, PT protoMessage[T]](ctx context.Context, st Storage, key string, msgs []PT) error { + pr, pw := io.Pipe() + go func() { + w := gogio.NewDelimitedWriter(pw) // varint-length-delimited + var err error + for _, m := range msgs { + if err = ctx.Err(); err != nil { + break + } + if err = w.WriteMsg(m); err != nil { + err = fmt.Errorf("write delimited: %w", err) + break + } + } + pw.CloseWithError(err) + }() + err := st.Put(ctx, UploadRequest{Key: key, Reader: pr}) + // Unblock the writer goroutine if Put stopped reading early. + pr.CloseWithError(err) + return err +} + // WriteGraphStream writes a list of GetTargetGraphResponse messages to the storage. // The messages are written as length-delimited protobuf, allowing streaming reads. // Typically this includes multiple OptimizedTargets chunks followed by Metadata. func WriteGraphStream(ctx context.Context, st Storage, key string, responses []*pb.GetTargetGraphResponse) error { - buf := &bytes.Buffer{} - w := gogio.NewDelimitedWriter(buf) // varint-length-delimited - for _, r := range responses { - if err := w.WriteMsg(r); err != nil { - return fmt.Errorf("write delimited: %w", err) - } - } - return st.Put(ctx, UploadRequest{Key: key, Reader: bytes.NewReader(buf.Bytes())}) + return writeStream[pb.GetTargetGraphResponse](ctx, st, key, responses) } From 7614cb3557adac1c8f277e9f302f1451a08b5b2e Mon Sep 17 00:00:00 2001 From: Sung Yoon Whang Date: Tue, 7 Jul 2026 15:32:22 -0700 Subject: [PATCH 2/3] wait for writer goroutine to finish --- core/storage/graphwriter.go | 17 ++++++++++++----- core/storage/storage_test.go | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/core/storage/graphwriter.go b/core/storage/graphwriter.go index 42c943f..6bd0a71 100644 --- a/core/storage/graphwriter.go +++ b/core/storage/graphwriter.go @@ -31,10 +31,12 @@ import ( // The writer goroutine checks ctx before each message so a cancellation // unwinds the encode loop promptly instead of waiting for Put to notice and // stop reading; the context error is propagated to the reader. If Put returns -// before draining the pipe, the reader is closed to unblock (and terminate) -// the writer goroutine either way. +// before draining the pipe, the reader is closed to unblock the writer. The +// goroutine is joined before returning, and its error is returned when Put +// succeeds. func writeStream[T any, PT protoMessage[T]](ctx context.Context, st Storage, key string, msgs []PT) error { pr, pw := io.Pipe() + writerErr := make(chan error, 1) go func() { w := gogio.NewDelimitedWriter(pw) // varint-length-delimited var err error @@ -48,11 +50,16 @@ func writeStream[T any, PT protoMessage[T]](ctx context.Context, st Storage, key } } pw.CloseWithError(err) + writerErr <- err }() - err := st.Put(ctx, UploadRequest{Key: key, Reader: pr}) + putErr := st.Put(ctx, UploadRequest{Key: key, Reader: pr}) // Unblock the writer goroutine if Put stopped reading early. - pr.CloseWithError(err) - return err + pr.CloseWithError(putErr) + writeErr := <-writerErr + if putErr != nil { + return putErr + } + return writeErr } // WriteGraphStream writes a list of GetTargetGraphResponse messages to the storage. diff --git a/core/storage/storage_test.go b/core/storage/storage_test.go index f1a5fd0..6598aa1 100644 --- a/core/storage/storage_test.go +++ b/core/storage/storage_test.go @@ -24,6 +24,40 @@ import ( "github.com/stretchr/testify/require" ) +var errMarshal = errors.New("marshal failed") + +type marshalErrorMessage struct{} + +func (*marshalErrorMessage) Reset() {} +func (*marshalErrorMessage) String() string { return "" } +func (*marshalErrorMessage) ProtoMessage() {} +func (*marshalErrorMessage) Marshal() ([]byte, error) { + return nil, errMarshal +} + +type discardStorage struct{} + +func (discardStorage) Get(context.Context, DownloadRequest) (DownloadResponse, error) { + return DownloadResponse{}, nil +} + +func (discardStorage) Put(context.Context, UploadRequest) error { return nil } + +func (discardStorage) Exists(context.Context, string) (bool, error) { return false, nil } + +func (discardStorage) List(context.Context, string) ([]string, error) { return nil, nil } + +func TestWriteStreamReturnsWriterError(t *testing.T) { + err := writeStream[marshalErrorMessage]( + context.Background(), + discardStorage{}, + "key", + []*marshalErrorMessage{{}}, + ) + + require.ErrorIs(t, err, errMarshal) +} + func TestMemoryStorage_List(t *testing.T) { ctx := context.Background() s := NewMemoryStorage() From b6c710f37c4efd4cd35e830fcb2e66d75081231c Mon Sep 17 00:00:00 2001 From: Sung Yoon Whang Date: Tue, 7 Jul 2026 17:52:03 -0700 Subject: [PATCH 3/3] fix mock that was relying on Put succeding without writes actually succeeding --- orchestrator/native_orchestrator_test.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/orchestrator/native_orchestrator_test.go b/orchestrator/native_orchestrator_test.go index 427d77c..a587052 100644 --- a/orchestrator/native_orchestrator_test.go +++ b/orchestrator/native_orchestrator_test.go @@ -95,7 +95,10 @@ func TestNative_GetTargetGraph_TreehashNotFound_NoError(t *testing.T) { // First attempt returns NotFound to trigger compute path. st.EXPECT().Get(gomock.Any(), gomock.Any()).Return(storage.DownloadResponse{}, &storage.NotFoundError{Path: "missing"}) // Expect writes (graph list and treehash cache mapping) - st.EXPECT().Put(gomock.Any(), gomock.Any()).Return(nil).MinTimes(2) + st.EXPECT().Put(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req storage.UploadRequest) error { + _, err := io.Copy(io.Discard, req.Reader) + return err + }).MinTimes(2) // After compute, second read returns a valid delimited stream with one message var buf bytes.Buffer _ = gogio.NewDelimitedWriter(&buf).WriteMsg(&pb.GetTargetGraphResponse{