core/storage: stream graph writes instead of buffering in full

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)
 }

core/storage/graphwriter.go

@@ -15,24 +15,56 @@
 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 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
+		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)
+		writerErr <- err
+	}()
+	putErr := st.Put(ctx, UploadRequest{Key: key, Reader: pr})
+	// Unblock the writer goroutine if Put stopped reading early.
+	pr.CloseWithError(putErr)
+	writeErr := <-writerErr
+	if putErr != nil {
+		return putErr
+	}
+	return writeErr
+}
+
 // 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)
 }

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()

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{