Add SetCapacity method to leveldb::Cache.

This change introduces a virtual SetCapacity method to the Cache interface, allowing the cache capacity to be adjusted dynamically. The LRUCache implementation evicts entries if the new capacity is smaller than the current usage. The ShardedLRUCache distributes the new capacity among its shards.

This functionality will be used in chromium to improve memory pressure handling. The
current Prune() functionality is not flexible enough to reduce the memory used by the
cache if necessary.

PiperOrigin-RevId: 974645811

include/leveldb/cache.h

@@ -93,6 +93,14 @@ class LEVELDB_EXPORT Cache {
   // leveldb may change Prune() to a pure abstract method.
   virtual void Prune() {}
 
+  // Set the capacity of the cache. Entries may be evicted if the new capacity
+  // is smaller than the current usage.
+  //
+  // This method is optional to implement and the default implementation does
+  // nothing.  Subclasses that do not support dynamic capacity sizing may ignore
+  // this call.
+  virtual void SetCapacity(size_t capacity) {}
+
   // Return an estimate of the combined charges of all elements stored in the
   // cache.
   virtual size_t TotalCharge() const = 0;

util/cache.cc

@@ -154,7 +154,7 @@ class LRUCache {
   ~LRUCache();
 
   // Separate from constructor so caller can easily make an array of LRUCache
-  void SetCapacity(size_t capacity) { capacity_ = capacity; }
+  void SetInitialCapacity(size_t capacity) { capacity_ = capacity; }
 
   // Like Cache methods, but with an extra "hash" parameter.
   Cache::Handle* Insert(const Slice& key, uint32_t hash, void* value,
@@ -164,6 +164,7 @@ class LRUCache {
   void Release(Cache::Handle* handle);
   void Erase(const Slice& key, uint32_t hash);
   void Prune();
+  void SetCapacity(size_t capacity);
   size_t TotalCharge() const {
     MutexLock l(&mutex_);
     return usage_;
@@ -333,6 +334,19 @@ void LRUCache::Prune() {
   }
 }
 
+void LRUCache::SetCapacity(size_t capacity) {
+  MutexLock l(&mutex_);
+  capacity_ = capacity;
+  while (usage_ > capacity_ && lru_.next != &lru_) {
+    LRUHandle* old = lru_.next;
+    assert(old->refs == 1);
+    bool erased = FinishErase(table_.Remove(old->key(), old->hash));
+    if (!erased) {  // to avoid unused variable when compiled NDEBUG
+      assert(erased);
+    }
+  }
+}
+
 static const int kNumShardBits = 4;
 static const int kNumShards = 1 << kNumShardBits;
 
@@ -352,7 +366,7 @@ class ShardedLRUCache : public Cache {
   explicit ShardedLRUCache(size_t capacity) : last_id_(0) {
     const size_t per_shard = (capacity + (kNumShards - 1)) / kNumShards;
     for (int s = 0; s < kNumShards; s++) {
-      shard_[s].SetCapacity(per_shard);
+      shard_[s].SetInitialCapacity(per_shard);
     }
   }
   ~ShardedLRUCache() override {}
@@ -385,6 +399,12 @@ class ShardedLRUCache : public Cache {
       shard_[s].Prune();
     }
   }
+  void SetCapacity(size_t capacity) override {
+    const size_t per_shard = (capacity + (kNumShards - 1)) / kNumShards;
+    for (int s = 0; s < kNumShards; s++) {
+      shard_[s].SetCapacity(per_shard);
+    }
+  }
   size_t TotalCharge() const override {
     size_t total = 0;
     for (int s = 0; s < kNumShards; s++) {

util/cache_test.cc

@@ -221,4 +221,24 @@ TEST_F(CacheTest, ZeroSizeCache) {
   ASSERT_EQ(-1, Lookup(1));
 }
 
+TEST_F(CacheTest, SetCapacity) {
+  // Insert many items.
+  for (int i = 0; i < 100; i++) {
+    Insert(i, 1000 + i);
+  }
+  ASSERT_EQ(100, cache_->TotalCharge());
+
+  // Reduce capacity to 10.
+  // Each shard gets (10 + 15) / 16 = 1 capacity.
+  // Total capacity = 16.
+  cache_->SetCapacity(10);
+
+  // Usage should be at most 16.
+  ASSERT_LE(cache_->TotalCharge(), 16);
+
+  // Reduce capacity to 0.
+  cache_->SetCapacity(0);
+  ASSERT_EQ(0, cache_->TotalCharge());
+}
+
 }  // namespace leveldb