Fix integration gap between blob direct write and blob_compression_opts (#14669)

Summary:
`blob_compression_opts` is dynamically changeable through `SetOptions()`, but
blob direct write cached its per-thread compressor by `CompressionType` only.
After updating `blob_compression_opts`, new blob records on an existing writer
thread could keep using stale compression settings.

This change threads `CompressionOptions` through the direct-write settings and
rebuilds the cached per-thread compressor whenever the published options for
that compression type change. When the options stay the same, the steady-state
cache behavior is unchanged.

It also adds a regression test that toggles the ZSTD checksum flag through
`SetOptions()` and inspects the raw blob records to confirm each direct-write
record used the latest compression options.

## Testing

- `timeout 60s ./db_blob_direct_write_test --gtest_filter='DBBlobDirectWriteTest.DirectWriteCompressionOptionsUpdateRebuildsCachedCompressor' --gtest_repeat=5`

Pull Request resolved: https://github.com/facebook/rocksdb/pull/14669

Reviewed By: jainraj91

Differential Revision: D102423376

Pulled By: xingbowang

fbshipit-source-id: 504243d72cf22e62097f7e88ac1cc36c5bf29eee

db/blob/blob_file_partition_manager.cc

@@ -54,13 +54,14 @@ class RoundRobinBlobFilePartitionStrategy : public BlobFilePartitionStrategy {
 };
 
 struct DirectWriteCompressionState {
+  CompressionOptions compression_opts;
   // `working_area` must be released before its owning compressor.
   std::unique_ptr<Compressor> compressor;
   Compressor::ManagedWorkingArea working_area;
 };
 
 DirectWriteCompressionState& GetDirectWriteCompressionState(
-    CompressionType compression) {
+    CompressionType compression, const CompressionOptions& compression_opts) {
   assert(compression <= kLastBuiltinCompression);
 
   static thread_local std::array<DirectWriteCompressionState,
@@ -71,12 +72,15 @@ DirectWriteCompressionState& GetDirectWriteCompressionState(
   auto& compression_state =
       compression_states[static_cast<size_t>(compression)];
   if (compression != kNoCompression &&
-      compression_state.compressor == nullptr) {
-    // BDW v1 mirrors BlobFileBuilder by using the built-in compressor with the
-    // default CompressionOptions only. Cache it per thread so repeated writes
-    // do not allocate a new compressor and working area every time.
+      (compression_state.compressor == nullptr ||
+       compression_state.compression_opts != compression_opts)) {
+    // BDW compression settings are mutable, so rebuild the per-thread cached
+    // compressor when the latest published opts for this type change.
+    compression_state.working_area = Compressor::ManagedWorkingArea{};
+    compression_state.compressor.reset();
+    compression_state.compression_opts = compression_opts;
     compression_state.compressor =
-        GetBuiltinV2CompressionManager()->GetCompressor(CompressionOptions{},
+        GetBuiltinV2CompressionManager()->GetCompressor(compression_opts,
                                                         compression);
     if (compression_state.compressor != nullptr) {
       compression_state.working_area =
@@ -390,15 +394,15 @@ bool BlobFilePartitionManager::MarkSealedFileGarbage(
 }
 
 Status BlobFilePartitionManager::MaybePrepopulateBlobCache(
-    const BlobDirectWriteSettings* settings, const Slice& original_value,
+    const BlobDirectWriteSettings& settings, const Slice& original_value,
     uint64_t blob_file_number, uint64_t blob_offset) {
-  if (settings == nullptr || settings->blob_cache == nullptr ||
-      settings->prepopulate_blob_cache != PrepopulateBlobCache::kFlushOnly) {
+  if (settings.blob_cache == nullptr ||
+      settings.prepopulate_blob_cache != PrepopulateBlobCache::kFlushOnly) {
     return Status::OK();
   }
 
   FullTypedCacheInterface<BlobContents, BlobContentsCreator> blob_cache{
-      settings->blob_cache};
+      settings.blob_cache};
   const OffsetableCacheKey base_cache_key(db_id_, db_session_id_,
                                           blob_file_number);
   const CacheKey cache_key = base_cache_key.WithOffset(blob_offset);
@@ -419,7 +423,7 @@ Status BlobFilePartitionManager::WriteBlob(
     const WriteOptions& write_options, uint32_t column_family_id,
     CompressionType compression, const Slice& key, const Slice& value,
     uint64_t* blob_file_number, uint64_t* blob_offset, uint64_t* blob_size,
-    const BlobDirectWriteSettings* settings, const uint32_t* partition_idx) {
+    const BlobDirectWriteSettings& settings, const uint32_t* partition_idx) {
   assert(blob_file_number != nullptr);
   assert(blob_offset != nullptr);
   assert(blob_size != nullptr);
@@ -431,7 +435,8 @@ Status BlobFilePartitionManager::WriteBlob(
   GrowableBuffer compressed_value;
   Slice write_value = value;
   if (compression != kNoCompression) {
-    auto& compression_state = GetDirectWriteCompressionState(compression);
+    auto& compression_state =
+        GetDirectWriteCompressionState(compression, settings.compression_opts);
     if (compression_state.compressor == nullptr) {
       return Status::NotSupported(
           "Blob direct write compression type not supported");

db/blob/blob_file_partition_manager.h

@@ -77,7 +77,7 @@ class BlobFilePartitionManager {
                    CompressionType compression, const Slice& key,
                    const Slice& value, uint64_t* blob_file_number,
                    uint64_t* blob_offset, uint64_t* blob_size,
-                   const BlobDirectWriteSettings* settings = nullptr,
+                   const BlobDirectWriteSettings& settings,
                    const uint32_t* partition_idx = nullptr);
 
   // Selects the partition to use for all blob-backed columns of one PutEntity
@@ -239,8 +239,8 @@ class BlobFilePartitionManager {
   static bool MarkSealedFileGarbage(std::vector<SealedFile>* sealed_files,
                                     uint64_t file_number, uint64_t blob_count,
                                     uint64_t blob_bytes);
-  // Optionally seeds the blob cache with the original uncompressed value.
-  Status MaybePrepopulateBlobCache(const BlobDirectWriteSettings* settings,
+  // Seeds the blob cache with the original uncompressed value when configured.
+  Status MaybePrepopulateBlobCache(const BlobDirectWriteSettings& settings,
                                    const Slice& original_value,
                                    uint64_t blob_file_number,
                                    uint64_t blob_offset);

db/blob/blob_write_batch_transformer.cc

@@ -107,7 +107,7 @@ Status BlobWriteBatchTransformer::MaybePreprocessWideColumns(
     uint64_t blob_size = 0;
     Status s = partition_mgr->WriteBlob(
         write_options, column_family_id, settings.compression_type, key,
-        column_value, &blob_file_number, &blob_offset, &blob_size, &settings,
+        column_value, &blob_file_number, &blob_offset, &blob_size, settings,
         &entity_partition);
     if (!s.ok()) {
       return s;
@@ -169,7 +169,7 @@ Status BlobWriteBatchTransformer::PutCF(uint32_t column_family_id,
 
   Status s = cached_partition_mgr_->WriteBlob(
       write_options_, column_family_id, settings.compression_type, key, value,
-      &blob_file_number, &blob_offset, &blob_size, &settings);
+      &blob_file_number, &blob_offset, &blob_size, settings);
   if (!s.ok()) {
     return s;
   }

db/blob/blob_write_batch_transformer.h

@@ -33,6 +33,8 @@ struct BlobDirectWriteSettings {
   uint64_t min_blob_size = 0;
   // Compression to use for newly written blob records.
   CompressionType compression_type = kNoCompression;
+  // Compression options for newly written blob records.
+  CompressionOptions compression_opts;
   // Raw pointer — the Cache is owned by ColumnFamilyOptions and outlives all
   // settings snapshots. Using raw avoids 2 atomic ref-count ops per Put().
   Cache* blob_cache = nullptr;

db/blob/db_blob_direct_write_test.cc

@@ -151,6 +151,48 @@ class DBBlobDirectWriteTest : public DBTestBase {
     return reader.ReadHeader(header);
   }
 
+  Status ReadBlobFileRecords(uint64_t blob_file_number, size_t num_records,
+                             std::vector<BlobLogRecord>* records) {
+    assert(records != nullptr);
+    records->clear();
+
+    std::unique_ptr<FSRandomAccessFile> file;
+    FileOptions file_options;
+    constexpr IODebugContext* dbg = nullptr;
+    const std::string blob_file_path = BlobFileName(dbname_, blob_file_number);
+    FileSystem* fs = env_->GetFileSystem().get();
+    SystemClock* clock = env_->GetSystemClock().get();
+
+    Status s =
+        fs->NewRandomAccessFile(blob_file_path, file_options, &file, dbg);
+    if (!s.ok()) {
+      return s;
+    }
+
+    std::unique_ptr<RandomAccessFileReader> file_reader(
+        new RandomAccessFileReader(std::move(file), blob_file_path, clock));
+    BlobLogSequentialReader reader(std::move(file_reader), clock,
+                                   /*statistics=*/nullptr);
+
+    BlobLogHeader header;
+    s = reader.ReadHeader(&header);
+    if (!s.ok()) {
+      return s;
+    }
+
+    records->reserve(num_records);
+    for (size_t i = 0; i < num_records; ++i) {
+      BlobLogRecord record;
+      s = reader.ReadRecord(&record,
+                            BlobLogSequentialReader::kReadHeaderKeyBlob);
+      if (!s.ok()) {
+        return s;
+      }
+      records->push_back(std::move(record));
+    }
+
+    return Status::OK();
+  }
   CompressionType GetSupportedCompressedBlobCompression() {
     static constexpr std::array<CompressionType, 6> kCandidates{
         kSnappyCompression, kLZ4Compression,   kZSTD,
@@ -944,6 +986,70 @@ TEST_F(DBBlobDirectWriteTest,
   ASSERT_EQ(Get(third_key), third_value);
 }
 
+TEST_F(DBBlobDirectWriteTest,
+       DirectWriteCompressionOptionsUpdateRebuildsCachedCompressor) {
+#ifndef ZSTD
+  ROCKSDB_GTEST_SKIP("This test requires ZSTD support");
+  return;
+#else
+  // Goal: verify BDW picks up dynamic blob_compression_opts updates instead of
+  // reusing a stale per-thread compressor keyed only by compression type.
+  // We toggle the ZSTD frame checksum flag on and off, then inspect the raw
+  // blob records in one blob file. With identical inputs and level, enabling
+  // checksum should add exactly 4 bytes to the stored compressed blob.
+  Options options = GetDirectWriteOptions();
+  options.blob_file_size = 1 << 20;
+  options.blob_compression_type = kZSTD;
+  options.blob_compression_opts.level = 1;
+  options.blob_compression_opts.checksum = false;
+
+  Reopen(options);
+
+  const std::string first_key = "checksum_off_before_update";
+  const std::string second_key = "checksum_on_after_update";
+  const std::string third_key = "checksum_off_after_second_update";
+  const std::string value(256, 'v');
+
+  ASSERT_OK(Put(first_key, value));
+  ASSERT_EQ(Get(first_key), value);
+
+  ASSERT_OK(
+      db_->SetOptions(db_->DefaultColumnFamily(),
+                      {{"blob_compression_opts", "{level=1;checksum=true}"}}));
+  ASSERT_OK(Put(second_key, value));
+  ASSERT_EQ(Get(second_key), value);
+
+  ASSERT_OK(
+      db_->SetOptions(db_->DefaultColumnFamily(),
+                      {{"blob_compression_opts", "{level=1;checksum=false}"}}));
+  ASSERT_OK(Put(third_key, value));
+  ASSERT_EQ(Get(third_key), value);
+
+  ASSERT_OK(Flush());
+  ASSERT_EQ(CountBlobFiles(), 1U);
+
+  const std::vector<uint64_t> blob_file_numbers = GetBlobFileNumbers();
+  ASSERT_EQ(blob_file_numbers.size(), 1U);
+
+  std::vector<BlobLogRecord> records;
+  ASSERT_OK(
+      ReadBlobFileRecords(blob_file_numbers[0], /*num_records=*/3, &records));
+  ASSERT_EQ(records.size(), 3U);
+  ASSERT_EQ(records[0].key.ToString(), first_key);
+  ASSERT_EQ(records[1].key.ToString(), second_key);
+  ASSERT_EQ(records[2].key.ToString(), third_key);
+  ASSERT_EQ(records[0].value.size(), records[2].value.size());
+  ASSERT_EQ(records[1].value.size(), records[0].value.size() + 4);
+
+  Close();
+  Reopen(options);
+
+  ASSERT_EQ(Get(first_key), value);
+  ASSERT_EQ(Get(second_key), value);
+  ASSERT_EQ(Get(third_key), value);
+#endif  // ZSTD
+}
+
 TEST_F(DBBlobDirectWriteTest, DirectWriteFailedBatchTrackedAsInitialGarbage) {
   Options options = GetDirectWriteOptions();
   options.blob_file_size = 1 << 20;

db/db_impl/db_impl_write.cc

@@ -509,6 +509,7 @@ struct DBImpl::BlobDirectWriteContext {
         sv->cfd->ioptions().enable_blob_direct_write;
     settings.min_blob_size = sv->mutable_cf_options.min_blob_size;
     settings.compression_type = sv->mutable_cf_options.blob_compression_type;
+    settings.compression_opts = sv->mutable_cf_options.blob_compression_opts;
     settings.blob_cache = sv->cfd->ioptions().blob_cache.get();
     settings.prepopulate_blob_cache =
         sv->mutable_cf_options.prepopulate_blob_cache;