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