SERVER-27896 Add support for enforcing single-thread context

9 files changed, 275 insertions, 5 deletions

diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index f0ef13a95e7..fdc017efb8b 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -164,6 +164,7 @@ baseEnv.Library(
         'util/system_clock_source.cpp',
         'util/system_tick_source.cpp',
         'util/text.cpp',
+        'util/thread_safety_context.cpp',
         'util/time_support.cpp',
         'util/timer.cpp',
         'util/uuid.cpp',
diff --git a/src/mongo/db/db.cpp b/src/mongo/db/db.cpp
index e79178ed28c..ac10b378700 100644
--- a/src/mongo/db/db.cpp
+++ b/src/mongo/db/db.cpp
@@ -1228,6 +1228,8 @@ void shutdownTask(const ShutdownTaskArgs& shutdownArgs) {
 }
 
 int mongoDbMain(int argc, char* argv[], char** envp) {
+    ThreadSafetyContext::getThreadSafetyContext()->forbidMultiThreading();
+
     registerShutdownTask(shutdownTask);
 
     setupSignalHandlers();
@@ -1277,6 +1279,9 @@ int mongoDbMain(int argc, char* argv[], char** envp) {
     if (!initializeServerSecurityGlobalState(service))
         quickExit(EXIT_FAILURE);
 
+    // There is no single-threaded guarantee beyond this point.
+    ThreadSafetyContext::getThreadSafetyContext()->allowMultiThreading();
+
     // Per SERVER-7434, startSignalProcessingThread must run after any forks (i.e.
     // initializeServerGlobalState) and before the creation of any other threads
     startSignalProcessingThread();
diff --git a/src/mongo/stdx/SConscript b/src/mongo/stdx/SConscript
index 02af02ec6f3..fc70228475e 100644
--- a/src/mongo/stdx/SConscript
+++ b/src/mongo/stdx/SConscript
@@ -10,6 +10,7 @@ env.Benchmark(
         'condition_variable_bm.cpp',
     ],
     LIBDEPS=[
+        '$BUILD_DIR/mongo/base',
     ],
 )
 
@@ -44,7 +45,7 @@ env.CppUnitTest(
         'sigaltstack_location_test.cpp',
     ],
     LIBDEPS=[
-        'stdx',
+        '$BUILD_DIR/mongo/base',
     ],
     UNITTEST_HAS_CUSTOM_MAINLINE=True,
 )
@@ -64,7 +65,7 @@ env.CppUnitTest(
         'set_terminate_dispatch_test.cpp',
     ],
     LIBDEPS=[
-        'stdx',
+        '$BUILD_DIR/mongo/base',
     ],
     UNITTEST_HAS_CUSTOM_MAINLINE=True,
 )
@@ -76,7 +77,7 @@ env.CppUnitTest(
         'set_terminate_from_main_die_in_thread_test.cpp',
     ],
     LIBDEPS=[
-        'stdx',
+        '$BUILD_DIR/mongo/base',
     ],
     UNITTEST_HAS_CUSTOM_MAINLINE=True,
 )
@@ -88,7 +89,7 @@ env.CppUnitTest(
         'set_terminate_from_thread_die_in_main_test.cpp',
     ],
     LIBDEPS=[
-        'stdx',
+        '$BUILD_DIR/mongo/base',
     ],
     UNITTEST_HAS_CUSTOM_MAINLINE=True,
 )
@@ -100,7 +101,7 @@ env.CppUnitTest(
         'set_terminate_from_thread_die_in_thread_test.cpp',
     ],
     LIBDEPS=[
-        'stdx',
+        '$BUILD_DIR/mongo/base',
     ],
     UNITTEST_HAS_CUSTOM_MAINLINE=True,
 )
diff --git a/src/mongo/stdx/thread.h b/src/mongo/stdx/thread.h
index e072de0cc72..f8058279e34 100644
--- a/src/mongo/stdx/thread.h
+++ b/src/mongo/stdx/thread.h
@@ -39,6 +39,7 @@
 #include <type_traits>
 
 #include "mongo/stdx/exception.h"
+#include "mongo/util/thread_safety_context.h"
 
 #if defined(__linux__) || defined(__FreeBSD__)
 #define MONGO_HAS_SIGALTSTACK 1
@@ -180,6 +181,7 @@ public:
                   ::std::set_terminate(  // NOLINT
                       ::mongo::stdx::TerminateHandlerDetailsInterface::dispatch);
 #endif
