coroutine stack: use boost::context::protected_fixedsize_stack

Author: xiexiaoy

src/mongo/db/kill_sessions_local.cpp

@@ -44,6 +44,8 @@
 #include "mongo/util/log.h"
 
 #include <boost/context/continuation_fcontext.hpp>
+#include <boost/context/protected_fixedsize_stack.hpp>
+
 
 namespace mongo {
 namespace {
@@ -71,6 +73,15 @@ SessionKiller::Result killSessionsLocal(OperationContext* opCtx,
     return {std::vector<HostAndPort>{}};
 }
 
+struct CoroCtx {
+    static constexpr size_t kCoroStackSize = 3200 * 1024;
+    boost::context::protected_fixedsize_stack salloc{kCoroStackSize};
+    boost::context::continuation source;
+    std::function<void()> yieldFunc;
+    std::function<void()> resumeFunc;
+    std::function<void()> longResumeFunc;
+};
+
 void killAllExpiredTransactions(OperationContext* opCtx) {
     RecoveryUnit* ru = opCtx->releaseRecoveryUnit();
     WriteUnitOfWork::RecoveryUnitState ruState = opCtx->getRecoveryUnitState();
@@ -84,38 +95,40 @@ void killAllExpiredTransactions(OperationContext* opCtx) {
             std::mutex mux;
             std::condition_variable cv;
 
+            std::shared_ptr<CoroCtx> coroCtx = std::make_shared<CoroCtx>();
+
             auto client = Client::releaseCurrent();
             dassert(client->coroutineFunctors() == CoroutineFunctors::Unavailable);
 
-            std::function<void()> yieldFunc, resumeFunc, longResumeFunc;
             client->setCoroutineFunctors(CoroutineFunctors{
-                &yieldFunc,
-                &resumeFunc,
-                &longResumeFunc,
+                &coroCtx->yieldFunc,
+                &coroCtx->resumeFunc,
+                &coroCtx->longResumeFunc,
                 nullptr,
             });
             transport::ServiceExecutor* serviceExecutor =
                 getGlobalServiceContext()->getServiceEntryPoint()->getServiceExecutor();
 
-            boost::context::continuation source;
-            std::function<void()> resumeTask = [&source, &client] {
-                log() << "abortArbitraryTransactionIfExpired call resume";
+            std::function<void()> resumeTask = [&source = coroCtx->source, &client] {
+                log() << "abortArbitraryTransactionIfExpired call resume.";
                 Client::setCurrent(std::move(client));
                 source = source.resume();
             };
-            resumeFunc =
+
+            coroCtx->resumeFunc =
                 serviceExecutor->coroutineResumeFunctor(session->ThreadGroupId(), resumeTask);
-            longResumeFunc =
+            coroCtx->longResumeFunc =
                 serviceExecutor->coroutineLongResumeFunctor(session->ThreadGroupId(), resumeTask);
 
-            auto task = [&finished, &mux, &cv, &source, &yieldFunc, opCtx, session, &client] {
+            auto task = [&finished, &mux, &cv, coroCtx, opCtx, session, &client] {
                 Client::setCurrent(std::move(client));
-
-                source = boost::context::callcc(
-                    [&finished, &mux, &cv, &yieldFunc, opCtx, session, &client](
+                coroCtx->source = boost::context::callcc(
+                    std::allocator_arg,
+                    coroCtx->salloc,
+                    [&finished, &mux, &cv, coroCtx, opCtx, session, &client](
                         boost::context::continuation&& sink) {
-                        yieldFunc = [&sink, &client]() {
-                            log() << "abortArbitraryTransactionIfExpired call yield";
+                        coroCtx->yieldFunc = [&sink, &client]() {
+                            log() << "abortArbitraryTransactionIfExpired call yield.";
                             client = Client::releaseCurrent();
                             sink = sink.resume();
                         };

src/mongo/transport/service_state_machine.cpp

@@ -545,22 +545,22 @@ void ServiceStateMachine::_runNextInGuard(ThreadGuard guard) {
                         _coroMigrateThreadGroup = std::bind(
                             &ServiceStateMachine::_migrateThreadGroup, this, std::placeholders::_1);
 
-                        boost::context::stack_context sc = _coroStackContext();
-                        boost::context::preallocated prealloc(sc.sp, sc.size, sc);
+                        std::weak_ptr<ServiceStateMachine> wssm = weak_from_this();
                         _source = boost::context::callcc(
                             std::allocator_arg,
-                            prealloc,
-                            NoopAllocator(),
-                            [ssm = shared_from_this(),
-                             &guard](boost::context::continuation&& sink) {
+                            _salloc,
+                            [wssm, &guard](boost::context::continuation&& sink) {
+                                auto ssm = wssm.lock();
+                                if (!ssm) {
+                                    return std::move(sink);
+                                }
                                 ssm->_coroYield = [ssm = ssm.get(), &sink]() {
                                     MONGO_LOG(3) << "call yield";
                                     ssm->_dbClient = Client::releaseCurrent();
-                                    ssm->_abortIfStackOverflow();
                                     sink = sink.resume();
                                 };
                                 ssm->_processMessage(std::move(guard));
-                                ssm->_abortIfStackOverflow();
+
                                 return std::move(sink);
                             });
 
@@ -578,15 +578,13 @@ void ServiceStateMachine::_runNextInGuard(ThreadGuard guard) {
                         //         _coroYield = [this, &sink]() {
                         //             MONGO_LOG(1) << "call yield";
                         //             _dbClient = Client::releaseCurrent();
-                        //             _abortIfStackOverflow();
                         //             sink = sink.resume();
                         //         };
                         //         _processMessage(std::move(guard));
                         //         return std::move(sink);
                         //     });
                     } else if (_coroStatus == CoroStatus::OnGoing) {
                         MONGO_LOG(1) << "coroutine ongoing";
-                        _abortIfStackOverflow();
                         _source = _source.resume();
                     }
                 }
@@ -618,7 +616,6 @@ void ServiceStateMachine::_runResumeProcess() {
     MONGO_LOG(3) << "ServiceStateMachine::_resumeRun";
     if (_coroStatus == CoroStatus::OnGoing) {
         MONGO_LOG(3) << "coroutine ongoing";
-        _abortIfStackOverflow();
         _source = _source.resume();
     }
 }

src/mongo/transport/service_state_machine.h

@@ -31,7 +31,7 @@
 #include <atomic>
 #include <boost/context/continuation.hpp>
 #include <boost/context/continuation_fcontext.hpp>
-#include <boost/context/stack_context.hpp>
+#include <boost/context/protected_fixedsize_stack.hpp>
 #include <functional>
 
 #include "boost/optional/optional.hpp"
@@ -261,64 +261,12 @@ class ServiceStateMachine : public std::enable_shared_from_this<ServiceStateMach
 #endif
     std::string _oldThreadName;
 
-    // Coroutine design
-    class NoopAllocator {
-    public:
-        NoopAllocator() = default;
-
-        boost::context::stack_context allocate() {
-            boost::context::stack_context sc;
-            return sc;
-        }
-
-        void deallocate(boost::context::stack_context& sc) {
-            // no-op
-        }
-    };
-
-    static constexpr size_t kCoroStackSize = 3200 * 1024;
-    static constexpr char kCanaryByte = 0xAB;
-    static constexpr size_t kCanarySize = 16;
-    static constexpr char kCanaryBytes[kCanarySize] = {
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-        kCanaryByte,
-    };
-
-    boost::context::stack_context _coroStackContext() {
-        boost::context::stack_context sc;
-        sc.size = kCoroStackSize;
-        // Because stack grows downwards from high address?
-        sc.sp = _coroStack + kCoroStackSize;
-        // Set canary bytes at the end of the stack to detect stack overflow.
-        std::memset(_coroStack, kCanaryByte, kCanarySize);
-        return sc;
-    }
-
-    void _abortIfStackOverflow() {
-        if (std::memcmp(_coroStack, kCanaryBytes, kCanarySize) != 0) {
-            std::abort();
-        }
-    }
-
     void _migrateThreadGroup(uint16_t threadGroupId);
 
 
+    static constexpr size_t kCoroStackSize = 3200 * 1024;
+    boost::context::protected_fixedsize_stack _salloc{kCoroStackSize};
     boost::context::continuation _source;
-    char _coroStack[kCoroStackSize];
 
     enum class CoroStatus { Empty = 0, OnGoing, Finished };
     CoroStatus _coroStatus{CoroStatus::Empty};