From 9e24ed82920b5ccdc24cc0c446ec86f8b5778479 Mon Sep 17 00:00:00 2001
From: Matt Jurik <exabytes18@gmail.com>
Date: Sun, 7 Jul 2024 18:50:39 -0700
Subject: [PATCH] Add support for enabling blob cache (#898)

* Add support for enabling blob cache

---------

Co-authored-by: Matt Jurik <matt.jurik@hulu.com>
---
 src/db_options.rs             | 18 ++++++++++++++++++
 tests/test_rocksdb_options.rs | 24 ++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/src/db_options.rs b/src/db_options.rs
index dabf86d28..efddc69bf 100644
--- a/src/db_options.rs
+++ b/src/db_options.rs
@@ -210,6 +210,7 @@ impl Cache {
 pub(crate) struct OptionsMustOutliveDB {
     env: Option<Env>,
     row_cache: Option<Cache>,
+    blob_cache: Option<Cache>,
     block_based: Option<BlockBasedOptionsMustOutliveDB>,
     write_buffer_manager: Option<WriteBufferManager>,
 }
@@ -219,6 +220,7 @@ impl OptionsMustOutliveDB {
         Self {
             env: self.env.clone(),
             row_cache: self.row_cache.clone(),
+            blob_cache: self.blob_cache.clone(),
             block_based: self
                 .block_based
                 .as_ref()
@@ -3329,6 +3331,22 @@ impl Options {
         }
     }
 
+    /// Sets the blob cache.
+    ///
+    /// Using a dedicated object for blobs and using the same object for the block and blob caches
+    /// are both supported. In the latter case, note that blobs are less valuable from a caching
+    /// perspective than SST blocks, and some cache implementations have configuration options that
+    /// can be used to prioritize items accordingly (see Cache::Priority and
+    /// LRUCacheOptions::{high,low}_pri_pool_ratio).
+    ///
+    /// Default: disabled
+    pub fn set_blob_cache(&mut self, cache: &Cache) {
+        unsafe {
+            ffi::rocksdb_options_set_blob_cache(self.inner, cache.0.inner.as_ptr());
+        }
+        self.outlive.blob_cache = Some(cache.clone());
+    }
+
     /// Set this option to true during creation of database if you want
     /// to be able to ingest behind (call IngestExternalFile() skipping keys
     /// that already exist, rather than overwriting matching keys).
diff --git a/tests/test_rocksdb_options.rs b/tests/test_rocksdb_options.rs
index 2cc062a70..96374eb19 100644
--- a/tests/test_rocksdb_options.rs
+++ b/tests/test_rocksdb_options.rs
@@ -313,3 +313,27 @@ fn test_set_ratelimiter() {
         assert_eq!(&*db.get(b"k2").unwrap().unwrap(), b"a");
     }
 }
+
+#[test]
+fn test_set_blob_cache() {
+    let path = DBPath::new("_set_blob_cache");
+    let cache = Cache::new_hyper_clock_cache(1024 * 1024, 4 * 1024);
+
+    let mut opts = Options::default();
+    opts.create_if_missing(true);
+    opts.set_enable_blob_files(true);
+    opts.set_min_blob_size(16);
+    opts.set_blob_cache(&cache);
+
+    let db = DB::open(&opts, &path).unwrap();
+
+    const KEY: &[u8] = b"k1";
+    const VALUE: &[u8] = b"01234567890123456789";
+    db.put(KEY, VALUE).unwrap();
+
+    // Cache miss
+    assert_eq!(&*db.get(KEY).unwrap().unwrap(), VALUE);
+
+    // Cache hit
+    assert_eq!(&*db.get(KEY).unwrap().unwrap(), VALUE);
+}