objectstorage: decode snapshot downloads frame-parallel by cheftako · Pull Request #2502 · agent-substrate/substrate

Walter Fender (cheftako) · 2026-10-12T15:44:59Z

A snapshot download decoded its zstd stream with one decoder on one core. On the counter-microvm demo that was ~95ms to decode a 71MB memory image, about the whole download phase of a resume. The upload side already compresses on every core.

parZstd writes independent frames of at most 8MiB that declare their size. parZstdReader splits the compressed stream at frame boundaries, decodes the frames concurrently and returns the output in order. A frame is decoded in parallel only when its decoded size, block count and window all fit in 16MiB. From the first frame that does not, the rest of the stream decodes serially, so any valid zstd stream still decodes.

Only downloads that GetObject splits into ranged requests (rangedReader) use parZstdReader. A single-request download arrives as one stream, where the serial decoder overlaps decoding with the download and is faster.

Readers share 4 decode tokens, and each has one of its own, so concurrent restores split the cores and a reader always makes progress. A reader holds at most workers+1 frames, and frame and output buffers come from the parZstd pools.

BenchmarkParZstdReader, 128MiB of compressible data to io.Discard (Intel Core Ultra 7 155H):

	ns/op	B/op
serial	89ms	10MB
parallel, 1 CPU	99ms	0.7MB
parallel, 2 CPUs	52ms	1MB
parallel, 8 CPUs	30ms	2.6MB

Decoding 512MiB into a file through decodeContent takes ~280ms at 2 CPUs and ~185ms at 4 or 8, with a peak heap of about 20MiB at 2 CPUs.

Written with Claude Code, including the changes for review feedback.

🤖 Generated with Claude Code

google-cla · 2026-10-12T15:45:32Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Benjamin Elder (BenTheElder) · 2026-10-14T00:07:56Z

cc ILLUM1N0X Anish Ramasekar (@aramase)

Anish Ramasekar (aramase) · 2026-10-14T04:13:58Z

