From 01e88dfeb4df9934e3933bdd61aafcf576c89627 Mon Sep 17 00:00:00 2001
From: Levi Tamasi <ltamasi@fb.com>
Date: Tue, 30 Aug 2022 17:03:45 -0700
Subject: [PATCH] Support using cache warming with the secondary blob cache
 (#10603)

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

Test Plan: `make check`

Reviewed By: riversand963

Differential Revision: D39117952

Pulled By: ltamasi

fbshipit-source-id: 5e956fa2fc18974876a5c87686acb50718e0edb7
---
 db/blob/blob_file_builder.cc  | 14 ++++++--
 db/blob/db_blob_basic_test.cc | 61 +++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/db/blob/blob_file_builder.cc b/db/blob/blob_file_builder.cc
index a8e94b6f4..84d00ab3b 100644
--- a/db/blob/blob_file_builder.cc
+++ b/db/blob/blob_file_builder.cc
@@ -420,9 +420,17 @@ Status BlobFileBuilder::PutBlobIntoCacheIfNeeded(const Slice& blob,
         BlobContents::GetCacheItemHelper();
     assert(cache_item_helper);
 
-    s = blob_cache->Insert(key, buf.get(), buf->ApproximateMemoryUsage(),
-                           cache_item_helper->del_cb,
-                           nullptr /* cache_handle */, priority);
+    if (immutable_options_->lowest_used_cache_tier ==
+        CacheTier::kNonVolatileBlockTier) {
+      s = blob_cache->Insert(key, buf.get(), cache_item_helper,
+                             buf->ApproximateMemoryUsage(),
+                             nullptr /* cache_handle */, priority);
+    } else {
+      s = blob_cache->Insert(key, buf.get(), buf->ApproximateMemoryUsage(),
+                             cache_item_helper->del_cb,
+                             nullptr /* cache_handle */, priority);
+    }
+
     if (s.ok()) {
       RecordTick(statistics, BLOB_DB_CACHE_ADD);
       RecordTick(statistics, BLOB_DB_CACHE_BYTES_WRITE, buf->size());
diff --git a/db/blob/db_blob_basic_test.cc b/db/blob/db_blob_basic_test.cc
index e3fa4a669..df009bffa 100644
--- a/db/blob/db_blob_basic_test.cc
+++ b/db/blob/db_blob_basic_test.cc
@@ -7,6 +7,7 @@
 #include <sstream>
 #include <string>
 
+#include "cache/compressed_secondary_cache.h"
 #include "db/blob/blob_index.h"
 #include "db/blob/blob_log_format.h"
 #include "db/db_test_util.h"
@@ -1701,6 +1702,66 @@ TEST_F(DBBlobBasicTest, DynamicallyWarmCacheDuringFlush) {
 }
 #endif  // !ROCKSDB_LITE
 
+TEST_F(DBBlobBasicTest, WarmCacheWithBlobsSecondary) {
+  CompressedSecondaryCacheOptions secondary_cache_opts;
+  secondary_cache_opts.capacity = 1 << 20;
+  secondary_cache_opts.num_shard_bits = 0;
+  secondary_cache_opts.metadata_charge_policy = kDontChargeCacheMetadata;
+  secondary_cache_opts.compression_type = kNoCompression;
+
+  LRUCacheOptions primary_cache_opts;
+  primary_cache_opts.capacity = 1024;
+  primary_cache_opts.num_shard_bits = 0;
+  primary_cache_opts.metadata_charge_policy = kDontChargeCacheMetadata;
+  primary_cache_opts.secondary_cache =
+      NewCompressedSecondaryCache(secondary_cache_opts);
+
+  Options options = GetDefaultOptions();
+  options.create_if_missing = true;
+  options.statistics = CreateDBStatistics();
+  options.enable_blob_files = true;
+  options.blob_cache = NewLRUCache(primary_cache_opts);
+  options.prepopulate_blob_cache = PrepopulateBlobCache::kFlushOnly;
+
+  DestroyAndReopen(options);
+
+  // Note: only one of the two blobs fit in the primary cache at any given time.
+  constexpr char first_key[] = "foo";
+  constexpr size_t first_blob_size = 512;
+  const std::string first_blob(first_blob_size, 'a');
+
+  constexpr char second_key[] = "bar";
+  constexpr size_t second_blob_size = 768;
+  const std::string second_blob(second_blob_size, 'b');
+
+  // First blob gets inserted into primary cache during flush
+  ASSERT_OK(Put(first_key, first_blob));
+  ASSERT_OK(Flush());
+  ASSERT_EQ(options.statistics->getAndResetTickerCount(BLOB_DB_CACHE_ADD), 1);
+
+  // Second blob gets inserted into primary cache during flush, first blob gets
+  // evicted to secondary cache
+  ASSERT_OK(Put(second_key, second_blob));
+  ASSERT_OK(Flush());
+  ASSERT_EQ(options.statistics->getAndResetTickerCount(BLOB_DB_CACHE_ADD), 1);
+
+  // First blob gets promoted back to primary cache b/c of lookup, second blob
+  // gets evicted to secondary cache
+  ASSERT_EQ(Get(first_key), first_blob);
+  ASSERT_EQ(options.statistics->getAndResetTickerCount(BLOB_DB_CACHE_MISS), 0);
+  ASSERT_EQ(options.statistics->getAndResetTickerCount(BLOB_DB_CACHE_HIT), 1);
+  ASSERT_EQ(options.statistics->getAndResetTickerCount(SECONDARY_CACHE_HITS),
+            1);
+
+  // Second blob gets promoted back to primary cache b/c of lookup, first blob
+  // gets evicted to secondary cache
+  ASSERT_EQ(Get(second_key), second_blob);
+  ASSERT_EQ(options.statistics->getAndResetTickerCount(BLOB_DB_CACHE_MISS), 0);
+  ASSERT_EQ(options.statistics->getAndResetTickerCount(BLOB_DB_CACHE_HIT), 1);
+  ASSERT_EQ(options.statistics->getAndResetTickerCount(SECONDARY_CACHE_HITS),
+            1);
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {