From 7d0d018f7661e1994c7f68c8b80483d80a96f9fd Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Mon, 3 Apr 2017 20:04:44 -0700
Subject: [PATCH] Don't allow singleton use w/o registrationComplete()

Summary: This is same as D4249032, but for now this is only enabled in dbg and dev builds. I'm planning to bump this to FATAL later, once we fix know problematic main()s.

Reviewed By: lbrandy, yfeldblum

Differential Revision: D4821763

fbshipit-source-id: 5f7930f8cbcb10275d23a89848f1ec8ee34a8020
---
 folly/Singleton-inl.h | 11 +++++++++++
 folly/Singleton.cpp   | 22 ++++++++++++++++++++++
 folly/Singleton.h     |  4 +++-
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/folly/Singleton-inl.h b/folly/Singleton-inl.h
index cacfec18..2e69d90e 100644
--- a/folly/Singleton-inl.h
+++ b/folly/Singleton-inl.h
@@ -230,6 +230,17 @@ void SingletonHolder<T>::createInstance() {
   creating_thread_.store(std::this_thread::get_id(), std::memory_order_release);
 
   auto state = vault_.state_.rlock();
+  if (vault_.type_ != SingletonVault::Type::Relaxed &&
+      !state->registrationComplete) {
+    auto stack_trace_getter = SingletonVault::stackTraceGetter().load();
+    auto stack_trace = stack_trace_getter ? stack_trace_getter() : "";
+    if (!stack_trace.empty()) {
+      stack_trace = "Stack trace:\n" + stack_trace;
+    }
+
+    LOG(DFATAL) << "Singleton " << type().name() << " requested before "
+                << "registrationComplete() call. " << stack_trace;
+  }
   if (state->state == SingletonVault::SingletonVaultState::Quiescing) {
     return;
   }
diff --git a/folly/Singleton.cpp b/folly/Singleton.cpp
index 67e6fd81..4f24328e 100644
--- a/folly/Singleton.cpp
+++ b/folly/Singleton.cpp
@@ -16,6 +16,10 @@
 
 #include <folly/Singleton.h>
 
+#ifndef _WIN32
+#include <dlfcn.h>
+#endif
+
 #include <atomic>
 #include <cstdio>
 #include <cstdlib>
@@ -24,8 +28,26 @@
 
 #include <folly/ScopeGuard.h>
 
+#if !defined(_WIN32) && !defined(__APPLE__)
+static void hs_init_weak(int* argc, char** argv[])
+    __attribute__((__weakref__("hs_init")));
+#endif
+
 namespace folly {
 
+SingletonVault::Type SingletonVault::defaultVaultType() {
+#if !defined(_WIN32) && !defined(__APPLE__)
+  bool isPython = dlsym(RTLD_DEFAULT, "Py_Main");
+  bool isHaskel = &::hs_init_weak || dlsym(RTLD_DEFAULT, "hs_init");
+  bool isJVM = dlsym(RTLD_DEFAULT, "JNI_GetCreatedJavaVMs");
+  bool isD = dlsym(RTLD_DEFAULT, "_d_run_main");
+
+  return isPython || isHaskel || isJVM || isD ? Type::Relaxed : Type::Strict;
+#else
+  return Type::Relaxed;
+#endif
+}
+
 namespace detail {
 
 [[noreturn]] void singletonWarnDoubleRegistrationAndAbort(
diff --git a/folly/Singleton.h b/folly/Singleton.h
index 2c699e47..e8344cef 100644
--- a/folly/Singleton.h
+++ b/folly/Singleton.h
@@ -358,7 +358,9 @@ class SingletonVault {
     }
   };
 
-  explicit SingletonVault(Type type = Type::Strict) : type_(type) {}
+  static Type defaultVaultType();
+
+  explicit SingletonVault(Type type = defaultVaultType()) : type_(type) {}
 
   // Destructor is only called by unit tests to check destroyInstances.
   ~SingletonVault();
-- 
2.34.1