From 8ff0d130807328323730f5ea974df3f978b1d499 Mon Sep 17 00:00:00 2001
From: Steve O'Brien <steveo@fb.com>
Date: Thu, 8 Oct 2015 15:03:22 -0700
Subject: [PATCH] aggregator: fix `doEagerInitVia`

Summary: Allow notification upon async eager initialization completion.

Reviewed By: @luciang

Differential Revision: D2518925

fb-gh-sync-id: 95f96da2afa6e5953630dd935d74819521716e5f
---
 folly/Singleton.cpp          | 24 +++++++++++++++++++++--
 folly/Singleton.h            | 37 +++++++++++++++++++++++++++++-------
 folly/test/SingletonTest.cpp |  6 ++++--
 3 files changed, 56 insertions(+), 11 deletions(-)

diff --git a/folly/Singleton.cpp b/folly/Singleton.cpp
index 52f3957f..c2a8afde 100644
--- a/folly/Singleton.cpp
+++ b/folly/Singleton.cpp
@@ -16,8 +16,11 @@
 
 #include <folly/Singleton.h>
 
+#include <atomic>
 #include <string>
 
+#include <folly/ScopeGuard.h>
+
 namespace folly {
 
 namespace detail {
@@ -128,7 +131,7 @@ void SingletonVault::doEagerInit() {
   }
 }
 
-void SingletonVault::doEagerInitVia(Executor* exe) {
+void SingletonVault::doEagerInitVia(Executor& exe, folly::Baton<>* done) {
   std::unordered_set<detail::SingletonHolderBase*> singletonSet;
   {
     RWSpinLock::ReadHolder rh(&stateMutex_);
@@ -139,8 +142,25 @@ void SingletonVault::doEagerInitVia(Executor* exe) {
     singletonSet = eagerInitSingletons_; // copy set of pointers
   }
 
+  auto countdown = std::make_shared<std::atomic<size_t>>(singletonSet.size());
   for (auto* single : singletonSet) {
-    exe->add([single] {
+    // countdown is retained by shared_ptr, and will be alive until last lambda
+    // is done.  notifyBaton is provided by the caller, and expected to remain
+    // present (if it's non-nullptr).  singletonSet can go out of scope but
+    // its values, which are SingletonHolderBase pointers, are alive as long as
+    // SingletonVault is not being destroyed.
+    exe.add([=] {
+      // decrement counter and notify if requested, whether initialization
+      // was successful, was skipped (already initialized), or exception thrown.
+      SCOPE_EXIT {
+        if (--(*countdown) == 0) {
+          if (done != nullptr) {
+            done->post();
+          }
+        }
+      };
+      // if initialization is in progress in another thread, don't try to init
+      // here.  Otherwise the current thread will block on 'createInstance'.
       if (!single->creationStarted()) {
         single->createInstance();
       }
diff --git a/folly/Singleton.h b/folly/Singleton.h
index ac5b8328..01e0bed5 100644
--- a/folly/Singleton.h
+++ b/folly/Singleton.h
@@ -112,16 +112,18 @@
 #include <folly/io/async/Request.h>
 
 #include <algorithm>
-#include <vector>
-#include <mutex>
-#include <thread>
+#include <atomic>
 #include <condition_variable>
+#include <functional>
+#include <memory>
+#include <mutex>
 #include <string>
+#include <thread>
+#include <typeindex>
+#include <typeinfo>
 #include <unordered_map>
 #include <unordered_set>
-#include <functional>
-#include <typeinfo>
-#include <typeindex>
+#include <vector>
 
 #include <glog/logging.h>
 
@@ -340,8 +342,23 @@ class SingletonVault {
 
   /**
    * Schedule eager singletons' initializations through the given executor.
+   * If baton ptr is not null, its `post` method is called after all
+   * early initialization has completed.
+   *
+   * If exceptions are thrown during initialization, this method will still
+   * `post` the baton to indicate completion.  The exception will not propagate
+   * and future attempts to `try_get` or `get_weak` the failed singleton will
+   * retry initialization.
+   *
+   * Sample usage:
+   *
+   *   wangle::IOThreadPoolExecutor executor(max_concurrency_level);
+   *   folly::Baton<> done;
+   *   doEagerInitVia(executor, &done);
+   *   done.wait();  // or 'timed_wait', or spin with 'try_wait'
+   *
    */
-  void doEagerInitVia(Executor* exe);
+  void doEagerInitVia(Executor& exe, folly::Baton<>* done = nullptr);
 
   // Destroy all singletons; when complete, the vault can't create
   // singletons once again until reenableInstances() is called.
@@ -357,6 +374,12 @@ class SingletonVault {
     return singletons_.size();
   }
 
+  /**
+   * Flips to true if eager initialization was used, and has completed.
+   * Never set to true if "doEagerInit()" or "doEagerInitVia" never called.
+   */
+  bool eagerInitComplete() const;
+
   size_t livingSingletonCount() const {
     RWSpinLock::ReadHolder rh(&mutex_);
 
diff --git a/folly/test/SingletonTest.cpp b/folly/test/SingletonTest.cpp
index 5eff5cc1..da62e18f 100644
--- a/folly/test/SingletonTest.cpp
+++ b/folly/test/SingletonTest.cpp
@@ -457,10 +457,12 @@ TEST(Singleton, SingletonEagerInitAsync) {
                   [&] {didEagerInit = true; return new std::string("foo"); })
               .shouldEagerInit();
   folly::EventBase eb;
+  folly::Baton<> done;
   vault.registrationComplete();
   EXPECT_FALSE(didEagerInit);
-  vault.doEagerInitVia(&eb);
+  vault.doEagerInitVia(eb, &done);
   eb.loop();
+  done.wait();
   EXPECT_TRUE(didEagerInit);
   sing.get_weak();  // (avoid compile error complaining about unused var 'sing')
 }
@@ -537,7 +539,7 @@ TEST(Singleton, SingletonEagerInitParallel) {
       for (size_t j = 0; j < kThreads; j++) {
         threads.push_back(std::make_shared<std::thread>([&] {
           barrier.wait();
-          vault.doEagerInitVia(&exe);
+          vault.doEagerInitVia(exe);
         }));
       }
 
-- 
2.34.1