[andygrove]

Which issue does this PR close?

Closes #7305.

Rationale for this change

An executor JVM with Native Memory Tracking enabled aborted with NMT corruption: ... footer canary broken (buffer overflow?) while materializing Comet's in-memory cache with the default zstd codec. The same job passes with lz4.

The cause is in Arrow 18.3.0's ZstdCompressionCodec.doCompress. It allocates compressBound(n) + 8 bytes, points zstd 8 bytes in (past the uncompressed-length prefix), but still tells zstd the capacity is compressBound(n) + 8. zstd guarantees that it never writes past the capacity it is given, not that it never writes past compressBound: while it entropy-codes the literals of a block that turns out not to compress, it fills the destination up to that capacity before falling back to a stored block. The attempt can therefore write up to 8 bytes past the buffer.

Two things make this rare. The allocator rounds requests up to a power of two, so the overrun lands in slack unless compressBound(n) + 8 is within 8 bytes below a power of two (n = 432..440 for 512, and so on; the block in #7305 was 437 bytes). And without NMT the extra bytes land in the allocator's padding or in whatever sits next, with no error at all; with NMT the footer canary sits right after the buffer and the JVM aborts when the discarded output is freed, which is why the crashed block had a zstd frame but no length prefix.

A standalone program that feeds Arrow's codec 437-byte inputs reproduces the exact abort in about a second under -XX:NativeMemoryTracking=detail; with the corrected capacity it runs clean. Arrow fixed this in apache/arrow-java#1119 (reported as apache/arrow-java#1116) for Arrow 20.0.0, which has not been released.

What changes are included in this PR?

• CometZstdCompressionCodec, a subclass of Arrow's codec whose doCompress passes compressBound(n) as the capacity, as the upstream fix does. Decompression is inherited unchanged.
• CachedBatchIpc.compressionCodec builds that subclass for zstd. It is the only place Comet compresses with Arrow's zstd codec.

The subclass can be removed once Comet moves to an Arrow release that contains the upstream fix.

How are these changes tested?

New CometZstdCompressionCodecSuite. It compresses through a BufferAllocator that hands the codec exactly the capacity it requests and keeps a sentinel-filled guard region behind every buffer, so an overrun shows up as a changed sentinel rather than as silent corruption. The inputs are the ones that provoke the bug: nearly incompressible bytes with scattered repeated values, in every length up to 8192 whose output buffer lands within 8 bytes below a power of two. The suite shows Arrow's codec overruns on those inputs, that Comet's codec does not, that both the stored-raw and the compressed outcomes round-trip, and that the guard itself detects a one-byte overrun.

Also verified with the standalone reproducer under -XX:NativeMemoryTracking=detail: Arrow's codec aborts the JVM, the fixed codec does not.

[comphead]

Thanks, the diagnosis checks out. I ran the codec's capacity math in a standalone Java program against zstd-jni 1.5.2-5, 1.5.5-4, 1.5.6-9, 1.5.7-9 and 1.5.7-23, using this suite's inputs and a guard right after the compressBound(n) + 8 buffer. On every version, 298 of the 1600 inputs wrote past the buffer with Arrow's capacity and none did with compressBound(n). So the "Arrow's zstd codec writes past" test should hold on every Spark profile. This was a local probe, not a run of the suite.

Two points outside the diff:

• CometInMemoryCacheSuite still constructs Arrow's new ZstdCompressionCodec(1) directly in the empty-batch test, in FailAfterCompressionCodec and in two other write-path tests. Those now exercise a codec the write path no longer uses. new CometZstdCompressionCodec(1) would keep them on the shipped one.
• Spark master has the same exposure. ArrowCachedBatchSerializer and ArrowCompressionUtils construct Arrow's ZstdCompressionCodec on Arrow 19.0.0, which does not have the upstream fix. It may be worth a SPARK JIRA.

[comphead]

The overrun is explained in full in the CometZstdCompressionCodec scaladoc, again here, and a third time on exposedLengths. A one-line pointer to the codec's scaladoc would keep a single copy.

[andygrove]

Done, the suite header now just points to the CometZstdCompressionCodec scaladoc, and exposedLengths is gone.

[comphead]

GuardedAllocator cuts every slice to the exact requested size, so the guard catches the overrun at any length, not only where power-of-two rounding leaves less than 8 bytes of slack. In the same local probe, Arrow's capacity overran on 19177 of 49550 inputs at the other lengths below 1024. A plain range such as 400 to 600 would drive both tests and drop the rounding math and its doc comment.

[andygrove]

Good point. Replaced it with 400 to 600 and dropped the rounding math and its doc comment. The Arrow test still sees overruns with the guard.

[comphead]

packed is counted but never read. Adding a few compressible inputs to this loop and asserting packed should be > 0 would cover the compressed outcome under the guard and let the next test go. Otherwise it can be dropped.

[andygrove]

Added compressible inputs to that loop and it now asserts packed > 0 alongside raw > 0.

[comphead]

This mostly repeats existing coverage. Empty input never reaches doCompress, because AbstractCompressionCodec.compress returns early when writerIndex() is 0. The compressed round trip through the config is covered by "Comet in-memory cache round-trips under every compression codec" in CometInMemoryCacheSuite, which now goes through this codec. The only new check is the guard on a compressed outcome, which fits in the previous test.

[andygrove]

Agreed, removed. The compressed outcome under the guard is covered by the previous test now, and the config round trip is in CometInMemoryCacheSuite.

[comphead]

The first test already shows that the guard detects an overrun, since it asserts overruns > 0 through the same GuardedAllocator. Once Arrow is bumped, that test, the codec and this suite go together, so this test never adds coverage beyond it.

[andygrove]

Agreed, removed this one too. The first test already shows the guard detects an overrun.

[andygrove]

Thanks for the review. I pushed 2c30c57 with the suite changes above. CometInMemoryCacheSuite now gets its zstd codec from CometCachedBatchHelper.compressionCodec("zstd"), so all four places that built Arrow's codec directly go through CometZstdCompressionCodec. I also rewrapped a comment in CachedBatchIpc that was one character over the scalastyle limit.

On Spark having the same exposure: agreed, I read the 19.0.0 bytecode and ArrowCompressionUtils and ArrowCachedBatchSerializer both build ZstdCompressionCodec directly (ArrowCompressionUtils is also on branch-4.1 with Arrow 18.3.0). I'll file the SPARK JIRA soon and will comment here once it's filed.