Skip to content

Commit 98693ec

Browse files
committed
fix(artifact-cas): stream uploads to object-store backends to bound CAS memory
The CAS bytestream Write RPC buffered the entire artifact in memory before flushing it to the storage backend, so a large upload grew CAS memory ~1:1 with artifact size and OOM-killed the pod mid-stream. Object-store backends (S3, S3 access point, Azure Blob) now advertise a StreamingUploader capability and are fed directly from the client stream through an io.Pipe, so peak CAS memory stays bounded regardless of artifact size. The OCI backend keeps the fully-buffered path because its push implementation needs the whole layer content up front. A backend upload failure is always surfaced as an internal error and never misclassified as a client disconnect, and the per-upload size cap is enforced on both paths. The download path, including digest verification, is unchanged. Assisted-by: Claude Code Signed-off-by: Javier Rodriguez <javier@chainloop.dev> Chainloop-Trace-Sessions: 884134e6-0e9d-43aa-9637-2c5e1cfac1e5
1 parent e77e078 commit 98693ec

12 files changed

Lines changed: 1226 additions & 50 deletions

File tree

‎app/artifact-cas/internal/service/bytestream.go‎

Lines changed: 228 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -115,29 +115,33 @@ func (s *ByteStreamService) Write(stream bytestream.ByteStream_WriteServer) erro
115115
}
116116

117117
s.log.Infow("msg", "artifact does not exist, uploading", "digest", req.resource.Digest, "name", req.resource.FileName)
118-
// Create a buffer that will be filled in the background before sending its content to the backend
119-
buffer := newStreamReader(info.MaxBytes)
120-
// Add data from first request
121-
if err = buffer.Write(req.GetData()); err != nil {
122-
if backend.IsUploadSizeExceeded(err) {
123-
return status.Error(codes.ResourceExhausted, err.Error())
124-
}
125-
return sl.LogAndMaskErr(err, s.log)
126-
}
127118

128-
// Start a goroutine that will fill the buffer in the background
129-
go bufferStream(ctx, stream, buffer, s.log)
119+
// Streaming-capable backends (object stores such as S3/Azure) are fed
120+
// directly from the client stream through an io.Pipe, so CAS memory stays
121+
// bounded by the chunk/pipe size regardless of artifact size (PFM-6923).
122+
// The OCI backend, whose push path needs the whole layer content up front,
123+
// does not advertise streaming and keeps the fully-buffered path.
124+
var committedSize int64
125+
if su, ok := storageBackend.(backend.StreamingUploader); ok && su.SupportsStreaming() {
126+
committedSize, err = s.streamUpload(ctx, stream, storageBackend, req, info.MaxBytes)
127+
} else {
128+
committedSize, err = s.bufferedUpload(ctx, stream, storageBackend, req, info.MaxBytes)
129+
}
130130

