Cache to have an option to fail Cache::Insert() when full

Summary:
Cache to have an option to fail Cache::Insert() when full. Update call sites to check status and handle error.

I totally have no idea what's correct behavior of all the call sites when they encounter error. Please let me know if you see something wrong or more unit test is needed.

Test Plan: make check -j32, see tests pass.

Reviewers: anthony, yhchiang, andrewkr, IslamAbdelRahman, kradhakrishnan, sdong

Reviewed By: sdong

Subscribers: andrewkr, dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D54705

CMakeLists.txt

@@ -338,6 +338,7 @@ set(TESTS
         db/db_tailing_iter_test.cc
         db/db_test.cc
         db/db_test2.cc
+        db/db_block_cache_test.cc
         db/db_universal_compaction_test.cc
         db/db_wal_test.cc
         db/dbformat_test.cc

HISTORY.md

@@ -2,6 +2,7 @@
 ## Unreleased
 ### Public API Changes
 * Change default of BlockBasedTableOptions.format_version to 2. It means default DB created by 4.6 or up cannot be opened by RocksDB version 3.9 or earlier.
+* Added strict_capacity_limit option to NewLRUCache. If the flag is set to true, insert to cache will fail if no enough capacity can be free. Signiture of Cache::Insert() is updated accordingly.
 ### New Features
 * Add CompactionPri::kMinOverlappingRatio, a compaction picking mode friendly to write amplification.
 * Deprecate Iterator::IsKeyPinned() and replace it with Iterator::GetProperty() with prop_name="rocksdb.iterator.is.key.pinned"

Makefile

@@ -246,6 +246,7 @@ BENCHTOOLOBJECTS = $(BENCH_SOURCES:.cc=.o) $(LIBOBJECTS) $(TESTUTIL)
 TESTS = \
 	db_test \
 	db_test2 \
+	db_block_cache_test \
 	db_iter_test \
 	db_log_iter_test \
 	db_compaction_filter_test \
@@ -794,6 +795,9 @@ db_test: db/db_test.o db/db_test_util.o $(LIBOBJECTS) $(TESTHARNESS)
 db_test2: db/db_test2.o db/db_test_util.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_LINK)
 
+db_block_cache_test: db/db_block_cache_test.o db/db_test_util.o $(LIBOBJECTS) $(TESTHARNESS)
+	$(AM_LINK)
+
 db_log_iter_test: db/db_log_iter_test.o db/db_test_util.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_LINK)
 

db/db_block_cache_test.cc

