@nxtedition/rocksdb 15.1.2 → 15.1.4

package/deps/rocksdb/rocksdb/table/block_based/block_based_table_builder.cc

@@ -110,6 +110,19 @@ FilterBlockBuilder* CreateFilterBlockBuilder(
   }
 }
 
+// A convenience function for populating the Compressor* fields; see ~Rep()
+Compressor* MaybeCloneSpecialized(
+    Compressor* compressor, CacheEntryRole block_type,
+    Compressor::DictSampleArgs&& dict_samples = {}) {
+  auto specialized =
+      compressor->MaybeCloneSpecialized(block_type, std::move(dict_samples));
+  if (specialized) {
+    // Caller is responsible for freeing when distinct
+    return specialized.release();
+  } else {
+    return compressor;
+  }
+}
 }  // namespace
 
 // kBlockBasedTableMagicNumber was picked by running
@@ -824,15 +837,17 @@ struct BlockBasedTableBuilder::Rep {
 
   // *** Compressors & decompressors - Yes, it seems like a lot here but ***
   // *** these are distinct fields to minimize extra conditionals and    ***
-  // *** field reads on hot code paths.                                  ***
+  // *** field reads on hot code paths. And to avoid interlocked         ***
+  // *** instructions associated with shared_ptr.                        ***
 
   // A compressor for blocks in general, without dictionary compression
   std::unique_ptr<Compressor> basic_compressor;
-  // A compressor using dictionary compression (when applicable)
-  std::unique_ptr<Compressor> compressor_with_dict;
-  // Once configured/determined, points to one of the above Compressors to
-  // use on data blocks.
-  Compressor* data_block_compressor = nullptr;
+  // A compressor for data blocks, which might be tuned differently and might
+  // use dictionary compression (when applicable). See ~Rep() for some details.
+  UnownedPtr<Compressor> data_block_compressor = nullptr;
+  // A compressor for index blocks, which might be tuned differently from
+  // basic_compressor. See ~Rep() for some details.
+  UnownedPtr<Compressor> index_block_compressor = nullptr;
   // A decompressor corresponding to basic_compressor (when non-nullptr).
   // Used for verification and cache warming.
   std::shared_ptr<Decompressor> basic_decompressor;
@@ -853,7 +868,7 @@ struct BlockBasedTableBuilder::Rep {
       compression_types_used;
 
   // Working area for basic_compressor when compression_parallel_threads==1
-  WorkingAreaPair basic_working_area;
+  WorkingAreaPair index_block_working_area;
   // Working area for data_block_compressor, for emit/compaction thread
   WorkingAreaPair data_block_working_area;
 
@@ -894,6 +909,7 @@ struct BlockBasedTableBuilder::Rep {
   std::unique_ptr<FilterBlockBuilder> filter_builder;
   OffsetableCacheKey base_cache_key;
   const TableFileCreationReason reason;
+  const bool target_file_size_is_upper_bound;
 
   BlockHandle pending_handle;  // Handle to add to index block
 
@@ -1041,6 +1057,8 @@ struct BlockBasedTableBuilder::Rep {
         use_delta_encoding_for_index_values(table_opt.format_version >= 4 &&
                                             !table_opt.block_align),
         reason(tbo.reason),
+        target_file_size_is_upper_bound(
+            tbo.moptions.target_file_size_is_upper_bound),
         flush_block_policy(
             table_options.flush_block_policy_factory->NewFlushBlockPolicy(
                 table_options, data_block)),
@@ -1096,7 +1114,10 @@ struct BlockBasedTableBuilder::Rep {
         filter_context, tbo.compression_opts, tbo.compression_type);
     if (basic_compressor) {
       if (table_options.enable_index_compression) {
-        basic_working_area.compress = basic_compressor->ObtainWorkingArea();
+        index_block_compressor = MaybeCloneSpecialized(
+            basic_compressor.get(), CacheEntryRole::kIndexBlock);
+        index_block_working_area.compress =
+            index_block_compressor->ObtainWorkingArea();
       }
       max_dict_sample_bytes = basic_compressor->GetMaxSampleSizeIfWantDict(
           CacheEntryRole::kDataBlock);
@@ -1111,8 +1132,10 @@ struct BlockBasedTableBuilder::Rep {
                                   tbo.compression_opts.max_dict_buffer_bytes);
         }
       } else {
-        // No distinct data block compressor using dictionary
-        data_block_compressor = basic_compressor.get();
+        // No distinct data block compressor using dictionary, but
+        // implementation might still want to specialize for data blocks
+        data_block_compressor = MaybeCloneSpecialized(
+            basic_compressor.get(), CacheEntryRole::kDataBlock);
         data_block_working_area.compress =
             data_block_compressor->ObtainWorkingArea();
       }
@@ -1126,8 +1149,9 @@ struct BlockBasedTableBuilder::Rep {
       if (table_options.verify_compression) {
         verify_decompressor = basic_decompressor.get();
         if (table_options.enable_index_compression) {
-          basic_working_area.verify = verify_decompressor->ObtainWorkingArea(
-              basic_compressor->GetPreferredCompressionType());
+          index_block_working_area.verify =
+              verify_decompressor->ObtainWorkingArea(
+                  index_block_compressor->GetPreferredCompressionType());
         }
         if (state == State::kUnbuffered) {
           assert(data_block_compressor);
@@ -1292,8 +1316,19 @@ struct BlockBasedTableBuilder::Rep {
   }
 
   ~Rep() {
+    // Delete working areas before their compressors.
+    index_block_working_area = {};
+    data_block_working_area = {};
     // Must have been cleaned up by StopParallelCompression
     assert(pc_rep == nullptr);
+    // Delete specialized compressors if they were distinct (avoiding extra
+    // fields and interlocked instructions with shared_ptr)
+    if (data_block_compressor.get() != basic_compressor.get()) {
+      delete data_block_compressor.get();
+    }
+    if (index_block_compressor.get() != basic_compressor.get()) {
+      delete index_block_compressor.get();
+    }
   }
 
   Rep(const Rep&) = delete;
@@ -1611,6 +1646,17 @@ void BlockBasedTableBuilder::Flush(const Slice* first_key_in_next_block) {
     rep_->data_begin_offset += uncompressed_block_data.size();
     MaybeEnterUnbuffered(first_key_in_next_block);
   } else {
+    // Increment num_data_blocks when a data block is finalized in the
+    // emit thread to avoid data races with write worker threads
+    ++r->props.num_data_blocks;
+
+    // Notify filter builder that a data block has been finalized
+    // This must happen on the emit thread before the block is added to the
+    // ring buffer to avoid race conditions with worker threads
+    if (r->filter_builder) {
+      r->filter_builder->OnDataBlockFinalized(r->props.num_data_blocks);
+    }
+
     if (r->IsParallelCompressionActive()) {
       EmitBlockForParallel(r->data_block.MutableBuffer(), r->last_ikey,
                            first_key_in_next_block);
@@ -1715,9 +1761,11 @@ void BlockBasedTableBuilder::WriteBlock(const Slice& uncompressed_block_data,
   assert(!r->IsParallelCompressionActive());
   CompressionType type;
   bool is_data_block = block_type == BlockType::kData;
+  // NOTE: only index and data blocks are currently compressed
+  assert(is_data_block || block_type == BlockType::kIndex);
   Status compress_status = CompressAndVerifyBlock(
       uncompressed_block_data, is_data_block,
-      is_data_block ? r->data_block_working_area : r->basic_working_area,
+      is_data_block ? r->data_block_working_area : r->index_block_working_area,
       &r->single_threaded_compressed_output, &type);
   r->SetStatus(compress_status);
   if (UNLIKELY(!ok())) {
@@ -1735,7 +1783,6 @@ void BlockBasedTableBuilder::WriteBlock(const Slice& uncompressed_block_data,
   if (is_data_block) {
     r->props.data_size = r->get_offset();
     r->props.uncompressed_data_size += uncompressed_block_data.size();
-    ++r->props.num_data_blocks;
   }
 }
 
@@ -1784,7 +1831,6 @@ void BlockBasedTableBuilder::BGWorker(WorkingAreaPair& working_area) {
       if (LIKELY(ios.ok())) {
         rep_->props.data_size = rep_->get_offset();
         rep_->props.uncompressed_data_size += block_rep->uncompressed.size();
-        ++rep_->props.num_data_blocks;
 
         rep_->index_builder->FinishIndexEntry(
             rep_->pending_handle, block_rep->prepared_index_entry.get(),
@@ -1833,13 +1879,13 @@ Status BlockBasedTableBuilder::CompressAndVerifyBlock(
   Rep* r = rep_.get();
   Status status;
 
-  Compressor* compressor = nullptr;
+  UnownedPtr<Compressor> compressor = nullptr;
   Decompressor* verify_decomp = nullptr;
   if (is_data_block) {
     compressor = r->data_block_compressor;
     verify_decomp = r->data_block_verify_decompressor.get();
   } else {
-    compressor = r->basic_compressor.get();
+    compressor = r->index_block_compressor;
     verify_decomp = r->verify_decompressor.get();
   }
 
@@ -1940,6 +1986,9 @@ void BlockBasedTableBuilder::WriteMaybeCompressedBlock(
     const Slice& block_contents, CompressionType comp_type, BlockHandle* handle,
     BlockType block_type, const Slice* uncompressed_block_data,
     bool* skip_delta_encoding) {
+  // Must have pre-checked status in single-threaded context
+  assert(status().ok());
+  assert(io_status().ok());
   rep_->SetIOStatus(WriteMaybeCompressedBlockImpl(
       block_contents, comp_type, handle, block_type, uncompressed_block_data,
       skip_delta_encoding));
@@ -2014,8 +2063,6 @@ IOStatus BlockBasedTableBuilder::WriteMaybeCompressedBlockImpl(
 
   handle->set_offset(offset);
   handle->set_size(block_contents.size());
-  assert(status().ok());
-  assert(io_status().ok());
   if (uncompressed_block_data == nullptr) {
     uncompressed_block_data = &block_contents;
     assert(comp_type == kNoCompression);
@@ -2103,7 +2150,7 @@ void BlockBasedTableBuilder::MaybeStartParallelCompression() {
   // that latency. So even with some optimizations, turning on the parallel
   // framework when compression is disabled just eats more CPU with little-to-no
   // improvement in throughput.
-  if (rep_->data_block_compressor == nullptr) {
+  if (!rep_->data_block_compressor) {
     // Force the generally best configuration for no compression: no parallelism
     return;
   }
@@ -2140,8 +2187,9 @@ void BlockBasedTableBuilder::StopParallelCompression(bool abort) {
     pc_rep.SetAbort(pc_rep.emit_thread_state);
   } else if (pc_rep.emit_thread_state !=
              ParallelCompressionRep::ThreadState::kEnd) {
-    // In case we didn't do a final flush with no next key
-    assert(rep_->props.num_data_blocks == 0);
+    // In case we didn't do a final flush with no next key, which might have
+    // been skipped if !ok() was set after the start of Finish()
+    assert(rep_->props.num_data_blocks == 0 || !ok());
     pc_rep.SetNoMoreToEmit(pc_rep.emit_thread_state, pc_rep.emit_slot);
   }
 #ifdef BBTB_PC_WATCHDOG
@@ -2449,8 +2497,8 @@ void BlockBasedTableBuilder::WritePropertiesBlock(
 void BlockBasedTableBuilder::WriteCompressionDictBlock(
     MetaIndexBuilder* meta_index_builder) {
   Slice compression_dict;
-  if (rep_->compressor_with_dict) {
-    compression_dict = rep_->compressor_with_dict->GetSerializedDict();
+  if (rep_->data_block_compressor) {
+    compression_dict = rep_->data_block_compressor->GetSerializedDict();
   }
   if (!compression_dict.empty()) {
     BlockHandle compression_dict_block_handle;
@@ -2545,6 +2593,7 @@ void BlockBasedTableBuilder::MaybeEnterUnbuffered(
     // The below code is neither safe nor necessary for handling zero data
     // blocks.
     // For PostPopulateCompressionProperties()
+    assert(!r->data_block_compressor);
     r->data_block_compressor = r->basic_compressor.get();
     return;
   }
@@ -2586,15 +2635,12 @@ void BlockBasedTableBuilder::MaybeEnterUnbuffered(
 
   assert(samples.sample_data.size() > 0);
 
-  // final sample data block flushed, now we can generate dictionary
-  r->compressor_with_dict = r->basic_compressor->MaybeCloneSpecialized(
-      CacheEntryRole::kDataBlock, std::move(samples));
+  // final sample data block flushed, now we can generate dictionary (or it
+  // might opt not to use a dictionary and that's ok)
+  r->data_block_compressor =
+      MaybeCloneSpecialized(r->basic_compressor.get(),
+                            CacheEntryRole::kDataBlock, std::move(samples));
 
-  // The compressor might opt not to use a dictionary, in which case we
-  // can use the same compressor as for e.g. index blocks.
-  r->data_block_compressor = r->compressor_with_dict
-                                 ? r->compressor_with_dict.get()
-                                 : r->basic_compressor.get();
   Slice serialized_dict = r->data_block_compressor->GetSerializedDict();
   if (r->verify_decompressor) {
     if (serialized_dict.empty()) {
@@ -2688,6 +2734,20 @@ void BlockBasedTableBuilder::MaybeEnterUnbuffered(
 Status BlockBasedTableBuilder::Finish() {
   Rep* r = rep_.get();
   assert(r->state != Rep::State::kClosed);
+
+#ifndef NDEBUG
+  {
+    // This sync point callback is a simple approximation of a failure detected
+    // in parallel compression after the start of calling Finish() but before
+    // Finish() calls Flush()
+    IOStatus s = rep_->GetIOStatus();
+    TEST_SYNC_POINT_CALLBACK("BlockBasedTableBuilder::Finish:ParallelIOStatus",
+                             &s);
+    if (!s.ok()) {
+      rep_->SetIOStatus(s);
+    }
+  }
+#endif  // !NDEBUG
   // To make sure properties block is able to keep the accurate size of index
   // block, we will finish writing all index entries first, in Flush().
   Flush(/*first_key_in_next_block=*/nullptr);
@@ -2701,6 +2761,8 @@ Status BlockBasedTableBuilder::Finish() {
 
   r->props.tail_start_offset = r->offset.LoadRelaxed();
 
+  uint64_t last_estimated_tail_size = EstimatedTailSize();
+
   // Write meta blocks, metaindex block and footer in the following order.
   //    1. [meta block: filter]
   //    2. [meta block: index]
@@ -2727,6 +2789,24 @@ Status BlockBasedTableBuilder::Finish() {
   r->state = Rep::State::kClosed;
   r->tail_size = r->offset.LoadRelaxed() - r->props.tail_start_offset;
 
+  // Assert tail size estimation is an overestimate only when tail size
+  // estimation option is enabled for compaction files with supported
+  // index/filter types:
+  // - Shortened indexes (kBinarySearch, kBinarySearchWithFirstKey)
+  // - Partitioned indexes (kTwoLevelIndexSearch)
+  // - Full filters
+  // - Partitioned filters
+  if (r->target_file_size_is_upper_bound &&
+      r->reason == TableFileCreationReason::kCompaction &&
+      r->table_options.index_type != BlockBasedTableOptions::kHashSearch) {
+    ROCKS_LOG_WARN(r->ioptions.info_log,
+                   "File number: %" PRIu64 ", Estimated tail size = %" PRIu64
+                   " bytes, Actual tail size = %" PRIu64 " bytes",
+                   r->props.orig_file_number, last_estimated_tail_size,
+                   r->tail_size);
+    assert(r->tail_size <= last_estimated_tail_size);
+  }
+
   return r->GetStatus();
 }
 
@@ -2764,6 +2844,49 @@ uint64_t BlockBasedTableBuilder::EstimatedFileSize() const {
   }
 }
 
+uint64_t BlockBasedTableBuilder::EstimatedTailSize() const {
+  uint64_t estimated_tail_size = 0;
+
+  // 1. Estimate index size
+  if (rep_->table_options.index_type ==
+      BlockBasedTableOptions::kTwoLevelIndexSearch) {
+    assert(rep_->p_index_builder_);
+    estimated_tail_size += rep_->p_index_builder_->CurrentIndexSizeEstimate();
+  } else {
+    assert(rep_->index_builder);
+    estimated_tail_size += rep_->index_builder->CurrentIndexSizeEstimate();
+  }
+
+  // 2. Estimate filter size
+  if (rep_->filter_builder) {
+    estimated_tail_size += rep_->filter_builder->CurrentFilterSizeEstimate();
+  }
+
+  // 3. Estimate compression dictionary size
+  if (rep_->data_block_compressor) {
+    Slice dict = rep_->data_block_compressor->GetSerializedDict();
+    if (!dict.empty()) {
+      estimated_tail_size += dict.size();
+    }
+  }
+
+  // 4. Estimate range deletion block size
+  if (!rep_->range_del_block.empty()) {
+    estimated_tail_size += rep_->range_del_block.CurrentSizeEstimate();
+  }
+
+  // 5. Estimate properties block size conservatively (~1-2KB)
+  estimated_tail_size += 2048;
+
+  // 6. Estimate meta-index block size conservatively (~1KB)
+  estimated_tail_size += 1024;
+
+  // 7. Add footer size
+  estimated_tail_size += Footer::kMaxEncodedLength;
+
+  return estimated_tail_size;
+}
+
 uint64_t BlockBasedTableBuilder::GetTailSize() const { return rep_->tail_size; }
 
 bool BlockBasedTableBuilder::NeedCompact() const {

package/deps/rocksdb/rocksdb/table/block_based/block_based_table_builder.h

@@ -89,11 +89,15 @@ class BlockBasedTableBuilder : public TableBuilder {
   // Finish() call, returns the size of the final generated file.
   uint64_t FileSize() const override;
 
-  // Estimated size of the file generated so far. This is used when
-  // FileSize() cannot estimate final SST size, e.g. parallel compression
-  // is enabled.
+  // Estimated size of the file generated so far (based on data blocks, this
+  // estimate does not include meta blocks). This is used when FileSize() cannot
+  // estimate final SST size, e.g. parallel compression is enabled.
   uint64_t EstimatedFileSize() const override;
 
+  // Estimated tail size of the SST file generated so far. The "tail" refers to
+  // all blocks written after data blocks (index + filter).
+  uint64_t EstimatedTailSize() const override;
+
   // Get the size of the "tail" part of a SST file. "Tail" refers to
   // all blocks after data blocks till the end of the SST file.
   uint64_t GetTailSize() const override;