131-
// Block until the buffer has been filled or the upload process has been canceled
132-
// This implementation is suboptimal since it requires the content to be uploaded in memory before pushing it
133-
// This is due to the fact that our OCI push implementation does not support streaming/chunks for uncompressed layers
134-
// We can not use stream.Layer since it only supports compressed layers, we want to store raw data and set custom mimetypes
135-
// https://github.com/google/go-containerregistry/blob/main/pkg/v1/stream/README.md
136-
// TODO: Split content in multiple layers and do concurrent uploads/downloads
137-
err = <-buffer.errorChan
138-
139-
// Now it's time to check if the data provider has sent an error
131+
// Classify the outcome. The error may come from two distinct stages, which
132+
// must be treated differently:
133+
// - A backend Upload failure (backendUploadError) is always masked as an
134+
// internal error. It must NOT be interpreted as a client disconnect even
135+
// when it wraps a network reset/cancellation originating backend-side —
136+
// doing so would falsely report success and silently drop the artifact.
137+
// - A stream-read (feed) error is classified: a client disconnect is not a
138+
// failure, an exceeded size cap maps to ResourceExhausted, anything else
139+
// is masked.
140140
if err != nil {
141+
var backendErr *backendUploadError
142+
if errors.As(err, &backendErr) {
143+
return sl.LogAndMaskErr(backendErr.err, s.log)
144+
}
141145
if isClientDisconnect(err) {
142146
s.log.Infow("msg", "upload canceled", "digest", req.resource.Digest, "name", req.resource.FileName)
143147
return nil
@@ -148,22 +152,198 @@ func (s *ByteStreamService) Write(stream bytestream.ByteStream_WriteServer) erro
148152
return sl.LogAndMaskErr(err, s.log)
149153
}
150154

151-
s.log.Infow("msg", "artifact received, uploading now to backend", "name", req.resource.FileName, "digest", req.resource.Digest, "size", buffer.size)
152-
if err := storageBackend.Upload(ctx, buffer, req.resource); err != nil {
153-
return sl.LogAndMaskErr(err, s.log)
154-
}
155-
156-
s.log.Infow("msg", "upload finished", "name", req.resource.FileName, "digest", req.resource.Digest, "size", buffer.size)
155+
s.log.Infow("msg", "upload finished", "name", req.resource.FileName, "digest", req.resource.Digest, "size", committedSize)
157156
s.audit.Dispatch(&events.CASArtifactUploaded{
158157
CASArtifactBase: &events.CASArtifactBase{
159158
Digest: req.resource.Digest,
160-
SizeBytes: buffer.size,
159+
SizeBytes: committedSize,
161160
FileName: req.resource.FileName,
162161
BackendType: info.BackendType,
163162
},
164163
}, info)
165164

166-
return stream.SendAndClose(&bytestream.WriteResponse{CommittedSize: buffer.size})
165+
return stream.SendAndClose(&bytestream.WriteResponse{CommittedSize: committedSize})
166+
}
167+
168+
// bufferedUpload accumulates the whole artifact in memory before handing it to
169+
// the backend. This is required by the OCI backend: its push implementation
170+
// does not support streaming/chunked uploads for uncompressed layers (we can not
171+
// use stream.Layer since it only supports compressed layers, and we want to
172+
// store raw data with custom mimetypes), so it needs the full content up front.
173+
// https://github.com/google/go-containerregistry/blob/main/pkg/v1/stream/README.md
174+
// It returns the total number of bytes committed to the backend. Feed errors are
175+
// returned unwrapped (classified by the caller); backend Upload failures are
176+
// wrapped in backendUploadError so the caller always masks them.
177+
func (s *ByteStreamService) bufferedUpload(ctx context.Context, stream bytestream.ByteStream_WriteServer, storageBackend backend.Uploader, req *writeRequest, maxBytes int64) (int64, error) {
178+
// Create a buffer that will be filled in the background before sending its content to the backend
179+
buffer := newStreamReader(maxBytes)
180+
// Add data from the first request
181+
if err := buffer.Write(req.GetData()); err != nil {
182+
return 0, err
183+
}
184+
185+
// Start a goroutine that will fill the buffer in the background
186+
go bufferStream(ctx, stream, buffer, s.log)
187+
188+
// Block until the buffer has been filled or the upload process has been canceled
189+
if err := <-buffer.errorChan; err != nil {
190+
return 0, err
191+
}
192+
193+
s.log.Infow("msg", "artifact received, uploading now to backend", "name", req.resource.FileName, "digest", req.resource.Digest, "size", buffer.size)
194+
if err := storageBackend.Upload(ctx, buffer, req.resource); err != nil {
195+
return 0, &backendUploadError{err}
196+
}
197+
198+
return buffer.size, nil
199+
}
200+
201+
// streamUpload pipes the client stream straight into the backend's Upload
202+
// without buffering the whole artifact in memory. A background goroutine feeds
203+
// received chunks into an io.Pipe while Upload consumes the other end, so the
204+
// two run concurrently and peak memory stays bounded (PFM-6923). It returns the
205+
// total number of bytes committed to the backend.
206+
//
207+
// Tradeoff vs. the buffered path: because bytes are forwarded as they arrive, an
208+
// over-cap upload has already streamed up to maxBytes to the backend by the time
209+
// the size guard trips (the client is still correctly rejected with
210+
// ResourceExhausted). For multipart backends the SDK aborts the in-flight upload
211+
// best-effort; operators should keep a bucket lifecycle rule to reap any
212+
// incomplete multipart uploads a failed abort could leave behind. This is
213+
// inherent to streaming — the alternative (buffer-then-validate) is exactly the
214+
// OOM this change removes.
215+
func (s *ByteStreamService) streamUpload(ctx context.Context, stream bytestream.ByteStream_WriteServer, storageBackend backend.Uploader, req *writeRequest, maxBytes int64) (int64, error) {
216+
pr, pw := io.Pipe()
217+
218+
var (
219+
uploadedSize int64
220+
feedErr error
221+
)
222+
done := make(chan struct{})
223+
go func() {
224+
defer close(done)
225+
uploadedSize, feedErr = feedPipe(ctx, stream, pw, req.GetData(), maxBytes, s.log, req.resource.Digest)
226+
// Closing with feedErr signals EOF to the reader when nil, or propagates
227+
// the failure so Upload stops reading.
228+
_ = pw.CloseWithError(feedErr)
229+
}()
230+
231+
uploadErr := storageBackend.Upload(ctx, streamingReader{pr}, req.resource)
232+
// If Upload returned without draining the pipe (a backend failure, or a
233+
// backend that reports success without reading to EOF), the feeding goroutine
234+
// may still be blocked on Write; closing the read end unblocks it. Then wait
235+
// for it so uploadedSize/feedErr are safe to read.
236+
_ = pr.CloseWithError(uploadErr)
237+
<-done
238+
239+
// errPipeConsumerGone means the feed only failed because the reader (Upload)
240+
// stopped consuming — a consequence of the upload outcome, not a genuine
241+
// stream-read failure, so the backend's own result is authoritative.
242+
if errors.Is(feedErr, errPipeConsumerGone) {
243+
feedErr = nil
244+
}
245+
246+
// A genuine feed-side error (client disconnect, exceeded size cap, stream
247+
// read failure) is the precise signal and takes precedence: when it occurs it
248+
// is what induced the backend error through the pipe. Returned unwrapped so
249+
// the caller classifies it (disconnect / ResourceExhausted / mask).
250+
if feedErr != nil {
251+
return 0, feedErr
252+
}
253+
// A backend failure is wrapped so the caller always masks it, never mistaking
254+
// a backend-side reset/cancellation for a client disconnect.
255+
if uploadErr != nil {
256+
return 0, &backendUploadError{uploadErr}
257+
}
258+
259+
return uploadedSize, nil
260+
}
261+
262+
// backendUploadError marks a failure returned by the storage backend's Upload,
263+
// as opposed to an error reading the client stream. Backend failures are always
264+
// masked as internal errors and are never interpreted as a client disconnect or
265+
// a size-cap violation, both of which only originate on the stream-read side.
266+
type backendUploadError struct{ err error }
267+
268+
func (e *backendUploadError) Error() string { return e.err.Error() }
269+
func (e *backendUploadError) Unwrap() error { return e.err }
270+
271+
// errPipeConsumerGone is returned by feedPipe when a write to the pipe fails,
272+
// which only happens once the reader (the backend Upload) has stopped consuming
273+
// — because Upload returned and streamUpload closed the read end, or because it
274+
// failed. It is not a genuine stream-read failure; streamUpload defers to the
275+
// backend's own error in that case.
276+
var errPipeConsumerGone = errors.New("pipe consumer stopped reading")
277+
278+
// streamingReader wraps the upload pipe reader with a stable string form. The
279+
// pipe is written to concurrently while the backend reads it; exposing the bare
280+
// *io.PipeReader lets a reflective consumer (a structured logger, a test's mock
281+
// matcher, etc.) walk the pipe's internal state and race with the writer. The
282+
// wrapper keeps io.Reader behaviour while presenting an opaque identity to fmt.
283+
type streamingReader struct {
284+
io.Reader
285+
}
286+
287+
func (streamingReader) String() string { return "cas-streaming-upload" }
288+
289+
// feedPipe forwards the artifact from the client stream into pw, enforcing the
290+
// max upload size as it goes. firstData is the payload already read from the
291+
// first request. It returns the total number of bytes forwarded.
292+
func feedPipe(ctx context.Context, stream bytestream.ByteStream_WriteServer, pw *io.PipeWriter, firstData []byte, maxSize int64, log *log.Helper, digest string) (int64, error) {
293+
var size int64
294+
write := func(data []byte) error {
295+
if len(data) == 0 {
296+
return nil
297+
}
298+
size += int64(len(data))
299+
if err := checkUploadSize(size, maxSize); err != nil {
300+
return err
301+
}
302+
if _, err := pw.Write(data); err != nil {
303+
// A write only fails once the reader has gone away; surface it as the
304+
// consumer-gone sentinel so streamUpload defers to the backend result
305+
// rather than treating this as a client-side stream failure.
306+
return errPipeConsumerGone
307+
}
308+
return nil
309+
}
310+
311+
// Forward the data from the first request.
312+
if err := write(firstData); err != nil {
313+
return size, err
314+
}
315+
316+
for {
317+
select {
318+
case <-ctx.Done():
319+
// DeadlineExceeded, or Canceled
320+
return size, ctx.Err()
321+
default:
322+
// Extract the next chunk of data from the stream request
323+
req, err := getWriteRequest(stream)
324+
if err != nil {
325+
// Finished reading the stream is not a real error
326+
if errors.Is(err, io.EOF) {
327+
return size, nil
328+
}
329+
return size, err
330+
}
331+
332+
// Forward this request's data first: a spec-compliant client may set
333+
// finish_write=true on the same message that carries the final chunk,
334+
// so the data must be written before the finish check or it is lost.
335+
if err := write(req.GetData()); err != nil {
336+
return size, err
337+
}
338+
339+
log.Debugw("msg", "upload chunk received (streaming)", "digest", digest, "currentSize", size, "maxSize", maxSize, "chunkSize", len(req.GetData()))
340+
341+
// Check if the client has finished sending data
342+
if req.GetFinishWrite() {
343+
return size, nil
344+
}
345+
}
346+
}
167347
}
168348

