Adding support for multiple apps in WDT

Summary:Adding support for multiple apps in WDT. Because of single app restriction, it
is hard to write tests involving multiple apps

Reviewed By: ldemailly

Differential Revision: D3069096

fb-gh-sync-id: 3cf069bc4d77d107d5c2afc3e4145bf4c665b204
shipit-source-id: 3cf069bc4d77d107d5c2afc3e4145bf4c665b204

Wdt.cpp

@@ -1,17 +1,16 @@
 #include <wdt/Wdt.h>
 #include <wdt/util/WdtFlags.h>
 
-#include <wdt/WdtResourceController.h>
-
 using std::string;
 
 namespace facebook {
 namespace wdt {
 
 // this must be called first and exactly once:
 Wdt &Wdt::initializeWdt(const std::string &appName) {
-  Wdt &res = getWdtInternal();
+  Wdt &res = getWdtInternal(appName);
   res.initializeWdtInternal(appName);
+  // TODO this should return the options
   WdtFlags::initializeFromFlags();
   // At fb we do this for services - that's floody for cmd line though
   // res.printWdtOptions(WLOG(INFO));
@@ -26,20 +25,18 @@ ErrorCode Wdt::initializeWdtInternal(const std::string &appName) {
   }
   appName_ = appName;
   initDone_ = true;
-  // TODO this should return the options
   return OK;
 }
 
 ErrorCode Wdt::applySettings() {
-  WdtResourceController::get()->setThrottler(
-      Throttler::makeThrottler(options_));
+  resourceController_.setThrottler(Throttler::makeThrottler(options_));
   settingsApplied_ = true;
   return OK;
 }
 
 // this can be called many times after initializeWdt()
-Wdt &Wdt::getWdt() {
-  Wdt &res = getWdtInternal();
+Wdt &Wdt::getWdt(const std::string &appName) {
+  Wdt &res = getWdtInternal(appName);
   if (!res.initDone_) {
     WLOG(ERROR) << "Called getWdt() before/without initializeWdt()";
     WDT_CHECK(false) << "Must call initializeWdt() once before getWdt()";
@@ -69,19 +66,19 @@ ErrorCode Wdt::wdtSend(const std::string &wdtNamespace,
 
   // try to create sender
   SenderPtr sender;
-  auto wdtController = WdtResourceController::get();
   // TODO should be using recoverid
   const std::string secondKey = req.hostName;
   ErrorCode errCode =
-      wdtController->createSender(wdtNamespace, secondKey, req, sender);
+      resourceController_.createSender(wdtNamespace, secondKey, req, sender);
   if (errCode == ALREADY_EXISTS && terminateExistingOne) {
     WLOG(WARNING) << "Found pre-existing sender for " << wdtNamespace << " "
                   << secondKey << " aborting it and making a new one";
     sender->abort(ABORTED_BY_APPLICATION);
     // This may log an error too
-    wdtController->releaseSender(wdtNamespace, secondKey);
+    resourceController_.releaseSender(wdtNamespace, secondKey);
     // Try#2
-    errCode = wdtController->createSender(wdtNamespace, secondKey, req, sender);
+    errCode =
+        resourceController_.createSender(wdtNamespace, secondKey, req, sender);
   }
   if (errCode != OK) {
     WLOG(ERROR) << "Failed to create sender " << errorCodeToStr(errCode) << " "
@@ -99,7 +96,7 @@ ErrorCode Wdt::wdtSend(const std::string &wdtNamespace,
   }
   auto transferReport = sender->transfer();
   ErrorCode ret = transferReport->getSummary().getErrorCode();
-  wdtController->releaseSender(wdtNamespace, secondKey);
+  resourceController_.releaseSender(wdtNamespace, secondKey);
   WLOG(INFO) << "wdtSend for " << wdtNamespace << " " << secondKey << " "
              << " ended with " << errorCodeToStr(ret);
   return ret;
@@ -119,9 +116,18 @@ WdtOptions &Wdt::getWdtOptions() {
 }
 
 // private version
-Wdt &Wdt::getWdtInternal() {
-  static Wdt s_wdtInstance{WdtOptions::getMutable()};
-  return s_wdtInstance;
+Wdt &Wdt::getWdtInternal(const std::string &appName) {
+  static std::unordered_map<std::string, std::unique_ptr<Wdt>> s_wdtMap;
+  static std::mutex mutex;
+  std::lock_guard<std::mutex> lock(mutex);
+  auto it = s_wdtMap.find(appName);
+  if (it != s_wdtMap.end()) {
+    return *(it->second);
+  }
+  Wdt *wdtPtr = new Wdt(WdtOptions::getMutable());
+  std::unique_ptr<Wdt> wdt(wdtPtr);
+  s_wdtMap.emplace(appName, std::move(wdt));
+  return *wdtPtr;
 }
 }
 }  // namespaces

Wdt.h

@@ -15,6 +15,7 @@
 #include <wdt/ErrorCodes.h>
 // For IAbortChecker and WdtTransferRequest - TODO: split out ?
 #include <wdt/WdtBase.h>
+#include <wdt/WdtResourceController.h>
 #include <wdt/util/EncryptionUtils.h>
 #include <ostream>
 
@@ -30,7 +31,7 @@ namespace wdt {
  *
  * Example of use:
  * // During the single threaded part of your service's initialization
- * Wdt &wdt = Wdt::initializeWdt();
+ * Wdt &wdt = Wdt::initializeWdt("app-name");
  * // Optionally: change the WdtOptions as you need, for instance:
  * wdt.getWdtOptions().overwrite = true;
  * // Will use the (possibly changed above) settings, to configure wdt,
@@ -57,7 +58,7 @@ class Wdt {
    * This is only needed if the caller code doesn't want to pass the Wdt
    * instance around.
    */
-  static Wdt &getWdt();
+  static Wdt &getWdt(const std::string &appName);
 
   /// High level APIs:
 
@@ -75,6 +76,10 @@ class Wdt {
 
   virtual ErrorCode printWdtOptions(std::ostream &out);
 
+  /// Virtual Destructor (for class hierarchy)
+  virtual ~Wdt() {
+  }
+
  protected:
   /// Set to true when the single instance is initialized
   bool initDone_{false};
@@ -94,6 +99,10 @@ class Wdt {
   /// responsible for initializing openssl
   WdtCryptoIntializer cryptoInitializer_;
 
+  // TODO: share resource controller across apps
+  /// wdt resource controller
+  WdtResourceController resourceController_;
+
   // Internal initialization so sub classes can share the code
   virtual ErrorCode initializeWdtInternal(const std::string &appName);
 
@@ -103,8 +112,8 @@ class Wdt {
       const WdtTransferRequest &req,
       std::shared_ptr<IAbortChecker> abortChecker);
 
-  // Internal singleton creator/holder
-  static Wdt &getWdtInternal();
+  // Internal wdt object creator/holder
+  static Wdt &getWdtInternal(const std::string &appName);
 
   /**
    * Apply the possibly changed settings, eg throttler.
@@ -121,10 +130,6 @@ class Wdt {
   /// Not copyable
   Wdt(const Wdt &) = delete;
   Wdt &operator=(const Wdt &) = delete;
-
-  /// Virtual Destructor (for class hierarchy)
-  virtual ~Wdt() {
-  }
 };
 }
 }  // namespaces

test/WdtUrlTest.cpp

@@ -6,8 +6,7 @@
  * LICENSE file in the root directory of this source tree. An additional grant
  * of patent rights can be found in the PATENTS file in the same directory.
  */
-#include <wdt/Receiver.h>
-#include <wdt/Sender.h>
+#include <wdt/Wdt.h>
 #include <wdt/test/TestCommon.h>
 
 #include <gflags/gflags.h>
@@ -306,10 +305,20 @@ TEST(TransferRequestTest, Encryption1) {
     EXPECT_EQ(req, unser);
   }
 }
+
+TEST(Wdt, WdtInstanceTest) {
+  Wdt& wdt1 = Wdt::initializeWdt("wdt");
+  Wdt& wdt2 = Wdt::getWdt("wdt");
+  Wdt& wdt3 = Wdt::getWdt("wdt");
+  Wdt& wdt4 = Wdt::initializeWdt("wdt1");
+  EXPECT_EQ(&wdt1, &wdt2);
+  EXPECT_EQ(&wdt2, &wdt3);
+  EXPECT_NE(&wdt1, &wdt4);
+}
 }
 }  // namespace end
 
-int main(int argc, char *argv[]) {
+int main(int argc, char* argv[]) {
   FLAGS_logtostderr = true;
   testing::InitGoogleTest(&argc, argv);
   google::ParseCommandLineFlags(&argc, &argv, true);