1 files changed, 84 insertions, 34 deletions

diff --git a/util/cache.cc b/util/cache.cc
index 0881bce..ce46886 100644
--- a/util/cache.cc
+++ b/util/cache.cc
@@ -19,6 +19,23 @@ Cache::~Cache() {
 namespace {
 
 // LRU cache implementation
+//
+// Cache entries have an "in_cache" boolean indicating whether the cache has a
+// reference on the entry.  The only ways that this can become false without the
+// entry being passed to its "deleter" are via Erase(), via Insert() when
+// an element with a duplicate key is inserted, or on destruction of the cache.
+//
+// The cache keeps two linked lists of items in the cache.  All items in the
+// cache are in one list or the other, and never both.  Items still referenced
+// by clients but erased from the cache are in neither list.  The lists are:
+// - in-use:  contains the items currently referenced by clients, in no
+//   particular order.  (This list is used for invariant checking.  If we
+//   removed the check, elements that would otherwise be on this list could be
+//   left as disconnected singleton lists.)
+// - LRU:  contains the items not currently referenced by clients, in LRU order
+// Elements are moved between these lists by the Ref() and Unref() methods,
+// when they detect an element in the cache acquiring or losing its only
+// external reference.
 
 // An entry is a variable length heap-allocated structure.  Entries
 // are kept in a circular doubly linked list ordered by access time.
@@ -30,7 +47,8 @@ struct LRUHandle {
   LRUHandle* prev;
   size_t charge;      // TODO(opt): Only allow uint32_t?
   size_t key_length;
-  uint32_t refs;
+  bool in_cache;      // Whether entry is in the cache.
+  uint32_t refs;      // References, including cache reference, if present.
   uint32_t hash;      // Hash of key(); used for fast sharding and comparisons
   char key_data[1];   // Beginning of key
 
@@ -155,8 +173,10 @@ class LRUCache {
 
  private:
   void LRU_Remove(LRUHandle* e);
-  void LRU_Append(LRUHandle* e);
+  void LRU_Append(LRUHandle*list, LRUHandle* e);
+  void Ref(LRUHandle* e);
   void Unref(LRUHandle* e);
+  bool FinishErase(LRUHandle* e);
 
   // Initialized before use.
   size_t capacity_;
@@ -167,34 +187,55 @@ class LRUCache {
 
   // Dummy head of LRU list.
   // lru.prev is newest entry, lru.next is oldest entry.
+  // Entries have refs==1 and in_cache==true.
   LRUHandle lru_;
 
+  // Dummy head of in-use list.
+  // Entries are in use by clients, and have refs >= 2 and in_cache==true.
+  LRUHandle in_use_;
+
   HandleTable table_;
 };
 
 LRUCache::LRUCache()
     : usage_(0) {
-  // Make empty circular linked list
+  // Make empty circular linked lists.
   lru_.next = &lru_;
   lru_.prev = &lru_;
+  in_use_.next = &in_use_;
+  in_use_.prev = &in_use_;
 }
 
 LRUCache::~LRUCache() {
+  assert(in_use_.next == &in_use_);  // Error if caller has an unreleased handle
   for (LRUHandle* e = lru_.next; e != &lru_; ) {
     LRUHandle* next = e->next;
-    assert(e->refs == 1);  // Error if caller has an unreleased handle
+    assert(e->in_cache);
+    e->in_cache = false;
+    assert(e->refs == 1);  // Invariant of lru_ list.
     Unref(e);
     e = next;
   }
 }
 
+void LRUCache::Ref(LRUHandle* e) {
+  if (e->refs == 1 && e->in_cache) {  // If on lru_ list, move to in_use_ list.
+    LRU_Remove(e);
+    LRU_Append(&in_use_, e);
+  }
+  e->refs++;
+}
+
 void LRUCache::Unref(LRUHandle* e) {
   assert(e->refs > 0);
   e->refs--;
-  if (e->refs <= 0) {
-    usage_ -= e->charge;
+  if (e->refs == 0) { // Deallocate.
+    assert(!e->in_cache);
     (*e->deleter)(e->key(), e->value);
     free(e);
+  } else if (e->in_cache && e->refs == 1) {  // No longer in use; move to lru_ list.
+    LRU_Remove(e);
+    LRU_Append(&lru_, e);
   }
 }
 
@@ -203,10 +244,10 @@ void LRUCache::LRU_Remove(LRUHandle* e) {
   e->prev->next = e->next;
 }
 
-void LRUCache::LRU_Append(LRUHandle* e) {
-  // Make "e" newest entry by inserting just before lru_
-  e->next = &lru_;
-  e->prev = lru_.prev;
+void LRUCache::LRU_Append(LRUHandle* list, LRUHandle* e) {
+  // Make "e" newest entry by inserting just before *list
+  e->next = list;
+  e->prev = list->prev;
   e->prev->next = e;
   e->next->prev = e;
 }
@@ -215,9 +256,7 @@ Cache::Handle* LRUCache::Lookup(const Slice& key, uint32_t hash) {
   MutexLock l(&mutex_);
   LRUHandle* e = table_.Lookup(key, hash);
   if (e != NULL) {
-    e->refs++;
-    LRU_Remove(e);
-    LRU_Append(e);
+    Ref(e);
   }
   return reinterpret_cast<Cache::Handle*>(e);
 }
@@ -239,46 +278,57 @@ Cache::Handle* LRUCache::Insert(
   e->charge = charge;
   e->key_length = key.size();
   e->hash = hash;
-  e->refs = 2;  // One from LRUCache, one for the returned handle
+  e->in_cache = false;
+  e->refs = 1;  // for the returned handle.
   memcpy(e->key_data, key.data(), key.size());
-  LRU_Append(e);
-  usage_ += charge;
 
-  LRUHandle* old = table_.Insert(e);
-  if (old != NULL) {
-    LRU_Remove(old);
-    Unref(old);
-  }
+  if (capacity_ > 0) {
+    e->refs++;  // for the cache's reference.
+    e->in_cache = true;
+    LRU_Append(&in_use_, e);
+    usage_ += charge;
+    FinishErase(table_.Insert(e));
+  } // else don't cache.  (Tests use capacity_==0 to turn off caching.)
 
   while (usage_ > capacity_ && lru_.next != &lru_) {
     LRUHandle* old = lru_.next;
-    LRU_Remove(old);
-    table_.Remove(old->key(), old->hash);
-    Unref(old);
+    assert(old->refs == 1);
+    bool erased = FinishErase(table_.Remove(old->key(), old->hash));
+    if (!erased) {  // to avoid unused variable when compiled NDEBUG
+      assert(erased);
+    }
   }
 
   return reinterpret_cast<Cache::Handle*>(e);
 }
 
-void LRUCache::Erase(const Slice& key, uint32_t hash) {
-  MutexLock l(&mutex_);
-  LRUHandle* e = table_.Remove(key, hash);
+// If e != NULL, finish removing *e from the cache; it has already been removed
+// from the hash table.  Return whether e != NULL.  Requires mutex_ held.
+bool LRUCache::FinishErase(LRUHandle* e) {
   if (e != NULL) {
+    assert(e->in_cache);
     LRU_Remove(e);
+    e->in_cache = false;
+    usage_ -= e->charge;
     Unref(e);
   }
+  return e != NULL;
+}
+
+void LRUCache::Erase(const Slice& key, uint32_t hash) {
+  MutexLock l(&mutex_);
+  FinishErase(table_.Remove(key, hash));
 }
 
 void LRUCache::Prune() {
   MutexLock l(&mutex_);
-  for (LRUHandle* e = lru_.next; e != &lru_; ) {
-    LRUHandle* next = e->next;
-    if (e->refs == 1) {
-      table_.Remove(e->key(), e->hash);
-      LRU_Remove(e);
-      Unref(e);
+  while (lru_.next != &lru_) {
+    LRUHandle* e = lru_.next;
+    assert(e->refs == 1);
+    bool erased = FinishErase(table_.Remove(e->key(), e->hash));
+    if (!erased) {  // to avoid unused variable when compiled NDEBUG
+      assert(erased);
     }
-    e = next;
   }
 }