169349
// Server-side streaming RPC for reading blobs, implements the bytestream interface
@@ -250,18 +430,20 @@ func bufferStream(ctx context.Context, stream bytestream.ByteStream_WriteServer,
250430
return
251431
}
252432

253-
// Check if the client has finished sending data
254-
if req.GetFinishWrite() {
255-
return
256-
}
257-
258-
// Write the data to the buffer
433+
// Write the data first: a spec-compliant client may set
434+
// finish_write=true on the same message that carries the final chunk,
435+
// so the data must be buffered before the finish check or it is lost.
259436
if err = buffer.Write(req.GetData()); err != nil {
260437
bufferErr = err
261438
return
262439
}
263440

264441
log.Debugw("msg", "upload chunk received", "digest", req.resource.Digest, "currentSize", buffer.size, "maxSize", buffer.maxSize, "chunkSize", len(req.GetData()))
442+
443+
// Check if the client has finished sending data
444+
if req.GetFinishWrite() {
445+
return
446+
}
265447
}
266448
}
267449
}
@@ -290,16 +472,24 @@ func newStreamReader(maxSize int64) *streamReader {
290472
func (r *streamReader) Write(data []byte) error {
291473
r.size += int64(len(data))
292474

293-
// Check if the size of the buffer has exceeded the maximum allowed size
294-
// if maxSize is 0, then there is no limit
295-
if r.maxSize != 0 && r.size > r.maxSize {
296-
return backend.NewErrUploadSizeExceeded(r.size, r.maxSize)
475+
if err := checkUploadSize(r.size, r.maxSize); err != nil {
476+
return err
297477
}
298478

299479
_, err := r.Buffer.Write(data)
300480
return err
301481
}
302482

483+
// checkUploadSize returns an ErrUploadSizeExceeded when total exceeds maxSize.
484+
// maxSize == 0 means no limit. It is shared by the buffered (streamReader) and
485+
// streaming (feedPipe) paths so their cap semantics cannot drift.
486+
func checkUploadSize(total, maxSize int64) error {
487+
if maxSize != 0 && total > maxSize {
488+
return backend.NewErrUploadSizeExceeded(total, maxSize)
489+
}
490+
return nil
491+
}
492+
303493
type writeRequest struct {
304494
*bytestream.WriteRequest
305495
resource *v1.CASResource

0 commit comments

Comments
 (0)