From 59620c802864d01a9417a39a0a9212bfbc359a49 Mon Sep 17 00:00:00 2001
From: Nitin Garg <gargn@meta.com>
Date: Sun, 23 Jun 2024 17:37:50 -0700
Subject: [PATCH] Fix TSAN deadlock in ThreadLocal's fork handler.

Summary:
The per ThreadEntrySet lock causes the number of locks to be tracked
by TSAN during fork to be too large. We cannot control order of fork handler
registration and it seems like TSAN's handlers are run earlier and lock down
some of its internal state. Later when a realloc triggered in TSAN's tracking
logic when we acquire all the ThreadEntrySet lock causes the fork to deadlock.

For now, we can simply skip locking the per ThreadEntrySet lock in fork
handlers. reset() and accessAllThreads() path are the only ones that end up
holding locks on them and both paths also hold the accessAllThreadsLock_ as
well and that is accounted for in the fork handler.

Reviewed By: chippermist

Differential Revision: D58920872

fbshipit-source-id: dbdc966f2586bb5479e0bf3274e54d3e9e5a2c6c
---
 folly/detail/ThreadLocalDetail.h | 27 ++++++++++-----------------
 1 file changed, 10 insertions(+), 17 deletions(-)

diff --git a/folly/detail/ThreadLocalDetail.h b/folly/detail/ThreadLocalDetail.h
index ea4660ae704..63b993f13fa 100644
--- a/folly/detail/ThreadLocalDetail.h
+++ b/folly/detail/ThreadLocalDetail.h
@@ -652,19 +652,14 @@ struct FOLLY_EXPORT StaticMeta final : StaticMetaBase {
       return false;
     }
     meta.lock_.lock();
-    size_t maxId = meta.allId2ThreadEntrySets_.size();
-    for (size_t id = 0; id < maxId; ++id) {
-      meta.allId2ThreadEntrySets_[id].unsafeGetMutex().lock();
-    }
+    // Okay to not lock each set in meta.allId2ThreadEntrySets
+    // as accessAllThreadsLock_ in held by calls to reset() and
+    // accessAllThreads.
     return true;
   }
 
   static void onForkParent() {
     auto& meta = instance();
-    size_t maxId = meta.allId2ThreadEntrySets_.size();
-    for (size_t id = 0; id < maxId; ++id) {
-      meta.allId2ThreadEntrySets_[id].unsafeGetMutex().unlock();
-    }
     meta.lock_.unlock();
     meta.accessAllThreadsLock_.unlock();
   }
@@ -676,14 +671,13 @@ struct FOLLY_EXPORT StaticMeta final : StaticMetaBase {
     // Loop through allId2ThreadEntrySets_; Only keep ThreadEntry* in the map
     // for ThreadEntry::elements that are still in use by the current thread.
     // Evict all of the ThreadEntry* from other threads.
-    size_t maxId = meta.allId2ThreadEntrySets_.size();
-    for (size_t id = 0; id < maxId; ++id) {
-      auto& teSet = meta.allId2ThreadEntrySets_[id];
-      auto& wlockedSet = teSet.unsafeGetUnlocked(); // locked in preFork
-      if (wlockedSet.contains(threadEntry)) {
+    uint32_t maxId = meta.nextId_.load();
+    for (uint32_t id = 0; id < maxId; ++id) {
+      auto wlockedSet = meta.allId2ThreadEntrySets_[id].wlock();
+      if (wlockedSet->contains(threadEntry)) {
         DCHECK(threadEntry->elements[id].isLinked);
-        wlockedSet.clear();
-        wlockedSet.insert(threadEntry);
+        wlockedSet->clear();
+        wlockedSet->insert(threadEntry);
       } else {
         // DCHECK if elements[id] is unlinked for this
         // threadEntry only (id < threadEntry->getElementsCapacity()).
@@ -691,9 +685,8 @@ struct FOLLY_EXPORT StaticMeta final : StaticMetaBase {
         if (id < threadEntry->getElementsCapacity()) {
           DCHECK(!threadEntry->elements[id].isLinked);
         }
-        wlockedSet.clear();
+        wlockedSet->clear();
       }
-      teSet.unsafeGetMutex().unlock();
     }
     meta.lock_.unlock();
     meta.accessAllThreadsLock_.unlock();