From 53bd909558f9a6db7710988f9f2134c8c8e4e253 Mon Sep 17 00:00:00 2001
From: Yedidya Feldblum <yfeldblum@meta.com>
Date: Mon, 5 Aug 2024 19:11:51 -0700
Subject: [PATCH] an overload of ThreadLocalPtr::reset taking std::shared_ptr

Summary:
A variation of `ThreadLocalPtr<std::shared_ptr<T>>` but with one less indirection to access the pointer.

Copies the `std::shared_ptr` into a custom deleter in order to keep the pointer alive.

Differential Revision: D59398514

fbshipit-source-id: 93d4df834b408dd5e86ea49c418c060ef7def4ef
---
 folly/ThreadLocal.h                |  4 +++
 folly/detail/ThreadLocalDetail.cpp | 10 ++++++
 folly/detail/ThreadLocalDetail.h   | 51 ++++++++++++++++++++++++++++--
 folly/test/ThreadLocalTest.cpp     | 12 +++++++
 4 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/folly/ThreadLocal.h b/folly/ThreadLocal.h
index 114b30f203d..34d20c21498 100644
--- a/folly/ThreadLocal.h
+++ b/folly/ThreadLocal.h
@@ -238,6 +238,10 @@ class ThreadLocalPtr {
     guard.dismiss();
   }
 
+  void reset(const std::shared_ptr<T>& newPtr) {
+    reset(newPtr.get(), threadlocal_detail::SharedPtrDeleter{newPtr});
+  }
+
   // Holds a global lock for iteration through all thread local child objects.
   // Can be used as an iterable container.
   // Use accessAllThreads() to obtain one.
diff --git a/folly/detail/ThreadLocalDetail.cpp b/folly/detail/ThreadLocalDetail.cpp
index 13cce1f9f67..c7f4376775e 100644
--- a/folly/detail/ThreadLocalDetail.cpp
+++ b/folly/detail/ThreadLocalDetail.cpp
@@ -31,6 +31,16 @@ constexpr auto kBigGrowthFactor = 1.7;
 namespace folly {
 namespace threadlocal_detail {
 
+SharedPtrDeleter::SharedPtrDeleter(std::shared_ptr<void> const& ts) noexcept
+    : ts_{ts} {}
+SharedPtrDeleter::SharedPtrDeleter(SharedPtrDeleter const& that) noexcept
+    : ts_{that.ts_} {}
+SharedPtrDeleter::~SharedPtrDeleter() = default;
+void SharedPtrDeleter::operator()(
+    void* /* ptr */, folly::TLPDestructionMode) const {
+  ts_.reset();
+}
+
 bool ThreadEntrySet::basicSanity() const {
   return //
       threadEntries.size() == entryToVectorSlot.size() &&
diff --git a/folly/detail/ThreadLocalDetail.h b/folly/detail/ThreadLocalDetail.h
index 477d69086cb..b8f1d9acf1d 100644
--- a/folly/detail/ThreadLocalDetail.h
+++ b/folly/detail/ThreadLocalDetail.h
@@ -57,6 +57,41 @@ namespace threadlocal_detail {
 
 constexpr uint32_t kEntryIDInvalid = std::numeric_limits<uint32_t>::max();
 
+//  as a memory-usage optimization, try to make this deleter fit in-situ in
+//  the deleter function storage rather than being heap-allocated separately
+//
+//  for libstdc++, specialization below of std::__is_location_invariant
+//
+//  TODO: ensure in-situ storage for other standard-library implementations
+struct SharedPtrDeleter {
+  mutable std::shared_ptr<void> ts_;
+  explicit SharedPtrDeleter(std::shared_ptr<void> const& ts) noexcept;
+  SharedPtrDeleter(SharedPtrDeleter const& that) noexcept;
+  void operator=(SharedPtrDeleter const& that) = delete;
+  ~SharedPtrDeleter();
+  void operator()(void* ptr, folly::TLPDestructionMode) const;
+};
+
+} // namespace threadlocal_detail
+
+} // namespace folly
+
+#if defined(__GLIBCXX__)
+
+namespace std {
+
+template <>
+struct __is_location_invariant<::folly::threadlocal_detail::SharedPtrDeleter>
+    : std::true_type {};
+
+} // namespace std
+
+#endif
+
+namespace folly {
+
+namespace threadlocal_detail {
+
 /**
  * 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.
@@ -118,6 +153,18 @@ struct ElementWrapper {
     ptr = p;
   }
 
+  template <typename Ptr, typename Deleter>
+  static auto makeDeleter(const Deleter& d) {
+    return [d](void* pt, TLPDestructionMode mode) {
+      d(static_cast<Ptr>(pt), mode);
+    };
+  }
+
+  template <typename Ptr>
+  static decltype(auto) makeDeleter(const SharedPtrDeleter& d) {
+    return d;
+  }
+
   template <class Ptr, class Deleter>
   void set(Ptr p, const Deleter& d) {
     DCHECK_EQ(static_cast<void*>(nullptr), ptr);
@@ -128,9 +175,7 @@ struct ElementWrapper {
     }
 
     auto guard = makeGuard([&] { d(p, TLPDestructionMode::THIS_THREAD); });
-    auto const obj = new DeleterObjType([d](void* pt, TLPDestructionMode mode) {
-      d(static_cast<Ptr>(pt), mode);
-    });
+    auto const obj = new DeleterObjType(makeDeleter<Ptr>(d));
     guard.dismiss();
     auto const raw = reinterpret_cast<uintptr_t>(obj);
     DCHECK_EQ(0, raw & deleter_all_mask);
diff --git a/folly/test/ThreadLocalTest.cpp b/folly/test/ThreadLocalTest.cpp
index f9482538d8b..5a5855c4c3c 100644
--- a/folly/test/ThreadLocalTest.cpp
+++ b/folly/test/ThreadLocalTest.cpp
@@ -255,6 +255,18 @@ TEST(ThreadLocalPtr, CustomDeleter2) {
   EXPECT_EQ(1010, Widget::totalVal_);
 }
 
+TEST(ThreadLocalPtr, SharedPtr) {
+  ThreadLocalPtr<int> tlp;
+  auto sp = std::make_shared<int>(7);
+  EXPECT_EQ(1, sp.use_count());
+  tlp.reset(sp);
+  EXPECT_EQ(2, sp.use_count());
+  EXPECT_EQ(sp.get(), tlp.get());
+  tlp.reset();
+  EXPECT_EQ(1, sp.use_count());
+  EXPECT_EQ(static_cast<void*>(nullptr), tlp.get());
+}
+
 TEST(ThreadLocal, NotDefaultConstructible) {
   struct Object {
     int value;