Add BlockBasedTablePinningTier enum and setter methods (#1024)

JackDrogon · web-flow · commit 8c47b014bc39 · 2025-07-31T08:35:02.000+01:00
Implements missing block cache pinning tier controls to align with C API.
Adds BlockBasedTablePinningTier enum with four variants (Fallback, None,
FlushAndSimilar, All) and corresponding setter methods in BlockBasedOptions:

- set_top_level_index_pinning_tier()
- set_partition_pinning_tier()
- set_unpartitioned_pinning_tier()

This provides fine-grained control over metadata block cache memory
management and pinning strategies for improved performance tuning.

Signed-off-by: Jack Drogon &lt;jack.xsuperman@gmail.com&gt;
diff --git a/src/db_options.rs b/src/db_options.rs
@@ -825,6 +825,80 @@ impl BlockBasedOptions {
             );
         }
     }
+
+    /// Set the top-level index pinning tier.
+    ///
+    /// Controls when top-level index blocks are pinned in block cache memory.
+    /// This affects memory usage and lookup performance for large databases with
+    /// multiple levels.
+    ///
+    /// Default: `BlockBasedTablePinningTier::Fallback`
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use rocksdb::{BlockBasedOptions, BlockBasedTablePinningTier};
+    ///
+    /// let mut opts = BlockBasedOptions::default();
+    /// opts.set_top_level_index_pinning_tier(BlockBasedTablePinningTier::FlushAndSimilar);
+    /// ```
+    pub fn set_top_level_index_pinning_tier(&mut self, pinning_tier: BlockBasedTablePinningTier) {
+        unsafe {
+            ffi::rocksdb_block_based_options_set_top_level_index_pinning_tier(
+                self.inner,
+                pinning_tier as c_int,
+            );
+        }
+    }
+
+    /// Set the partition pinning tier.
+    ///
+    /// Controls when partition blocks (used in partitioned indexes and filters)
+    /// are pinned in block cache memory. This affects performance for databases
+    /// using partitioned metadata.
+    ///
+    /// Default: `BlockBasedTablePinningTier::Fallback`
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use rocksdb::{BlockBasedOptions, BlockBasedTablePinningTier};
+    ///
+    /// let mut opts = BlockBasedOptions::default();
+    /// opts.set_partition_pinning_tier(BlockBasedTablePinningTier::All);
+    /// ```
+    pub fn set_partition_pinning_tier(&mut self, pinning_tier: BlockBasedTablePinningTier) {
+        unsafe {
+            ffi::rocksdb_block_based_options_set_partition_pinning_tier(
+                self.inner,
+                pinning_tier as c_int,
+            );
+        }
+    }
+
+    /// Set the unpartitioned pinning tier.
+    ///
+    /// Controls when unpartitioned metadata blocks (index and filter blocks that
+    /// are not partitioned) are pinned in block cache memory.
+    ///
+    /// Default: `BlockBasedTablePinningTier::Fallback`
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use rocksdb::{BlockBasedOptions, BlockBasedTablePinningTier};
+    ///
+    /// let mut opts = BlockBasedOptions::default();
+    /// opts.set_unpartitioned_pinning_tier(BlockBasedTablePinningTier::None);
+    /// ```
+    pub fn set_unpartitioned_pinning_tier(&mut self, pinning_tier: BlockBasedTablePinningTier) {
+        unsafe {
+            ffi::rocksdb_block_based_options_set_unpartitioned_pinning_tier(
+                self.inner,
+                pinning_tier as c_int,
+            );
+        }
+    }
 }
 
 impl Default for BlockBasedOptions {
@@ -4232,6 +4306,20 @@ pub enum DataBlockIndexType {
     BinaryAndHash = 1,
 }
 
+/// Used by BlockBasedOptions for setting metadata cache pinning tiers.
+/// Controls how metadata blocks (index, filter, etc.) are pinned in block cache.
+#[repr(C)]
+pub enum BlockBasedTablePinningTier {
+    /// Use fallback pinning tier (context-dependent)
+    Fallback = ffi::rocksdb_block_based_k_fallback_pinning_tier as isize,
+    /// No pinning - blocks can be evicted at any time
+    None = ffi::rocksdb_block_based_k_none_pinning_tier as isize,
+    /// Pin blocks for flushed files and similar scenarios
+    FlushAndSimilar = ffi::rocksdb_block_based_k_flush_and_similar_pinning_tier as isize,
+    /// Pin all blocks (highest priority)
+    All = ffi::rocksdb_block_based_k_all_pinning_tier as isize,
+}
+
 /// Defines the underlying memtable implementation.
 /// See official [wiki](https://github.com/facebook/rocksdb/wiki/MemTable) for more information.
 pub enum MemtableFactory {
diff --git a/src/lib.rs b/src/lib.rs
@@ -121,11 +121,12 @@ pub use crate::{
         DBWALIterator, Direction, IteratorMode,
     },
     db_options::{
-        BlockBasedIndexType, BlockBasedOptions, BottommostLevelCompaction, Cache, ChecksumType,
-        CompactOptions, CompactionPri, CuckooTableOptions, DBCompactionStyle, DBCompressionType,
-        DBPath, DBRecoveryMode, DataBlockIndexType, FifoCompactOptions, FlushOptions,
-        IngestExternalFileOptions, KeyEncodingType, LogLevel, LruCacheOptions, MemtableFactory,
-        Options, PlainTableFactoryOptions, ReadOptions, ReadTier, UniversalCompactOptions,
+        BlockBasedIndexType, BlockBasedOptions, BlockBasedTablePinningTier,
+        BottommostLevelCompaction, Cache, ChecksumType, CompactOptions, CompactionPri,
+        CuckooTableOptions, DBCompactionStyle, DBCompressionType, DBPath, DBRecoveryMode,
+        DataBlockIndexType, FifoCompactOptions, FlushOptions, IngestExternalFileOptions,
+        KeyEncodingType, LogLevel, LruCacheOptions, MemtableFactory, Options,
+        PlainTableFactoryOptions, ReadOptions, ReadTier, UniversalCompactOptions,
         UniversalCompactionStopStyle, WaitForCompactOptions, WriteBufferManager, WriteOptions,
     },
     db_pinnable_slice::DBPinnableSlice,
diff --git a/tests/test_rocksdb_options.rs b/tests/test_rocksdb_options.rs
@@ -18,8 +18,8 @@ use std::{fs, io::Read as _};
 
 use rocksdb::checkpoint::Checkpoint;
 use rocksdb::{
-    BlockBasedOptions, Cache, DBCompressionType, DataBlockIndexType, Env, LruCacheOptions, Options,
-    ReadOptions, DB,
+    BlockBasedOptions, BlockBasedTablePinningTier, Cache, DBCompressionType, DataBlockIndexType,
+    Env, LruCacheOptions, Options, ReadOptions, DB,
 };
 use util::DBPath;
 
@@ -423,3 +423,35 @@ fn test_set_write_dbid_to_manifest() {
         String::from_utf8_lossy(&db_checkpoint_id)
     );
 }
+
+#[test]
+fn test_block_based_table_pinning_tier() {
+    let path = DBPath::new("_block_based_table_pinning_tier");
+
+    // Test that we can create all enum variants
+    let _fallback = BlockBasedTablePinningTier::Fallback;
+    let _none = BlockBasedTablePinningTier::None;
+    let _flush_and_similar = BlockBasedTablePinningTier::FlushAndSimilar;
+    let _all = BlockBasedTablePinningTier::All;
+
+    // Test that enum values match expected C constants
+    assert_eq!(BlockBasedTablePinningTier::Fallback as i32, 0);
+    assert_eq!(BlockBasedTablePinningTier::None as i32, 1);
+    assert_eq!(BlockBasedTablePinningTier::FlushAndSimilar as i32, 2);
+    assert_eq!(BlockBasedTablePinningTier::All as i32, 3);
+
+    // Test that we can use the setter methods with BlockBasedOptions
+    let mut block_opts = BlockBasedOptions::default();
+    block_opts.set_top_level_index_pinning_tier(BlockBasedTablePinningTier::FlushAndSimilar);
+    block_opts.set_partition_pinning_tier(BlockBasedTablePinningTier::All);
+    block_opts.set_unpartitioned_pinning_tier(BlockBasedTablePinningTier::None);
+
+    // Test that we can create a database with these options
+    let mut opts = Options::default();
+    opts.create_if_missing(true);
+    opts.set_block_based_table_factory(&block_opts);
+
+    let db = DB::open(&opts, &path).unwrap();
+    db.put(b"test_key", b"test_value").unwrap();
+    assert_eq!(&*db.get(b"test_key").unwrap().unwrap(), b"test_value");
+}
