zstd (zstandard) decompression stream heap allocates

[andrewrk]

Decompression should use O(1) memory.

#18923 removed the requirement of an allocator from decompressing deflate and gzip streams. This is great, but zstd is left behind, leaving this sad situation:

 switch (res.request.compression) {
 .none => {},
 .deflate => {},
 .gzip => {},
 .zstd => |*zstd| zstd.deinit(),
 }

and

 .deflate => res.request.compression = .{
 .deflate = std.compress.zlib.decompressStream(res.transferReader()),
 },
 .gzip, .@"x-gzip" => res.request.compression = .{
 .gzip = std.compress.gzip.decompress(res.transferReader()),
 },
 .zstd => res.request.compression = .{
 .zstd = std.compress.zstd.decompressStream(res.allocator, res.transferReader()),
 },

zstd is holding us back.

Here's a diff to get started on this issue:

diff --git a/lib/std/compress/zstandard.zig b/lib/std/compress/zstandard.zig
index 4d9421acac..a3dcd37c2a 100644
--- a/lib/std/compress/zstandard.zig
+++ b/lib/std/compress/zstandard.zig
@@ -20,7 +20,6 @@ pub fn DecompressStream(
 return struct {
 const Self = @This();
 
- allocator: Allocator,
 source: std.io.CountingReader(ReaderType),
 state: enum { NewFrame, InFrame, LastBlock },
 decode_state: decompress.block.DecodeState,
@@ -44,9 +43,8 @@ pub fn DecompressStream(
 
 pub const Reader = std.io.Reader(*Self, Error, read);
 
- pub fn init(allocator: Allocator, source: ReaderType) Self {
+ pub fn init(source: ReaderType) Self {
 return Self{
- .allocator = allocator,
 .source = std.io.countingReader(source),
 .state = .NewFrame,
 .decode_state = undefined,
@@ -126,16 +124,6 @@ pub fn DecompressStream(
 }
 }
 
- pub fn deinit(self: *Self) void {
- if (self.state == .NewFrame) return;
- self.allocator.free(self.decode_state.literal_fse_buffer);
- self.allocator.free(self.decode_state.match_fse_buffer);
- self.allocator.free(self.decode_state.offset_fse_buffer);
- self.allocator.free(self.literals_buffer);
- self.allocator.free(self.sequence_buffer);
- self.buffer.deinit(self.allocator);
- }
-
 pub fn reader(self: *Self) Reader {
 return .{ .context = self };
 }
@@ -235,18 +223,16 @@ pub fn DecompressStream(
 }
 
 pub fn decompressStreamOptions(
- allocator: Allocator,
 reader: anytype,
 comptime options: DecompressStreamOptions,
 ) DecompressStream(@TypeOf(reader, options)) {
- return DecompressStream(@TypeOf(reader), options).init(allocator, reader);
+ return DecompressStream(@TypeOf(reader), options).init(reader);
 }
 
 pub fn decompressStream(
- allocator: Allocator,
 reader: anytype,
 ) DecompressStream(@TypeOf(reader), .{}) {
- return DecompressStream(@TypeOf(reader), .{}).init(allocator, reader);
+ return DecompressStream(@TypeOf(reader), .{}).init(reader);
 }
 
 fn testDecompress(data: []const u8) ![]u8 {

[dweiller]

This could be a straightforward change to make in exchange for several MiB of stack space as the Zstandard decompressor already uses O(1) memory. The required space for internal buffers will be a couple hundred kilobytes plus the maximum supported window size, which currently defaults to 8 MiB (at least 8 MiB is recommended by the spec for improved interoperability).

[truemedian]

Given std.http was used as the example where zstd causes a problem: I must point out that the decompressors are stored by value in Response, which is stored by value in Request, which is returned by value from Client.open, so it is highly likely that whatever memory is placed into a decompressor struct will end up on the stack (at least in the current architecture).

If zstd needs 8 MiB of space, putting it on the stack is untenable, a single copy will use all 16 MiB of stack space allocated on linux; however I see no reason why the decompressors (or zstd specifically) can't just ask for a buffer of some defined constant length from the user.

If this is untenable, then std.http will need to change to allocate the Client.Request and Server.Response structures (both of the types returned from the api, that both eventually hold the decompressor).

[dweiller]

This could be a straightforward change to make in exchange for several MiB of stack space...

And this is too much stack space to use. The best fix as far as Zstandard is concerned would be to change the decompressStream API to accept a buffer to use for the window.

[dweiller]

however I see no reason why the decompressors (or zstd specifically) can't just ask for a buffer of some defined constant length from the user.

Yes, this I think this is the right approach to take, but std.http will still need to be changed to allocate the window buffer and pass it to the decompressor. I would have initially thought to do this with the decompressor init function decompressStream, but a quick look through std.http.Client makes it look like this won't be suitable.

[andrewrk]

a quick look through std.http.Client makes it look like this won't be suitable.

Give me an hour to wrap up some local changes and then it might be worth re-evaluating this :-)

[dweiller]

I'm not actually familiar with std.http and "suitable" may have been the wrong word. The point is that the wait() function would need to allocate the window buffer requiring the addition of error.OutOfMemory to wait()'s error set. However, it is simple to remove allocation from the zstandard decompressor API as long as a mechanism to pass a buffer to it is added.

[andrewrk]

Are there any other valid values for window size? Do I understand correctly that any buffer size can be used for the window, but it will cause a failure to decompress if the particular stream required a bigger window size than was allocated to the decompression stream?

[truemedian]

From the zstd spec:

The minimum Window_Size is 1 KB. The maximum Window_Size is (1<<41) + 7*(1<<38) bytes, which is 3.75 TB.

In general, larger Window_Size tend to improve compression ratio, but at the cost of memory usage.

To properly decode compressed data, a decoder will need to allocate a buffer of at least Window_Size bytes.

...

For improved interoperability, it's recommended for decoders to support Window_Size of up to 8 MB, and it's recommended for encoders to not generate frame requiring Window_Size larger than 8 MB. It's merely a recommendation though, decoders are free to support larger or lower limits, depending on local limitations.

It seems 8 MiB is the recommended size if we want it to almost always work.

[dweiller]

... but it will cause a failure to decompress if the particular stream required a bigger window size than was allocated to the decompression stream?

That is correct. My impression is that If a compressed stream used a larger window than the decompressor supports it is not possible to do the decompression. It looks like the reference implementation uses a limit of 8MiB in the zstd program by default, but depending on CLI parameters can make it higher.

[andrewrk]

Give me an hour to wrap up some local changes and then it might be worth re-evaluating this :-)

Following up: I'm super close to having this work done but I gotta go, so this is coming tomorrow rather than this evening.