Support Zstd compression level in Leveldb

PiperOrigin-RevId: 520556840

6 files changed, 30 insertions, 6 deletions

diff --git a/benchmarks/db_bench.cc b/benchmarks/db_bench.cc
index db866f5..8e3f4e7 100644
--- a/benchmarks/db_bench.cc
+++ b/benchmarks/db_bench.cc
@@ -126,6 +126,9 @@ static bool FLAGS_compression = true;
 // Use the db with the following name.
 static const char* FLAGS_db = nullptr;
 
+// ZSTD compression level to try out
+static int FLAGS_zstd_compression_level = 1;
+
 namespace leveldb {
 
 namespace {
@@ -779,11 +782,21 @@ class Benchmark {
   }
 
   void ZstdCompress(ThreadState* thread) {
-    Compress(thread, "zstd", &port::Zstd_Compress);
+    Compress(thread, "zstd",
+             [](const char* input, size_t length, std::string* output) {
+               return port::Zstd_Compress(FLAGS_zstd_compression_level, input,
+                                          length, output);
+             });
   }
 
   void ZstdUncompress(ThreadState* thread) {
-    Uncompress(thread, "zstd", &port::Zstd_Compress, &port::Zstd_Uncompress);
+    Uncompress(
+        thread, "zstd",
+        [](const char* input, size_t length, std::string* output) {
+          return port::Zstd_Compress(FLAGS_zstd_compression_level, input,
+                                     length, output);
+        },
+        &port::Zstd_Uncompress);
   }
 
   void Open() {
diff --git a/include/leveldb/options.h b/include/leveldb/options.h
index 79bcdbb..d755f46 100644
--- a/include/leveldb/options.h
+++ b/include/leveldb/options.h
@@ -131,6 +131,10 @@ struct LEVELDB_EXPORT Options {
   // efficiently detect that and will switch to uncompressed mode.
   CompressionType compression = kSnappyCompression;
 
+  // Compression level for zstd.
+  // Currently only the range [-5,22] is supported. Default is 1.
+  int zstd_compression_level = 1;
+
   // EXPERIMENTAL: If true, append to existing MANIFEST and log files
   // when a database is opened.  This can significantly speed up open.
   //
diff --git a/port/port_example.h b/port/port_example.h
index b1a1c32..5d50ffa 100644
--- a/port/port_example.h
+++ b/port/port_example.h
@@ -83,7 +83,8 @@ bool Snappy_Uncompress(const char* input_data, size_t input_length,
 
 // Store the zstd compression of "input[0,input_length-1]" in *output.
 // Returns false if zstd is not supported by this port.
-bool Zstd_Compress(const char* input, size_t input_length, std::string* output);
+bool Zstd_Compress(int level, const char* input, size_t input_length,
+                   std::string* output);
 
 // If input[0,input_length-1] looks like a valid zstd compressed
 // buffer, store the size of the uncompressed data in *result and
diff --git a/port/port_stdcxx.h b/port/port_stdcxx.h
index ca961e6..6f503f6 100644
--- a/port/port_stdcxx.h
+++ b/port/port_stdcxx.h
@@ -29,6 +29,7 @@
 #include <snappy.h>
 #endif  // HAVE_SNAPPY
 #if HAVE_ZSTD
+#define ZSTD_STATIC_LINKING_ONLY  // For ZSTD_compressionParameters.
 #include <zstd.h>
 #endif  // HAVE_ZSTD
 
@@ -129,7 +130,7 @@ inline bool Snappy_Uncompress(const char* input, size_t length, char* output) {
 #endif  // HAVE_SNAPPY
 }
 
-inline bool Zstd_Compress(const char* input, size_t length,
+inline bool Zstd_Compress(int level, const char* input, size_t length,
                           std::string* output) {
 #if HAVE_ZSTD
   // Get the MaxCompressedLength.
@@ -139,6 +140,9 @@ inline bool Zstd_Compress(const char* input, size_t length,
   }
   output->resize(outlen);
   ZSTD_CCtx* ctx = ZSTD_createCCtx();
+  ZSTD_compressionParameters parameters =
+      ZSTD_getCParams(level, std::max(length, size_t{1}), /*dictSize=*/0);
+  ZSTD_CCtx_setCParams(ctx, parameters);
   outlen = ZSTD_compress2(ctx, &(*output)[0], output->size(), input, length);
   ZSTD_freeCCtx(ctx);
   if (ZSTD_isError(outlen)) {
@@ -148,6 +152,7 @@ inline bool Zstd_Compress(const char* input, size_t length,
   return true;
 #else
   // Silence compiler warnings about unused arguments.
+  (void)level;
   (void)input;
   (void)length;
   (void)output;
diff --git a/table/table_builder.cc b/table/table_builder.cc
index ba3df9e..0932c94 100644
--- a/table/table_builder.cc
+++ b/table/table_builder.cc
@@ -171,7 +171,8 @@ void TableBuilder::WriteBlock(BlockBuilder* block, BlockHandle* handle) {
 
     case kZstdCompression: {
       std::string* compressed = &r->compressed_output;
-      if (port::Zstd_Compress(raw.data(), raw.size(), compressed) &&
+      if (port::Zstd_Compress(r->options.zstd_compression_level, raw.data(),
+                              raw.size(), compressed) &&
           compressed->size() < raw.size() - (raw.size() / 8u)) {
         block_contents = *compressed;
       } else {
diff --git a/table/table_test.cc b/table/table_test.cc
index b3baf95..aea0697 100644
--- a/table/table_test.cc
+++ b/table/table_test.cc
@@ -791,7 +791,7 @@ static bool CompressionSupported(CompressionType type) {
   if (type == kSnappyCompression) {
     return port::Snappy_Compress(in.data(), in.size(), &out);
   } else if (type == kZstdCompression) {
-    return port::Zstd_Compress(in.data(), in.size(), &out);
+    return port::Zstd_Compress(/*level=*/1, in.data(), in.size(), &out);
   }
   return false;
 }