Add CompressionContext for reusing compression working memory Today, every Compress()/RawCompress() call allocates and frees an internal WorkingMemory. Callers that compress frequently, or that run with custom or per-thread allocators sensitive to large contiguous allocations, currently have no way to avoid this per-call allocation, and other mainstream codecs offer reusable contexts for exactly this purpose. Add an opaque CompressionContext holding a reusable WorkingMemory sized for kBlockSize (usable with inputs of any size and both compression levels), plus Compress()/RawCompress() overloads that use it instead of allocating internally.
diff --git a/snappy.cc b/snappy.cc index 7f6fa44..efca3af 100644 --- a/snappy.cc +++ b/snappy.cc
@@ -1945,6 +1945,33 @@ return InternalCompress(reader, writer, options, &wmem); } +size_t Compress(Source* reader, Sink* writer, CompressionOptions options, + CompressionContext* ctx) { + assert(ctx != nullptr); + assert(ctx->working_memory_ != nullptr); + return InternalCompress(reader, writer, options, ctx->working_memory_); +} + +CompressionContext::CompressionContext() + : working_memory_(new internal::WorkingMemory(kBlockSize)) {} + +CompressionContext::~CompressionContext() { delete working_memory_; } + +CompressionContext::CompressionContext(CompressionContext&& other) noexcept + : working_memory_(other.working_memory_) { + other.working_memory_ = nullptr; +} + +CompressionContext& CompressionContext::operator=( + CompressionContext&& other) noexcept { + if (this != &other) { + delete working_memory_; + working_memory_ = other.working_memory_; + other.working_memory_ = nullptr; + } + return *this; +} + // ----------------------------------------------------------------------- // IOVec interfaces // ----------------------------------------------------------------------- @@ -2390,6 +2417,17 @@ *compressed_length = (writer.CurrentDestination() - compressed); } +void RawCompress(const char* input, size_t input_length, char* compressed, + size_t* compressed_length, CompressionOptions options, + CompressionContext* ctx) { + ByteArraySource reader(input, input_length); + UncheckedByteArraySink writer(compressed); + Compress(&reader, &writer, options, ctx); + + // Compute how many bytes were added + *compressed_length = (writer.CurrentDestination() - compressed); +} + void RawCompressFromIOVec(const struct iovec* iov, size_t uncompressed_length, char* compressed, size_t* compressed_length) { RawCompressFromIOVec(iov, uncompressed_length, compressed, compressed_length,
diff --git a/snappy.h b/snappy.h index ecd83d7..423142b 100644 --- a/snappy.h +++ b/snappy.h
@@ -50,6 +50,10 @@ class Source; class Sink; + namespace internal { + class WorkingMemory; + } // end namespace internal + struct CompressionOptions { // Compression level. // Level 1 is the fastest @@ -73,6 +77,33 @@ static constexpr int DefaultCompressionLevel() { return 1; } }; + // Scratch memory for compression, reusable across compressions. Callers that + // compress frequently, or that need to avoid large heap allocations can + // allocate a CompressionContext once and pass it to Compress()/RawCompress() + // to reuse the working memory across calls. + // + // The context is sized for the largest block and works for inputs of any + // size. A context may be used by any number of sequential compressions, but + // must not be used from multiple threads concurrently. A moved-from context + // may only be destroyed or assigned to. + class CompressionContext { + public: + CompressionContext(); + ~CompressionContext(); + + CompressionContext(CompressionContext&& other) noexcept; + CompressionContext& operator=(CompressionContext&& other) noexcept; + + CompressionContext(const CompressionContext&) = delete; + CompressionContext& operator=(const CompressionContext&) = delete; + + private: + friend size_t Compress(Source* reader, Sink* writer, + CompressionOptions options, CompressionContext* ctx); + + internal::WorkingMemory* working_memory_; + }; + // ------------------------------------------------------------------------ // Generic compression/decompression routines. // ------------------------------------------------------------------------ @@ -84,6 +115,11 @@ size_t Compress(Source* reader, Sink* writer, CompressionOptions options); + // Same as the above, but uses the working memory of "*ctx" instead of + // allocating it internally. See CompressionContext. + size_t Compress(Source* reader, Sink* writer, CompressionOptions options, + CompressionContext* ctx); + // Find the uncompressed length of the given stream, as given by the header. // Note that the true length could deviate from this; the stream could e.g. // be truncated. @@ -165,6 +201,12 @@ void RawCompress(const char* input, size_t input_length, char* compressed, size_t* compressed_length, CompressionOptions options); + // Same as the above, but uses the working memory of "*ctx" instead of + // allocating it internally. See CompressionContext. + void RawCompress(const char* input, size_t input_length, char* compressed, + size_t* compressed_length, CompressionOptions options, + CompressionContext* ctx); + // Same as `RawCompress` above but taking an `iovec` array as input. Note that // `uncompressed_length` is the total number of bytes to be read from the // elements of `iov` (_not_ the number of elements in `iov`).
diff --git a/snappy_unittest.cc b/snappy_unittest.cc index 2d04200..2ee6e22 100644 --- a/snappy_unittest.cc +++ b/snappy_unittest.cc
@@ -543,6 +543,58 @@ } } +TEST(Snappy, CompressionContext) { + std::minstd_rand0 rng(snappy::GetFlag(FLAGS_test_random_seed)); + std::uniform_int_distribution<int> uniform_byte(0, 255); + + // A single context, reused across every compression below. + CompressionContext ctx; + + const size_t sizes[] = {0, + 1, + 100, + kBlockSize - 1, + kBlockSize, + kBlockSize + 1, + 2 * kBlockSize, + (1 << 20) + 17}; + for (int level = CompressionOptions::MinCompressionLevel(); + level <= CompressionOptions::MaxCompressionLevel(); ++level) { + CompressionOptions options(level); + for (size_t len : sizes) { + for (bool compressible : {true, false}) { + std::string input; + input.reserve(len); + while (input.size() < len) { + input.push_back(compressible + ? static_cast<char>('a' + input.size() % 4) + : static_cast<char>(uniform_byte(rng))); + } + + std::string plain(MaxCompressedLength(len), '\0'); + size_t plain_len = 0; + RawCompress(input.data(), input.size(), &plain[0], &plain_len, options); + plain.resize(plain_len); + + std::string with_context(MaxCompressedLength(len), '\0'); + size_t with_context_len = 0; + RawCompress(input.data(), input.size(), &with_context[0], + &with_context_len, options, &ctx); + with_context.resize(with_context_len); + + // Compressing with a reused context must produce output identical to + // the context-free API. + EXPECT_EQ(plain, with_context) << "level=" << level << " len=" << len + << " compressible=" << compressible; + + std::string uncompressed; + EXPECT_TRUE(Uncompress(with_context, &uncompressed)); + EXPECT_EQ(input, uncompressed); + } + } + } +} + TEST(Snappy, FourByteOffset) { // The new compressor cannot generate four-byte offsets since // it chops up the input into 32KB pieces. So we hand-emit the