Add TerminateCancellationToken utility

Summary: `getTerminateCancellationToken()` is the public API that can be used to get a cancellation token that'll be cancelled when any of the registered signal (SIGTERM or SIGINT) is received.

Reviewed By: andriigrynenko

Differential Revision: D29883069

fbshipit-source-id: 4270c42316b80afb2c75f54de7e720dcb7a083ca

folly/io/async/TerminateCancellationToken.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/io/async/TerminateCancellationToken.h>
+
+#include <csignal>
+
+#include <folly/Singleton.h>
+#include <folly/io/async/AsyncSignalHandler.h>
+#include <folly/io/async/ScopedEventBaseThread.h>
+
+namespace folly {
+namespace {
+/**
+ * A helper class that starts a new thread to listen for signals. It can issue
+ * CancellationToken that can be used to schedule callback to execute on
+ * receiving the signals.
+ */
+class ScopedTerminateSignalHandler : private AsyncSignalHandler {
+ public:
+  ScopedTerminateSignalHandler() : AsyncSignalHandler(nullptr) {
+    attachEventBase(scopedEventBase_.getEventBase());
+    // To prevent data race with unregisterSignalHandler
+    scopedEventBase_.getEventBase()->runInEventBaseThreadAndWait([this]() {
+      registerSignalHandler(SIGTERM);
+      registerSignalHandler(SIGINT);
+    });
+  }
+
+  CancellationToken getCancellationToken() {
+    return cancellationSource_.getToken();
+  }
+
+ private:
+  void signalReceived(int) noexcept override {
+    // Unregister signal handler as we want to handle the signal only once
+    unregisterSignalHandler(SIGTERM);
+    unregisterSignalHandler(SIGINT);
+    cancellationSource_.requestCancellation();
+  }
+
+  ScopedEventBaseThread scopedEventBase_;
+  CancellationSource cancellationSource_;
+};
+
+folly::Singleton<ScopedTerminateSignalHandler> singletonSignalHandler;
+
+} // namespace
+
+CancellationToken getTerminateCancellationToken() {
+  auto signalHandler = singletonSignalHandler.try_get();
+  if (signalHandler == nullptr) {
+    CancellationSource cs;
+    cs.requestCancellation();
+    return cs.getToken();
+  }
+  return signalHandler->getCancellationToken();
+}
+
+} // namespace folly

folly/io/async/TerminateCancellationToken.h

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <folly/CancellationToken.h>
+
+namespace folly {
+
+/**
+ * Returns a CancellationToken that can be used to schedule callbacks. The
+ * CancellationToken is cancelled when any of SIGTERM and SIGINT
+ * signal is received.
+ */
+CancellationToken getTerminateCancellationToken();
+
+} // namespace folly

folly/io/async/test/TerminateCancellationTokenTest.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/io/async/TerminateCancellationToken.h>
+
+#include <chrono>
+
+#include <folly/portability/GTest.h>
+#include <folly/synchronization/Baton.h>
+
+using namespace folly;
+using namespace std::chrono;
+
+TEST(TerminateCancellationTokenTest, addCallbacks) {
+  bool called1 = false, called2 = false;
+  Baton<> done1, done2;
+  auto cb1 = CancellationCallback(getTerminateCancellationToken(), [&]() {
+    called1 = true;
+    done1.post();
+  });
+  auto cb2 = CancellationCallback(getTerminateCancellationToken(), [&]() {
+    called2 = true;
+    done2.post();
+  });
+  kill(getpid(), SIGTERM);
+  EXPECT_TRUE(done1.try_wait_for(seconds(2)));
+  EXPECT_TRUE(called1);
+  EXPECT_TRUE(done2.try_wait_for(seconds(2)));
+  EXPECT_TRUE(called2);
+}