VirtualExecutor join should wait for function to be destroyed

Reviewed By: aary

Differential Revision: D8393345

fbshipit-source-id: 5fc9bbfb3fa1dc63077cd406f4335c2a7967a53d

folly/Makefile.am

@@ -519,7 +519,8 @@ nobase_follyinclude_HEADERS = \
 	Uri.h \
 	Uri-inl.h \
 	Utility.h \
-	Varint.h
+	Varint.h \
+	VirtualExecutor.h
 
 libfollybasesse42_la_SOURCES = \
 	detail/RangeSse42.cpp \

folly/VirtualExecutor.h

@@ -30,6 +30,24 @@ namespace folly {
  * backed by it are released.
  */
 class VirtualExecutor : public DefaultKeepAliveExecutor {
+  auto wrapFunc(Func f) {
+    class FuncAndKeepAlive {
+     public:
+      FuncAndKeepAlive(Func&& f, VirtualExecutor* executor)
+          : keepAlive_(getKeepAliveToken(executor)), f_(std::move(f)) {}
+
+      void operator()() {
+        f_();
+      }
+
+     private:
+      Executor::KeepAlive<VirtualExecutor> keepAlive_;
+      Func f_;
+    };
+
+    return FuncAndKeepAlive(std::move(f), this);
+  }
+
  public:
   explicit VirtualExecutor(KeepAlive<> executor)
       : executor_(std::move(executor)) {
@@ -50,16 +68,11 @@ class VirtualExecutor : public DefaultKeepAliveExecutor {
   }
 
   void add(Func f) override {
-    executor_->add([func = std::move(f),
-                    me = getKeepAliveToken(this)]() mutable { func(); });
+    executor_->add(wrapFunc(std::move(f)));
   }
 
   void addWithPriority(Func f, int8_t priority) override {
-    executor_->addWithPriority(
-        [func = std::move(f), me = getKeepAliveToken(this)]() mutable {
-          func();
-        },
-        priority);
+    executor_->addWithPriority(wrapFunc(std::move(f)), priority);
   }
 
   ~VirtualExecutor() override {

folly/executors/test/ThreadPoolExecutorTest.cpp

@@ -775,6 +775,21 @@ static void virtualExecutorTest() {
               .semi();
     }
     EXPECT_EQ(1, counter);
+
+    bool functionDestroyed{false};
+    bool functionCalled{false};
+    {
+      VirtualExecutor ve(fe);
+      auto guard = makeGuard([&functionDestroyed] {
+        std::this_thread::sleep_for(100ms);
+        functionDestroyed = true;
+      });
+      ve.add([&functionCalled, guard = std::move(guard)] {
+        functionCalled = true;
+      });
+    }
+    EXPECT_TRUE(functionCalled);
+    EXPECT_TRUE(functionDestroyed);
   }
   EXPECT_TRUE(f->isReady());
   EXPECT_NO_THROW(std::move(*f).get());