From af8129a20218967099a90fc419762f3770574377 Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Thu, 12 May 2016 16:20:40 -0700
Subject: [PATCH] Register custom stack for signal handlers

Summary: This updates FiberManager to automatically register custom stack for signal handlers. This ensures signal handlers will work correctly not only in case of stack overflow, but also if fiber stack size is just not enough for them.

Reviewed By: yfeldblum

Differential Revision: D3289912

fbshipit-source-id: 44959b727f70e1f1748d837b1a06c7414433e5ec
---
 folly/experimental/fibers/FiberManager-inl.h  |  4 ++
 folly/experimental/fibers/FiberManager.cpp    | 60 +++++++++++++++++++
 folly/experimental/fibers/FiberManager.h      |  4 ++
 .../fibers/test/StackOverflow.cpp             | 39 ++++++++++++
 4 files changed, 107 insertions(+)
 create mode 100644 folly/experimental/fibers/test/StackOverflow.cpp

diff --git a/folly/experimental/fibers/FiberManager-inl.h b/folly/experimental/fibers/FiberManager-inl.h
index 52c283f5..697f14ba 100644
--- a/folly/experimental/fibers/FiberManager-inl.h
+++ b/folly/experimental/fibers/FiberManager-inl.h
@@ -164,6 +164,10 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
 }
 
 inline bool FiberManager::loopUntilNoReady() {
+  if (UNLIKELY(!alternateSignalStackRegistered_)) {
+    registerAlternateSignalStack();
+  }
+
   // Support nested FiberManagers
   auto originalFiberManager = this;
   std::swap(currentFiberManager_, originalFiberManager);
diff --git a/folly/experimental/fibers/FiberManager.cpp b/folly/experimental/fibers/FiberManager.cpp
index 77239537..2592ee1c 100644
--- a/folly/experimental/fibers/FiberManager.cpp
+++ b/folly/experimental/fibers/FiberManager.cpp
@@ -15,6 +15,7 @@
  */
 #include "FiberManager.h"
 
+#include <signal.h>
 #include <sys/syscall.h>
 #include <unistd.h>
 
@@ -26,6 +27,8 @@
 #include <folly/experimental/fibers/Fiber.h>
 #include <folly/experimental/fibers/LoopController.h>
 
+#include <folly/SingletonThreadLocal.h>
+
 #ifdef FOLLY_SANITIZE_ADDRESS
 
 #include <dlfcn.h>
@@ -273,5 +276,62 @@ static AsanUnpoisonMemoryRegionFuncPtr getUnpoisonMemoryRegionFunc() {
 }
 
 #endif // FOLLY_SANITIZE_ADDRESS
+
+namespace {
+
+// SIGSTKSZ (8 kB on our architectures) isn't always enough for
+// folly::symbolizer, so allocate 32 kB.
+constexpr size_t kAltStackSize = folly::constexpr_max(SIGSTKSZ, 32 * 1024);
+
+bool hasAlternateStack() {
+  stack_t ss;
+  sigaltstack(nullptr, &ss);
+  return !(ss.ss_flags & SS_DISABLE);
+}
+
+int setAlternateStack(char* sp, size_t size) {
+  CHECK(sp);
+  stack_t ss{};
+  ss.ss_sp = sp;
+  ss.ss_size = size;
+  return sigaltstack(&ss, nullptr);
+}
+
+int unsetAlternateStack() {
+  stack_t ss{};
+  ss.ss_flags = SS_DISABLE;
+  return sigaltstack(&ss, nullptr);
+}
+
+class ScopedAlternateSignalStack {
+ public:
+  ScopedAlternateSignalStack() {
+    if (hasAlternateStack()) {
+      return;
+    }
+
+    stack_ = folly::make_unique<AltStackBuffer>();
+
+    setAlternateStack(stack_->data(), stack_->size());
+  }
+
+  ~ScopedAlternateSignalStack() {
+    if (stack_) {
+      unsetAlternateStack();
+    }
+  }
+
+ private:
+  using AltStackBuffer = std::array<char, kAltStackSize>;
+  std::unique_ptr<AltStackBuffer> stack_;
+};
+}
+
+void FiberManager::registerAlternateSignalStack() {
+  static folly::SingletonThreadLocal<ScopedAlternateSignalStack> singleton;
+  singleton.get();
+
+  alternateSignalStackRegistered_ = true;
+}
 }
 }
diff --git a/folly/experimental/fibers/FiberManager.h b/folly/experimental/fibers/FiberManager.h
index 405b2d19..aa250a74 100644
--- a/folly/experimental/fibers/FiberManager.h
+++ b/folly/experimental/fibers/FiberManager.h
@@ -467,6 +467,10 @@ class FiberManager : public ::folly::Executor {
   void unpoisonFiberStack(const Fiber* fiber);
 
 #endif // FOLLY_SANITIZE_ADDRESS
+
+  bool alternateSignalStackRegistered_{false};
+
+  void registerAlternateSignalStack();
 };
 
 /**
diff --git a/folly/experimental/fibers/test/StackOverflow.cpp b/folly/experimental/fibers/test/StackOverflow.cpp
new file mode 100644
index 00000000..d8462e1e
--- /dev/null
+++ b/folly/experimental/fibers/test/StackOverflow.cpp
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2016 Facebook, Inc.
+ *
+ * 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/experimental/fibers/FiberManagerMap.h>
+#include <folly/experimental/symbolizer/SignalHandler.h>
+#include <folly/init/Init.h>
+
+void f(int* p) {
+  // Make sure recursion is not optimized out
+  int a[100];
+  for (size_t i = 0; i < 100; ++i) {
+    a[i] = i;
+    ++(a[i]);
+    if (p) {
+      a[i] += p[i];
+    }
+  }
+  f(a);
+}
+
+int main(int argc, char* argv[]) {
+  folly::init(&argc, &argv);
+
+  folly::EventBase evb;
+  folly::fibers::getFiberManager(evb).addTask([&]() { f(nullptr); });
+  evb.loop();
+}
-- 
2.34.1