add optional thread finalizer to InitThreadFactory

Summary: Optional thread finalizer can help do cleanup work before thread exit.

Reviewed By: yfeldblum

Differential Revision: D8576224

fbshipit-source-id: b9ecd1b790b1935616fbe71c870dad82ec379b4e

folly/executors/test/ThreadPoolExecutorTest.cpp

@@ -463,6 +463,42 @@ TEST(InitThreadFactoryTest, InitializerCalled) {
   EXPECT_EQ(initializerCalledCount, 1);
 }
 
+TEST(InitThreadFactoryTest, InitializerAndFinalizerCalled) {
+  bool initializerCalled = false;
+  bool taskBodyCalled = false;
+  bool finalizerCalled = false;
+
+  InitThreadFactory factory(
+      std::make_shared<NamedThreadFactory>("test"),
+      [&] {
+        // thread initializer
+        EXPECT_FALSE(initializerCalled);
+        EXPECT_FALSE(taskBodyCalled);
+        EXPECT_FALSE(finalizerCalled);
+        initializerCalled = true;
+      },
+      [&] {
+        // thread finalizer
+        EXPECT_TRUE(initializerCalled);
+        EXPECT_TRUE(taskBodyCalled);
+        EXPECT_FALSE(finalizerCalled);
+        finalizerCalled = true;
+      });
+
+  factory
+      .newThread([&]() {
+        EXPECT_TRUE(initializerCalled);
+        EXPECT_FALSE(taskBodyCalled);
+        EXPECT_FALSE(finalizerCalled);
+        taskBodyCalled = true;
+      })
+      .join();
+
+  EXPECT_TRUE(initializerCalled);
+  EXPECT_TRUE(taskBodyCalled);
+  EXPECT_TRUE(finalizerCalled);
+}
+
 class TestData : public folly::RequestData {
  public:
   explicit TestData(int data) : data_(data) {}

folly/executors/thread_factory/InitThreadFactory.h

@@ -19,6 +19,7 @@
 #include <memory>
 #include <thread>
 
+#include <folly/ScopeGuard.h>
 #include <folly/executors/thread_factory/ThreadFactory.h>
 
 namespace folly {
@@ -27,22 +28,34 @@ class InitThreadFactory : public ThreadFactory {
  public:
   explicit InitThreadFactory(
       std::shared_ptr<ThreadFactory> threadFactory,
-      Func&& threadInitializer)
+      Func&& threadInitializer,
+      Func&& threadFinializer = [] {})
       : threadFactory_(std::move(threadFactory)),
-        threadInitializer_(
-            std::make_shared<Func>(std::move(threadInitializer))) {}
+        threadInitFini_(std::make_shared<ThreadInitFini>(
+            std::move(threadInitializer),
+            std::move(threadFinializer))) {}
 
   std::thread newThread(Func&& func) override {
     return threadFactory_->newThread(
-        [func = std::move(func), initializer = threadInitializer_]() mutable {
-          (*initializer)();
+        [func = std::move(func), threadInitFini = threadInitFini_]() mutable {
+          threadInitFini->initializer();
+          SCOPE_EXIT {
+            threadInitFini->finalizer();
+          };
           func();
         });
   }
 
  private:
   std::shared_ptr<ThreadFactory> threadFactory_;
-  std::shared_ptr<Func> threadInitializer_;
+  struct ThreadInitFini {
+    ThreadInitFini(Func&& init, Func&& fini)
+        : initializer(std::move(init)), finalizer(std::move(fini)) {}
+
+    Func initializer;
+    Func finalizer;
+  };
+  std::shared_ptr<ThreadInitFini> threadInitFini_;
 };
 
 } // namespace folly