[CORE-16902] compression: allocate and re-use a single WorkingMemory space for snappy compression

MODULE.bazel

@@ -61,7 +61,15 @@ bazel_dep(name = "rules_pkg", version = "1.0.1")
 bazel_dep(name = "rules_proto", version = "7.1.0")
 bazel_dep(name = "rules_python", version = "1.7.0")
 bazel_dep(name = "rules_rust", version = "0.70.0")
-bazel_dep(name = "snappy", version = "1.2.1")
+bazel_dep(name = "snappy", version = "1.2.2.bcr.3")
+single_version_override(
+    module_name = "snappy",
+    patch_strip = 1,
+    patches = [
+        "//bazel/thirdparty:snappy-export-internal-h.patch",
+    ],
+)
+
 bazel_dep(name = "utf8proc", version = "2.11.0")
 bazel_dep(name = "yaml-cpp", version = "0.8.0")
 bazel_dep(name = "zlib", version = "1.3.1.bcr.8")

bazel/thirdparty/snappy-export-internal-h.patch

@@ -0,0 +1,19 @@
+Export snappy-internal.h so dependents can drive CompressFragment with a
+caller-owned WorkingMemory (see src/v/compression/internal/
+snappy_java_compressor.cc).
+
+--- a/BUILD.bazel
++++ b/BUILD.bazel
+@@ -74,11 +74,11 @@ cc_library(
+     name = "snappy",
+     srcs = [
+         "snappy.cc",
+-        "snappy-internal.h",
+         "snappy-sinksource.cc",
+     ],
+     hdrs = [
+         "snappy.h",
++        "snappy-internal.h",
+         "snappy-sinksource.h",
+     ],
+     copts = select({

src/v/compression/BUILD

@@ -39,6 +39,7 @@ redpanda_cc_library(
         "@lz4//:lz4_frame",
         "@seastar",
         "@snappy",
+        "@snappy//:snappy-stubs-internal",
         "@zlib",
         "@zstd",
     ],

src/v/compression/internal/snappy_java_compressor.cc

@@ -12,22 +12,98 @@
 #include "base/likely.h"
 #include "bytes/details/io_iterator_consumer.h"
 #include "compression/snappy_standard_compressor.h"
-#include "snappy.h"
+
+#include <seastar/core/temporary_buffer.hh>
 
 #include <cstring>
+#include <snappy-internal.h>
+#include <snappy-stubs-internal.h>
+#include <snappy.h>
 
 namespace compression::internal {
 
-size_t find_max_size_in_frags(const iobuf& x) {
-    size_t ret = 0;
-    for (const auto& f : x) {
-        if (f.size() > ret) {
-            ret = f.size();
-        }
+// We drive snappy's internal block compressor directly instead of calling
+// snappy::RawCompress per block: the public one-shot API allocates a ~`200KiB`
+// `WorkingMemory` object internally on every call, tripping the
+// large-allocation warning threshold on hot paths like produce-time
+// The public `snappy` API currently offers no way to reuse this working memory:
+// https://github.com/google/snappy/blob/1.2.2/snappy.cc#L1813.
+// So, instead of using those public APIs, we reach into `snappy` internals,
+// allocate our own thread local `WorkingMemory` object at startup, and reuse it
+// across all compression calls. This also involves hand-rolling parts of
+// `snappy::RawCompress()`.
+//
+// There are several contracts we must uphold for CompressFragment:
+// https://github.com/google/snappy/blob/1.2.2/snappy-internal.h#L142-L157.
+//   - input_length <= kBlockSize
+//   - output buffer >= MaxCompressedLength(input_length)
+//   - hash table zeroed, size a power of two (GetHashTable guarantees both)
+static_assert(
+  SNAPPY_MAJOR == 1 && SNAPPY_MINOR == 2 && SNAPPY_PATCHLEVEL == 2,
+  "snappy version changed: re-audit the internal-header usage in this file "
+  "(WorkingMemory, CompressFragment and Varint contracts) and the BUILD.bazel "
+  "hunks in bazel/thirdparty/snappy-export-internal-h.patch before bumping");
+
+namespace {
+
+/// Scratch space for compression, reused across calls on this shard. Holds
+/// snappy's working memory (hash table + scratch buffers, ~203KiB) and a
+/// staging buffer for one compressed block (~75KiB). Allocated once at
+/// startup via init_workspace().
+class snappy_workspace {
+private:
+    /// Upper bound on one framed block in the staging buffer: a varint32
+    /// uncompressed-length prefix, followed by the compressed block data for
+    /// which CompressFragment's contract demands room for, which is
+    /// `MaxCompressedLength(input)` bytes.
+    /// `snappy::Compress()` composes the same bound, with the varint32 prefix
+    /// in a separate stack buffer instead:
+    /// https://github.com/google/snappy/blob/1.2.2/snappy.cc#L1808-L1810).
+    static size_t max_framed_block_len() {
+        return snappy::Varint::kMax32
+               + snappy::MaxCompressedLength(snappy::kBlockSize);
+    }
+
+public:
+    snappy_workspace()
+      : _staging(max_framed_block_len()) {}
+
+    /// Compress one block of at most `snappy::kBlockSize` bytes into the
+    /// staging buffer. This mirrors the prologue and per-block body of
+    /// `snappy::Compress()`, with the WorkingMemory hoisted out into this
+    /// class: https://github.com/google/snappy/blob/1.2.2/snappy.cc#L1801-L1877
+    ///
+    /// Returns the framed size.
+    size_t compress_block(const char* input, size_t len) {
+        char* dst = _staging.get_write();
+        char* op = snappy::Varint::Encode32(dst, static_cast<uint32_t>(len));
+        int table_size = 0;
+        uint16_t* table = _wmem.GetHashTable(len, &table_size);
+        char* end = snappy::internal::CompressFragment(
+          input, len, op, table, table_size);
+        return end - dst;
     }
-    return snappy::MaxCompressedLength(ret);
+
+    const char* data() const { return _staging.get(); }
+
+private:
+    snappy::internal::WorkingMemory _wmem{snappy::kBlockSize};
+    ss::temporary_buffer<char> _staging;
+};
+
+thread_local std::unique_ptr<snappy_workspace> shard_workspace;
+
+snappy_workspace& get_workspace() {
+    if (unlikely(!shard_workspace)) {
+        shard_workspace = std::make_unique<snappy_workspace>();
+    }
+    return *shard_workspace;
 }
 
+} // namespace
+
+void snappy_java_compressor::init_workspace() { get_workspace(); }
+
 template<typename T, typename = std::enable_if_t<std::is_integral_v<T>, T>>
 void append_be(iobuf& o, T t) {
     auto x = ss::cpu_to_be(t);
@@ -41,22 +117,28 @@ void append_le(iobuf& o, T t) {
     o.append((const char*)&x, sizeof(x));
 }
 iobuf snappy_java_compressor::compress(const iobuf& x) {
+    auto& wksp = get_workspace();
     iobuf ret;
     ret.append(
       snappy_magic::java_magic.data(), snappy_magic::java_magic.size());
     // versions in header are big-endian. See:
     // https://github.com/xerial/snappy-java/blob/65e1ec3de1a0d447b137c6dd6393629aa3d75b8b/src/main/java/org/xerial/snappy/SnappyCodec.java#L78-L81
     append_be(ret, snappy_magic::default_version);
     append_be(ret, snappy_magic::min_compatible_version);
-    // staging buffer
-    ss::temporary_buffer<char> obuf(find_max_size_in_frags(x));
     for (const auto& f : x) {
-        // do compression
-        size_t omax = obuf.size();
-        snappy::RawCompress(f.get(), f.size(), obuf.get_write(), &omax);
-        // must be int32 to be compatible && in big endian
-        append_be(ret, int32_t(omax));
-        ret.append(obuf.get(), omax);
+        // chop fragments at kBlockSize, as snappy::Compress does internally
+        // (`CompressFragment()` rejects anything larger):
+        // https://github.com/google/snappy/blob/1.2.2/snappy.cc#L1815-L1877
+        for (size_t offset = 0; offset < f.size();
+             offset += snappy::kBlockSize) {
+            const size_t block_len = std::min(
+              snappy::kBlockSize, f.size() - offset);
+            const size_t framed = wksp.compress_block(
+              f.get() + offset, block_len);
+            // must be int32 to be compatible && in big endian
+            append_be(ret, int32_t(framed));
+            ret.append(wksp.data(), framed);
+        }
     }
     return ret;
 }

src/v/compression/internal/snappy_java_compressor.h

@@ -32,6 +32,11 @@ struct snappy_java_compressor {
 
     static iobuf compress(const iobuf&);
     static iobuf uncompress(const iobuf&);
+
+    /// Preallocate this shard's compression workspace (~278KiB). Called
+    /// during broker startup, before the large-allocation warning threshold
+    /// is armed. Otherwise the workspace is created lazily on first use.
+    static void init_workspace();
 };
 
 } // namespace compression::internal

src/v/compression/tests/snappy_tests.cc

@@ -7,6 +7,7 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0
 
+#include "base/units.h"
 #include "bytes/iobuf.h"
 #include "compression/internal/snappy_java_compressor.h"
 #include "compression/snappy_standard_compressor.h"
@@ -139,6 +140,109 @@ TEST(SnappyTest, CompressedVersionHeadersSnappyJavaTest) {
     compressed_frag->trim_front(sizeof(decompressed_size));
 }
 
+namespace {
+
+iobuf gen_iobuf(size_t size, bool compressible) {
+    iobuf ret;
+    if (compressible) {
+        const auto data = random_generators::gen_alphanum_string(512);
+        while (ret.size_bytes() < size) {
+            ret.append(
+              data.data(), std::min(data.size(), size - ret.size_bytes()));
+        }
+    } else {
+        while (ret.size_bytes() < size) {
+            const auto data = random_generators::gen_alphanum_string(
+              std::min<size_t>(4096, size - ret.size_bytes()));
+            ret.append(data.data(), data.size());
+        }
+    }
+    return ret;
+}
+
+// input sizes covering the snappy::kBlockSize (64KiB) and iobuf max fragment
+// (128KiB) boundaries
+constexpr std::array<size_t, 9> boundary_sizes = {
+  0,
+  1,
+  100,
+  64_KiB - 1,
+  64_KiB,
+  64_KiB + 1,
+  128_KiB,
+  128_KiB + 1,
+  1_MiB,
+};
+
+} // namespace
+
+TEST(SnappyTest, RoundTripBlockBoundariesSnappyJavaTest) {
+    for (size_t size : boundary_sizes) {
+        for (bool compressible : {true, false}) {
+            auto buf = gen_iobuf(size, compressible);
+            auto compressed
+              = compression::internal::snappy_java_compressor::compress(buf);
+            auto decompressed
+              = compression::internal::snappy_java_compressor::uncompress(
+                compressed);
+            EXPECT_EQ(buf, decompressed)
+              << "size=" << size << " compressible=" << compressible;
+        }
+    }
+}
+
+// Differential check of the compressor's snappy-internal.h based
+// implementation (a caller-owned WorkingMemory driving CompressFragment)
+// against the public one-shot API: each java-framed block must be
+// byte-identical to what snappy::RawCompress produces for the same input
+// block, since both emit a varint32 uncompressed-length prefix followed by
+// the same block compression. If this fails after a snappy version bump, the
+// internal contracts drifted; see the static_assert in
+// snappy_java_compressor.cc.
+TEST(SnappyTest, DifferentialVsPublicApiSnappyJavaTest) {
+    using compressor = compression::internal::snappy_java_compressor;
+    for (size_t size : boundary_sizes) {
+        for (bool compressible : {true, false}) {
+            auto buf = gen_iobuf(size, compressible);
+
+            iobuf expected;
+            expected.append(
+              compressor::snappy_magic::java_magic.data(),
+              compressor::snappy_magic::java_magic.size());
+            auto be_version = ss::cpu_to_be(
+              compressor::snappy_magic::default_version);
+            expected.append(
+              reinterpret_cast<const char*>(&be_version), sizeof(be_version));
+            auto be_compat = ss::cpu_to_be(
+              compressor::snappy_magic::min_compatible_version);
+            expected.append(
+              reinterpret_cast<const char*>(&be_compat), sizeof(be_compat));
+            ss::temporary_buffer<char> obuf(
+              snappy::MaxCompressedLength(snappy::kBlockSize));
+            for (const auto& frag : buf) {
+                for (size_t offset = 0; offset < frag.size();
+                     offset += snappy::kBlockSize) {
+                    const size_t block_len = std::min(
+                      snappy::kBlockSize, frag.size() - offset);
+                    size_t compressed_len = obuf.size();
+                    snappy::RawCompress(
+                      frag.get() + offset,
+                      block_len,
+                      obuf.get_write(),
+                      &compressed_len);
+                    auto be_len = ss::cpu_to_be(int32_t(compressed_len));
+                    expected.append(
+                      reinterpret_cast<const char*>(&be_len), sizeof(be_len));
+                    expected.append(obuf.get(), compressed_len);
+                }
+            }
+
+            EXPECT_EQ(compressor::compress(buf), expected)
+              << "size=" << size << " compressible=" << compressible;
+        }
+    }
+}
+
 TEST(SnappyTest, LittleEndianHeadersBackwardsCompatibilitySnappyJavaTest) {
     // Previously, version fields were erroneously written with
     // little-endian encoding. They are now corrected to be written and decoded

src/v/redpanda/application.cc

@@ -22,6 +22,7 @@
 #include "cluster/node_isolation_watcher.h"
 #include "cluster/topic_recovery_service.h"
 #include "compression/async_stream_zstd.h"
+#include "compression/internal/snappy_java_compressor.h"
 #include "compression/lz4_decompression_buffers.h"
 #include "compression/stream_zstd.h"
 #include "config/configuration.h"
@@ -465,6 +466,8 @@ void application::initialize(
           compression::lz4_decompression_buffers::bufsize,
           compression::lz4_decompression_buffers::min_threshold,
           config::shard_local_cfg().lz4_decompress_reusable_buffers_disabled());
+
+        compression::internal::snappy_java_compressor::init_workspace();
     }).get();
 
     if (config::shard_local_cfg().enable_pid_file()) {