Add ThreadLocalRegistry (apache#8623)

Yuhta · facebook-github-bot · commit 12ed56bfe18f · 2024-02-01T15:50:41.000-08:00
Summary: Pull Request resolved: facebookincubator/velox#8623 Add this utility class to keep static thread local objects while being able to iterate over all of them from one single thread. Reviewed By: oerling Differential Revision: D53277766 fbshipit-source-id: f36d38934358080ed7210bb36c4619b859df7758
diff --git a/velox/common/process/ThreadLocalRegistry.h b/velox/common/process/ThreadLocalRegistry.h
@@ -0,0 +1,88 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <list>
+#include <memory>
+#include <mutex>
+
+namespace facebook::velox::process {
+
+/// A registry for keeping static thread local objects of type T.  Similar to
+/// folly::ThreadLocal but a little bit more efficient in terms of performance
+/// and memory usage, because we do not support thread local with lexical scope.
+///
+/// NOTE: only one instance of ThreadLocalRegistry<T> can be created with each
+/// T.
+template <typename T>
+class ThreadLocalRegistry {
+ public:
+  class Reference;
+
+  /// Access values from all threads.  Takes a global lock and should be used
+  /// with caution.
+  template <typename F>
+  void forAllValues(F f) {
+    std::lock_guard<std::mutex> entriesLock(entriesMutex_);
+    for (auto& entry : entries_) {
+      std::lock_guard<std::mutex> lk(entry->mutex);
+      f(entry->value);
+    }
+  }
+
+ private:
+  struct Entry {
+    std::mutex mutex;
+    T value;
+  };
+
+  std::list<std::unique_ptr<Entry>> entries_;
+  std::mutex entriesMutex_;
+};
+
+/// Reference to one thread local value.  Should be stored in thread local
+/// memory.
+template <typename T>
+class ThreadLocalRegistry<T>::Reference {
+ public:
+  explicit Reference(const std::shared_ptr<ThreadLocalRegistry>& registry)
+      : registry_(registry) {
+    auto entry = std::make_unique<Entry>();
+    std::lock_guard<std::mutex> lk(registry_->entriesMutex_);
+    iterator_ =
+        registry_->entries_.insert(registry_->entries_.end(), std::move(entry));
+  }
+
+  ~Reference() {
+    std::lock_guard<std::mutex> lk(registry_->entriesMutex_);
+    registry_->entries_.erase(iterator_);
+  }
+
+  /// Obtain the thread local value and process it with the functor `f'.
+  template <typename F>
+  auto withValue(F f) {
+    auto* entry = iterator_->get();
+    std::lock_guard<std::mutex> lk(entry->mutex);
+    return f(entry->value);
+  }
+
+ private:
+  std::shared_ptr<ThreadLocalRegistry> const registry_;
+  typename std::list<std::unique_ptr<Entry>>::iterator iterator_;
+};
+
+} // namespace facebook::velox::process
diff --git a/velox/common/process/tests/CMakeLists.txt b/velox/common/process/tests/CMakeLists.txt
@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-add_executable(velox_process_test TraceContextTest.cpp)
+add_executable(velox_process_test TraceContextTest.cpp
+                                  ThreadLocalRegistryTest.cpp)
 
 add_test(velox_process_test velox_process_test)
 
diff --git a/velox/common/process/tests/ThreadLocalRegistryTest.cpp b/velox/common/process/tests/ThreadLocalRegistryTest.cpp
@@ -0,0 +1,112 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "velox/common/process/ThreadLocalRegistry.h"
+
+#include <folly/synchronization/Baton.h>
+#include <folly/synchronization/Latch.h>
+#include <gtest/gtest.h>
+
+#include <atomic>
+
+namespace facebook::velox::process {
+namespace {
+
+template <typename Tag>
+class TestObject {
+ public:
+  static std::atomic_int& count() {
+    static std::atomic_int value;
+    return value;
+  }
+
+  TestObject() : threadId_(std::this_thread::get_id()) {
+    ++count();
+  }
+
+  ~TestObject() {
+    --count();
+  }
+
+  std::thread::id threadId() const {
+    return threadId_;
+  }
+
+ private:
+  const std::thread::id threadId_;
+};
+
+TEST(ThreadLocalRegistryTest, basic) {
+  struct Tag {};
+  using T = TestObject<Tag>;
+  ASSERT_EQ(T::count(), 0);
+  auto registry = std::make_shared<ThreadLocalRegistry<T>>();
+  registry->forAllValues([](const T&) { FAIL(); });
+  thread_local ThreadLocalRegistry<T>::Reference ref(registry);
+  const T* object = ref.withValue([](const T& x) {
+    EXPECT_EQ(T::count(), 1);
+    return &x;
+  });
+  ASSERT_EQ(object->threadId(), std::this_thread::get_id());
+  ref.withValue([&](const T& x) { ASSERT_EQ(&x, object); });
+  int count = 0;
+  registry->forAllValues([&](const T& x) {
+    ++count;
+    ASSERT_EQ(x.threadId(), std::this_thread::get_id());
+  });
+  ASSERT_EQ(count, 1);
+  ASSERT_EQ(T::count(), 1);
+}
+
+TEST(ThreadLocalRegistryTest, multiThread) {
+  struct Tag {};
+  using T = TestObject<Tag>;
+  ASSERT_EQ(T::count(), 0);
+  auto registry = std::make_shared<ThreadLocalRegistry<T>>();
+  constexpr int kNumThreads = 7;
+  std::vector<std::thread> threads;
+  folly::Latch latch(kNumThreads);
+  folly::Baton<> batons[kNumThreads];
+  const T* objects[kNumThreads];
+  for (int i = 0; i < kNumThreads; ++i) {
+    threads.emplace_back([&, i] {
+      thread_local ThreadLocalRegistry<T>::Reference ref(registry);
+      objects[i] = ref.withValue([](const T& x) { return &x; });
+      latch.count_down();
+      batons[i].wait();
+    });
+  }
+  latch.wait();
+  std::vector<int> indices;
+  registry->forAllValues([&](const T& x) {
+    auto it = std::find(std::begin(objects), std::end(objects), &x);
+    indices.push_back(it - std::begin(objects));
+  });
+  ASSERT_EQ(indices.size(), kNumThreads);
+  std::sort(indices.begin(), indices.end());
+  for (int i = 0; i < kNumThreads; ++i) {
+    ASSERT_EQ(indices[i], i);
+    ASSERT_EQ(objects[i]->threadId(), threads[i].get_id());
+    ASSERT_EQ(T::count(), kNumThreads - i);
+    batons[i].post();
+    threads[i].join();
+  }
+  ASSERT_EQ(T::count(), 0);
+  registry->forAllValues([](const T&) { FAIL(); });
+}
+
+} // namespace
+} // namespace facebook::velox::process
