From e73dcca89ea17363d5fd5a11aaa9514a95cfcd9f Mon Sep 17 00:00:00 2001 From: Chenyi Wang Date: Tue, 20 Oct 2026 00:29:10 +0000 Subject: [PATCH] objectstorage: size small snapshot upload buffers to actual data On GCS, streamed snapshot files smaller than uploadCompositeMin (64 MiB) allocated and reserved 64 MiB for PutObject's lookahead buffer, 64 MiB for putSingle's Writer.ChunkSize, and 64 MiB for parZstd's 8 MiB chunk buffers. Each small file reserved 192 MiB against the 512 MiB default upload budget, serializing a 7-file actor snapshot into two-file waves. Wrap bounded upload streams in boundedReader so PutObject and putSingle size their single-request buffers to the next 256 KiB multiple above the stream bound. Scale parZstd's frame chunk size down to a power of two in [1 MiB, 8 MiB] for sparse files with less than 8 MiB of populated data, and update planStream and streamBufferMemory to reserve the corresponding memory. --- cmd/objectstore-plugin/README.md | 12 +- internal/objectstoreplugin/upload_test.go | 15 +- pkg/objectstorage/gcs.go | 70 ++++++-- pkg/objectstorage/objects.go | 23 ++- pkg/objectstorage/parzstd.go | 61 ++++++- pkg/objectstorage/sparsezstd.go | 50 +++--- pkg/objectstorage/upload.go | 42 +++-- pkg/objectstorage/upload_bench_test.go | 190 +++++++++++++++++++++- pkg/objectstorage/upload_gcs_test.go | 178 ++++++++++++++++++++ pkg/objectstorage/upload_test.go | 16 +- 10 files changed, 561 insertions(+), 96 deletions(-) diff --git a/cmd/objectstore-plugin/README.md b/cmd/objectstore-plugin/README.md index a81ca91f2e..ce02d4c5c1 100644 --- a/cmd/objectstore-plugin/README.md +++ b/cmd/objectstore-plugin/README.md @@ -68,11 +68,13 @@ in memory. The sparse format and download readers are unchanged. With the current buffer sizes, the minimum budget is `384Mi` for GCS and `224Mi` for S3. Startup rejects smaller values. Each upload reserves for the path it -takes: on GCS, a file split into range parts reserves only those parts, and an -object whose sparse stream is bounded below one request (`64Mi`) reserves that -request's buffers instead of the composite-upload buffers. Other streamed -objects still reserve for a composite upload, even when they ultimately fit in -one request. Raise the budget to trade memory for more simultaneous uploads. +takes: on GCS, a file split into range parts reserves only those parts, sparse +files with less than `8Mi` of populated data scale their compression chunk +buffers down to `1Mi`, and an object whose stream is bounded below one request +(`64Mi`) sizes both its lookahead and single-request chunk buffer to that +bound instead of the composite-upload buffers. Unbounded streamed objects still +reserve for a composite upload, even when they ultimately fit in one request. +Raise the budget to trade memory for more simultaneous uploads. The budget is **not a process memory limit**. The Go runtime, garbage awaiting collection, connections, downloads, and asset transfers need additional memory. diff --git a/internal/objectstoreplugin/upload_test.go b/internal/objectstoreplugin/upload_test.go index cbec1a68d8..2fb8f19151 100644 --- a/internal/objectstoreplugin/upload_test.go +++ b/internal/objectstoreplugin/upload_test.go @@ -91,8 +91,9 @@ func TestUploadBudgetSharedAcrossSnapshots(t *testing.T) { synctest.Test(t, func(t *testing.T) { gate := make(chan struct{}) store := &slowSnapshotStore{memObjects: newMemObjects(), release: gate} - // This backend has no SDK buffers. Two compression reservations fit. - node, err := NewNodePlugin(store, root, 128<<20) + // This backend has no SDK buffers. Eight small-file compression + // reservations (8 MiB each) fit in the 64 MiB minimum budget. + node, err := NewNodePlugin(store, root, 64<<20) if err != nil { t.Fatal(err) } @@ -101,8 +102,8 @@ func TestUploadBudgetSharedAcrossSnapshots(t *testing.T) { go func() { _, err := node.UploadSnapshot(t.Context(), req); done <- err }() } synctest.Wait() - if got := store.active.Load(); got != 2 { - t.Errorf("active uploads = %d, want 2 across all requests", got) + if got := store.active.Load(); got != 8 { + t.Errorf("active uploads = %d, want 8 across all requests", got) } if got := store.manifests.Load(); got != 0 { t.Errorf("committed %d manifests before data finished", got) @@ -114,8 +115,8 @@ func TestUploadBudgetSharedAcrossSnapshots(t *testing.T) { t.Fatal(err) } } - if got := store.peak.Load(); got != 2 { - t.Errorf("peak uploads = %d, want 2", got) + if got := store.peak.Load(); got != 8 { + t.Errorf("peak uploads = %d, want 8", got) } if got := len(store.keys()); got != snapshots*(filesPerSnapshot+1) { t.Fatalf("uploaded %d objects", got) @@ -170,7 +171,7 @@ func TestUploadFailureDoesNotCommitManifest(t *testing.T) { func TestUploadCanceledWaiterReturnsCanceled(t *testing.T) { root := t.TempDir() dir := filepath.Join(root, "src") - files := snapshotUploadFiles(t, dir, 1) + files := snapshotUploadFiles(t, dir, 8) synctest.Test(t, func(t *testing.T) { gate := make(chan struct{}) store := &slowSnapshotStore{memObjects: newMemObjects(), release: gate} diff --git a/pkg/objectstorage/gcs.go b/pkg/objectstorage/gcs.go index 327af424f5..1445f406a1 100644 --- a/pkg/objectstorage/gcs.go +++ b/pkg/objectstorage/gcs.go @@ -105,38 +105,74 @@ func setRetry(c *storage.Client) { // the signal. func (g *gcsClient) supportsStreamingPut() {} -// uploadChunkSize is how much of a streamed object the GCS client buffers before -// starting a request. Each chunk costs a round trip, so a snapshot that fits in one -// chunk pays only one: measured on a GKE worker node, a 24 MiB object took 425-530ms -// at the 16 MiB default versus 258-314ms at 64 MiB. The buffer is capped by the object -// size, so small objects still cost only their own bytes. -const uploadChunkSize = 64 << 20 +// boundedReader wraps an io.Reader whose remaining length is at most max bytes. +type boundedReader struct { + io.Reader + max int64 +} + +func (b boundedReader) sizeBound() int64 { return b.max } +func (b boundedReader) Unwrap() io.Reader { return b.Reader } + +// readerSizeBound returns the upper bound on r's remaining length if r or an +// unwrapped inner reader reports one via sizeBound(). +func readerSizeBound(r io.Reader) (int64, bool) { + for r != nil { + if b, ok := r.(interface{ sizeBound() int64 }); ok { + return b.sizeBound(), true + } + u, ok := r.(interface{ Unwrap() io.Reader }) + if !ok { + return 0, false + } + r = u.Unwrap() + } + return 0, false +} + +// singleUploadChunkSize returns the Writer.ChunkSize that sends an object of at +// most size bytes in one request. googleapi rounds ChunkSize up to a multiple of +// MinUploadChunkSize (256 KiB), allocates a retry buffer of that size, and sends +// media strictly smaller than ChunkSize in a single request. +func singleUploadChunkSize(size int64) int64 { + const unit = int64(googleapi.MinUploadChunkSize) + return ((max(size, 0) + unit) / unit) * unit +} + +// singleUploadLookahead returns the PutObject lookahead size for reader. +func singleUploadLookahead(reader io.Reader) int64 { + if bound, ok := readerSizeBound(reader); ok && bound >= 0 && bound < uploadCompositeMin { + return singleUploadChunkSize(bound) + } + return uploadCompositeMin +} // PutObject writes reader to the object, in a single request below -// uploadCompositeMin and as parallel parts above it. The size is not known up front, -// so this reads that many bytes to find out which case it is, then hands them on. +// uploadCompositeMin and as parallel parts above it. When reader carries an +// upper bound below uploadCompositeMin, only that request's chunk size is +// allocated for the lookahead instead of the full uploadCompositeMin. func (g *gcsClient) PutObject(ctx context.Context, bucket, object string, reader io.Reader) error { - head := make([]byte, uploadCompositeMin) + head := make([]byte, singleUploadLookahead(reader)) n, err := io.ReadFull(reader, head) switch { case errors.Is(err, io.EOF), errors.Is(err, io.ErrUnexpectedEOF): // The whole object is in hand and fits in one request. - return g.putSingle(ctx, bucket, object, bytes.NewReader(head[:n])) + return g.putSingle(ctx, bucket, object, head[:n]) case err != nil: return fmt.Errorf("while reading object body: %w", err) } return g.putComposite(ctx, bucket, object, bytes.NewReader(head[:n]), reader) } -// putSingle writes the whole body in one resumable request, on a pooled client so -// the small objects a node uploads concurrently do not queue on one connection. -func (g *gcsClient) putSingle(ctx context.Context, bucket, object string, reader io.Reader) error { +// putSingle writes body in one request, on a pooled client so the small objects +// a node uploads concurrently do not queue on one connection. +func (g *gcsClient) putSingle(ctx context.Context, bucket, object string, body []byte) error { wc := g.nextClient().Bucket(bucket).Object(object).NewWriter(ctx) - wc.ChunkSize = uploadChunkSize - // io.Copy reports local read errors; wc.Close() reports the actual - // GCS upload (auth, permissions, transient). Join both so the caller + wc.ChunkSize = int(singleUploadChunkSize(int64(len(body)))) + // Write reports local pipe errors; wc.Close() reports the actual GCS + // upload (auth, permissions, transient). Join both so the caller // doesn't lose either. - _, copyErr := io.Copy(wc, reader) + _, copyErr := wc.Write(body) closeErr := wc.Close() if err := errors.Join(copyErr, closeErr); err != nil { return fmt.Errorf("while putting GCS object: %w", err) diff --git a/pkg/objectstorage/objects.go b/pkg/objectstorage/objects.go index c7c696df97..e1827b32b3 100644 --- a/pkg/objectstorage/objects.go +++ b/pkg/objectstorage/objects.go @@ -184,6 +184,13 @@ type writeContentResult struct { // plain zstd stream otherwise. It touches only io, so it is unit-testable without // an object store, and is shared by the buffered and streaming upload paths. func writeContent(ctx context.Context, out io.Writer, content io.Reader, workers int) (writeContentResult, error) { + if sf, ok := content.(*sparseFile); ok { + logical, populated, err := writeSparseExtents(ctx, out, sf.File, sf.size, sf.exts, sf.populated, workers) + if err != nil { + return writeContentResult{}, err + } + return writeContentResult{logicalBytes: logical, populatedBytes: populated, sparse: true}, nil + } if f, ok := content.(*os.File); ok { logical, populated, err := writeSparseZstd(ctx, out, f, workers) if err != nil { @@ -223,7 +230,7 @@ func (u *Uploader) sendZstd(ctx context.Context, gsURL string, content io.Reader if err != nil { return fmt.Errorf("while parsing URL: %w", err) } - plan, err := u.plan(content) + plan, streamMax, plannedContent, err := u.planStream(content) if err != nil { return err } @@ -232,7 +239,7 @@ func (u *Uploader) sendZstd(ctx context.Context, gsURL string, content io.Reader } defer u.budget.reserved.Release(plan.memory) tStart := time.Now() - if sp, ok := client.(sparseFilePutter); ok { + if sp, ok := client.(sparseFilePutter); ok && plan.parts >= 2 { if f, isFile := content.(*os.File); isFile { res, err := sp.PutSparseFile(ctx, bucket, object, f, plan.parts) switch { @@ -254,9 +261,9 @@ func (u *Uploader) sendZstd(ctx context.Context, gsURL string, content io.Reader } } if _, ok := client.(streamingPutter); ok { - return sendStreamingZstd(ctx, client, bucket, object, content, tStart, plan.workers) + return sendStreamingZstd(ctx, client, bucket, object, plannedContent, tStart, plan.workers, streamMax) } - return sendBufferedZstd(ctx, client, bucket, object, content, tStart, plan.workers) + return sendBufferedZstd(ctx, client, bucket, object, plannedContent, tStart, plan.workers) } // sendBufferedZstd compresses content to a seekable temp file, then uploads it. @@ -295,7 +302,7 @@ func sendBufferedZstd(ctx context.Context, client ObjectStorage, bucket, object // PutObject streams the read end to the object store. No seekable temp file, and // the compress runs concurrently with the network PUT. Used only for streaming // backends (GCS and S3); see sendZstd. -func sendStreamingZstd(ctx context.Context, client ObjectStorage, bucket, object string, content io.Reader, tStart time.Time, workers int) error { +func sendStreamingZstd(ctx context.Context, client ObjectStorage, bucket, object string, content io.Reader, tStart time.Time, workers int, streamMax int64) error { ctx, cancel := context.WithCancel(ctx) defer cancel() type result struct { @@ -311,7 +318,11 @@ func sendStreamingZstd(ctx context.Context, client ObjectStorage, bucket, object ch <- result{res: res, err: err} }() - putErr := client.PutObject(ctx, bucket, object, pr) + var body io.Reader = pr + if streamMax >= 0 { + body = boundedReader{Reader: pr, max: streamMax} + } + putErr := client.PutObject(ctx, bucket, object, body) if putErr != nil { cancel() // PutObject bailed (e.g. mid-stream); unblock the compressor goroutine so it diff --git a/pkg/objectstorage/parzstd.go b/pkg/objectstorage/parzstd.go index 7802703d36..4a0157b683 100644 --- a/pkg/objectstorage/parzstd.go +++ b/pkg/objectstorage/parzstd.go @@ -35,8 +35,11 @@ import ( // Independent chunks cost a little ratio (no matches across a boundary), which // parZstdChunk is sized to keep negligible. const ( - // parZstdChunk is the per-frame chunk size. + // parZstdChunk is the maximum per-frame chunk size. parZstdChunk = 8 << 20 + // parZstdMinChunk is the smallest per-frame chunk size, covering a fast + // encoder's fixed tables, block buffers, and minimum history. + parZstdMinChunk = 1 << 20 // parZstdQueue is how many chunks may be outstanding per worker, so the producer // can run ahead while workers compress. parZstdQueue = 2 @@ -57,18 +60,43 @@ var ( parZstdFramePool = sync.Pool{New: func() any { b := make([]byte, 0, parZstdFrameCap); return &b }} ) +// parZstdChunkSize returns the per-frame chunk size for a sparse file with +// populated data bytes, as a power of two between parZstdMinChunk and parZstdChunk. +func parZstdChunkSize(populated int64) int64 { + chunk := int64(parZstdMinChunk) + for chunk < parZstdChunk && chunk < populated { + chunk <<= 1 + } + return chunk +} + +// parZstdWorkerMemory returns the memory reservation per worker for a given +// frame chunk size. +func parZstdWorkerMemory(chunk int64) int64 { + return (parZstdQueue*3 + 2) * chunk +} + // newChunkEncoder returns an encoder for EncodeAll use. Single-segment frames always // declare their content size, even the short last frame of a stream (zstd otherwise // leaves it out under 256 bytes), so a reader can size every frame without having // to bound it; a chunk is at most parZstdChunk, so the frame never needs a larger // window than that. func newChunkEncoder() *zstd.Encoder { - enc, err := zstd.NewWriter(nil, + return newChunkEncoderSized(parZstdChunk) +} + +func newChunkEncoderSized(chunk int64) *zstd.Encoder { + opts := []zstd.EOption{ zstd.WithEncoderLevel(zstd.SpeedFastest), zstd.WithEncoderConcurrency(1), - zstd.WithSingleSegment(true)) + zstd.WithSingleSegment(true), + } + if chunk > 0 && chunk < 4<<20 { + opts = append(opts, zstd.WithWindowSize(int(max(chunk, parZstdMinChunk)))) + } + enc, err := zstd.NewWriter(nil, opts...) if err != nil { - // NewWriter only fails on bad options, which are compile-time constants here. + // NewWriter only fails on bad options, which are bounded powers of two here. panic(err) } return enc @@ -77,6 +105,13 @@ func newChunkEncoder() *zstd.Encoder { func getChunk() []byte { return (*parZstdChunkPool.Get().(*[]byte))[:0] } func getFrame() []byte { return (*parZstdFramePool.Get().(*[]byte))[:0] } +func getChunkSized(chunk int64) []byte { + if chunk == parZstdChunk { + return getChunk() + } + return make([]byte, 0, chunk) +} + func putChunk(b []byte) { if cap(b) == parZstdChunk { b = b[:0] @@ -107,6 +142,7 @@ func putFrame(b []byte) { type parZstd struct { dst io.Writer workers int + chunk int64 buf []byte free chan []byte @@ -131,28 +167,37 @@ type parZstdJob struct { // One worker, after the cap, means inline mode with no goroutines, so with // GOMAXPROCS=1 every caller compresses inline. func newParZstd(dst io.Writer, workers int) *parZstd { + return newParZstdChunk(dst, workers, parZstdChunk) +} + +func newParZstdChunk(dst io.Writer, workers int, chunk int64) *parZstd { if workers <= 0 || workers > runtime.GOMAXPROCS(0) { workers = runtime.GOMAXPROCS(0) } + if chunk <= 0 || chunk > parZstdChunk { + chunk = parZstdChunk + } if workers == 1 { return &parZstd{ dst: dst, workers: 1, + chunk: chunk, inline: true, - enc: newChunkEncoder(), - buf: getChunk(), + enc: newChunkEncoderSized(chunk), + buf: getChunkSized(chunk), } } p := &parZstd{ dst: dst, workers: workers, + chunk: chunk, free: make(chan []byte, workers*parZstdQueue), jobs: make(chan parZstdJob, workers), ordered: make(chan chan []byte, workers*parZstdQueue), done: make(chan struct{}), } for range workers * parZstdQueue { - p.free <- getChunk() + p.free <- getChunkSized(chunk) } for range workers { p.workerWG.Go(p.worker) @@ -165,7 +210,7 @@ func newParZstd(dst io.Writer, workers int) *parZstd { // worker compresses whole chunks. Each holds its own encoder: the encoders are // single-shot EncodeAll users, so one per worker keeps their state private. func (p *parZstd) worker() { - enc := newChunkEncoder() + enc := newChunkEncoderSized(p.chunk) defer enc.Close() for j := range p.jobs { // Only a full chunk takes a pooled frame. A partial one (a small file, or diff --git a/pkg/objectstorage/sparsezstd.go b/pkg/objectstorage/sparsezstd.go index 05b18643e0..d5e59f0ae6 100644 --- a/pkg/objectstorage/sparsezstd.go +++ b/pkg/objectstorage/sparsezstd.go @@ -23,7 +23,6 @@ import ( "os" "github.com/klauspost/compress/zstd" - "golang.org/x/sys/unix" ) // sparseMagic marks the sparse-extent snapshot format (see writeSparseZstd). It is @@ -44,6 +43,16 @@ const sparseVersion uint32 = 2 // and the reader stops when it sees the sentinel. const sparseEndOffset int64 = -1 +// sparseFile holds a pre-scanned sparse file's logical size, extents, and +// populated byte count so planning and compression share one SEEK_DATA/SEEK_HOLE +// pass and use the same chunk size. +type sparseFile struct { + *os.File + size int64 + exts []extentRange + populated int64 +} + // writeSparseZstd encodes a sparse file src to dst in the sparse-extent format: // // magic[8] | version:u32 | zstd( totalSize:i64 | (off:i64, len:i64, data[len])* | -1:i64 ) @@ -66,7 +75,17 @@ func writeSparseZstd(ctx context.Context, dst io.Writer, src *os.File, workers i return 0, 0, err } size := fi.Size() + exts, populated, err := sparseExtents(src, size) + if err != nil { + return 0, 0, err + } + if _, err := src.Seek(0, io.SeekStart); err != nil { + return 0, 0, err + } + return writeSparseExtents(ctx, dst, src, size, exts, populated, workers) +} +func writeSparseExtents(ctx context.Context, dst io.Writer, src *os.File, size int64, exts []extentRange, populated int64, workers int) (logical, dataBytes int64, err error) { // magic + version in the clear (buffered: a couple of tiny writes). bw := bufio.NewWriter(dst) if _, err := bw.WriteString(sparseMagic); err != nil { @@ -79,7 +98,7 @@ func writeSparseZstd(ctx context.Context, dst io.Writer, src *os.File, workers i return 0, 0, err } - zw := newParZstd(dst, workers) + zw := newParZstdChunk(dst, workers, parZstdChunkSize(populated)) cw := contextWriter{ctx: ctx, dst: zw} // fail closes the encoder before returning err (Close flushes/frees state). fail := func(e error) (int64, int64, error) { @@ -90,36 +109,21 @@ func writeSparseZstd(ctx context.Context, dst io.Writer, src *os.File, workers i return fail(err) } - fd := int(src.Fd()) - off := int64(0) - for off < size { - ds, serr := unix.Seek(fd, off, unix.SEEK_DATA) - if serr != nil { - if serr == unix.ENXIO { // no more data: the rest is a hole - break - } - return fail(fmt.Errorf("SEEK_DATA: %w", serr)) - } - de, serr := unix.Seek(fd, ds, unix.SEEK_HOLE) - if serr != nil { - return fail(fmt.Errorf("SEEK_HOLE: %w", serr)) - } - length := de - ds - if err := binary.Write(cw, binary.LittleEndian, ds); err != nil { + for _, e := range exts { + if err := binary.Write(cw, binary.LittleEndian, e.off); err != nil { return fail(err) } - if err := binary.Write(cw, binary.LittleEndian, length); err != nil { + if err := binary.Write(cw, binary.LittleEndian, e.length); err != nil { return fail(err) } - if _, err := src.Seek(ds, io.SeekStart); err != nil { + if _, err := src.Seek(e.off, io.SeekStart); err != nil { return fail(err) } - n, cerr := io.CopyN(cw, src, length) + n, cerr := io.CopyN(cw, src, e.length) dataBytes += n if cerr != nil { - return fail(fmt.Errorf("reading extent @%d+%d: %w", ds, length, cerr)) + return fail(fmt.Errorf("reading extent @%d+%d: %w", e.off, e.length, cerr)) } - off = de } if err := binary.Write(cw, binary.LittleEndian, sparseEndOffset); err != nil { return fail(err) diff --git a/pkg/objectstorage/upload.go b/pkg/objectstorage/upload.go index bca454f9d1..247d1e1572 100644 --- a/pkg/objectstorage/upload.go +++ b/pkg/objectstorage/upload.go @@ -97,25 +97,32 @@ type uploadPlan struct { } func (u *Uploader) plan(content io.Reader) (uploadPlan, error) { + plan, _, _, err := u.planStream(content) + return plan, err +} + +func (u *Uploader) planStream(content io.Reader) (uploadPlan, int64, io.Reader, error) { buffers := uploadBufferMemory(u.client) workers := min(int64(runtime.GOMAXPROCS(0)), (u.budget.capacity-buffers)/parallelWorkerMemory) + workerMemory := parallelWorkerMemory parts := 1 + streamMax := int64(-1) if f, ok := content.(*os.File); ok { info, err := f.Stat() if err != nil { - return uploadPlan{}, err + return uploadPlan{}, 0, nil, err } position, err := f.Seek(0, io.SeekCurrent) if err != nil { - return uploadPlan{}, err + return uploadPlan{}, 0, nil, err } exts, populated, err := sparseExtents(f, info.Size()) _, seekErr := f.Seek(position, io.SeekStart) if err != nil { - return uploadPlan{}, err + return uploadPlan{}, 0, nil, err } if seekErr != nil { - return uploadPlan{}, seekErr + return uploadPlan{}, 0, nil, seekErr } // Holes never enter the compressor or range uploads. workers = min(workers, populated/parZstdChunk+1) @@ -124,17 +131,15 @@ func (u *Uploader) plan(content io.Reader) (uploadPlan, error) { // PutSparseFile splits a file into parts whenever it has two or // more (see errSparseTooSmall), so the streaming fallback never // runs: reserve only what the parts hold. - return uploadPlan{workers: int(workers), parts: parts, memory: int64(parts) * sparseUploadMemory}, nil + return uploadPlan{workers: int(workers), parts: parts, memory: int64(parts) * sparseUploadMemory}, -1, content, nil } - buffers = streamBufferMemory(u.client, sparseEncodedBound(populated, len(exts))) - } - memory := buffers + workers*parallelWorkerMemory - if _, ok := u.client.(sparseFilePutter); ok { - // The range-parallel path and streaming fallback are alternatives, - // so reserve the larger requirement rather than adding both. - memory = max(memory, int64(parts)*sparseUploadMemory) + streamMax = sparseEncodedBound(populated, len(exts)) + buffers = streamBufferMemory(u.client, streamMax) + workerMemory = parZstdWorkerMemory(parZstdChunkSize(populated)) + content = &sparseFile{File: f, size: info.Size(), exts: exts, populated: populated} } - return uploadPlan{workers: int(workers), parts: parts, memory: memory}, nil + memory := buffers + workers*workerMemory + return uploadPlan{workers: int(workers), parts: parts, memory: memory}, streamMax, content, nil } // sparseEncodedBound bounds the size of the sparse-extent stream for a file @@ -147,11 +152,12 @@ func sparseEncodedBound(populated int64, extents int) int64 { // streamBufferMemory is the backend buffering for a streamed object of at // most size bytes. GCS sends an object smaller than uploadCompositeMin in one -// request: it holds PutObject's lookahead and the client's chunk buffer for -// that request (see putSingle), never the composite parts. +// request: PutObject allocates singleUploadChunkSize(size) for the lookahead +// and putSingle sets the same ChunkSize on the writer, never allocating the +// composite parts. func streamBufferMemory(client ObjectStorage, size int64) int64 { if _, ok := client.(*gcsClient); ok && size < uploadCompositeMin { - return uploadCompositeMin + uploadChunkSize + return 2 * singleUploadChunkSize(size) } return uploadBufferMemory(client) } @@ -179,7 +185,9 @@ func (u *Uploader) Send(ctx context.Context, url string, content io.Reader) erro // A manifest is a few KiB: size its reservation by the file rather // than by the largest object the backend could buffer. if info, err := f.Stat(); err == nil && info.Mode().IsRegular() { - memory = streamBufferMemory(u.client, info.Size()) + size := info.Size() + memory = streamBufferMemory(u.client, size) + content = boundedReader{Reader: content, max: size} } } if err := u.budget.reserved.Acquire(ctx, memory); err != nil { diff --git a/pkg/objectstorage/upload_bench_test.go b/pkg/objectstorage/upload_bench_test.go index 40beb8ae93..c1ef431041 100644 --- a/pkg/objectstorage/upload_bench_test.go +++ b/pkg/objectstorage/upload_bench_test.go @@ -15,17 +15,24 @@ package objectstorage import ( + "bytes" "context" + "encoding/json" "fmt" "io" "log/slog" + "net/http" + "net/url" "os" "path/filepath" "runtime" + "strings" "sync" "sync/atomic" "testing" "time" + + "google.golang.org/api/option" ) type delayedSnapshotStore struct{} @@ -49,9 +56,54 @@ func (delayedSnapshotStore) PutObject(ctx context.Context, _, _ string, r io.Rea return err } +type delayedGCSTransport struct{} + +func (delayedGCSTransport) RoundTrip(r *http.Request) (*http.Response, error) { + reply := func(body string) *http.Response { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + Request: r, + } + } + switch { + case r.URL.Query().Get("uploadType") == "resumable": + var metadata struct { + Name string `json:"name"` + } + if err := json.NewDecoder(r.Body).Decode(&metadata); err != nil { + return nil, err + } + response := reply("") + response.Header.Set("Location", "https://upload.test/session/"+url.PathEscape(metadata.Name)) + return response, nil + case strings.HasPrefix(r.URL.Path, "/session/"), r.URL.Query().Get("uploadType") == "multipart": + name := strings.TrimPrefix(r.URL.Path, "/session/") + if r.URL.Query().Get("uploadType") == "multipart" { + name = r.URL.Query().Get("name") + } + if _, err := io.Copy(io.Discard, r.Body); err != nil { + return nil, err + } + select { + case <-time.After(20 * time.Millisecond): + case <-r.Context().Done(): + return nil, r.Context().Err() + } + return reply(fmt.Sprintf(`{"name":%q,"bucket":"bucket"}`, name)), nil + case strings.HasSuffix(r.URL.Path, "/compose"): + return reply(`{"name":"composed","bucket":"bucket"}`), nil + case r.Method == http.MethodDelete: + return reply(""), nil + default: + return nil, fmt.Errorf("unexpected GCS request: %s %s", r.Method, r.URL) + } +} + // BenchmarkConcurrentSnapshotUploads compares the shared compression budget with -// admitting all 80 uploads at once. It samples peak heap, not cumulative allocation. -// SDK buffering is covered separately by the slow GCS and S3 tests. +// admitting all 80 uploads at once, and measures the full GCS single-upload path +// under the default 512 MiB budget. It samples peak heap, not cumulative allocation. func BenchmarkConcurrentSnapshotUploads(b *testing.B) { const files = 80 path := filepath.Join(b.TempDir(), "metadata") @@ -61,15 +113,32 @@ func BenchmarkConcurrentSnapshotUploads(b *testing.B) { logger := slog.Default() slog.SetDefault(slog.New(slog.NewTextHandler(io.Discard, nil))) defer slog.SetDefault(logger) + + gcsStore, err := NewGCSClient(b.Context(), + option.WithoutAuthentication(), + option.WithHTTPClient(&http.Client{Transport: delayedGCSTransport{}})) + if err != nil { + b.Fatal(err) + } + g := gcsStore.(*gcsClient) + defer func() { + g.client.Close() + for _, c := range g.pool { + c.Close() + } + }() + for _, tc := range []struct { name string + client ObjectStorage capacity int64 }{ - {"shared_512MiB", DefaultUploadMemory}, - {"all_80_admitted", files * parallelWorkerMemory}, + {"shared_512MiB", delayedSnapshotStore{}, DefaultUploadMemory}, + {"all_80_admitted", delayedSnapshotStore{}, files * parallelWorkerMemory}, + {"gcs_shared_512MiB", gcsStore, DefaultUploadMemory}, } { b.Run(tc.name, func(b *testing.B) { - u, err := NewUploader(delayedSnapshotStore{}, tc.capacity) + u, err := NewUploader(tc.client, tc.capacity) if err != nil { b.Fatal(err) } @@ -122,3 +191,114 @@ func BenchmarkConcurrentSnapshotUploads(b *testing.B) { }) } } + +// BenchmarkLargeSparseUpload verifies that large sparse memory images (both +// multi-chunk streaming and multi-part GCS ranged uploads) keep their full +// chunk sizes and throughput, and measures a full 7-file gVisor snapshot +// (6 small files + 1 large pages.img) under the 512 MiB default budget. +func BenchmarkLargeSparseUpload(b *testing.B) { + dir := b.TempDir() + largePath := filepath.Join(dir, "pages.img") + file, err := os.Create(largePath) + if err != nil { + b.Fatal(err) + } + block := bytes.Repeat([]byte{0x5a}, 1<<20) + const populatedMiB = 256 + for range populatedMiB { + if _, err := file.Write(block); err != nil { + b.Fatal(err) + } + } + if err := file.Truncate(2 << 30); err != nil { + b.Fatal(err) + } + if err := file.Close(); err != nil { + b.Fatal(err) + } + + var smallPaths []string + for i := range 6 { + p := filepath.Join(dir, fmt.Sprintf("small-%d.img", i)) + if err := os.WriteFile(p, bytes.Repeat([]byte{byte(i + 1)}, 64<<10), 0o600); err != nil { + b.Fatal(err) + } + smallPaths = append(smallPaths, p) + } + + logger := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(io.Discard, nil))) + defer slog.SetDefault(logger) + + gcsStore, err := NewGCSClient(b.Context(), + option.WithoutAuthentication(), + option.WithHTTPClient(&http.Client{Transport: delayedGCSTransport{}})) + if err != nil { + b.Fatal(err) + } + g := gcsStore.(*gcsClient) + defer func() { + g.client.Close() + for _, c := range g.pool { + c.Close() + } + }() + + b.Run("large_256MiB_only", func(b *testing.B) { + u, err := NewUploader(gcsStore, DefaultUploadMemory) + if err != nil { + b.Fatal(err) + } + b.SetBytes(populatedMiB << 20) + b.ResetTimer() + for range b.N { + f, err := os.Open(largePath) + if err != nil { + b.Fatal(err) + } + err = u.SendFile(b.Context(), "gs://bucket/pages.img", f) + f.Close() + if err != nil { + b.Fatal(err) + } + } + }) + + b.Run("4_actors_gvisor_7files_each", func(b *testing.B) { + u, err := NewUploader(gcsStore, DefaultUploadMemory) + if err != nil { + b.Fatal(err) + } + b.ResetTimer() + for range b.N { + var wg sync.WaitGroup + for actor := range 4 { + for i, p := range smallPaths { + wg.Go(func() { + f, err := os.Open(p) + if err != nil { + b.Error(err) + return + } + defer f.Close() + if err := u.SendFile(b.Context(), fmt.Sprintf("gs://bucket/a%d-small-%d.img", actor, i), f); err != nil { + b.Error(err) + } + }) + } + wg.Go(func() { + f, err := os.Open(largePath) + if err != nil { + b.Error(err) + return + } + defer f.Close() + if err := u.SendFile(b.Context(), fmt.Sprintf("gs://bucket/a%d-pages.img", actor), f); err != nil { + b.Error(err) + } + }) + } + wg.Wait() + } + }) +} diff --git a/pkg/objectstorage/upload_gcs_test.go b/pkg/objectstorage/upload_gcs_test.go index ea865f978b..830706c16d 100644 --- a/pkg/objectstorage/upload_gcs_test.go +++ b/pkg/objectstorage/upload_gcs_test.go @@ -29,6 +29,7 @@ import ( "testing" "testing/synctest" + "google.golang.org/api/googleapi" "google.golang.org/api/option" ) @@ -176,3 +177,180 @@ func TestGCSSparseUploadHonorsReservedParts(t *testing.T) { } }) } + +func TestUploaderGCSSnapshotBurstFitsDefaultBudget(t *testing.T) { + dir := t.TempDir() + var smallFiles []*os.File + for i := range 6 { + path := filepath.Join(dir, fmt.Sprintf("small-%d.img", i)) + writeSparseSource(t, path, 2<<30, []region{{0, 64 << 10, byte(i + 1)}}) + f, err := os.Open(path) + if err != nil { + t.Fatal(err) + } + defer f.Close() + smallFiles = append(smallFiles, f) + } + pagesPath := filepath.Join(dir, "pages.img") + writeSparseSource(t, pagesPath, 2<<30, []region{{0, 2 * sparsePartTarget, 0x7f}}) + pages, err := os.Open(pagesPath) + if err != nil { + t.Fatal(err) + } + defer pages.Close() + + synctest.Test(t, func(t *testing.T) { + gate := make(chan struct{}) + transport := &slowGCSTransport{parts: make(chan string, 16), release: gate} + client, err := NewGCSClient(t.Context(), option.WithoutAuthentication(), option.WithHTTPClient(&http.Client{Transport: transport})) + if err != nil { + t.Fatal(err) + } + g := client.(*gcsClient) + defer func() { + g.client.Close() + for _, c := range g.pool { + c.Close() + } + }() + u, err := NewUploader(client, DefaultUploadMemory) + if err != nil { + t.Fatal(err) + } + done := make(chan error, len(smallFiles)+1) + for i, f := range smallFiles { + go func() { + done <- u.SendFile(t.Context(), fmt.Sprintf("gs://bucket/small-%d.img", i), f) + }() + } + go func() { + done <- u.SendFile(t.Context(), "gs://bucket/pages.img", pages) + }() + synctest.Wait() + // All 6 small files and both ranged parts of pages.img must fit in + // the 512 MiB default budget simultaneously without queueing. + if got, want := len(transport.parts), len(smallFiles)+2; got != want { + t.Errorf("started %d concurrent GCS uploads, want %d", got, want) + } + close(gate) + synctest.Wait() + for range len(smallFiles) + 1 { + if err := <-done; err != nil { + t.Fatal(err) + } + } + }) +} + +func TestSingleUploadChunkSize(t *testing.T) { + const unit = int64(googleapi.MinUploadChunkSize) + for _, tc := range []struct { + size int64 + want int64 + }{ + {-1, unit}, + {0, unit}, + {1, unit}, + {unit - 1, unit}, + {unit, 2 * unit}, + {unit + 1, 2 * unit}, + {uploadCompositeMin - 1, uploadCompositeMin}, + } { + got := singleUploadChunkSize(tc.size) + if got != tc.want { + t.Errorf("singleUploadChunkSize(%d) = %d, want %d", tc.size, got, tc.want) + } + if got <= max(tc.size, 0) { + t.Errorf("singleUploadChunkSize(%d) = %d, must be strictly greater than size", tc.size, got) + } + if got%unit != 0 { + t.Errorf("singleUploadChunkSize(%d) = %d, must be a multiple of %d", tc.size, got, unit) + } + } + + // An object whose size is an exact multiple of MinUploadChunkSize (256 KiB) + // must still hit EOF inside PutObject's lookahead and take the single-request + // path rather than spilling into composite parts. + closed := make(chan struct{}) + close(closed) + transport := &slowGCSTransport{parts: make(chan string, 4), release: closed} + client, err := NewGCSClient(t.Context(), option.WithoutAuthentication(), option.WithHTTPClient(&http.Client{Transport: transport})) + if err != nil { + t.Fatal(err) + } + g := client.(*gcsClient) + defer func() { + g.client.Close() + for _, c := range g.pool { + c.Close() + } + }() + payload := bytes.Repeat([]byte{0xab}, int(unit)) + body := boundedReader{Reader: bytes.NewReader(payload), max: unit} + if err := client.PutObject(t.Context(), "bucket", "exact-256k", body); err != nil { + t.Fatal(err) + } + if got := len(transport.parts); got != 1 { + t.Errorf("256 KiB object sent %d requests, want 1", got) + } +} + +type wrappedReader struct { + r io.Reader +} + +func (w wrappedReader) Read(p []byte) (int, error) { return w.r.Read(p) } +func (w wrappedReader) Unwrap() io.Reader { return w.r } + +type lookaheadCaptureStore struct { + lookaheads []int64 +} + +func (*lookaheadCaptureStore) supportsStreamingPut() {} +func (*lookaheadCaptureStore) GetObject(context.Context, string, string) (io.ReadCloser, error) { + return nil, fmt.Errorf("unused") +} +func (s *lookaheadCaptureStore) PutObject(_ context.Context, _, _ string, r io.Reader) error { + // Wrap the reader before checking singleUploadLookahead to verify the size + // bound survives decorator unwrapping through the public upload paths. + s.lookaheads = append(s.lookaheads, singleUploadLookahead(wrappedReader{r: r})) + _, err := io.Copy(io.Discard, r) + return err +} + +func TestSingleUploadLookaheadThroughPublicUploaders(t *testing.T) { + dir := t.TempDir() + sparsePath := filepath.Join(dir, "small.img") + writeSparseSource(t, sparsePath, 2<<30, []region{{0, 64 << 10, 0x33}}) + sparseFile, err := os.Open(sparsePath) + if err != nil { + t.Fatal(err) + } + defer sparseFile.Close() + + manifestPath := filepath.Join(dir, "manifest.json") + if err := os.WriteFile(manifestPath, []byte(`{"version":1}`), 0o600); err != nil { + t.Fatal(err) + } + manifestFile, err := os.Open(manifestPath) + if err != nil { + t.Fatal(err) + } + defer manifestFile.Close() + + store := &lookaheadCaptureStore{} + u, err := NewUploader(store, DefaultUploadMemory) + if err != nil { + t.Fatal(err) + } + if err := u.SendFile(t.Context(), "gs://bucket/small.img", sparseFile); err != nil { + t.Fatal(err) + } + if err := u.Send(t.Context(), "gs://bucket/manifest.json", manifestFile); err != nil { + t.Fatal(err) + } + const want = int64(googleapi.MinUploadChunkSize) + if len(store.lookaheads) != 2 || store.lookaheads[0] != want || store.lookaheads[1] != want { + t.Errorf("captured lookaheads = %v, want [%d %d]", store.lookaheads, want, want) + } +} diff --git a/pkg/objectstorage/upload_test.go b/pkg/objectstorage/upload_test.go index 199dc8b366..84d8356b03 100644 --- a/pkg/objectstorage/upload_test.go +++ b/pkg/objectstorage/upload_test.go @@ -60,7 +60,7 @@ func TestUploaderBudgetFitsBothPaths(t *testing.T) { continue } // An all-holes file streams a few bytes. - if plan.memory < streamBufferMemory(client, sparseEncodedBound(0, 0))+int64(plan.workers)*parallelWorkerMemory { + if plan.memory < streamBufferMemory(client, sparseEncodedBound(0, 0))+int64(plan.workers)*parZstdWorkerMemory(parZstdMinChunk) { t.Fatalf("%T: streaming pipeline under-reserved: %+v", client, plan) } } @@ -75,9 +75,9 @@ func TestUploaderPlansFromPopulatedBytes(t *testing.T) { populatedMiB int gcs, s3 uploadPlan }{ - {"all_holes", 0, uploadPlan{1, 1, 192 << 20}, uploadPlan{1, 1, 224 << 20}}, - {"small", 1, uploadPlan{1, 1, 192 << 20}, uploadPlan{1, 1, 224 << 20}}, - {"two_chunks", 9, uploadPlan{2, 1, 256 << 20}, uploadPlan{2, 1, 288 << 20}}, + {"all_holes", 0, uploadPlan{1, 1, 8<<20 + 512<<10}, uploadPlan{1, 1, 168 << 20}}, + {"small", 1, uploadPlan{1, 1, 8<<20 + 2560<<10}, uploadPlan{1, 1, 168 << 20}}, + {"two_chunks", 9, uploadPlan{2, 1, 146<<20 + 512<<10}, uploadPlan{2, 1, 288 << 20}}, {"multiple_ranges", 192, uploadPlan{3, 3, 144 << 20}, uploadPlan{5, 3, 480 << 20}}, } { t.Run(tc.name, func(t *testing.T) { @@ -140,13 +140,13 @@ func TestUploaderReservesOnlyTheChosenPath(t *testing.T) { if got.parts != 4 || got.memory != 4*sparseUploadMemory { t.Errorf("capacity %d: ranged plan = %+v, want 4 parts reserving %d", capacity, got, 4*sparseUploadMemory) } - // 1 MiB populated: one request, so PutObject's lookahead, the client's - // chunk buffer, and one worker. + // 1 MiB populated: one request, sized to the bounded stream and 1 MiB + // compression chunks. got, err = u.plan(small) if err != nil { t.Fatal(err) } - if want := uploadCompositeMin + uploadChunkSize + parallelWorkerMemory; got.memory != want { + if want := streamBufferMemory(&gcsClient{}, sparseEncodedBound(1<<20, 1)) + parZstdWorkerMemory(parZstdMinChunk); got.memory != want { t.Errorf("capacity %d: small plan = %+v, want %d", capacity, got, want) } } @@ -162,7 +162,7 @@ func TestUploaderSmallSparseFilesShareBudget(t *testing.T) { synctest.Test(t, func(t *testing.T) { gate := make(chan struct{}) store := &gatedUploadStore{started: make(chan string, len(files)), release: gate} - u, err := NewUploader(store, 2*parallelWorkerMemory) + u, err := NewUploader(store, parallelWorkerMemory) if err != nil { t.Fatal(err) }