Bug 1433351 - Add nsBaseHashtable::EntryPtr::OrRemove method to abort nsBaseHashtable::LookupForAdd on miss. r=froydnj

In SourceSurfaceImage::GetTextureClient, we use LookupForAdd. This is
because we typically will create a new TextureClient if there isn't
already one created. This creation can fail because the size is too big,
or we don't have the memory available for it. Unfortunately LookupForAdd
is an infallible operation; it is expected we will always add something
to the hashtable if we don't find an entry. This patch adds an OrRemove
method to cover the corner case where we are unable to complete the
insertion.

gfx/layers/ImageContainer.cpp

@@ -865,7 +865,7 @@ SourceSurfaceImage::GetTextureClient(KnowsCompositor* aForwarder)
   }
 
   // Remove the speculatively added entry.
-  mTextureClients.Remove(aForwarder->GetSerial());
+  entry.OrRemove();
   return nullptr;
 }
 

xpcom/ds/nsBaseHashtable.h

@@ -265,29 +265,29 @@ public:
 
   struct EntryPtr {
   private:
-    EntryType& mEntry;
+    EntryType* mEntry;
     bool mExistingEntry;
+    nsBaseHashtable& mTable;
     // For debugging purposes
 #ifdef DEBUG
-    nsBaseHashtable& mTable;
     uint32_t mTableGeneration;
     bool mDidInitNewEntry;
 #endif
 
   public:
     EntryPtr(nsBaseHashtable& aTable, EntryType* aEntry, bool aExistingEntry)
-      : mEntry(*aEntry)
+      : mEntry(aEntry)
       , mExistingEntry(aExistingEntry)
-#ifdef DEBUG
       , mTable(aTable)
+#ifdef DEBUG
       , mTableGeneration(aTable.GetGeneration())
       , mDidInitNewEntry(false)
 #endif
     {}
     ~EntryPtr()
     {
-      MOZ_ASSERT(mExistingEntry || mDidInitNewEntry,
-                 "Forgot to call OrInsert() on a new entry");
+      MOZ_ASSERT(mExistingEntry || mDidInitNewEntry || !mEntry,
+                 "Forgot to call OrInsert() or OrRemove() on a new entry");
     }
 
     // Is there something stored in the table already?
@@ -301,19 +301,29 @@ public:
     UserDataType OrInsert(F func)
     {
       MOZ_ASSERT(mTableGeneration == mTable.GetGeneration());
+      MOZ_ASSERT(mEntry);
       if (!mExistingEntry) {
-        mEntry.mData = func();
+        mEntry->mData = func();
 #ifdef DEBUG
         mDidInitNewEntry = true;
 #endif
       }
-      return mEntry.mData;
+      return mEntry->mData;
+    }
+
+    void OrRemove()
+    {
+      MOZ_ASSERT(mTableGeneration == mTable.GetGeneration());
+      MOZ_ASSERT(mEntry);
+      mTable.RemoveEntry(mEntry);
+      mEntry = nullptr;
     }
 
     MOZ_MUST_USE DataType& Data()
     {
       MOZ_ASSERT(mTableGeneration == mTable.GetGeneration());
-      return mEntry.mData;
+      MOZ_ASSERT(mEntry);
+      return mEntry->mData;
     }
   };
 

xpcom/tests/gtest/TestHashtables.cpp

@@ -474,6 +474,28 @@ TEST(Hashtables, DataHashtable_LookupForAdd)
     }
   }
   ASSERT_TRUE(0 == UniToEntity.Count());
+
+  // Remove newly added entries via OrRemove.
+  for (auto& entity : gEntities) {
+    auto entry = UniToEntity.LookupForAdd(entity.mUnicode);
+    ASSERT_FALSE(entry);
+    entry.OrRemove();
+  }
+  ASSERT_TRUE(0 == UniToEntity.Count());
+
+  // Remove existing entries via OrRemove.
+  for (auto& entity : gEntities) {
+    auto entry = UniToEntity.LookupForAdd(entity.mUnicode);
+    const char* val = entry.OrInsert([&entity] () { return entity.mStr; });
+    ASSERT_FALSE(entry);
+    ASSERT_TRUE(val == entity.mStr);
+    ASSERT_TRUE(entry.Data() == entity.mStr);
+
+    auto entry2 = UniToEntity.LookupForAdd(entity.mUnicode);
+    ASSERT_TRUE(entry2);
+    entry2.OrRemove();
+  }
+  ASSERT_TRUE(0 == UniToEntity.Count());
 }
 
 TEST(Hashtables, ClassHashtable_LookupForAdd)
@@ -519,4 +541,26 @@ TEST(Hashtables, ClassHashtable_LookupForAdd)
     }
   }
   ASSERT_TRUE(0 == EntToUniClass.Count());
+
+  // Remove newly added entries via OrRemove.
+  for (auto& entity : gEntities) {
+    auto entry = EntToUniClass.LookupForAdd(nsDependentCString(entity.mStr));
+    ASSERT_FALSE(entry);
+    entry.OrRemove();
+  }
+  ASSERT_TRUE(0 == EntToUniClass.Count());
+
+  // Remove existing entries via OrRemove.
+  for (auto& entity : gEntities) {
+    auto entry = EntToUniClass.LookupForAdd(nsDependentCString(entity.mStr));
+    const TestUniChar* val = entry.OrInsert([] () { return nullptr; });
+    ASSERT_FALSE(entry);
+    ASSERT_TRUE(val == nullptr);
+    ASSERT_TRUE(entry.Data() == nullptr);
+
+    auto entry2 = EntToUniClass.LookupForAdd(nsDependentCString(entity.mStr));
+    ASSERT_TRUE(entry2);
+    entry2.OrRemove();
+  }
+  ASSERT_TRUE(0 == EntToUniClass.Count());
 }