Iterate only through the threads that have in use entries

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

Reviewed By: djwatson

Differential Revision: D7931363

fbshipit-source-id: 577063f920fb432e255a59e90faa32ab72107124

folly/detail/ThreadLocalDetail.cpp

@@ -21,6 +21,46 @@
 
 namespace folly { namespace threadlocal_detail {
 
+void ThreadEntryNode::initIfZero() {
+  if (UNLIKELY(!next)) {
+    next = prev = parent;
+    parent->meta->pushBackLocked(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 +110,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_EACH_RANGE (i, 0, threadEntry->elementsCapacity) {
+        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.
     }
@@ -115,6 +160,10 @@ void StaticMetaBase::onThreadExit(void* ptr) {
           if (tmp->elements[i].dispose(TLPDestructionMode::THIS_THREAD)) {
             shouldRunInner = true;
             shouldRunOuter = true;
+
+            // now delete the entry if not zero
+            std::lock_guard<std::mutex> g(meta.lock_);
+            tmp->elements[i].node.eraseZero();
           }
         }
       }
@@ -164,6 +213,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 +246,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]);
 
@@ -209,11 +267,13 @@ void StaticMetaBase::destroy(EntryID* ent) {
              * it's illegal to call get on a thread local that's
              * destructing.
              */
+
             e->elements[id].ptr = nullptr;
             e->elements[id].deleter1 = nullptr;
             e->elements[id].ownsDeleter = false;
           }
         }
+
         meta.freeIds_.push_back(id);
       }
     }
@@ -226,22 +286,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* FOLLY_NULLABLE 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 +343,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 +387,44 @@ 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);
+    }
+    head_.elementsCapacity = newCapacity;
+    free(reallocated);
+  }
+
+  head_.elements[id].node.init(&head_, id);
+}
+
+void StaticMetaBase::pushBackLocked(ThreadEntry* t, uint32_t id) {
+  if (LIKELY(!t->removed_)) {
+    auto* node = &t->elements[id].node;
+    std::lock_guard<std::mutex> g(lock_);
+    node->push_back(&head_);
+  }
+}
 
 FOLLY_STATIC_CTOR_PRIORITY_MAX
 PthreadKeyUnregister PthreadKeyUnregister::instance_;

folly/detail/ThreadLocalDetail.h

@@ -58,6 +58,49 @@ 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();
+
+  void init(ThreadEntry* entry, uint32_t newId) {
+    id = newId;
+    prev = next = parent = 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);
+  }
+
+  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 +136,7 @@ struct ElementWrapper {
     DCHECK(deleter1 == nullptr);
 
     if (p) {
+      node.initIfZero();
       ptr = p;
       deleter1 = [](void* pt, TLPDestructionMode) {
         delete static_cast<Ptr>(pt);
@@ -112,6 +156,7 @@ struct ElementWrapper {
     DCHECK(ptr == nullptr);
     DCHECK(deleter2 == nullptr);
     if (p) {
+      node.initIfZero();
       ptr = p;
       auto d2 = d; // gcc-4.8 doesn't decay types correctly in lambda captures
       deleter2 = new std::function<DeleterFunType>(
@@ -138,6 +183,7 @@ struct ElementWrapper {
     std::function<DeleterFunType>* deleter2;
   };
   bool ownsDeleter;
+  ThreadEntryNode node;
 };
 
 struct StaticMetaBase;
@@ -157,6 +203,7 @@ struct ThreadEntry {
   ThreadEntryList* list{nullptr};
   ThreadEntry* listNext{nullptr};
   StaticMetaBase* meta{nullptr};
+  bool removed_{false};
 };
 
 struct ThreadEntryList {
@@ -164,8 +211,6 @@ struct ThreadEntryList {
   size_t count{0};
 };
 
-constexpr uint32_t kEntryIDInvalid = std::numeric_limits<uint32_t>::max();
-
 struct PthreadKeyUnregisterTester;
 
 /**
@@ -303,6 +348,21 @@ 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);
+
+  // 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_;