@@ -0,0 +1,237 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under the BSD-style license found in the
+//  LICENSE file in the root directory of this source tree. An additional grant
+//  of patent rights can be found in the PATENTS file in the same directory.
+//
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+#include <cstdlib>
+#include "db/db_test_util.h"
+#include "port/stack_trace.h"
+
+namespace rocksdb {
+
+static uint64_t TestGetTickerCount(const Options& options,
+                                   Tickers ticker_type) {
+  return options.statistics->getTickerCount(ticker_type);
+}
+
+class DBBlockCacheTest : public DBTestBase {
+ private:
+  size_t miss_count_ = 0;
+  size_t hit_count_ = 0;
+  size_t insert_count_ = 0;
+  size_t failure_count_ = 0;
+  size_t compressed_miss_count_ = 0;
+  size_t compressed_hit_count_ = 0;
+  size_t compressed_insert_count_ = 0;
+  size_t compressed_failure_count_ = 0;
+
+ public:
+  const size_t kNumBlocks = 10;
+  const size_t kValueSize = 100;
+
+  DBBlockCacheTest() : DBTestBase("/db_block_cache_test") {}
+
+  BlockBasedTableOptions GetTableOptions() {
+    BlockBasedTableOptions table_options;
+    // Set a small enough block size so that each key-value get its own block.
+    table_options.block_size = 1;
+    return table_options;
+  }
+
+  Options GetOptions(const BlockBasedTableOptions& table_options) {
+    Options options = CurrentOptions();
+    options.create_if_missing = true;
+    // options.compression = kNoCompression;
+    options.statistics = rocksdb::CreateDBStatistics();
+    options.table_factory.reset(new BlockBasedTableFactory(table_options));
+    return options;
+  }
+
+  void InitTable(const Options& options) {
+    std::string value(kValueSize, 'a');
+    for (size_t i = 0; i < kNumBlocks; i++) {
+      ASSERT_OK(Put(ToString(i), value.c_str()));
+    }
+  }
+
+  void RecordCacheCounters(const Options& options) {
+    miss_count_ = TestGetTickerCount(options, BLOCK_CACHE_MISS);
+    hit_count_ = TestGetTickerCount(options, BLOCK_CACHE_HIT);
+    insert_count_ = TestGetTickerCount(options, BLOCK_CACHE_ADD);
+    failure_count_ = TestGetTickerCount(options, BLOCK_CACHE_ADD_FAILURES);
+    compressed_miss_count_ =
+        TestGetTickerCount(options, BLOCK_CACHE_COMPRESSED_MISS);
+    compressed_hit_count_ =
+        TestGetTickerCount(options, BLOCK_CACHE_COMPRESSED_HIT);
+    compressed_insert_count_ =
+        TestGetTickerCount(options, BLOCK_CACHE_COMPRESSED_ADD);
+    compressed_failure_count_ =
+        TestGetTickerCount(options, BLOCK_CACHE_COMPRESSED_ADD_FAILURES);
+  }
+
+  void CheckCacheCounters(const Options& options, size_t expected_misses,
+                          size_t expected_hits, size_t expected_inserts,
+                          size_t expected_failures) {
+    size_t new_miss_count = TestGetTickerCount(options, BLOCK_CACHE_MISS);
+    size_t new_hit_count = TestGetTickerCount(options, BLOCK_CACHE_HIT);
+    size_t new_insert_count = TestGetTickerCount(options, BLOCK_CACHE_ADD);
+    size_t new_failure_count =
+        TestGetTickerCount(options, BLOCK_CACHE_ADD_FAILURES);
+    ASSERT_EQ(miss_count_ + expected_misses, new_miss_count);
+    ASSERT_EQ(hit_count_ + expected_hits, new_hit_count);
+    ASSERT_EQ(insert_count_ + expected_inserts, new_insert_count);
+    ASSERT_EQ(failure_count_ + expected_failures, new_failure_count);
+    miss_count_ = new_miss_count;
+    hit_count_ = new_hit_count;
+    insert_count_ = new_insert_count;
+    failure_count_ = new_failure_count;
+  }
+
+  void CheckCompressedCacheCounters(const Options& options,
+                                    size_t expected_misses,
+                                    size_t expected_hits,
+                                    size_t expected_inserts,
+                                    size_t expected_failures) {
+    size_t new_miss_count =
+        TestGetTickerCount(options, BLOCK_CACHE_COMPRESSED_MISS);
+    size_t new_hit_count =
+        TestGetTickerCount(options, BLOCK_CACHE_COMPRESSED_HIT);
+    size_t new_insert_count =
+        TestGetTickerCount(options, BLOCK_CACHE_COMPRESSED_ADD);
+    size_t new_failure_count =
+        TestGetTickerCount(options, BLOCK_CACHE_COMPRESSED_ADD_FAILURES);
+    ASSERT_EQ(compressed_miss_count_ + expected_misses, new_miss_count);
+    ASSERT_EQ(compressed_hit_count_ + expected_hits, new_hit_count);
+    ASSERT_EQ(compressed_insert_count_ + expected_inserts, new_insert_count);
+    ASSERT_EQ(compressed_failure_count_ + expected_failures, new_failure_count);
+    compressed_miss_count_ = new_miss_count;
+    compressed_hit_count_ = new_hit_count;
+    compressed_insert_count_ = new_insert_count;
+    compressed_failure_count_ = new_failure_count;
+  }
+};
+
+TEST_F(DBBlockCacheTest, TestWithoutCompressedBlockCache) {
+  ReadOptions read_options;
+  auto table_options = GetTableOptions();
+  auto options = GetOptions(table_options);
+  InitTable(options);
+
+  std::shared_ptr<Cache> cache = NewLRUCache(0, 0, false);
+  table_options.block_cache = cache;
+  options.table_factory.reset(new BlockBasedTableFactory(table_options));
+  Reopen(options);
+  RecordCacheCounters(options);
+
+  std::vector<std::unique_ptr<Iterator>> iterators(kNumBlocks - 1);
+  Iterator* iter = nullptr;
+
+  // Load blocks into cache.
+  for (size_t i = 0; i < kNumBlocks - 1; i++) {
+    iter = db_->NewIterator(read_options);
+    iter->Seek(ToString(i));
+    ASSERT_OK(iter->status());
+    CheckCacheCounters(options, 1, 0, 1, 0);
+    iterators[i].reset(iter);
+  }
+  size_t usage = cache->GetUsage();
+  ASSERT_LT(0, usage);
+  cache->SetCapacity(usage);
+  ASSERT_EQ(usage, cache->GetPinnedUsage());
+
+  // Test with strict capacity limit.
+  cache->SetStrictCapacityLimit(true);
+  iter = db_->NewIterator(read_options);
+  iter->Seek(ToString(kNumBlocks - 1));
+  ASSERT_TRUE(iter->status().IsIncomplete());
+  CheckCacheCounters(options, 1, 0, 0, 1);
+  delete iter;
+  iter = nullptr;
+
+  // Release interators and access cache again.
+  for (size_t i = 0; i < kNumBlocks - 1; i++) {
+    iterators[i].reset();
+    CheckCacheCounters(options, 0, 0, 0, 0);
+  }
+  ASSERT_EQ(0, cache->GetPinnedUsage());
+  for (size_t i = 0; i < kNumBlocks - 1; i++) {
+    iter = db_->NewIterator(read_options);
+    iter->Seek(ToString(i));
+    ASSERT_OK(iter->status());
+    CheckCacheCounters(options, 0, 1, 0, 0);
+    iterators[i].reset(iter);
+  }
+}
+
+TEST_F(DBBlockCacheTest, TestWithCompressedBlockCache) {
+  ReadOptions read_options;
+  auto table_options = GetTableOptions();
+  auto options = GetOptions(table_options);
+  InitTable(options);
+
+  std::shared_ptr<Cache> cache = NewLRUCache(0, 0, false);
+  std::shared_ptr<Cache> compressed_cache = NewLRUCache(0, 0, false);
+  table_options.block_cache = cache;
+  table_options.block_cache_compressed = compressed_cache;
+  options.table_factory.reset(new BlockBasedTableFactory(table_options));
+  Reopen(options);
+  RecordCacheCounters(options);
+
+  std::vector<std::unique_ptr<Iterator>> iterators(kNumBlocks - 1);
+  Iterator* iter = nullptr;
+
+  // Load blocks into cache.
+  for (size_t i = 0; i < kNumBlocks - 1; i++) {
+    iter = db_->NewIterator(read_options);
+    iter->Seek(ToString(i));
+    ASSERT_OK(iter->status());
+    CheckCacheCounters(options, 1, 0, 1, 0);
+    CheckCompressedCacheCounters(options, 1, 0, 1, 0);
+    iterators[i].reset(iter);
+  }
+  size_t usage = cache->GetUsage();
+  ASSERT_LT(0, usage);
+  ASSERT_EQ(usage, cache->GetPinnedUsage());
+  size_t compressed_usage = compressed_cache->GetUsage();
+  ASSERT_LT(0, compressed_usage);
+  // Compressed block cache cannot be pinned.
+  ASSERT_EQ(0, compressed_cache->GetPinnedUsage());
+
+  // Set strict capacity limit flag. Now block will only load into compressed
+  // block cache.
+  cache->SetCapacity(usage);
+  cache->SetStrictCapacityLimit(true);
+  ASSERT_EQ(usage, cache->GetPinnedUsage());
+  // compressed_cache->SetCapacity(compressed_usage);
+  compressed_cache->SetCapacity(0);
+  // compressed_cache->SetStrictCapacityLimit(true);
+  iter = db_->NewIterator(read_options);
+  iter->Seek(ToString(kNumBlocks - 1));
+  ASSERT_TRUE(iter->status().IsIncomplete());
+  CheckCacheCounters(options, 1, 0, 0, 1);
+  CheckCompressedCacheCounters(options, 1, 0, 1, 0);
+  delete iter;
+  iter = nullptr;
+
+  // Clear strict capacity limit flag. This time we shall hit compressed block
+  // cache.
+  cache->SetStrictCapacityLimit(false);
+  iter = db_->NewIterator(read_options);
+  iter->Seek(ToString(kNumBlocks - 1));
+  ASSERT_OK(iter->status());
+  CheckCacheCounters(options, 1, 0, 1, 0);
+  CheckCompressedCacheCounters(options, 0, 1, 0, 0);
+  delete iter;
+  iter = nullptr;
+}
+
+}  // namespace rocksdb
+
+int main(int argc, char** argv) {
+  rocksdb::port::InstallStackTraceHandler();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

db/table_cache.cc

@@ -143,8 +143,12 @@ Status TableCache::FindTable(const EnvOptions& env_options,
       // We do not cache error results so that if the error is transient,
       // or somebody repairs the file, we recover automatically.
     } else {
-      *handle = cache_->Insert(key, table_reader.release(), 1,
-                               &DeleteEntry<TableReader>);
+      s = cache_->Insert(key, table_reader.get(), 1, &DeleteEntry<TableReader>,
+                         handle);
+      if (s.ok()) {
+        // Release ownership of table reader.
+        table_reader.release();
+      }
     }
   }
   return s;