+                  ThreadSafetyContext::getThreadSafetyContext()->onThreadCreate();
                   auto sigAltStackGuard = sigAltStackController.makeInstallGuard();
                   return std::apply(std::move(f), std::move(pack));
               }) {
diff --git a/src/mongo/transport/service_entry_point_utils.cpp b/src/mongo/transport/service_entry_point_utils.cpp
index f4e894e1168..aab92e246c3 100644
--- a/src/mongo/transport/service_entry_point_utils.cpp
+++ b/src/mongo/transport/service_entry_point_utils.cpp
@@ -40,6 +40,7 @@
 #include "mongo/stdx/thread.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/debug_util.h"
+#include "mongo/util/thread_safety_context.h"
 
 #if !defined(_WIN32)
 #include <sys/resource.h>
@@ -101,6 +102,7 @@ Status launchServiceWorkerThread(std::function<void()> task) {
 
         pthread_t thread;
         auto ctx = std::make_unique<std::function<void()>>(std::move(task));
+        ThreadSafetyContext::getThreadSafetyContext()->onThreadCreate();
         int failed = pthread_create(&thread, &attrs, runFunc, ctx.get());
 
         pthread_attr_destroy(&attrs);
diff --git a/src/mongo/util/SConscript b/src/mongo/util/SConscript
index 50b413173dd..46fa8aec9b1 100644
--- a/src/mongo/util/SConscript
+++ b/src/mongo/util/SConscript
@@ -213,6 +213,16 @@ env.Library(
     ]
 )
 
+env.CppUnitTest(
+    target='thread_safety_context_test',
+    source=[
+        'thread_safety_context_test.cpp',
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/base',
+    ],
+)
+
 if env['MONGO_ALLOCATOR'] in ['tcmalloc', 'tcmalloc-experimental']:
     tcmspEnv = env.Clone()
 
diff --git a/src/mongo/util/thread_safety_context.cpp b/src/mongo/util/thread_safety_context.cpp
new file mode 100644
index 00000000000..1b40c2746c5
--- /dev/null
+++ b/src/mongo/util/thread_safety_context.cpp
@@ -0,0 +1,59 @@
+/**
+ *    Copyright (C) 2020-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/util/thread_safety_context.h"
+
+#include "mongo/platform/compiler.h"
+#include "mongo/util/assert_util.h"
+
+namespace mongo {
+
+ThreadSafetyContext* ThreadSafetyContext::getThreadSafetyContext() noexcept {
+    static auto safetyContext = new ThreadSafetyContext();  // Intentionally leaked
+    return safetyContext;
+}
+
+void ThreadSafetyContext::forbidMultiThreading() noexcept {
+    invariant(_isSingleThreaded.load());
+    invariant(_safeToCreateThreads.swap(false));
+}
+
+void ThreadSafetyContext::allowMultiThreading() noexcept {
+    invariant(_isSingleThreaded.load());
+    invariant(!_safeToCreateThreads.swap(true));
+}
+
+void ThreadSafetyContext::onThreadCreate() noexcept {
+    invariant(_safeToCreateThreads.load());
+    if (MONGO_unlikely(_isSingleThreaded.load())) {
+        _isSingleThreaded.store(false);
+    }
+}
+
+}  // namespace mongo
diff --git a/src/mongo/util/thread_safety_context.h b/src/mongo/util/thread_safety_context.h
new file mode 100644
index 00000000000..cbf3f153efd
--- /dev/null
+++ b/src/mongo/util/thread_safety_context.h
@@ -0,0 +1,87 @@
+/**
+ *    Copyright (C) 2020-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/platform/atomic_word.h"
+
+namespace mongo {
+
+class ThreadSafetyContextTest;
+
+/**
+ * Provides a singleton object that allows enforcing a single-threaded context at startup.
+ *
+ * This class, currently, intercepts `stdx::thread` creations and aborts the process if a single
+ * threaded context is expected.
+ *
+ * You should avoid using the interfaces that specify single-threaded context (i.e.,
+ * `forbidMultiThreading()` and `allowMultiThreading()`) outside the main function.
+ *
+ * If using other APIs for spawning threads (e.g., `pthread_create()`), make sure to precede the
+ * API call with `onThreadCreate()`.
+ */
+class ThreadSafetyContext final {
+public:
+    ThreadSafetyContext(const ThreadSafetyContext&) = delete;
+
+    ThreadSafetyContext(ThreadSafetyContext&&) noexcept = delete;
+
+    ThreadSafetyContext& operator=(const ThreadSafetyContext&) = delete;
+
+    ThreadSafetyContext& operator=(ThreadSafetyContext&&) = delete;
+
+    static ThreadSafetyContext* getThreadSafetyContext() noexcept;
+
+    // Prevents a multi-threaded context -- aborts the process on thread creation.
+    // If the program is already multi-threaded, it will abort the program.
+    void forbidMultiThreading() noexcept;
+
+    // Allows a multi-threaded context by lifting the limit enforced by `forbidMultiThreading()`.
+    void allowMultiThreading() noexcept;
+
+    // Allows detecting thread creation, and thus a multi-threaded context.
+    // If not using `stdx::thread`, you must always call this method before spawning a new thread.
+    void onThreadCreate() noexcept;
+
+    // Restricts the ability of resetting the context to the unit-test fixture.
+    friend class ThreadSafetyContextTest;
+
+private:
+    // Prevents creating any instance of the object (other than the singleton).
+    ThreadSafetyContext() = default;
+
+    // An indicator on whether the program is single-threaded.
+    AtomicWord<bool> _isSingleThreaded{true};
+
+    // If set to "false", it will prevent creation of new threads.
+    AtomicWord<bool> _safeToCreateThreads{true};
+};
+
+}  // namespace mongo
diff --git a/src/mongo/util/thread_safety_context_test.cpp b/src/mongo/util/thread_safety_context_test.cpp
new file mode 100644
index 00000000000..4e27614b514
--- /dev/null
+++ b/src/mongo/util/thread_safety_context_test.cpp
@@ -0,0 +1,103 @@
+/**
+ *    Copyright (C) 2020-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+#include <vector>
+
+#include "mongo/stdx/thread.h"
+#include "mongo/unittest/death_test.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/thread_safety_context.h"
+#include "mongo/util/time_support.h"
+
+namespace mongo {
+
+class ThreadSafetyContextTest : public unittest::Test {
+public:
+    void setUp() override {
+        // Always start tests with a clean context
+        resetContext();
+    }
+
+    void tearDown() override {
+        // Clear the context for the next test
+        resetContext();
+    }
+
+private:
+    void resetContext() {
+        auto context = ThreadSafetyContext::getThreadSafetyContext();
+        context->_isSingleThreaded.store(true);
+        context->_safeToCreateThreads.store(true);
+    }
+};
+
+TEST_F(ThreadSafetyContextTest, CreateThreadsWithNoSafetyContext) {
+    constexpr auto threadCount = 16;
+    std::vector<stdx::thread> threads;
+
+    for (auto i = 0; i < threadCount; i++) {
+        threads.emplace_back([] { sleepFor(Milliseconds(10)); });
+    }
+
+    for (auto i = 0; i < threadCount; i++) {
+        threads[i].join();
+    }
+}
+
+DEATH_TEST_F(ThreadSafetyContextTest, CreateThreadsAfterForbidingMultiThreading, "invariant") {
+    ThreadSafetyContext::getThreadSafetyContext()->forbidMultiThreading();
+    auto thread = stdx::thread([] {});
+    // Must never reach here or the test fails
+}
+
+DEATH_TEST_F(ThreadSafetyContextTest, ForbidMultiThreadingAfterCreatingThreads, "invariant") {
+    auto thread = stdx::thread([] { sleepFor(Milliseconds(50)); });
+
+    ThreadSafetyContext::getThreadSafetyContext()->forbidMultiThreading();
+    // Must never reach here or the test fails
+}
+
+TEST_F(ThreadSafetyContextTest, CreateThreadsAfterSafetyContext) {
+    ThreadSafetyContext::getThreadSafetyContext()->forbidMultiThreading();
+    // Do something inside the single-threaded context
+    sleepFor(Milliseconds(10));
+    ThreadSafetyContext::getThreadSafetyContext()->allowMultiThreading();
+
+    constexpr auto threadCount = 16;
+    std::vector<stdx::thread> threads;
+
+    for (auto i = 0; i < threadCount; i++) {
+        threads.emplace_back([] { sleepFor(Milliseconds(10)); });
+    }
+
+    for (auto i = 0; i < threadCount; i++) {
+        threads[i].join();
+    }
+}
+
+}  // namespace mongo