Overload getKeepAliveToken

Summary: getKeepAlive can get a KeepAlive object and return a copy of itself

Reviewed By: yfeldblum

Differential Revision: D14095402

fbshipit-source-id: 9d00ebecddc938e8cf02bcb7cee0642a666db349

folly/Executor.h

@@ -243,4 +243,10 @@ Executor::KeepAlive<ExecutorT> getKeepAliveToken(ExecutorT& executor) {
   return getKeepAliveToken(&executor);
 }
 
+template <typename ExecutorT>
+Executor::KeepAlive<ExecutorT> getKeepAliveToken(
+    Executor::KeepAlive<ExecutorT>& ka) {
+  return ka.copy();
+}
+
 } // namespace folly

folly/test/ExecutorTest.cpp

@@ -196,4 +196,23 @@ TEST(ExecutorTest, KeepAliveCopy) {
   EXPECT_EQ(0, exec.refCount);
 }
 
+TEST(ExecutorTest, GetKeepAliveTokenFromToken) {
+  KeepAliveTestExecutor exec;
+
+  {
+    auto ka = getKeepAliveToken(exec);
+    EXPECT_TRUE(ka);
+    EXPECT_EQ(&exec, ka.get());
+    EXPECT_EQ(1, exec.refCount);
+
+    auto ka2 = getKeepAliveToken(ka);
+    EXPECT_TRUE(ka);
+    EXPECT_TRUE(ka2);
+    EXPECT_EQ(&exec, ka2.get());
+    EXPECT_EQ(2, exec.refCount);
+  }
+
+  EXPECT_EQ(0, exec.refCount);
+}
+
 } // namespace folly