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..6bd0a71 100644 --- a/core/storage/graphwriter.go +++ b/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) } 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() 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{