objectstorage: size small snapshot upload buffers to actual data

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.

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}

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)

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

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