From ed5ad8e1d385fa579f0ce553351d366a285c9c98 Mon Sep 17 00:00:00 2001
From: Julian Zhao <jzzhao@meta.com>
Date: Fri, 9 Feb 2024 13:43:46 -0800
Subject: [PATCH] Fix potential deadlock inside unregisterCallback()

Summary:
Special thanks to andriigrynenko for pointing out the potential dead lock problem when implementing a dynamic counter inside another dynamic counter.

Dead lock example explained:

a,b are both dynamic counters (callback counter) where b's lambda calls getCounter(a)

1.  thread 1 called getCounter(b)
2. thread 1 locked the map, got pointer to a counter, unlocked the map
3. thread 1 grabbed the callback lock
4. thread 1 started running the lambda (delay then getCounter(a))
5. thread 2 started unregistering dynamic counter b
6. thread 2 lock the whole map
7. thread 2 waits here (Waiting on cb lock)
8. thread 1 calls getCounter(a), which is blocked on the map lock.

Reviewed By: andriigrynenko

Differential Revision: D53455516

fbshipit-source-id: db172661d50fe568111f6b0f82b98b1f69b388ba
---
 fb303/CallbackValuesMap-inl.h        |  6 +++++-
 fb303/test/CallbackValuesMapTest.cpp | 28 ++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/fb303/CallbackValuesMap-inl.h b/fb303/CallbackValuesMap-inl.h
index 799b19a42..caae0b13d 100644
--- a/fb303/CallbackValuesMap-inl.h
+++ b/fb303/CallbackValuesMap-inl.h
@@ -104,7 +104,7 @@ bool CallbackValuesMap<T>::unregisterCallback(folly::StringPiece name) {
   if (entry == wlock->map.end()) {
     return false;
   }
-  entry->second->clear();
+  auto callbackCopy = std::move(entry->second);
 
   // avoid fetch_add() to avoid extra fences, since we hold the lock already
   uint64_t epoch = wlock->mapEpoch.load();
@@ -112,6 +112,10 @@ bool CallbackValuesMap<T>::unregisterCallback(folly::StringPiece name) {
 
   wlock->map.erase(entry);
   VLOG(5) << "Unregistered  callback: " << name;
+
+  // clear the callback after releasing the lock
+  wlock.unlock();
+  callbackCopy->clear();
   return true;
 }
 
diff --git a/fb303/test/CallbackValuesMapTest.cpp b/fb303/test/CallbackValuesMapTest.cpp
index 5e0644fcb..848273c71 100644
--- a/fb303/test/CallbackValuesMapTest.cpp
+++ b/fb303/test/CallbackValuesMapTest.cpp
@@ -20,6 +20,7 @@
 #include <thread>
 
 #include <boost/bind.hpp>
+#include <folly/synchronization/Baton.h>
 #include <gtest/gtest.h>
 
 using boost::bind;
@@ -119,3 +120,30 @@ TEST(CallbackValuesMapTest, GetCallback) {
   key2Cb->getValue(&val2);
   EXPECT_EQ(val2, 321);
 }
+
+TEST(CallbackValuesMapTest, DoubleDynamicCounterDeadlock) {
+  TestCallbackValuesMap callbackMap;
+  callbackMap.registerCallback("a", []() { return 42; });
+  callbackMap.registerCallback("b", [&callbackMap]() {
+    sleep(2);
+    int val;
+    callbackMap.getValue("a", &val);
+    return val;
+  });
+
+  std::thread t1([&callbackMap]() {
+    int val;
+    callbackMap.getValue("b", &val);
+  });
+
+  folly::Baton<> baton;
+  std::thread t2([&callbackMap, &baton, &t1]() {
+    sleep(1);
+    callbackMap.unregisterCallback("b");
+    t1.join();
+    baton.post();
+  });
+  ASSERT_TRUE(baton.try_wait_for(std::chrono::seconds(10)));
+  t2.join();
+  SUCCEED();
+}