Fix mocking to support multiple overrides

Summary: When installing the mock, we should make sure to remove singleton from the creation_order list.

Differential Revision: D3580725

fbshipit-source-id: dfb489de1be860ab639380644eab0b45a07a1450

folly/Singleton-inl.h

@@ -71,6 +71,16 @@ void SingletonHolder<T>::registerSingletonMock(CreateFunc c, TeardownFunc t) {
   }
   destroyInstance();
 
+  {
+    RWSpinLock::WriteHolder wh(&vault_.mutex_);
+
+    auto it = std::find(
+        vault_.creation_order_.begin(), vault_.creation_order_.end(), type());
+    if (it != vault_.creation_order_.end()) {
+      vault_.creation_order_.erase(it);
+    }
+  }
+
   std::lock_guard<std::mutex> entry_lock(mutex_);
 
   create_ = std::move(c);

folly/test/SingletonTest.cpp

@@ -582,6 +582,18 @@ TEST(Singleton, MockTest) {
 
   // If serial_count value is the same, then singleton was not replaced.
   EXPECT_NE(serial_count_first, serial_count_mock);
+
+  // Override existing mock using make_mock one more time
+  SingletonMock<Watchdog>::make_mock();
+
+  EXPECT_EQ(vault.registeredSingletonCount(), 1);
+  int serial_count_mock2 = SingletonMock<Watchdog>::try_get()->serial_number;
+
+  // If serial_count value is the same, then singleton was not replaced.
+  EXPECT_NE(serial_count_first, serial_count_mock2);
+  EXPECT_NE(serial_count_mock, serial_count_mock2);
+
+  vault.destroyInstances();
 }
 
 TEST(Singleton, DoubleRegistrationLogging) {