Moving predictor to WorkingArea to make it thread safe

tools/db_crashtest.py

@@ -1011,9 +1011,7 @@ def finalize_and_sanitize(src_params):
         dest_params["compression_type"] = "zstd"
         dest_params["bottommost_compression_type"] = "zstd"
     elif dest_params.get("compression_manager") == "autoskip":
-        # disabling compression parallel threads if auto skip manager is being used as the predictor is not thread safe
-        dest_params["compression_parallel_threads"] = 1
-        # esuring the compression is being used
+        # ensuring the compression is being used
         if dest_params.get("compression_type") == "none":
             dest_params["compression_type"] = random.choice(
                 ["snappy", "zlib", "lz4", "lz4hc", "xpress", "zstd"]

util/auto_skip_compressor.cc

@@ -65,27 +65,44 @@ Status AutoSkipCompressorWrapper::CompressBlock(
     return CompressBlockAndRecord(uncompressed_data, compressed_output,
                                   out_compression_type, wa);
   } else {
-    auto prediction = predictor_->Predict();
+    auto predictor_ptr =
+        static_cast<AutoSkipCompressionContext*>(wa->get())->predictor;
+    auto prediction = predictor_ptr->Predict();
     if (prediction <= kProbabilityCutOff) {
       // decide to compress
       return CompressBlockAndRecord(uncompressed_data, compressed_output,
                                     out_compression_type, wa);
     } else {
-      // decide to bypass compression
+      // bypassed compression
       *out_compression_type = kNoCompression;
       return Status::OK();
     }
   }
   return Status::OK();
 }
 
+Compressor::ManagedWorkingArea AutoSkipCompressorWrapper::ObtainWorkingArea() {
+  auto wrap_wa = wrapped_->ObtainWorkingArea();
+  return ManagedWorkingArea(
+      static_cast<WorkingArea*>(
+          new AutoSkipCompressionContext(std::move(wrap_wa))),
+      this);
+}
+void AutoSkipCompressorWrapper::ReleaseWorkingArea(WorkingArea* wa) {
+  delete static_cast<AutoSkipCompressionContext*>(wa);
+}
+
 Status AutoSkipCompressorWrapper::CompressBlockAndRecord(
     Slice uncompressed_data, std::string* compressed_output,
     CompressionType* out_compression_type, ManagedWorkingArea* wa) {
+  auto wrap_wa =
+      &(static_cast<AutoSkipCompressionContext*>(wa->get())->wrapped);
   Status status = wrapped_->CompressBlock(uncompressed_data, compressed_output,
-                                          out_compression_type, wa);
+                                          out_compression_type, wrap_wa);
   // determine if it was rejected or compressed
-  predictor_->Record(uncompressed_data, compressed_output, kOpts);
+  auto predictor_ptr =
+      static_cast<AutoSkipCompressionContext*>(wa->get())->predictor;
+  predictor_ptr->Record(uncompressed_data, compressed_output, kOpts);
   return status;
 }
 

util/auto_skip_compressor.h

@@ -12,8 +12,8 @@
 #include "rocksdb/advanced_compression.h"
 
 namespace ROCKSDB_NAMESPACE {
-// Predict rejection probability using a moving window approach
-// This class is not thread safe
+// Predicts Rejection Probability using previous using past data of certain
+// window size
 class CompressionRejectionProbabilityPredictor {
  public:
   CompressionRejectionProbabilityPredictor(int window_size)
@@ -33,6 +33,20 @@ class CompressionRejectionProbabilityPredictor {
   size_t window_size_;
 };
 
+class AutoSkipCompressionContext : public Compressor::WorkingArea {
+ public:
+  explicit AutoSkipCompressionContext(Compressor::ManagedWorkingArea&& wa)
+      : wrapped(std::move(wa)),
+        predictor(
+            std::make_shared<CompressionRejectionProbabilityPredictor>(10)) {}
+  ~AutoSkipCompressionContext() {}
+  AutoSkipCompressionContext(const AutoSkipCompressionContext&) = delete;
+  AutoSkipCompressionContext& operator=(const AutoSkipCompressionContext&) =
+      delete;
+  Compressor::ManagedWorkingArea wrapped;
+  std::shared_ptr<CompressionRejectionProbabilityPredictor> predictor;
+};
+
 class AutoSkipCompressorWrapper : public CompressorWrapper {
  public:
   const char* Name() const override;
@@ -42,6 +56,8 @@ class AutoSkipCompressorWrapper : public CompressorWrapper {
   Status CompressBlock(Slice uncompressed_data, std::string* compressed_output,
                        CompressionType* out_compression_type,
                        ManagedWorkingArea* wa) override;
+  ManagedWorkingArea ObtainWorkingArea() override;
+  void ReleaseWorkingArea(WorkingArea* wa) override;
 
  private:
   Status CompressBlockAndRecord(Slice uncompressed_data,