+			return nil, 0, 0, false, fmt.Errorf("reading zstd block: %w", err)
+		}
+		maxSize += zstdMaxBlockOut
+		if size < 0 && maxSize > parZstdMaxFrame && !last {


With && !last, an unsized frame whose last block pushes maxSize past 16MiB still goes to a worker, and the worker's 16MiB MaxMemory then rejects it ("decompressed size exceeds configured limit"). A 129x128KiB raw-block frame decodes on main and fails here. Could this fall back whenever maxSize > parZstdMaxFrame, last block or not?

Fixed. Any frame whose block count times 128KiB goes over parZstdMaxFrame now falls back, last block or not. The block header is checked before its contents are read. TestParZstdReaderFallsBackForLargeFrames has your 129-block case.

Anish Ramasekar (aramase) · 2026-10-14T04:13:58Z

+
+func (r *parZstdReader) decodeWorker(jobs <-chan parZstdDecodeJob) {
+	defer r.wg.Done()
+	dec, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1),


WithDecoderMaxMemory(16MiB) also clamps the decoder's max window to 16MiB, and readZstdFrame never looks at the window descriptor. So zstd --long=27 -c < small_file (unsized, 128MiB window) now fails with "window size exceeded" where main decoded it. Could the splitter parse the window size and send those frames to the serial path?

Fixed. readZstdFrame parses the window descriptor and sends frames with a window over parZstdMaxFrame to the serial path. A single-segment frame's window is its declared size, which is already checked. I also decoded zstd --long=27 -c < small_file output from the CLI, and it matches.

Anish Ramasekar (aramase) · 2026-10-14T04:13:58Z

+		case 3:
+			return nil, 0, 0, false, errors.New("reserved zstd block type")
+		}
+		if n > parZstdMaxFrame {


n is a 21-bit field, so this check never fires, and a sized frame has no cap on block count. A frame declaring 100 bytes followed by 256MiB of raw blocks gets buffered whole (~960MiB heap here) before the worker rejects it. Checking n > zstdMaxBlockOut and capping a sized frame's buffered bytes near parZstdMaxFrame would keep the bound in the comment above true.

Fixed. A block over 128KiB is now an error, as in klauspost's decoder. The block-count bound applies to sized frames too, so readZstdFrame buffers at most about parZstdMaxFrame before it falls back. Your 100-byte frame with 256MiB of blocks is a row in TestParZstdReaderFallsBackForLargeFrames.

Anish Ramasekar (aramase) · 2026-10-14T04:13:58Z

+	parZstdMaxFrame = 2 * parZstdChunk
+	// parZstdMaxDecoders caps the decode workers, and with them the memory held in
+	// decoded frames (about 3*parZstdMaxDecoders*parZstdMaxFrame at most).
+	parZstdMaxDecoders = 8


Decoding 512MiB into a file on an 8-core Linux box, this peaks at ~325-360MiB heap (~180MiB live, ~380MiB RSS) per reader at 8 procs vs 8MiB on main, and it takes ~360ms at both 2 and 8 procs (main ~520ms), so the extra workers cost memory without saving time there. FetchSnapshot decodes every file at once and the plugin requests 128Mi with no limit, so would a package-level budget shared across readers work here, letting concurrent restores split the cores instead of each holding up to 8 frames? Pooling the frame/output buffers would also cut the ~770MiB allocated per 512MiB decoded.

Done, roughly as you suggested:

All readers share 4 decode tokens (parZstdDecodeTokens), and each reader has one of its own so it always makes progress. Concurrent downloads hold at most 4 frames plus one per reader.

A reader holds at most workers+1 frames.

Frame and output buffers come from the parZstd pools.

I matched what you saw: more than 4 workers didn't decode a file any faster. Decoding 512MiB into a file now takes ~280ms at 2 procs and ~185ms at 4 or 8. Peak heap is ~20MiB at 2 procs. Allocation dropped from ~200MB to ~3MB per 128MiB decoded.

Anish Ramasekar (aramase) · 2026-10-14T04:13:58Z

+	// Sized frames rather than one streaming frame per part, so the download can
+	// decode them in parallel (parZstdReader). One worker: the parts already run
+	// one per core.
+	zw := newParZstd(dst, 1)


newParZstd(dst, 1) still preallocates two 8MiB chunks per part. I see ~70MiB peak heap per 64MiB part vs ~12MiB before (and ~80ms vs ~63ms), times up to 12 parts. For the single-worker case, could this encode each chunk synchronously into a reused buffer?

main already has this since #2835: with one worker, parZstd encodes each chunk inline into pooled buffers. This PR now only changes the comment here.

Anish Ramasekar (aramase) · 2026-10-14T04:13:58Z

+	return io.ReadAll(r)
+}
+
+func TestParZstdReaderRoundTrip(t *testing.T) {


Could you add a small Benchmark here comparing this to a serial zstd.NewReader on the same stream (SetBytes + ReportAllocs)? That would back the speedup in the description and show the allocation cost.

Added BenchmarkParZstdReader. On 128MiB to io.Discard: serial 89ms and 10MB/op, parallel 99ms at 1 proc, 52ms at 2, ~30ms at 8, with 1-3MB/op.

Anish Ramasekar (aramase) · 2026-10-14T04:13:58Z

+	"github.com/klauspost/compress/zstd"
+)
+
+// zstd decodes one frame on one core however many it is given, which capped a


nit: the measurement story (95ms, 10ms fetch) would fit better in the PR description. The comment can just say what parZstdReader does and when it falls back.

Done. The numbers are in the PR description now.

