gh-155460: Pre-allocate the output buffer in compression.zstd decompression - #155464
Open
ThomasWaldmann wants to merge 1 commit into
Open
gh-155460: Pre-allocate the output buffer in compression.zstd decompression#155464ThomasWaldmann wants to merge 1 commit into
ThomasWaldmann wants to merge 1 commit into
Conversation
…ecompression Frames produced by the one-shot compression APIs record the decompressed size in the frame header. When the input starts at such a frame header, read that size with ZSTD_getFrameContentSize() and allocate the output buffer at its exact size, instead of growing it progressively and shrinking it on finish. An exactly filled buffer is then returned without a copy. This is only a sizing hint: decompression does not rely on it, so a hand-crafted header recording a wrong size produces exactly the same results and exceptions as before. Two guards bound the allocation: the recorded size is trusted only up to 1 GiB, and recorded sizes claiming more than a 32768x expansion of the available input -- more than the format can produce -- are ignored. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
compression.zstddecompression always builds its output through the growingoutput buffer (initial 32 KiB, then progressively larger growth steps, finished
with a resize to the actual size), even though most zstd frames record the exact
decompressed size in the frame header: the one-shot compression APIs
(
zstd.compress(),ZSTD_compress2()) write it by default.When the input starts at such a frame header,
decompress_lock_held()now readsthe recorded size with
ZSTD_getFrameContentSize()and allocates the outputbuffer at its exact size via
_OutputBuffer_InitWithSize(). An exactly filledbuffer is returned without any resize or copy. Everything else falls back to the
existing
_OutputBuffer_InitAndGrow()path: frames without a recorded size,continuation calls in the middle of a frame (
ZSTD_CONTENTSIZE_ERROR), andskippable frames (recorded size 0).
This is only a sizing hint — decompression does not rely on it:
and exceptions as before (libzstd validates the recorded size itself; the
buffer grows further if the hint was too small, or is shrunk to the actual
size on finish);
OUTPUT_PREALLOC_MAX), so acrafted header cannot request an arbitrarily large allocation;
input, so recorded sizes claiming more than a 32768x expansion
(
OUTPUT_MAX_EXPANSION) of the available input are unfulfillable and areignored;
_OutputBuffer_InitWithSize()clamps tomax_length, so themax_lengthargument keeps working unchanged.
This covers both
ZstdDecompressor.decompress()and the module-levelzstd.decompress().Tests
Six tests are added to
DecompressorTestCase, including a helper that rewritesthe
Frame_Content_Sizefield of a frame per RFC 8878 §3.1.1.1 so that wrongand absurd recorded sizes can be tested: sizes that are too big, too small or
zero,
0xFFFFFFFFdeclared on a 130 KiB frame (hits the 1 GiB guard) and on atiny RLE frame (hits the expansion-ratio guard), round-trips at block
boundaries, the unknown-size fallback path,
max_length, and split input.They all pass on unpatched
mainas well — deliberately, since they pinbehaviour rather than the implementation.
Benchmarks
zstd.decompress()of level-1 compressed frames of moderately compressibledata; 10 timed runs after an untimed warmup, median run. Full tables, per-size
results and the profiling analysis are in gh-155460.
A mixed-size workload, 1000 chunks with sizes drawn from a normal distribution
in log2(size) centred on 2 MiB spanning 512 KiB to 8 MiB (2269 MiB total):
Single-size frames, same two machines: on Linux the patch is faster at every
size measured (+5–8.6% up to 5 MiB, +41–55% for 6–16 MiB, +0.7–2.7% at
32–256 MiB); on macOS +10% at 1 MiB, +18–19% at 6–8 MiB, +4.6–12% at 9–16 MiB,
and a 2–8.5% regression in the 3–5 MiB band, where the unpatched path needs no
growth step and its cache-resident internal window beats direct writes spanning
the whole output.
Prior art:
python-zstandard'sdecompress()allocates the output buffer fromthe frame's recorded content size, and pyzstd (from which this module descends)
did the same before the stdlib port.
Disclosure: this change was prepared with AI assistance (Claude Code); the
benchmarks and the analysis above were run and reviewed by me.