Detect singletons that are created in parent process and used in child process

Summary: folly::Singleton is used for singletons that manage resources (thread, fds etc.) and it's generally not safe to use previously created folly::Singletons in child process after fork.

Reviewed By: yfeldblum

Differential Revision: D25292667

fbshipit-source-id: 5343d26c72bb9203429a918a3bb5d458c44e9f3b

folly/Singleton-inl.h

@@ -74,7 +74,8 @@ void SingletonHolder<T>::registerSingletonMock(CreateFunc c, TeardownFunc t) {
   if (state_ == SingletonHolderState::NotRegistered) {
     detail::singletonWarnRegisterMockEarlyAndAbort(type());
   }
-  if (state_ == SingletonHolderState::Living) {
+  if (state_ == SingletonHolderState::Living ||
+      state_ == SingletonHolderState::LivingInChildAfterFork) {
     destroyInstance();
   }
 
@@ -171,6 +172,11 @@ void SingletonHolder<T>::preDestroyInstance(
 
 template <typename T>
 void SingletonHolder<T>::destroyInstance() {
+  if (state_.load(std::memory_order_relaxed) ==
+      SingletonHolderState::LivingInChildAfterFork) {
+    LOG(DFATAL) << "Attempting to destroy singleton " << type().name()
+                << " in child process after fork";
+  }
   state_ = SingletonHolderState::Dead;
   instance_.reset();
   instance_copy_.reset();
@@ -191,6 +197,16 @@ void SingletonHolder<T>::destroyInstance() {
   }
 }
 
+template <typename T>
+void SingletonHolder<T>::inChildAfterFork() {
+  auto expected = SingletonHolderState::Living;
+  state_.compare_exchange_strong(
+      expected,
+      SingletonHolderState::LivingInChildAfterFork,
+      std::memory_order_relaxed,
+      std::memory_order_relaxed);
+}
+
 template <typename T>
 SingletonHolder<T>::SingletonHolder(
     TypeDescriptor typeDesc,
@@ -224,6 +240,18 @@ void SingletonHolder<T>::createInstance() {
   if (state_.load(std::memory_order_acquire) == SingletonHolderState::Living) {
     return;
   }
+  if (state_.load(std::memory_order_relaxed) ==
+      SingletonHolderState::LivingInChildAfterFork) {
+    LOG(DFATAL) << "Attempting to use singleton " << type().name()
+                << " in child process after fork";
+    auto expected = SingletonHolderState::LivingInChildAfterFork;
+    state_.compare_exchange_strong(
+        expected,
+        SingletonHolderState::Living,
+        std::memory_order_relaxed,
+        std::memory_order_relaxed);
+    return;
+  }
   if (state_.load(std::memory_order_acquire) ==
       SingletonHolderState::NotRegistered) {
     detail::singletonWarnCreateUnregisteredAndAbort(type());

folly/Singleton.cpp

@@ -204,7 +204,34 @@ FatalHelper __attribute__((__init_priority__(101))) fatalHelper;
 
 } // namespace
 
+SingletonVault::SingletonVault(Type type) noexcept : type_(type) {
+  detail::AtFork::registerHandler(
+      this,
+      /*prepare*/
+      [this]() {
+        const auto& singletons = singletons_.unsafeGetUnlocked();
+        const auto& creationOrder = creationOrder_.unsafeGetUnlocked();
+
+        CHECK_GE(singletons.size(), creationOrder.size());
+
+        for (const auto& singletonType : creationOrder) {
+          liveSingletonsPreFork_.insert(singletons.at(singletonType));
+        }
+
+        return true;
+      },
+      /*parent*/ [this]() { liveSingletonsPreFork_.clear(); },
+      /*child*/
+      [this]() {
+        for (auto singleton : liveSingletonsPreFork_) {
+          singleton->inChildAfterFork();
+        }
+        liveSingletonsPreFork_.clear();
+      });
+}
+
 SingletonVault::~SingletonVault() {
+  detail::AtFork::unregisterHandler(this);
   destroyInstances();
 }
 

folly/Singleton.h

@@ -291,6 +291,7 @@ class SingletonHolderBase {
   virtual bool creationStarted() = 0;
   virtual void preDestroyInstance(ReadMostlyMainPtrDeleter<>&) = 0;
   virtual void destroyInstance() = 0;
+  virtual void inChildAfterFork() = 0;
 
  private:
   TypeDescriptor type_;
@@ -323,6 +324,7 @@ struct SingletonHolder : public SingletonHolderBase {
   bool creationStarted() override;
   void preDestroyInstance(ReadMostlyMainPtrDeleter<>&) override;
   void destroyInstance() override;
+  void inChildAfterFork() override;
 
  private:
   template <typename Tag, typename VaultTag>
@@ -334,6 +336,7 @@ struct SingletonHolder : public SingletonHolderBase {
     NotRegistered,
     Dead,
     Living,
+    LivingInChildAfterFork,
   };
 
   SingletonVault& vault_;
@@ -408,8 +411,7 @@ class SingletonVault {
 
   static Type defaultVaultType();
 
-  explicit SingletonVault(Type type = defaultVaultType()) noexcept
-      : type_(type) {}
+  explicit SingletonVault(Type type = defaultVaultType()) noexcept;
 
   // Destructor is only called by unit tests to check destroyInstances.
   ~SingletonVault();
@@ -553,6 +555,7 @@ class SingletonVault {
       eagerInitSingletons_;
   Synchronized<std::vector<detail::TypeDescriptor>, SharedMutexSuppressTSAN>
       creationOrder_;
+  std::unordered_set<detail::SingletonHolderBase*> liveSingletonsPreFork_;
 
   // Using SharedMutexReadPriority is important here, because we want to make
   // sure we don't block nested singleton creation happening concurrently with

folly/test/SingletonTest.cpp

@@ -1061,3 +1061,38 @@ TEST(Singleton, ShutdownTimerDisable) {
   vault.startShutdownTimer();
   vault.destroyInstances();
 }
+
+TEST(Singleton, ForkInChild) {
+  struct VaultTag {};
+  struct PrivateTag {};
+  struct ForkObject {};
+  using SingletonObject = Singleton<ForkObject, PrivateTag, VaultTag>;
+
+  auto& vault = *SingletonVault::singleton<VaultTag>();
+  SingletonObject object;
+  vault.registrationComplete();
+
+  // We use EXPECT_DEATH here to run code in the child process.
+  EXPECT_DEATH(
+      [&]() {
+        object.try_get();
+        vault.destroyInstances();
+
+        LOG(FATAL) << "Finished successfully";
+      }(),
+      "Finished successfully");
+
+  object.try_get();
+
+  if (!folly::kIsDebug) {
+    return;
+  }
+
+  EXPECT_DEATH(
+      [&]() { object.try_get(); }(),
+      "Attempting to use singleton .*ForkObject.* in child process");
+
+  EXPECT_DEATH(
+      [&]() { vault.destroyInstances(); }(),
+      "Attempting to destroy singleton .*ForkObject.* in child process");
+}