From 2fa98577f4b73f8a2180fb88b5f075cce6d650e6 Mon Sep 17 00:00:00 2001
From: Anton Likhtarov <alikhtarov@fb.com>
Date: Thu, 19 Nov 2015 11:32:26 -0800
Subject: [PATCH] Fix invalid DCHECK

Summary: There's a race between insert() and erase(): as soon as insert()
releases the lock (swaps kLockedKey_ with the actual key), an erase() might
jump in and invalidate the key.

As far as I can tell, this bug existed since the beginning.

Reviewed By: nbronson

Differential Revision: D2673099

fb-gh-sync-id: 4721893d2ad4836e11acc0fb4ecb0dd7b2b69be1
---
 folly/AtomicHashArray-inl.h      |  4 +++-
 folly/test/AtomicHashMapTest.cpp | 28 ++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/folly/AtomicHashArray-inl.h b/folly/AtomicHashArray-inl.h
index e3b32a04..1125bb98 100644
--- a/folly/AtomicHashArray-inl.h
+++ b/folly/AtomicHashArray-inl.h
@@ -144,9 +144,11 @@ insertInternal(KeyT key_in, ArgTs&&... vCtorArgs) {
             --numPendingEntries_;
             throw;
           }
+          // An erase() can race here and delete right after our insertion
           // Direct comparison rather than EqualFcn ok here
           // (we just inserted it)
-          DCHECK(relaxedLoadKey(*cell) == key_in);
+          DCHECK(relaxedLoadKey(*cell) == key_in ||
+                 relaxedLoadKey(*cell) == kErasedKey_);
           --numPendingEntries_;
           ++numEntries_;  // This is a thread cached atomic increment :)
           if (numEntries_.readFast() >= maxEntries_) {
diff --git a/folly/test/AtomicHashMapTest.cpp b/folly/test/AtomicHashMapTest.cpp
index 46223107..63523eda 100644
--- a/folly/test/AtomicHashMapTest.cpp
+++ b/folly/test/AtomicHashMapTest.cpp
@@ -667,6 +667,34 @@ TEST(Ahm, erase_find_race) {
   write_thread.join();
 }
 
+// Erase right after insert race bug repro (t9130653)
+TEST(Ahm, erase_after_insert_race) {
+  const uint64_t limit = 10000;
+  const size_t num_threads = 100;
+  const size_t num_iters = 500;
+  AtomicHashMap<uint64_t, uint64_t> map(limit + 10);
+
+  std::atomic<bool> go{false};
+  std::vector<std::thread> ts;
+  for (size_t i = 0; i < num_threads; ++i) {
+    ts.emplace_back([&]() {
+        while (!go) {
+          continue;
+        }
+        for (size_t n = 0; n < num_iters; ++n) {
+          map.erase(1);
+          map.insert(1, 1);
+        }
+      });
+  }
+
+  go = true;
+
+  for (auto& t : ts) {
+    t.join();
+  }
+}
+
 // Repro for a bug when iterator didn't skip empty submaps.
 TEST(Ahm, iterator_skips_empty_submaps) {
   AtomicHashMap<uint64_t, uint64_t>::Config config;
-- 
2.34.1