Make ThreadEntry::elementsCapacity atomic

Summary: Make ThreadEntry::elementsCapacity atomic

Reviewed By: yfeldblum

Differential Revision: D8963304

fbshipit-source-id: 0b37bbe97475a400e9b574b8285bde8f38a479f5

folly/detail/ThreadLocalDetail.cpp

@@ -125,7 +125,8 @@ void StaticMetaBase::onThreadExit(void* ptr) {
       // mark it as removed
       threadEntry->removed_ = true;
       meta.erase(&(*threadEntry));
-      for (size_t i = 0u; i < threadEntry->elementsCapacity; ++i) {
+      auto elementsCapacity = threadEntry->getElementsCapacity();
+      for (size_t i = 0u; i < elementsCapacity; ++i) {
         threadEntry->elements[i].node.eraseZero();
       }
       // No need to hold the lock any longer; the ThreadEntry is private to this
@@ -137,7 +138,8 @@ void StaticMetaBase::onThreadExit(void* ptr) {
     // may be required.
     for (bool shouldRun = true; shouldRun;) {
       shouldRun = false;
-      FOR_EACH_RANGE (i, 0, threadEntry->elementsCapacity) {
+      auto elementsCapacity = threadEntry->getElementsCapacity();
+      FOR_EACH_RANGE (i, 0, elementsCapacity) {
         if (threadEntry->elements[i].dispose(TLPDestructionMode::THIS_THREAD)) {
           threadEntry->elements[i].cleanup();
           shouldRun = true;
@@ -169,7 +171,8 @@ void StaticMetaBase::onThreadExit(void* ptr) {
       }
       for (bool shouldRunInner = true; shouldRunInner;) {
         shouldRunInner = false;
-        FOR_EACH_RANGE (i, 0, tmp->elementsCapacity) {
+        auto elementsCapacity = tmp->getElementsCapacity();
+        FOR_EACH_RANGE (i, 0, elementsCapacity) {
           if (tmp->elements[i].dispose(TLPDestructionMode::THIS_THREAD)) {
             tmp->elements[i].cleanup();
             shouldRunInner = true;
@@ -191,7 +194,7 @@ void StaticMetaBase::onThreadExit(void* ptr) {
     if (tmp->elements) {
       free(tmp->elements);
       tmp->elements = nullptr;
-      tmp->elementsCapacity = 0;
+      tmp->setElementsCapacity(0);
     }
 
 #ifndef FOLLY_TLD_USE_FOLLY_TLS
@@ -207,7 +210,7 @@ void StaticMetaBase::onThreadExit(void* ptr) {
 uint32_t StaticMetaBase::elementsCapacity() const {
   ThreadEntry* threadEntry = (*threadEntry_)();
 
-  return FOLLY_LIKELY(!!threadEntry) ? threadEntry->elementsCapacity : 0;
+  return FOLLY_LIKELY(!!threadEntry) ? threadEntry->getElementsCapacity() : 0;
 }
 
 uint32_t StaticMetaBase::allocate(EntryID* ent) {
@@ -268,8 +271,8 @@ void StaticMetaBase::destroy(EntryID* ent) {
           next->eraseZero();
 
           ThreadEntry* e = next->parent;
-
-          if (id < e->elementsCapacity && e->elements[id].ptr) {
+          auto elementsCapacity = e->getElementsCapacity();
+          if (id < elementsCapacity && e->elements[id].ptr) {
             elements.push_back(e->elements[id]);
 
             /*
@@ -306,15 +309,16 @@ ElementWrapper* StaticMetaBase::reallocate(
     ThreadEntry* threadEntry,
     uint32_t idval,
     size_t& newCapacity) {
-  size_t prevCapacity = threadEntry->elementsCapacity;
+  size_t prevCapacity = threadEntry->getElementsCapacity();
 
   // Growth factor < 2, see folly/docs/FBVector.md; + 5 to prevent
   // very slow start.
   auto smallCapacity = static_cast<size_t>((idval + 5) * kSmallGrowthFactor);
   auto bigCapacity = static_cast<size_t>((idval + 5) * kBigGrowthFactor);
 
-  newCapacity = (threadEntry->meta &&
-                 (bigCapacity <= threadEntry->meta->head_.elementsCapacity))
+  newCapacity =
+      (threadEntry->meta &&
+       (bigCapacity <= threadEntry->meta->head_.getElementsCapacity()))
       ? bigCapacity
       : smallCapacity;
 
@@ -377,7 +381,7 @@ ElementWrapper* StaticMetaBase::reallocate(
 void StaticMetaBase::reserve(EntryID* id) {
   auto& meta = *this;
   ThreadEntry* threadEntry = (*threadEntry_)();
-  size_t prevCapacity = threadEntry->elementsCapacity;
+  size_t prevCapacity = threadEntry->getElementsCapacity();
 
   uint32_t idval = id->getOrAllocate(meta);
   if (prevCapacity > idval) {
@@ -415,15 +419,15 @@ void StaticMetaBase::reserve(EntryID* id) {
       threadEntry->elements[i].node.initZero(threadEntry, i);
     }
 
-    threadEntry->elementsCapacity = newCapacity;
+    threadEntry->setElementsCapacity(newCapacity);
   }
 
   free(reallocated);
 }
 
 void StaticMetaBase::reserveHeadUnlocked(uint32_t id) {
-  if (head_.elementsCapacity <= id) {
-    size_t prevCapacity = head_.elementsCapacity;
+  if (head_.getElementsCapacity() <= id) {
+    size_t prevCapacity = head_.getElementsCapacity();
     size_t newCapacity;
     ElementWrapper* reallocated = reallocate(&head_, id, newCapacity);
 
@@ -439,7 +443,7 @@ void StaticMetaBase::reserveHeadUnlocked(uint32_t id) {
       head_.elements[i].node.init(&head_, i);
     }
 
-    head_.elementsCapacity = newCapacity;
+    head_.setElementsCapacity(newCapacity);
     free(reallocated);
   }
 }

folly/detail/ThreadLocalDetail.h

@@ -203,16 +203,26 @@ struct ThreadEntryList;
  * This is written from the owning thread only (under the lock), read
  * from the owning thread (no lock necessary), and read from other threads
  * (under the lock).
+ * StaticMetaBase::head_ elementsCapacity can be read from any thread on
+ * reallocate (no lock)
  */
 struct ThreadEntry {
   ElementWrapper* elements{nullptr};
-  size_t elementsCapacity{0};
+  std::atomic<size_t> elementsCapacity{0};
   ThreadEntry* next{nullptr};
   ThreadEntry* prev{nullptr};
   ThreadEntryList* list{nullptr};
   ThreadEntry* listNext{nullptr};
   StaticMetaBase* meta{nullptr};
   bool removed_{false};
+
+  size_t getElementsCapacity() const noexcept {
+    return elementsCapacity.load(std::memory_order_relaxed);
+  }
+
+  void setElementsCapacity(size_t capacity) noexcept {
+    elementsCapacity.store(capacity, std::memory_order_relaxed);
+  }
 };
 
 struct ThreadEntryList {
@@ -448,11 +458,11 @@ struct StaticMeta : StaticMetaBase {
       size_t& capacity) {
     auto& inst = instance();
     threadEntry = inst.threadEntry_();
-    if (UNLIKELY(threadEntry->elementsCapacity <= id)) {
+    if (UNLIKELY(threadEntry->getElementsCapacity() <= id)) {
       inst.reserve(ent);
       id = ent->getOrInvalid();
     }
-    capacity = threadEntry->elementsCapacity;
+    capacity = threadEntry->getElementsCapacity();
     assert(capacity > id);
   }
 
@@ -505,12 +515,14 @@ struct StaticMeta : StaticMetaBase {
     // init the head list
     head.next = head.prev = &head;
     // init the circular lists
-    for (size_t i = 0u; i < head.elementsCapacity; ++i) {
+    auto elementsCapacity = head.getElementsCapacity();
+    for (size_t i = 0u; i < 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) {
+    elementsCapacity = threadEntry->getElementsCapacity();
+    for (size_t i = 0u; i < elementsCapacity; ++i) {
       if (!threadEntry->elements[i].node.zero()) {
         threadEntry->elements[i].node.initZero(
             threadEntry, static_cast<uint32_t>(i));
@@ -519,7 +531,7 @@ struct StaticMeta : StaticMetaBase {
     }
 
     // If this thread was in the list before the fork, add it back.
-    if (threadEntry->elementsCapacity != 0) {
+    if (elementsCapacity != 0) {
       instance().push_back(threadEntry);
     }
     instance().lock_.unlock();