Moving predictor to WorkingArea to make it thread safe (#13706)

Summary:
**Summary:**

We need to move the Predictor to WorkingArea so that it is local to each thread and thus is thread safe.

Pull Request resolved: #13706

Test Plan: It should pass the test case written in ./compression_test.

Reviewed By: pdillinger

Differential Revision: D76836846

Pulled By: shubhajeet

fbshipit-source-id: 0d0170baf65f4bb95ba107fec77151e66b8a4449

Author: shubhajeet

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

@@ -56,20 +56,30 @@ const char* AutoSkipCompressorWrapper::Name() const {
 Status AutoSkipCompressorWrapper::CompressBlock(
     Slice uncompressed_data, std::string* compressed_output,
     CompressionType* out_compression_type, ManagedWorkingArea* wa) {
+  // Check if the managed working area is provided or owned by this object.
+  // If not, bypass auto-skip logic since the working area lacks a predictor to
+  // record or make necessary decisions to compress or bypass compression of the
+  // block
+  if (wa == nullptr || wa->owner() != this) {
+    return wrapped_->CompressBlock(uncompressed_data, compressed_output,
+                                   out_compression_type, wa);
+  }
   bool exploration =
       Random::GetTLSInstance()->PercentTrue(kExplorationPercentage);
   TEST_SYNC_POINT_CALLBACK(
       "AutoSkipCompressorWrapper::CompressBlock::exploitOrExplore",
       &exploration);
+  auto autoskip_wa = static_cast<AutoSkipWorkingArea*>(wa->get());
   if (exploration) {
     return CompressBlockAndRecord(uncompressed_data, compressed_output,
-                                  out_compression_type, wa);
+                                  out_compression_type, autoskip_wa);
   } else {
-    auto prediction = predictor_->Predict();
+    auto predictor_ptr = autoskip_wa->predictor;
+    auto prediction = predictor_ptr->Predict();
     if (prediction <= kProbabilityCutOff) {
       // decide to compress
       return CompressBlockAndRecord(uncompressed_data, compressed_output,
-                                    out_compression_type, wa);
+                                    out_compression_type, autoskip_wa);
     } else {
       // decide to bypass compression
       *out_compression_type = kNoCompression;
@@ -79,13 +89,22 @@ Status AutoSkipCompressorWrapper::CompressBlock(
   return Status::OK();
 }
 
+Compressor::ManagedWorkingArea AutoSkipCompressorWrapper::ObtainWorkingArea() {
+  auto wrap_wa = wrapped_->ObtainWorkingArea();
+  return ManagedWorkingArea(new AutoSkipWorkingArea(std::move(wrap_wa)), this);
+}
+void AutoSkipCompressorWrapper::ReleaseWorkingArea(WorkingArea* wa) {
+  delete static_cast<AutoSkipWorkingArea*>(wa);
+}
+
 Status AutoSkipCompressorWrapper::CompressBlockAndRecord(
     Slice uncompressed_data, std::string* compressed_output,
-    CompressionType* out_compression_type, ManagedWorkingArea* wa) {
+    CompressionType* out_compression_type, AutoSkipWorkingArea* wa) {
   Status status = wrapped_->CompressBlock(uncompressed_data, compressed_output,
-                                          out_compression_type, wa);
+                                          out_compression_type, &(wa->wrapped));
   // determine if it was rejected or compressed
-  predictor_->Record(uncompressed_data, compressed_output, kOpts);
+  auto predictor_ptr = wa->predictor;
+  predictor_ptr->Record(uncompressed_data, compressed_output, kOpts);
   return status;
 }
 

util/auto_skip_compressor.h

@@ -13,7 +13,6 @@
 
 namespace ROCKSDB_NAMESPACE {
 // Predict rejection probability using a moving window approach
-// This class is not thread safe
 class CompressionRejectionProbabilityPredictor {
  public:
   CompressionRejectionProbabilityPredictor(int window_size)
@@ -33,6 +32,30 @@ class CompressionRejectionProbabilityPredictor {
   size_t window_size_;
 };
 
+class AutoSkipWorkingArea : public Compressor::WorkingArea {
+ public:
+  explicit AutoSkipWorkingArea(Compressor::ManagedWorkingArea&& wa)
+      : wrapped(std::move(wa)),
+        predictor(
+            std::make_shared<CompressionRejectionProbabilityPredictor>(10)) {}
+  ~AutoSkipWorkingArea() {}
+  AutoSkipWorkingArea(const AutoSkipWorkingArea&) = delete;
+  AutoSkipWorkingArea& operator=(const AutoSkipWorkingArea&) = delete;
+  AutoSkipWorkingArea(AutoSkipWorkingArea&& other) noexcept
+      : wrapped(std::move(other.wrapped)),
+        predictor(std::move(other.predictor)) {}
+
+  AutoSkipWorkingArea& operator=(AutoSkipWorkingArea&& other) noexcept {
+    if (this != &other) {
+      wrapped = std::move(other.wrapped);
+      predictor = std::move(other.predictor);
+    }
+    return *this;
+  }
+  Compressor::ManagedWorkingArea wrapped;
+  std::shared_ptr<CompressionRejectionProbabilityPredictor> predictor;
+};
+
 class AutoSkipCompressorWrapper : public CompressorWrapper {
  public:
   const char* Name() const override;
@@ -42,12 +65,14 @@ 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,
                                 std::string* compressed_output,
                                 CompressionType* out_compression_type,
-                                ManagedWorkingArea* wa);
+                                AutoSkipWorkingArea* wa);
   static constexpr int kExplorationPercentage = 10;
   static constexpr int kProbabilityCutOff = 50;
   const CompressionOptions kOpts;