Add CompressionContext for reusing compression working memory

snappy-internal.h

@@ -137,8 +137,16 @@ inline V128 V128_DupChar(char c) { return vdupq_n_u8(c); }
 class WorkingMemory {
  public:
   explicit WorkingMemory(size_t input_size);
+
+  // Non-allocating: lays out the scratch space in the caller-provided
+  // buffer, which must be at least RequiredSize(input_size) bytes, aligned
+  // at least as strictly as uint16_t, and must outlive "*this".
+  WorkingMemory(size_t input_size, char* buffer);
   ~WorkingMemory();
 
+  // The buffer size required by the non-allocating constructor above.
+  static size_t RequiredSize(size_t input_size);
+
   // Allocates and clears a hash table using memory in "*this",
   // stores the number of buckets in "*table_size" and returns a pointer to
   // the base of the hash table.
@@ -149,6 +157,7 @@ class WorkingMemory {
  private:
   char* mem_;        // the allocated memory, never nullptr
   size_t size_;      // the size of the allocated memory, never 0
+  bool owns_mem_;    // whether the destructor should free mem_
   uint16_t* table_;  // the pointer to the hashtable
   char* input_;      // the pointer to the input scratch buffer
   char* output_;     // the pointer to the output scratch buffer

snappy.cc

@@ -76,6 +76,7 @@
 #include <cstring>
 #include <limits>
 #include <memory>
+#include <new>
 #include <string>
 #include <utility>
 #include <vector>
@@ -757,19 +758,36 @@ uint32_t CalculateTableSize(uint32_t input_size) {
 }  // namespace
 
 namespace internal {
-WorkingMemory::WorkingMemory(size_t input_size) {
+size_t WorkingMemory::RequiredSize(size_t input_size) {
   const size_t max_fragment_size = std::min(input_size, kBlockSize);
   const size_t table_size = CalculateTableSize(max_fragment_size);
-  size_ = table_size * sizeof(*table_) + max_fragment_size +
-          MaxCompressedLength(max_fragment_size);
-  mem_ = std::allocator<char>().allocate(size_);
+  return table_size * sizeof(uint16_t) + max_fragment_size +
+         MaxCompressedLength(max_fragment_size);
+}
+
+WorkingMemory::WorkingMemory(size_t input_size)
+    : WorkingMemory(input_size,
+                    std::allocator<char>().allocate(RequiredSize(input_size))) {
+  owns_mem_ = true;
+}
+
+WorkingMemory::WorkingMemory(size_t input_size, char* buffer) {
+  assert(buffer != nullptr);
+  assert(reinterpret_cast<uintptr_t>(buffer) % alignof(uint16_t) == 0);
+  const size_t max_fragment_size = std::min(input_size, kBlockSize);
+  const size_t table_size = CalculateTableSize(max_fragment_size);
+  mem_ = buffer;
+  size_ = RequiredSize(input_size);
+  owns_mem_ = false;
   table_ = reinterpret_cast<uint16_t*>(mem_);
   input_ = mem_ + table_size * sizeof(*table_);
   output_ = input_ + max_fragment_size;
 }
 
 WorkingMemory::~WorkingMemory() {
-  std::allocator<char>().deallocate(mem_, size_);
+  if (owns_mem_) {
+    std::allocator<char>().deallocate(mem_, size_);
+  }
 }
 
 uint16_t* WorkingMemory::GetHashTable(size_t fragment_size,
@@ -1865,7 +1883,9 @@ size_t Compress(Source* reader, Sink* writer) {
   return Compress(reader, writer, CompressionOptions{});
 }
 
-size_t Compress(Source* reader, Sink* writer, CompressionOptions options) {
+static size_t InternalCompress(Source* reader, Sink* writer,
+                               CompressionOptions options,
+                               internal::WorkingMemory* wmem) {
   assert(options.level == 1 || options.level == 2);
   size_t written = 0;
   size_t N = reader->Available();
@@ -1875,8 +1895,6 @@ size_t Compress(Source* reader, Sink* writer, CompressionOptions options) {
   writer->Append(ulength, p - ulength);
   written += (p - ulength);
 
-  internal::WorkingMemory wmem(N);
-
   while (N > 0) {
     // Get next block to compress (without copying if possible)
     size_t fragment_size;
@@ -1891,7 +1909,7 @@ size_t Compress(Source* reader, Sink* writer, CompressionOptions options) {
       pending_advance = num_to_read;
       fragment_size = num_to_read;
     } else {
-      char* scratch = wmem.GetScratchInput();
+      char* scratch = wmem->GetScratchInput();
       std::memcpy(scratch, fragment, bytes_read);
       reader->Skip(bytes_read);
 
@@ -1910,7 +1928,7 @@ size_t Compress(Source* reader, Sink* writer, CompressionOptions options) {
 
     // Get encoding table for compression
     int table_size;
-    uint16_t* table = wmem.GetHashTable(num_to_read, &table_size);
+    uint16_t* table = wmem->GetHashTable(num_to_read, &table_size);
 
     // Compress input_fragment and append to dest
     int max_output = MaxCompressedLength(num_to_read);
@@ -1920,7 +1938,7 @@ size_t Compress(Source* reader, Sink* writer, CompressionOptions options) {
     // scratch_output[] region is big enough for this iteration.
     // Need a scratch buffer for the output, in case the byte sink doesn't
     // have room for us directly.
-    char* dest = writer->GetAppendBuffer(max_output, wmem.GetScratchOutput());
+    char* dest = writer->GetAppendBuffer(max_output, wmem->GetScratchOutput());
     char* end = nullptr;
       if (options.level == 1) {
         end = internal::CompressFragment(fragment, fragment_size, dest, table,
@@ -1940,6 +1958,69 @@ size_t Compress(Source* reader, Sink* writer, CompressionOptions options) {
   return written;
 }
 
+size_t Compress(Source* reader, Sink* writer, CompressionOptions options) {
+  internal::WorkingMemory wmem(reader->Available());
+  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)),
+      owns_working_memory_(true) {}
+
+size_t CompressionContext::WorkspaceSize() {
+  return sizeof(internal::WorkingMemory) +
+         internal::WorkingMemory::RequiredSize(kBlockSize);
+}
+
+CompressionContext::CompressionContext(void* workspace, size_t workspace_size)
+    : owns_working_memory_(false) {
+  assert(workspace != nullptr);
+  assert(workspace_size >= WorkspaceSize());
+  assert(reinterpret_cast<uintptr_t>(workspace) %
+             alignof(internal::WorkingMemory) ==
+         0);
+  (void)workspace_size;
+  char* base = static_cast<char*>(workspace);
+  working_memory_ = new (base) internal::WorkingMemory(
+      kBlockSize, base + sizeof(internal::WorkingMemory));
+}
+
+void CompressionContext::Reset() {
+  if (working_memory_ == nullptr) return;
+  if (owns_working_memory_) {
+    delete working_memory_;
+  } else {
+    working_memory_->~WorkingMemory();
+  }
+  working_memory_ = nullptr;
+}
+
+CompressionContext::~CompressionContext() { Reset(); }
+
+CompressionContext::CompressionContext(CompressionContext&& other) noexcept
+    : working_memory_(other.working_memory_),
+      owns_working_memory_(other.owns_working_memory_) {
+  other.working_memory_ = nullptr;
+}
+
+CompressionContext& CompressionContext::operator=(
+    CompressionContext&& other) noexcept {
+  if (this != &other) {
+    Reset();
+    working_memory_ = other.working_memory_;
+    owns_working_memory_ = other.owns_working_memory_;
+    other.working_memory_ = nullptr;
+  }
+  return *this;
+}
+
 // -----------------------------------------------------------------------
 // IOVec interfaces
 // -----------------------------------------------------------------------
@@ -2385,6 +2466,17 @@ void RawCompress(const char* input, size_t input_length, char* compressed,
   *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,

snappy.h

@@ -50,6 +50,10 @@ namespace snappy {
   class Source;
   class Sink;
 
+  namespace internal {
+  class WorkingMemory;
+  }  // end namespace internal
+
   struct CompressionOptions {
     // Compression level.
     // Level 1 is the fastest
@@ -73,6 +77,53 @@ namespace snappy {
     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:
+    // Allocates the working memory on the heap.
+    CompressionContext();
+
+    // Constructs a context whose working memory is placed in the
+    // caller-provided "workspace" instead of being heap-allocated; the
+    // library performs no allocation at all.
+    //
+    // REQUIRES: "workspace" points to at least "workspace_size" bytes with
+    // "workspace_size >= WorkspaceSize()", is suitably aligned for any
+    // object type (as if returned by malloc), and outlives "*this".
+    CompressionContext(void* workspace, size_t workspace_size);
+
+    ~CompressionContext();
+
+    CompressionContext(CompressionContext&& other) noexcept;
+    CompressionContext& operator=(CompressionContext&& other) noexcept;
+
+    CompressionContext(const CompressionContext&) = delete;
+    CompressionContext& operator=(const CompressionContext&) = delete;
+
+    // The workspace size required by the non-allocating constructor above.
+    static size_t WorkspaceSize();
+
+   private:
+    friend size_t Compress(Source* reader, Sink* writer,
+                           CompressionOptions options, CompressionContext* ctx);
+
+    // Destroys the working memory as appropriate for how it was created
+    // (delete if heap-allocated, in-place destruction if placement-constructed
+    // in a caller-provided workspace).
+    void Reset();
+
+    internal::WorkingMemory* working_memory_;
+    bool owns_working_memory_;
+  };
+
   // ------------------------------------------------------------------------
   // Generic compression/decompression routines.
   // ------------------------------------------------------------------------
@@ -84,6 +135,11 @@ namespace snappy {
   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 +221,12 @@ namespace snappy {
   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`).