@@ -285,9 +289,8 @@ Status TableCache::Get(const ReadOptions& options,
     size_t charge =
         row_cache_key.Size() + row_cache_entry->size() + sizeof(std::string);
     void* row_ptr = new std::string(std::move(*row_cache_entry));
-    auto row_handle = ioptions_.row_cache->Insert(
-        row_cache_key.GetKey(), row_ptr, charge, &DeleteEntry<std::string>);
-    ioptions_.row_cache->Release(row_handle);
+    ioptions_.row_cache->Insert(row_cache_key.GetKey(), row_ptr, charge,
+                                &DeleteEntry<std::string>);
   }
 #endif  // ROCKSDB_LITE
 

include/rocksdb/cache.h

@@ -25,6 +25,7 @@
 #include <memory>
 #include <stdint.h>
 #include "rocksdb/slice.h"
+#include "rocksdb/status.h"
 
 namespace rocksdb {
 
@@ -33,12 +34,15 @@ using std::shared_ptr;
 class Cache;
 
 // Create a new cache with a fixed size capacity. The cache is sharded
-// to 2^numShardBits shards, by hash of the key. The total capacity
+// to 2^num_shard_bits shards, by hash of the key. The total capacity
 // is divided and evenly assigned to each shard.
 //
-// The functions without parameter numShardBits uses default value, which is 4
+// The parameter num_shard_bits defaults to 4, and strict_capacity_limit
+// defaults to false.
 extern shared_ptr<Cache> NewLRUCache(size_t capacity);
-extern shared_ptr<Cache> NewLRUCache(size_t capacity, int numShardBits);
+extern shared_ptr<Cache> NewLRUCache(size_t capacity, int num_shard_bits);
+extern shared_ptr<Cache> NewLRUCache(size_t capacity, int num_shard_bits,
+                                     bool strict_capacity_limit);
 
 class Cache {
  public:
@@ -55,15 +59,22 @@ class Cache {
 
   // Insert a mapping from key->value into the cache and assign it
   // the specified charge against the total cache capacity.
+  // If strict_capacity_limit is true and cache reaches its full capacity,
+  // return Status::Incomplete.
   //
-  // Returns a handle that corresponds to the mapping.  The caller
-  // must call this->Release(handle) when the returned mapping is no
-  // longer needed.
+  // If handle is not nullptr, returns a handle that corresponds to the
+  // mapping. The caller must call this->Release(handle) when the returned
+  // mapping is no longer needed. In case of error caller is responsible to
+  // cleanup the value (i.e. calling "deleter").
+  //
+  // If handle is nullptr, it is as if Release is called immediately after
+  // insert. In case of error value will be cleanup.
   //
   // When the inserted entry is no longer needed, the key and
   // value will be passed to "deleter".
-  virtual Handle* Insert(const Slice& key, void* value, size_t charge,
-                         void (*deleter)(const Slice& key, void* value)) = 0;
+  virtual Status Insert(const Slice& key, void* value, size_t charge,
+                        void (*deleter)(const Slice& key, void* value),
+                        Handle** handle = nullptr) = 0;
 
   // If the cache has no mapping for "key", returns nullptr.
   //
@@ -100,6 +111,14 @@ class Cache {
   // purge the released entries from the cache in order to lower the usage
   virtual void SetCapacity(size_t capacity) = 0;
 
+  // Set whether to return error on insertion when cache reaches its full
+  // capacity.
+  virtual void SetStrictCapacityLimit(bool strict_capacity_limit) = 0;
+
+  // Set whether to return error on insertion when cache reaches its full
+  // capacity.
+  virtual bool HasStrictCapacityLimit() const = 0;
+
   // returns the maximum configured capacity of the cache
   virtual size_t GetCapacity() const = 0;