Iterate only through the threads that have in use entries - v2

Summary: Iterate only through the threads that have in use entries - v2

Reviewed By: djwatson

Differential Revision: D8143429

fbshipit-source-id: 51098ff8e1fef7aaf9f1c0f2ddbf7ebb579cbeb4

folly/detail/ThreadLocalDetail.cpp

@@ -21,6 +21,49 @@
 
 namespace folly { namespace threadlocal_detail {
 
+void ThreadEntryNode::initIfZero(bool locked) {
+  if (UNLIKELY(!next)) {
+    if (LIKELY(locked)) {
+      parent->meta->pushBackLocked(parent, id);
+    } else {
+      parent->meta->pushBackUnlocked(parent, id);
+    }
+  }
+}
+
+ThreadEntryNode* ThreadEntryNode::getNext() {
+  return &next->elements[id].node;
+}
+
+void ThreadEntryNode::push_back(ThreadEntry* head) {
+  // get the head prev and next nodes
+  ThreadEntryNode* hnode = &head->elements[id].node;
+
+  // update current
+  next = head;
+  prev = hnode->prev;
+
+  // hprev
+  ThreadEntryNode* hprev = &hnode->prev->elements[id].node;
+  hprev->next = parent;
+  hnode->prev = parent;
+}
+
+void ThreadEntryNode::eraseZero() {
+  if (LIKELY(prev != nullptr)) {
+    // get the prev and next nodes
+    ThreadEntryNode* nprev = &prev->elements[id].node;
+    ThreadEntryNode* nnext = &next->elements[id].node;
+
+    // update the prev and next
+    nnext->prev = prev;
+    nprev->next = next;
+
+    // set the prev and next to nullptr
+    next = prev = nullptr;
+  }
+}
+
 StaticMetaBase::StaticMetaBase(ThreadEntry* (*threadEntry)(), bool strict)
     : nextId_(1), threadEntry_(threadEntry), strict_(strict) {
   head_.next = head_.prev = &head_;
@@ -70,7 +113,12 @@ void StaticMetaBase::onThreadExit(void* ptr) {
     }
     {
       std::lock_guard<std::mutex> g(meta.lock_);
+      // mark it as removed
+      threadEntry->removed_ = true;
       meta.erase(&(*threadEntry));
+      for (size_t i = 0u; i < threadEntry->elementsCapacity; ++i) {
+        threadEntry->elements[i].node.eraseZero();
+      }
       // No need to hold the lock any longer; the ThreadEntry is private to this
       // thread now that it's been removed from meta.
     }
@@ -135,9 +183,7 @@ void StaticMetaBase::onThreadExit(void* ptr) {
       tmp->elementsCapacity = 0;
     }
 
-#ifndef FOLLY_TLD_USE_FOLLY_TLS
     delete tmp;
-#endif
   }
 
 #ifndef FOLLY_TLD_USE_FOLLY_TLS
@@ -164,6 +210,9 @@ uint32_t StaticMetaBase::allocate(EntryID* ent) {
 
   uint32_t old_id = ent->value.exchange(id);
   DCHECK_EQ(old_id, kEntryIDInvalid);
+
+  reserveHeadUnlocked(id);
+
   return id;
 }
 
@@ -194,7 +243,13 @@ void StaticMetaBase::destroy(EntryID* ent) {
           return;
         }
 
-        for (ThreadEntry* e = meta.head_.next; e != &meta.head_; e = e->next) {
+        auto& node = meta.head_.elements[id].node;
+        while (!node.empty()) {
+          auto* next = node.getNext();
+          next->eraseZero();
+
+          ThreadEntry* e = next->parent;
+
           if (id < e->elementsCapacity && e->elements[id].ptr) {
             elements.push_back(e->elements[id]);
 
@@ -226,22 +281,15 @@ void StaticMetaBase::destroy(EntryID* ent) {
   }
 }
 
-/**
- * Reserve enough space in the ThreadEntry::elements for the item
- * @id to fit in.
- */
-void StaticMetaBase::reserve(EntryID* id) {
-  auto& meta = *this;
-  ThreadEntry* threadEntry = (*threadEntry_)();
+ElementWrapper* StaticMetaBase::reallocate(
+    ThreadEntry* threadEntry,
+    uint32_t idval,
+    size_t& newCapacity) {
   size_t prevCapacity = threadEntry->elementsCapacity;
 
-  uint32_t idval = id->getOrAllocate(meta);
-  if (prevCapacity > idval) {
-    return;
-  }
   // Growth factor < 2, see folly/docs/FBVector.md; + 5 to prevent
   // very slow start.
-  size_t newCapacity = static_cast<size_t>((idval + 5) * 1.7);
+  newCapacity = static_cast<size_t>((idval + 5) * 1.7);
   assert(newCapacity > prevCapacity);
   ElementWrapper* reallocated = nullptr;
 
@@ -290,6 +338,27 @@ void StaticMetaBase::reserve(EntryID* id) {
     }
   }
 
+  return reallocated;
+}
+
+/**
+ * Reserve enough space in the ThreadEntry::elements for the item
+ * @id to fit in.
+ */
+
+void StaticMetaBase::reserve(EntryID* id) {
+  auto& meta = *this;
+  ThreadEntry* threadEntry = (*threadEntry_)();
+  size_t prevCapacity = threadEntry->elementsCapacity;
+
+  uint32_t idval = id->getOrAllocate(meta);
+  if (prevCapacity > idval) {
+    return;
+  }
+
+  size_t newCapacity;
+  ElementWrapper* reallocated = reallocate(threadEntry, idval, newCapacity);
+
   // Success, update the entry
   {
     std::lock_guard<std::mutex> g(meta.lock_);
@@ -313,12 +382,54 @@ void StaticMetaBase::reserve(EntryID* id) {
       }
       std::swap(reallocated, threadEntry->elements);
     }
+
+    for (size_t i = prevCapacity; i < newCapacity; i++) {
+      threadEntry->elements[i].node.initZero(threadEntry, i);
+    }
+
     threadEntry->elementsCapacity = newCapacity;
   }
 
   free(reallocated);
 }
 
+void StaticMetaBase::reserveHeadUnlocked(uint32_t id) {
+  if (head_.elementsCapacity <= id) {
+    size_t prevCapacity = head_.elementsCapacity;
+    size_t newCapacity;
+    ElementWrapper* reallocated = reallocate(&head_, id, newCapacity);
+
+    if (reallocated) {
+      if (prevCapacity != 0) {
+        memcpy(
+            reallocated, head_.elements, sizeof(*reallocated) * prevCapacity);
+      }
+      std::swap(reallocated, head_.elements);
+    }
+
+    for (size_t i = prevCapacity; i < newCapacity; i++) {
+      head_.elements[i].node.init(&head_, i);
+    }
+
+    head_.elementsCapacity = newCapacity;
+    free(reallocated);
+  }
+}
+
+void StaticMetaBase::pushBackLocked(ThreadEntry* t, uint32_t id) {
+  if (LIKELY(!t->removed_)) {
+    std::lock_guard<std::mutex> g(lock_);
+    auto* node = &t->elements[id].node;
+    node->push_back(&head_);
+  }
+}
+
+void StaticMetaBase::pushBackUnlocked(ThreadEntry* t, uint32_t id) {
+  if (LIKELY(!t->removed_)) {
+    auto* node = &t->elements[id].node;
+    node->push_back(&head_);
+  }
+}
 
 FOLLY_STATIC_CTOR_PRIORITY_MAX
 PthreadKeyUnregister PthreadKeyUnregister::instance_;

folly/detail/ThreadLocalDetail.h

@@ -58,6 +58,53 @@ struct AccessModeStrict {};
 
 namespace threadlocal_detail {
 
+constexpr uint32_t kEntryIDInvalid = std::numeric_limits<uint32_t>::max();
+
+struct ThreadEntry;
+/* This represents a node in doubly linked list where all the nodes
+ * are part of an ElementWrapper struct that has the same id.
+ * we cannot use prev and next as ThreadEntryNode pointers since the
+ * ThreadEntry::elements can be reallocated and the pointers will change
+ * in this case. So we keep a pointer to the parent ThreadEntry struct
+ * one for the prev and next and also the id.
+ * We will traverse and update the list only when holding the
+ * StaticMetaBase::lock_
+ */
+struct ThreadEntryNode {
+  uint32_t id;
+  ThreadEntry* parent;
+  ThreadEntry* prev;
+  ThreadEntry* next;
+
+  void initIfZero(bool locked);
+
+  void init(ThreadEntry* entry, uint32_t newId) {
+    id = newId;
+    parent = prev = next = entry;
+  }
+
+  void initZero(ThreadEntry* entry, uint32_t newId) {
+    id = newId;
+    parent = entry;
+    prev = next = nullptr;
+  }
+
+  // if the list this node is part of is empty
+  bool empty() const {
+    return (next == parent);
+  }
+
+  bool zero() const {
+    return (!prev);
+  }
+
+  ThreadEntryNode* getNext();
+
+  void push_back(ThreadEntry* head);
+
+  void eraseZero();
+};
+
 /**
  * POD wrapper around an element (a void*) and an associated deleter.
  * This must be POD, as we memset() it to 0 and memcpy() it around.
@@ -93,6 +140,7 @@ struct ElementWrapper {
     DCHECK(deleter1 == nullptr);
 
     if (p) {
+      node.initIfZero(true /*locked*/);
       ptr = p;
       deleter1 = [](void* pt, TLPDestructionMode) {
         delete static_cast<Ptr>(pt);
@@ -112,6 +160,7 @@ struct ElementWrapper {
     DCHECK(ptr == nullptr);
     DCHECK(deleter2 == nullptr);
     if (p) {
+      node.initIfZero(true /*locked*/);
       ptr = p;
       auto d2 = d; // gcc-4.8 doesn't decay types correctly in lambda captures
       deleter2 = new std::function<DeleterFunType>(
@@ -138,6 +187,7 @@ struct ElementWrapper {
     std::function<DeleterFunType>* deleter2;
   };
   bool ownsDeleter;
+  ThreadEntryNode node;
 };
 
 struct StaticMetaBase;
@@ -157,6 +207,7 @@ struct ThreadEntry {
   ThreadEntryList* list{nullptr};
   ThreadEntry* listNext{nullptr};
   StaticMetaBase* meta{nullptr};
+  bool removed_{false};
 };
 
 struct ThreadEntryList {
@@ -164,8 +215,6 @@ struct ThreadEntryList {
   size_t count{0};
 };
 
-constexpr uint32_t kEntryIDInvalid = std::numeric_limits<uint32_t>::max();
-
 struct PthreadKeyUnregisterTester;
 
 /**
@@ -303,6 +352,22 @@ struct StaticMetaBase {
 
   ElementWrapper& getElement(EntryID* ent);
 
+  // reserve an id in the head_ ThreadEntry->elements
+  // array if not already there
+  void reserveHeadUnlocked(uint32_t id);
+
+  // push back an entry in the doubly linked list
+  // that corresponds to idx id
+  void pushBackLocked(ThreadEntry* t, uint32_t id);
+  void pushBackUnlocked(ThreadEntry* t, uint32_t id);
+
+  // static helper method to reallocate the ThreadEntry::elements
+  // returns != nullptr if the ThreadEntry::elements was reallocated
+  // nullptr if the ThreadEntry::elements was just extended
+  // and throws stdd:bad_alloc if memory cannot be allocated
+  static ElementWrapper*
+  reallocate(ThreadEntry* threadEntry, uint32_t idval, size_t& newCapacity);
+
   uint32_t nextId_;
   std::vector<uint32_t> freeIds_;
   std::mutex lock_;
@@ -374,7 +439,7 @@ struct StaticMeta : StaticMetaBase {
     assert(capacity > id);
   }
 
-  static ThreadEntry* getThreadEntrySlow() {
+  FOLLY_EXPORT FOLLY_NOINLINE static ThreadEntry* getThreadEntrySlow() {
     auto& meta = instance();
     auto key = meta.pthreadKey_;
     ThreadEntry* threadEntry =
@@ -382,8 +447,8 @@ struct StaticMeta : StaticMetaBase {
     if (!threadEntry) {
       ThreadEntryList* threadEntryList = StaticMeta::getThreadEntryList();
 #ifdef FOLLY_TLD_USE_FOLLY_TLS
-      static FOLLY_TLS ThreadEntry threadEntrySingleton;
-      threadEntry = &threadEntrySingleton;
+      static FOLLY_TLS ThreadEntry* one;
+      threadEntry = FOLLY_LIKELY(!!one) ? one : (one = new ThreadEntry());
 #else
       threadEntry = new ThreadEntry();
 #endif
@@ -419,8 +484,23 @@ struct StaticMeta : StaticMetaBase {
 
   static void onForkChild() {
     // only the current thread survives
-    instance().head_.next = instance().head_.prev = &instance().head_;
+    auto& head = instance().head_;
+    // init the head list
+    head.next = head.prev = &head;
+    // init the circular lists
+    for (size_t i = 0u; i < head.elementsCapacity; ++i) {
+      head.elements[i].node.init(&head, static_cast<uint32_t>(i));
+    }
+    // init the thread entry
     ThreadEntry* threadEntry = instance().threadEntry_();
+    for (size_t i = 0u; i < threadEntry->elementsCapacity; ++i) {
+      if (!threadEntry->elements[i].node.zero()) {
+        threadEntry->elements[i].node.initZero(
+            threadEntry, static_cast<uint32_t>(i));
+        threadEntry->elements[i].node.initIfZero(false /*locked*/);
+      }
+    }
+
     // If this thread was in the list before the fork, add it back.
     if (threadEntry->elementsCapacity != 0) {
       instance().push_back(threadEntry);