diff --git a/include/rocksdb/cache.h b/include/rocksdb/cache.h index 201d82f19..c8542f072 100644 --- a/include/rocksdb/cache.h +++ b/include/rocksdb/cache.h @@ -104,6 +104,9 @@ class Cache { // returns the maximum configured capacity of the cache virtual size_t GetCapacity() const = 0; + // returns the memory size for the entries residing in the cache. + virtual size_t GetUsage() const = 0; + private: void LRU_Remove(Handle* e); void LRU_Append(Handle* e); diff --git a/util/cache.cc b/util/cache.cc index b9d41be49..8fceefc82 100644 --- a/util/cache.cc +++ b/util/cache.cc @@ -156,6 +156,7 @@ class LRUCache { Cache::Handle* Lookup(const Slice& key, uint32_t hash); void Release(Cache::Handle* handle); void Erase(const Slice& key, uint32_t hash); + size_t GetUsage() const { return usage_.load(std::memory_order_relaxed); } private: void LRU_Remove(LRUHandle* e); @@ -172,7 +173,7 @@ class LRUCache { // mutex_ protects the following state. port::Mutex mutex_; - size_t usage_; + std::atomic_size_t usage_; // Dummy head of LRU list. // lru.prev is newest entry, lru.next is oldest entry. @@ -214,7 +215,7 @@ void LRUCache::FreeEntry(LRUHandle* e) { void LRUCache::LRU_Remove(LRUHandle* e) { e->next->prev = e->prev; e->prev->next = e->next; - usage_ -= e->charge; + usage_.fetch_sub(e->charge, std::memory_order_relaxed); } void LRUCache::LRU_Append(LRUHandle* e) { @@ -223,7 +224,7 @@ void LRUCache::LRU_Append(LRUHandle* e) { e->prev = lru_.prev; e->prev->next = e; e->next->prev = e; - usage_ += e->charge; + usage_.fetch_add(e->charge, std::memory_order_relaxed); } Cache::Handle* LRUCache::Lookup(const Slice& key, uint32_t hash) { @@ -282,7 +283,7 @@ Cache::Handle* LRUCache::Insert( // referenced by the cache first. LRUHandle* cur = lru_.next; for (unsigned int scanCount = 0; - usage_ > capacity_ && cur != &lru_ + GetUsage() > capacity_ && cur != &lru_ && scanCount < remove_scan_count_limit_; scanCount++) { LRUHandle* next = cur->next; if (cur->refs <= 1) { @@ -298,7 +299,7 @@ Cache::Handle* LRUCache::Insert( // Free the space following strict LRU policy until enough space // is freed. - while (usage_ > capacity_ && lru_.next != &lru_) { + while (GetUsage() > capacity_ && lru_.next != &lru_) { LRUHandle* old = lru_.next; LRU_Remove(old); table_.Remove(old->key(), old->hash); @@ -340,10 +341,10 @@ static int kRemoveScanCountLimit = 0; // default values, can be overridden class ShardedLRUCache : public Cache { private: - LRUCache* shard_; + LRUCache* shards_; port::Mutex id_mutex_; uint64_t last_id_; - int numShardBits; + int num_shard_bits_; size_t capacity_; static inline uint32_t HashSlice(const Slice& s) { @@ -352,18 +353,18 @@ class ShardedLRUCache : public Cache { uint32_t Shard(uint32_t hash) { // Note, hash >> 32 yields hash in gcc, not the zero we expect! - return (numShardBits > 0) ? (hash >> (32 - numShardBits)) : 0; + return (num_shard_bits_ > 0) ? (hash >> (32 - num_shard_bits_)) : 0; } void init(size_t capacity, int numbits, int removeScanCountLimit) { - numShardBits = numbits; + num_shard_bits_ = numbits; capacity_ = capacity; - int numShards = 1 << numShardBits; - shard_ = new LRUCache[numShards]; - const size_t per_shard = (capacity + (numShards - 1)) / numShards; - for (int s = 0; s < numShards; s++) { - shard_[s].SetCapacity(per_shard); - shard_[s].SetRemoveScanCountLimit(removeScanCountLimit); + int num_shards = 1 << num_shard_bits_; + shards_ = new LRUCache[num_shards]; + const size_t per_shard = (capacity + (num_shards - 1)) / num_shards; + for (int s = 0; s < num_shards; s++) { + shards_[s].SetCapacity(per_shard); + shards_[s].SetRemoveScanCountLimit(removeScanCountLimit); } } @@ -372,30 +373,30 @@ class ShardedLRUCache : public Cache { : last_id_(0) { init(capacity, kNumShardBits, kRemoveScanCountLimit); } - ShardedLRUCache(size_t capacity, int numShardBits, + ShardedLRUCache(size_t capacity, int num_shard_bits, int removeScanCountLimit) : last_id_(0) { - init(capacity, numShardBits, removeScanCountLimit); + init(capacity, num_shard_bits, removeScanCountLimit); } virtual ~ShardedLRUCache() { - delete[] shard_; + delete[] shards_; } virtual Handle* Insert(const Slice& key, void* value, size_t charge, void (*deleter)(const Slice& key, void* value)) { const uint32_t hash = HashSlice(key); - return shard_[Shard(hash)].Insert(key, hash, value, charge, deleter); + return shards_[Shard(hash)].Insert(key, hash, value, charge, deleter); } virtual Handle* Lookup(const Slice& key) { const uint32_t hash = HashSlice(key); - return shard_[Shard(hash)].Lookup(key, hash); + return shards_[Shard(hash)].Lookup(key, hash); } virtual void Release(Handle* handle) { LRUHandle* h = reinterpret_cast(handle); - shard_[Shard(h->hash)].Release(handle); + shards_[Shard(h->hash)].Release(handle); } virtual void Erase(const Slice& key) { const uint32_t hash = HashSlice(key); - shard_[Shard(hash)].Erase(key, hash); + shards_[Shard(hash)].Erase(key, hash); } virtual void* Value(Handle* handle) { return reinterpret_cast(handle)->value; @@ -407,6 +408,16 @@ class ShardedLRUCache : public Cache { virtual size_t GetCapacity() const { return capacity_; } + virtual size_t GetUsage() const { + // We will not lock the cache when getting the usage from shards. + // for (size_t i = 0; i < num_shard_bits_; ++i) + int num_shards = 1 << num_shard_bits_; + size_t usage = 0; + for (int s = 0; s < num_shards; s++) { + usage += shards_[s].GetUsage(); + } + return usage; + } }; } // end anonymous namespace @@ -415,17 +426,17 @@ shared_ptr NewLRUCache(size_t capacity) { return NewLRUCache(capacity, kNumShardBits); } -shared_ptr NewLRUCache(size_t capacity, int numShardBits) { - return NewLRUCache(capacity, numShardBits, kRemoveScanCountLimit); +shared_ptr NewLRUCache(size_t capacity, int num_shard_bits) { + return NewLRUCache(capacity, num_shard_bits, kRemoveScanCountLimit); } -shared_ptr NewLRUCache(size_t capacity, int numShardBits, +shared_ptr NewLRUCache(size_t capacity, int num_shard_bits, int removeScanCountLimit) { - if (numShardBits >= 20) { + if (num_shard_bits >= 20) { return nullptr; // the cache cannot be sharded into too many fine pieces } return std::make_shared(capacity, - numShardBits, + num_shard_bits, removeScanCountLimit); } diff --git a/util/cache_test.cc b/util/cache_test.cc index 87ab91389..2e10bdc3a 100644 --- a/util/cache_test.cc +++ b/util/cache_test.cc @@ -107,6 +107,39 @@ class CacheTest { }; CacheTest* CacheTest::current_; +void dumbDeleter(const Slice& key, void* value) { } + +TEST(CacheTest, UsageTest) { + // cache is shared_ptr and will be automatically cleaned up. + const uint64_t kCapacity = 100000; + auto cache = NewLRUCache(kCapacity, 8, 200); + + size_t usage = 0; + const char* value = "abcdef"; + // make sure everything will be cached + for (int i = 1; i < 100; ++i) { + std::string key(i, 'a'); + auto kv_size = key.size() + 5; + cache->Release( + cache->Insert(key, (void*)value, kv_size, dumbDeleter) + ); + usage += kv_size; + ASSERT_EQ(usage, cache->GetUsage()); + } + + // make sure the cache will be overloaded + for (int i = 1; i < kCapacity; ++i) { + auto key = std::to_string(i); + cache->Release( + cache->Insert(key, (void*)value, key.size() + 5, dumbDeleter) + ); + } + + // the usage should be close to the capacity + ASSERT_GT(kCapacity, cache->GetUsage()); + ASSERT_LT(kCapacity * 0.95, cache->GetUsage()); +} + TEST(CacheTest, HitAndMiss) { ASSERT_EQ(-1, Lookup(100)); @@ -353,7 +386,6 @@ void deleter(const Slice& key, void* value) { delete (Value *)value; } - TEST(CacheTest, BadEviction) { int n = 10;