From d70b7399f39ea6c183d1070a654a24026985cb10 Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Thu, 23 Mar 2017 11:17:04 -0700
Subject: [PATCH] Make sure singleton destructor is always called from main
 thread.

Reviewed By: yfeldblum, ot

Differential Revision: D4742262

fbshipit-source-id: 7137bb8e1cb39c74e8ba726e76e9a8a6f3e1984c
---
 folly/Singleton-inl.h        | 48 +++++++++++++++++++-----------------
 folly/test/SingletonTest.cpp | 43 ++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 23 deletions(-)

diff --git a/folly/Singleton-inl.h b/folly/Singleton-inl.h
index 207a0f7d..40c7a27a 100644
--- a/folly/Singleton-inl.h
+++ b/folly/Singleton-inl.h
@@ -153,14 +153,18 @@ void SingletonHolder<T>::destroyInstance() {
   instance_copy_.reset();
   if (destroy_baton_) {
     constexpr std::chrono::seconds kDestroyWaitTime{5};
-    auto wait_result = destroy_baton_->timed_wait(
-      std::chrono::steady_clock::now() + kDestroyWaitTime);
-    if (!wait_result) {
+    auto last_reference_released = destroy_baton_->timed_wait(
+        std::chrono::steady_clock::now() + kDestroyWaitTime);
+    if (last_reference_released) {
+      teardown_(instance_ptr_);
+    } else {
       print_destructor_stack_trace_->store(true);
       LOG(ERROR) << "Singleton of type " << type().name() << " has a "
                  << "living reference at destroyInstances time; beware! Raw "
                  << "pointer is " << instance_ptr_ << ". It is very likely "
                  << "that some other singleton is holding a shared_ptr to it. "
+                 << "This singleton will be leaked (even if a shared_ptr to it "
+                 << "is eventually released)."
                  << "Make sure dependencies between these singletons are "
                  << "properly defined.";
     }
@@ -242,30 +246,28 @@ void SingletonHolder<T>::createInstance() {
   auto destroy_baton = std::make_shared<folly::Baton<>>();
   auto print_destructor_stack_trace =
     std::make_shared<std::atomic<bool>>(false);
-  auto teardown = teardown_;
 
   // Can't use make_shared -- no support for a custom deleter, sadly.
   std::shared_ptr<T> instance(
-    create_(),
-    [destroy_baton, print_destructor_stack_trace, teardown, type = type()]
-    (T* instance_ptr) mutable {
-      teardown(instance_ptr);
-      destroy_baton->post();
-      if (print_destructor_stack_trace->load()) {
-        std::string output = "Singleton " + type.name() + " was destroyed.\n";
-
-        auto stack_trace_getter = SingletonVault::stackTraceGetter().load();
-        auto stack_trace = stack_trace_getter ? stack_trace_getter() : "";
-        if (stack_trace.empty()) {
-          output += "Failed to get destructor stack trace.";
-        } else {
-          output += "Destructor stack trace:\n";
-          output += stack_trace;
+      create_(),
+      [ destroy_baton, print_destructor_stack_trace, type = type() ](
+          T*) mutable {
+        destroy_baton->post();
+        if (print_destructor_stack_trace->load()) {
+          std::string output = "Singleton " + type.name() + " was released.\n";
+
+          auto stack_trace_getter = SingletonVault::stackTraceGetter().load();
+          auto stack_trace = stack_trace_getter ? stack_trace_getter() : "";
+          if (stack_trace.empty()) {
+            output += "Failed to get release stack trace.";
+          } else {
+            output += "Release stack trace:\n";
+            output += stack_trace;
+          }
+
+          LOG(ERROR) << output;
         }
-
-        LOG(ERROR) << output;
-      }
-    });
+      });
 
   // We should schedule destroyInstances() only after the singleton was
   // created. This will ensure it will be destroyed before singletons,
diff --git a/folly/test/SingletonTest.cpp b/folly/test/SingletonTest.cpp
index 98ca0d31..36aeb4e1 100644
--- a/folly/test/SingletonTest.cpp
+++ b/folly/test/SingletonTest.cpp
@@ -674,3 +674,46 @@ TEST(Singleton, ConcurrentCreationDestruction) {
 
   needyThread.join();
 }
+
+struct MainThreadDestructorTag {};
+template <typename T, typename Tag = detail::DefaultTag>
+using SingletonMainThreadDestructor =
+    Singleton<T, Tag, MainThreadDestructorTag>;
+
+struct ThreadLoggingSingleton {
+  ThreadLoggingSingleton() {
+    initThread = std::this_thread::get_id();
+  }
+
+  ~ThreadLoggingSingleton() {
+    destroyThread = std::this_thread::get_id();
+  }
+
+  static std::thread::id initThread;
+  static std::thread::id destroyThread;
+};
+std::thread::id ThreadLoggingSingleton::initThread{};
+std::thread::id ThreadLoggingSingleton::destroyThread{};
+
+TEST(Singleton, MainThreadDestructor) {
+  auto& vault = *SingletonVault::singleton<MainThreadDestructorTag>();
+  SingletonMainThreadDestructor<ThreadLoggingSingleton> singleton;
+
+  vault.registrationComplete();
+  EXPECT_EQ(std::thread::id(), ThreadLoggingSingleton::initThread);
+
+  singleton.try_get();
+  EXPECT_EQ(std::this_thread::get_id(), ThreadLoggingSingleton::initThread);
+
+  std::thread t([instance = singleton.try_get()] {
+    /* sleep override */ std::this_thread::sleep_for(
+        std::chrono::milliseconds{100});
+  });
+
+  EXPECT_EQ(std::thread::id(), ThreadLoggingSingleton::destroyThread);
+
+  vault.destroyInstances();
+  EXPECT_EQ(std::this_thread::get_id(), ThreadLoggingSingleton::destroyThread);
+
+  t.join();
+}
-- 
2.34.1