…oad buffers (agent-substrate#2835) Part of agent-substrate#2553. Related to agent-substrate#2502 and agent-substrate#2708. Large GCS snapshots (`PutSparseFile` → `writeSparsePart`) previously compressed each 64 MiB part as one unsized streaming zstd frame with freshly allocated buffers. `agent-substrate#2502`'s `parZstdReader` only parallelizes frames whose size is declared or bounded at `≤ 16 MiB`; a part written as one streaming frame of 64 MiB or more is neither, so restore goes serial from the first part. Meanwhile, switching parts to the worker-pool `parZstd` with one worker uses ~77 MiB peak per part (2.4× over `agent-substrate#2708`'s 32 MiB `serialWorkerMemory` budget). This PR: - Adds an inline mode to `parZstd` for `workers == 1`: no goroutines; each `≤ 8 MiB` chunk is compressed on the caller's goroutine into a frame that declares its size. `writeSparsePart` uses `newParZstd(dst, 1)`, and small files get it too. - Pools the 8 MiB chunk, 8.5 MiB frame, and 1 MiB sparse copy buffers, and removes unused `newSerialZstd`. - Makes `Write` and a second `Close` after `Close` return `os.ErrClosed`. - Adds `BenchmarkWriteSparsePart`, `BenchmarkDecodeSparse`, and format tests, including one that pins the output byte for byte to per-chunk `EncodeAll`. ### Results (`-count=10`, 512 MiB populated / 2 GiB sparse) | Benchmark | `main` | `agent-substrate#2502` as-is | **this PR** | |---|---|---|---| | **Write 64 MiB part:** time / peak heap (32 MiB budget) | 148 ms / 9.8 MiB | 133 ms / **76.9 MiB** ❌ | 144 ms / **26.4 MiB** ✅ | | **Write 64 MiB part:** frames with content size | 0 / 1 | 8 / 9 | **9 / 9** | | **Read 512 MiB:** time (this PR + `agent-substrate#2502`: 254 ms, 3.8×) | 971 ms | 285 ms (3.4×) | **552 ms (1.76×)** | | **`ConcurrentSnapshotUploads` (`shared_512MiB`):** `B/op` / peak heap¹ | 1,326 / 272 MiB | 1,326 / 286 MiB | **70 / 196 MiB** | - **Serial restore is 1.76× faster even without `agent-substrate#2502`:** sized frames decode straight into a buffer of the declared size (`runtime.memmove` drops from 40% to 9% of decode CPU). - **Stress case (`all_80_admitted`)¹:** batch time roughly halves (`57.6 → ~29 ms`) and peak heap drops about a third (`1,994 → 1,330 MiB`). - **Cluster A/B (`sweperf`, 150u on `c3-highcpu-192-metal`, with `agent-substrate#2502`):** microVM `restore_golden_download` −31% at p50 (`1,095 → 752 ms`) and −63% at p95 (`3,419 → 1,255 ms`); `restore_p50_s` −29% (`945 → 669 ms`). ¹ This PR's upload numbers are from a later session on the same machine than the `main` and `agent-substrate#2502` columns, so treat the deltas as approximate. ### Compatibility & follow-up - Snapshot format is unchanged; old and new snapshots decode byte-exact on both readers. - `agent-substrate#2502`'s `newParZstd(dst, 1)` in `writeSparsePart` lands as-is; only `TestReadZstdFrameBoundsUnsizedSmallFrames` needs to build its unsized frame with `zstd.NewWriter(nil).EncodeAll`. - Follow-up (now that `agent-substrate#2848` is merged): 1-worker uploads could reserve `serialWorkerMemory` instead of `parallelWorkerMemory`. - [x] Tests pass - [x] Appropriate changes to documentation are included in the PR (none needed)

A snapshot download decoded its zstd stream with one decoder on one core. On the counter-microvm demo that was ~95ms to decode a 71MB memory image, essentially the whole download phase of a resume, while the upload side already compresses on every core. The uploader writes independent zstd frames of at most 8MiB (parZstd), so the download can split the compressed stream at frame boundaries, decode the frames concurrently, and hand the output back in order: parZstdReader. A frame is decoded in parallel only when its decoded size is known to be at most 16MiB, from its header or from its block count; from the first frame that is not, the rest of the stream is decoded serially, so any valid zstd stream still decodes. The GCS sparse-part upload compressed each part as one streaming frame with no declared size, which the reader would decode serially. Write each part through parZstd with one worker instead, so its frames are sized too. newSerialZstd has no callers left and is removed.

parZstd now writes single-segment frames, which always declare their content size, so it no longer produces the sub-256-byte frames with no declared size that TestReadZstdFrameBoundsUnsizedSmallFrames covers. Build them with a multi-segment encoder so the block-count bound is still exercised. Rebased with help from Claude.

Send a frame to the serial decoder when its window, its declared size, or its block count times 128KiB is over 16MiB. Such frames used to reach a decode worker and fail its memory limit, and a sized frame with many blocks was buffered whole. A block over 128KiB is now an error. Readers share 4 decode tokens and each has one of its own, so concurrent downloads hold at most 4 frames plus one per reader. A reader holds at most workers+1 frames. Frame and output buffers come from the parZstd pools. Add BenchmarkParZstdReader to compare with a serial decoder.

Nishanth Kotla (Nishanth29) · 2026-10-20T23:00:16Z

LGTM! Tested this locally and the updates look clean.

Benjamin Elder (BenTheElder)

I noticed the upstream library has some parallelism support so I asked my agent to probe that a bit while I read through this, the code looks pretty good but I'm not sure we want to compete downstream on the WithDecoderConcurrency functionality vs patching upstream?

Couple findings attached.

Benjamin Elder (BenTheElder) · 2026-10-20T23:23:48Z

+// free, else a shared one. It returns nil once the reader is closed.
+func (r *parZstdReader) acquire() chan struct{} {
+	select {
+	case r.own <- struct{}{}:


🤖 nit 🟢 – Close can wait for more than the frame the splitter is on. When own is free, this takes it without checking stop, and the select at line 200 picks between ordered and stop at random. So after Close the splitter can still read a few more frames from src. Checking stop first here, and before the send to ordered, would match the Close comment.

Fixed. acquire and the new send check stop first, so after Close the splitter only finishes the frame it is on.

Benjamin Elder (BenTheElder) · 2026-10-20T23:23:48Z

+			// Close discards the output.
+		default:
+			res.data, res.err = dec.DecodeAll(j.frame, getBuf(j.maxSize, parZstdChunk, getChunk))
+			if res.err == nil && j.size >= 0 && len(res.data) != j.size {


🤖 nit 🟢 – I don't think this check can fire. DecodeAll already returns ErrFrameSizeMismatch when the output doesn't match a declared size, so we can drop this and the size field on the job.

https://github.com/klauspost/compress/blob/2602f4afea09fe72f2b58d4ed04d43a6047a0131/zstd/framedec.go#L399-L400

Right, runDecoder checks it. Removed the check, the size field on the job, and the extra return value from readZstdFrame.

Benjamin Elder (BenTheElder) · 2026-10-20T23:23:48Z

+		b.SetBytes(int64(len(data)))
+		b.ReportAllocs()
+		for b.Loop() {
+			zr, err := zstd.NewReader(bytes.NewReader(stream), zstd.WithDecoderConcurrency(1))


🤖 nit 🟢 – I measured zstd's default async stream decoder as well. At 2-8 CPUs it's only 5-7% faster than WithDecoderConcurrency(1), where parZstdReader is 1.9x at 2 CPUs and 3.6-3.8x at 4 and 8 (64MiB and 512MiB of compressibleData, M4 Pro). So it doesn't come close and the extra code seems worth it. Might be worth adding as a row so the benchmark shows that.

We ... might keep it for that speed, but I have some concern that upstream will just outpace us later.

Added a default row. On an Intel Core Ultra 7 155H it matches serial at every CPU count (~90ms per 128MiB), while parallel is ~50ms at 2 CPUs and ~30ms at 4 and 8.

Fair concern. The parallelism upstream has is for its stream decoder, and the default row shows that doesn't help here. Frame-parallel decoding depends on the independent sized frames parZstd writes. If upstream adds a frame-parallel stream decoder later, we can drop parZstdReader and call it from newZstdReader.

We should probably file an issue to ask? Not blocking for merge, but seems like the right thing to do.

klauspost/compress#1274

Benjamin Elder (BenTheElder) · 2026-10-21T00:01:51Z

-	if err != nil {
-		return 0, err
-	}
+	zr := newParZstdReader(src)


🤖 should-fix 🟡 – Objects up to downloadChunkSize (16MiB compressed) skip rangedReader and come back from GetObject as one plain stream. For those, this is slower than the serial decoder whenever the download is slower than decode. A frame can't start decoding until all of it has arrived, so the last frame's decode never overlaps the download. Under contention it gets worse: a reader with only its own token stops reading src until the caller drains the previous frame.

I simulated the download at 100-300MB/s per stream with 30ms to first byte, decoding 32MiB of compressibleData (10.8MiB compressed) at GOMAXPROCS 4 and 8 on an M4 Pro. Alone it's 4-9% slower. With a 256KiB receive buffer it's up to 44% slower at 4-16 concurrent readers, and with a 4MiB buffer 4-18% slower. At 1GB/s it's 14% faster alone. Larger objects, through rangedReader, are 23-38% faster alone and no slower under contention.

Maybe skip parZstdReader for objects at or under downloadChunkSize? GetObject's readers already report ObjectSize.

Reproduced it: on a single throttled stream, parallel was 5-10% slower at 100-300MB/s and ~12% faster at 1GB/s. fetchFromGCSWithZstd now uses parZstdReader only when ObjectSize reports more than downloadChunkSize. Smaller or unsized objects use the serial decoder, as on main. Through rustfs on kind, a 512MiB snapshot decoded in ~390ms against ~630ms serial.

An object of at most downloadChunkSize arrives as one stream. The serial decoder overlaps decoding with the download, but parZstdReader cannot start a frame until all of it has arrived, so it was slower there. fetchFromGCSWithZstd now uses parZstdReader only for larger objects. Close no longer reads further frames from src once called. Drop the decoded size check, since DecodeAll already returns ErrFrameSizeMismatch. Add zstd's default stream decoder to BenchmarkParZstdReader.

Benjamin Elder (BenTheElder)

nothing blocking for now, and given the pending freeze, inclined to merge and iterate.

one high level thought: we should probably update the sparse extant format to include a de-compressed memory size (and maybe just to disambiguate both that and the full size including holes), so we can more reliably detect if it's a small highly compressed object or not and select the parallel decoder instead of having to use the compressed footprint as a proxy without knowing the compression ratio. we'd need to make sure we DO NOT use this for allocating buffers etc (seems like a trivial attack) but we could record it on upload and use it as a hint for download.

that's a format change, but we could do it non-breaking post-GA because the format is already versioned

Benjamin Elder (BenTheElder) · 2026-10-21T10:54:48Z

+	size, ok := ObjectSize(rc)
 	t0 := time.Now()
-	res, err := decodeContent(out, rc)
+	res, err := decodeContent(out, rc, ok && size > downloadChunkSize)


🤖 nit 🟢 – This repeats GetObject's threshold, and it still picks parZstdReader for one stream when an S3 server ignores Range, since s3Body then reports the whole object size (s3ranged.go:59). Checking for a *rangedReader directly would match the comment above in both cases.

Done. fetchFromGCSWithZstd now uses parZstdReader only when GetObject returned a *rangedReader, so an S3 server that ignores Range gets the serial decoder. Through rustfs on kind, a 512MiB object still took the parallel path and a 1MiB one the serial path, and both round-tripped.

Benjamin Elder (BenTheElder) · 2026-10-21T10:54:48Z

+		if err != nil || !ok || frame == nil {
+			t.Fatalf("n=%d: ok=%v err=%v, want a parallel-decodable frame", n, ok, err)
+		}
+		if maxSize < n || maxSize > parZstdMaxFrame {


🤖 nit 🟢 – With size gone, this no longer checks that the frames are unsized, so a sized frame would pass too. Failing on maxSize <= n would keep that.

Done. The test now fails when maxSize <= n. I checked that it fails with a single-segment (sized) encoder.

Benjamin Elder (BenTheElder) · 2026-10-21T11:04:30Z

would also like #2502 (comment) but not blocking

fetchFromGCSWithZstd used parZstdReader for objects over downloadChunkSize. That repeated GetObject's threshold, and an S3 server that ignores Range returns the whole object as one stream with its full size. Use parZstdReader only when GetObject returned a rangedReader. TestReadZstdFrameBoundsUnsizedSmallFrames now fails on a sized frame.

Walter Fender (cheftako) force-pushed the resume-parallel-zstd-decode branch from 757f55e to d9bcce6 Compare October 13, 2026 06:10

Anish Ramasekar (aramase) reviewed Oct 14, 2026

View reviewed changes

Nishanth Kotla (Nishanth29) mentioned this pull request Oct 18, 2026

objectstorage: write snapshot parts as sized zstd frames and pool upload buffers #2835

Merged

2 tasks done

Bowei Du (bowei) added area/storage kind/feature An enhancement / feature request or implementation labels Oct 19, 2026

Walter Fender (cheftako) and others added 2 commits October 20, 2026 13:06

Walter Fender (cheftako) force-pushed the resume-parallel-zstd-decode branch from d9bcce6 to 18ac297 Compare October 20, 2026 20:23

Benjamin Elder (BenTheElder) reviewed Oct 21, 2026

View reviewed changes

Walter Fender (cheftako) mentioned this pull request Oct 21, 2026

objectstorage: cancelled ranged download ends in io.EOF instead of an error #3047

Open

Benjamin Elder (BenTheElder) previously approved these changes Oct 21, 2026

View reviewed changes

Benjamin Elder (BenTheElder) added this pull request to the merge queue Oct 21, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Oct 21, 2026

Benjamin Elder (BenTheElder) mentioned this pull request Oct 21, 2026

e2e: use a port-forward per router CONNECT tunnel #3051

Open

2 tasks done

Walter Fender (cheftako) mentioned this pull request Oct 21, 2026

Add support for parrallel decode on frame boundaries. klauspost/compress#1274

Open

Walter Fender (cheftako) dismissed Benjamin Elder (BenTheElder)’s stale review via f7ea562 October 21, 2026 21:55

Conversation

Walter Fender (cheftako) commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

google-cla Bot commented Oct 12, 2026

Uh oh!

Benjamin Elder (BenTheElder) commented Oct 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Nishanth Kotla (Nishanth29) commented Oct 20, 2026

Uh oh!

Benjamin Elder (BenTheElder) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Benjamin Elder (BenTheElder) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Benjamin Elder (BenTheElder) commented Oct 21, 2026

Uh oh!

Uh oh!

Reviewers

Walter Fender (cheftako) commented Oct 12, 2026 •

edited

Loading