From 0fae325d76abb49c4a74e73ebe4bcccfcde9f88f Mon Sep 17 00:00:00 2001
From: Nikolai Kochetov <nik-kochetov@yandex-team.ru>
Date: Wed, 2 Dec 2020 14:18:46 +0300
Subject: [PATCH] Add FiberStack

---
 src/Common/FiberStack.h            | 45 +++++++++++++++++++
 src/Interpreters/tests/context.cpp | 72 ++++++++++++++++++++++++++++--
 2 files changed, 113 insertions(+), 4 deletions(-)
 create mode 100644 src/Common/FiberStack.h

diff --git a/src/Common/FiberStack.h b/src/Common/FiberStack.h
new file mode 100644
index 00000000000..c133cb6250c
--- /dev/null
+++ b/src/Common/FiberStack.h
@@ -0,0 +1,45 @@
+#include <boost/context/stack_context.hpp>
+#include <Common/Allocator.h>
+
+#if defined(BOOST_USE_VALGRIND)
+#include <valgrind/valgrind.h>
+#endif
+
+/// This is an implementation of allocator for fiber stack.
+/// It uses internal allocator, so we track memory usage. It is the main reason why this class is needed.
+/// The reference implementations are pooled_fixedsize_stack and protected_fixedsize_stack from boost::context.
+template <typename TAllocator = Allocator<false, false>>
+class FiberStack
+{
+private:
+    size_t stack_size;
+public:
+    /// 8MB of memory per fiber stack may seem too expensive. It is indeed.
+    /// The reason is that current (patched) libunwind needs > 4MB of stack memory to unwind stack.
+    /// If we allocate less memory, any thrown exception inside fiber will cause segfault.
+    static constexpr size_t default_stack_size = 8 * 1024 * 1024;
+
+    explicit FiberStack(size_t stack_size_ = default_stack_size) : stack_size(stack_size_) {}
+
+    boost::context::stack_context allocate()
+    {
+        void * vp = TAllocator().alloc(stack_size);
+
+        boost::context::stack_context sctx;
+        sctx.size = stack_size;
+        sctx.sp = static_cast< char * >(vp) + sctx.size;
+#if defined(BOOST_USE_VALGRIND)
+        sctx.valgrind_stack_id = VALGRIND_STACK_REGISTER(sctx.sp, vp);
+#endif
+        return sctx;
+    }
+
+    void deallocate(boost::context::stack_context & sctx)
+    {
+#if defined(BOOST_USE_VALGRIND)
+        VALGRIND_STACK_DEREGISTER( sctx.valgrind_stack_id);
+#endif
+        void * vp = static_cast< char * >(sctx.sp) - sctx.size;
+        TAllocator().free(vp, stack_size);
+    }
+};
diff --git a/src/Interpreters/tests/context.cpp b/src/Interpreters/tests/context.cpp
index eed7ca60790..407a7a9ba3b 100644
--- a/src/Interpreters/tests/context.cpp
+++ b/src/Interpreters/tests/context.cpp
@@ -1,26 +1,90 @@
 #include <iostream>
+/// #define BOOST_USE_UCONTEXT
 #include <boost/context/fiber.hpp>
+#include <boost/context/pooled_fixedsize_stack.hpp>
+#include <boost/context/segmented_stack.hpp>
+#include <Common/Exception.h>
+#include <Common/FiberStack.h>
+
+void __attribute__((__noinline__)) foo(std::exception_ptr exception)
+{
+    if (exception)
+        std::rethrow_exception(exception);
+}
+
+void __attribute__((__noinline__)) bar(int a)
+{
+    std::cout << StackTrace().toString() << std::endl;
+
+    if (a > 0)
+        throw DB::Exception(0, "hello");
+}
+
+void __attribute__((__noinline__)) gar(int a)
+{
+    char buf[1024];
+    buf[1023] = a & 255;
+    if (a > 2)
+        return gar(a - 1);
+    else
+        bar(a);
+}
 
 int main(int, char **)
-{
+try {
     namespace ctx=boost::context;
     int a;
-    ctx::fiber source{[&a](ctx::fiber&& sink)
+    std::exception_ptr exception;
+    // ctx::protected_fixedsize allocator
+    // ctx::pooled_fixedsize_stack(1024 * 64 + 2 * 2 * 1024 * 1024 * 16, 1)
+    ctx::fiber source{std::allocator_arg_t(), FiberStack(), [&](ctx::fiber&& sink)
     {
         a=0;
         int b=1;
-        while (true)
+        for (size_t i = 0; i < 9; ++i)
         {
             sink=std::move(sink).resume();
             int next=a+b;
             a=b;
             b=next;
         }
+        try
+        {
+            gar(1024);
+        }
+        catch (...)
+        {
+            std::cout << "Saving exception\n";
+            exception = std::current_exception();
+        }
         return std::move(sink);
     }};
+
     for (int j=0;j<10;++j)
     {
-        source=std::move(source).resume();
+        try
+        {
+            source=std::move(source).resume();
+        }
+        catch (DB::Exception & e)
+        {
+            std::cout << "Caught exception in resume " << e.getStackTraceString() << std::endl;
+        }
         std::cout << a << " ";
     }
+
+    std::cout << std::endl;
+
+    try
+    {
+        foo(exception);
+    }
+    catch (const DB::Exception & e)
+    {
+        std::cout << e.getStackTraceString() << std::endl;
+    }
+}
+catch (...)
+{
+    std::cerr << "Uncaught exception\n";
 }