SERVER-52822 Introduced ThreadContext class

7 files changed, 495 insertions, 0 deletions

diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index b9027f0ded0..18bbee84283 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -241,6 +241,7 @@ baseEnv.Library(
         'util/testing_proctor.cpp',
         'util/text.cpp',
         'util/thread_safety_context.cpp',
+        'util/thread_context.cpp',
         'util/time_support.cpp',
         'util/timer.cpp',
         'util/uuid.cpp',
diff --git a/src/mongo/platform/process_id.cpp b/src/mongo/platform/process_id.cpp
index 48bc0354316..cff6af7e2db 100644
--- a/src/mongo/platform/process_id.cpp
+++ b/src/mongo/platform/process_id.cpp
@@ -31,11 +31,21 @@
 
 #include "mongo/platform/process_id.h"
 
+#ifndef _WIN32
+#include <pthread.h>
+#endif
+
+#if defined(__linux__)
+#include <sys/syscall.h>
+#include <sys/types.h>
+#endif
+
 #include <iostream>
 #include <limits>
 #include <sstream>
 
 #include "mongo/base/static_assert.h"
+#include "mongo/util/assert_util.h"
 
 namespace mongo {
 
@@ -51,12 +61,33 @@ inline NativeProcessId getCurrentNativeProcessId() {
     return getpid();
 }
 #endif
+
+#ifdef _WIN32
+inline NativeProcessId getCurrentNativeThreadId() {
+    return GetCurrentThreadId();
+}
+#elif __APPLE__
+inline NativeProcessId getCurrentNativeThreadId() {
+    // macOS deprecated syscall in 10.12.
+    uint64_t tid;
+    invariant(::pthread_threadid_np(NULL, &tid) == 0);
+    return tid;
+}
+#else
+inline NativeProcessId getCurrentNativeThreadId() {
+    return ::syscall(SYS_gettid);
+}
+#endif
 }  // namespace
 
 ProcessId ProcessId::getCurrent() {
     return fromNative(getCurrentNativeProcessId());
 }
 
+ProcessId ProcessId::getCurrentThreadId() {
+    return fromNative(getCurrentNativeThreadId());
+}
+
 int64_t ProcessId::asInt64() const {
     typedef std::numeric_limits<NativeProcessId> limits;
     if (limits::is_signed)
diff --git a/src/mongo/platform/process_id.h b/src/mongo/platform/process_id.h
index 2e9d1f721f7..2e2e6396de9 100644
--- a/src/mongo/platform/process_id.h
+++ b/src/mongo/platform/process_id.h
@@ -56,6 +56,11 @@ public:
     static ProcessId getCurrent();
 
     /**
+     * Gets the thread id for the currently executing process.
+     */
+    static ProcessId getCurrentThreadId();
+
+    /**
      * Constructs a ProcessId from a NativeProcessId.
      */
     static inline ProcessId fromNative(NativeProcessId npid) {
diff --git a/src/mongo/util/SConscript b/src/mongo/util/SConscript
index 13f9ec6ab3e..df6e25513b5 100644
--- a/src/mongo/util/SConscript
+++ b/src/mongo/util/SConscript
@@ -666,6 +666,7 @@ icuEnv.CppUnitTest(
         'strong_weak_finish_line_test.cpp',
         'summation_test.cpp',
         'text_test.cpp',
+        'thread_context_test.cpp',
         'tick_source_test.cpp',
         'time_support_test.cpp',
         'unique_function_test.cpp',
diff --git a/src/mongo/util/thread_context.cpp b/src/mongo/util/thread_context.cpp
new file mode 100644
index 00000000000..5d0e23647eb
--- /dev/null
+++ b/src/mongo/util/thread_context.cpp
@@ -0,0 +1,85 @@
+/**
+ *    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/platform/basic.h"
+
+#include "mongo/util/thread_context.h"
+
+#include "mongo/base/init.h"
+#include "mongo/stdx/thread.h"
+
+namespace mongo {
+namespace {
+const auto kMainThreadId = stdx::this_thread::get_id();
+AtomicWord<bool> gHasInitializedMain{false};
+}  // namespace
+
+MONGO_INITIALIZER(ThreadContextsInitialized)(InitializerContext*) {
+    ThreadContext::initializeMain();
+    return Status::OK();
+}
+
+void ThreadContext::initializeMain() {
+    invariant(stdx::this_thread::get_id() == kMainThreadId,
+              "initializeMain() must be called on the main thread");
+
+    if (gHasInitializedMain.swap(true)) {
+        return;
+    }
+
+    _handle.init();
+}
+
+ThreadContext::Handle::Handle() {
+    if (!gHasInitializedMain.loadRelaxed()) {
+        // If we have not initialized main, then we delay until the MONGO_INITIALIZER runs.
+        return;
+    }
+
+    init();
+}
+
+void ThreadContext::Handle::init() {
+    // Note that construction happens before assignment to the thread_local.
+    instance = make_intrusive<ThreadContext>();
+}
+
+ThreadContext::Handle::~Handle() {
+    if (!instance) {
+        // If we don't have an instance, just skip. This is mostly going to be pre-main failures.
+        return;
+    }
+
+    // Remove from the thread local access, then destroy our pointer to it.
+    auto localInstance = std::exchange(instance, {});
+    localInstance->_isAlive.store(false);
+    localInstance.reset();
+}
+
+}  // namespace mongo
diff --git a/src/mongo/util/thread_context.h b/src/mongo/util/thread_context.h
new file mode 100644
index 00000000000..7c4b9bbfbc8
--- /dev/null
+++ b/src/mongo/util/thread_context.h
@@ -0,0 +1,124 @@
+/**
+ *    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"
+#include "mongo/platform/process_id.h"
+#include "mongo/util/assert_util.h"
+#include "mongo/util/decorable.h"
+#include "mongo/util/intrusive_counter.h"
+
+namespace mongo {
+
+/**
+ * A ThreadContext is a simple decorable that has an explicit one-to-one relationship with threads.
+ *
+ * There are three lifetime tricks that this class does:
+ * 1. It only exists on the main thread after we run MONGO_INITIALIZERS (a.k.a. post-init).
+ * 2. It constructs and destructs on the local stack and thus cannot reference itself in its
+ *    constructor or destructor.
+ * 3. It can be persisted past the death of its original thread. This means that decorations may
+ *    be destructed on a thread with a different ThreadContext attached. It also means that
+ *    ThreadContexts can be safely tracked in data structures without complicated lifetime logic.
+ *
+ * There may be situations where you want to do one thing to a decoration when the thread dies and
+ * another when the decoration is destructed. If this comes up, we will need a graph of ordered
+ * actions like we have for ServiceContext.
+ */
+class ThreadContext final : public Decorable<ThreadContext>, public RefCountable {
+public:
+    ThreadContext() = default;
+    virtual ~ThreadContext() = default;
+
+    /**
+     * This initializes the main thread in a MONGO_INITIALIZER.
+     *
+     * This is invalid to invoke on other threads.
+     */
+    static void initializeMain();
+
+    /**
+     * Get the ThreadContext for the current thread.
+     *
+     * If you are in pre-init, this will return an empty pointer. If you want to access a given
+     * ThreadContext during thread death and beyond, store a copy of this intrusive_ptr somewhere.
+     */
+    static const boost::intrusive_ptr<ThreadContext>& get() {
+        return _handle.instance;
+    }
+
+    /**
+     * Get the thread id for the current thread.
+     */
+    const auto& threadId() const {
+        return _threadId;
+    }
+
+    /**
+     * Get if the current thread is still running.
+     *
+     * This value only transitions from true to false.
+     */
+    bool isAlive() const {
+        return _isAlive.load();
+    }
+
+private:
+    /**
+     * This handle class is the actual thread_local variable.
+     *
+     * Its functions manage lifetime and set _isAlive.
+     */
+    struct Handle {
+        /**
+         * Create a new ThreadContext post-init and do nothing pre-init.
+         */
+        Handle();
+
+        /**
+         * Move the ThreadContext to the local stack and destroy it.
+         */
+        ~Handle();
+
+        /**
+         * Create and bind a new ThreadContext.
+         */
+        void init();
+
+        boost::intrusive_ptr<ThreadContext> instance;
+    };
+
+    inline static thread_local auto _handle = Handle{};
+
+    const ProcessId _threadId = ProcessId::getCurrentThreadId();
+    AtomicWord<bool> _isAlive{true};
+};
+
+}  // namespace mongo
diff --git a/src/mongo/util/thread_context_test.cpp b/src/mongo/util/thread_context_test.cpp
new file mode 100644
index 00000000000..8d777dbe02b
--- /dev/null
+++ b/src/mongo/util/thread_context_test.cpp
@@ -0,0 +1,248 @@
+/**
+ *    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/platform/basic.h"
+
+#include "mongo/util/thread_context.h"
+
+#include "mongo/stdx/thread.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/time_support.h"
+
+namespace mongo {
+namespace {
+
+/**
+ * This decoration increments a set of global counters on creation and destruction.
+ */
+class TestDecoration {
+public:
+    TestDecoration() {
+        ASSERT(!ThreadContext::get())
+            << "ThreadContext decorations should be created before the ThreadContext is set";
+
+        _created.fetchAndAdd(1);
+    };
+
+    ~TestDecoration() {
+        _destroyed.fetchAndAdd(1);
+
+        if (ThreadContext::get()) {
+            // We should only be able to reference a ThreadContext in our destructor if our
+            // lifetime was extended to be off thread.
+            _destroyedOffThread.fetchAndAdd(1);
+        }
+    }
+
+    static auto created() {
+        return _created.load();
+    }
+
+    static auto destroyed() {
+        return _destroyed.load();
+    }
+
+    static auto destroyedOffThread() {
+        return _destroyedOffThread.load();
+    }
+
+private:
+    static inline AtomicWord<size_t> _created{0};
+    static inline AtomicWord<size_t> _destroyed{0};
+    static inline AtomicWord<size_t> _destroyedOffThread{0};
+};
+
+const auto getThreadTestDecoration = ThreadContext::declareDecoration<TestDecoration>();
+
+class ThreadContextTest : public unittest::Test {
+public:
+    void setUp() override;
+    void tearDown() override;
+
+    /**
+     * Get the ThreadContext for the current thread and assert that it is valid and alive.
+     */
+    auto getThreadContext() {
+        auto context = ThreadContext::get();
+
+        ASSERT(context);
+        ASSERT(context->isAlive());
+
+        return context;
+    }
+
+    /**
+     * Verify that the given ThreadContext is valid but not alive.
+     */
+    void assertNotAlive(boost::intrusive_ptr<ThreadContext> context) {
+        ASSERT(context);
+        ASSERT(!context->isAlive());
+    }
+
+    /**
+     * Launch a thread and then immediately join it.
+     */
+    template <typename F>
+    void launchAndJoinThread(F&& f) {
+        auto t = stdx::thread(std::forward<F>(f));
+        t.join();
+    }
+
+    /**
+     * Get the amount of TestDecoration instances created since the start of this test.
+     */
+    auto decorationsCreated() const {
+        return TestDecoration::created() - _created;
+    }
+
+    /**
+     * Get the amount of TestDecoration instances destroyed since the start of this test.
+     */
+    auto decorationsDestroyed() const {
+        return TestDecoration::destroyed() - _destroyed;
+    }
+
+    /**
+     * Get the amount of TestDecoration instances destroyed off thread since the start of this test.
+     */
+    auto decorationsDestroyedOffThread() const {
+        return TestDecoration::destroyedOffThread() - _destroyedOffThread;
+    }
+
+private:
+    size_t _created;
+    size_t _destroyed;
+    size_t _destroyedOffThread;
+};
+
+void ThreadContextTest::setUp() {
+    _created = TestDecoration::created();
+    _destroyed = TestDecoration::destroyed();
+    _destroyedOffThread = TestDecoration::destroyedOffThread();
+}
+
+void ThreadContextTest::tearDown() {
+    ASSERT_EQ(decorationsCreated(), decorationsDestroyed())
+        << "Each created decoration should also be destroyed";
+    ASSERT_GTE(decorationsDestroyed(), decorationsDestroyedOffThread())
+        << "We can never have more decorations destroyed off thread than we have made in total";
+}
+
+TEST_F(ThreadContextTest, HasLocalThreadContext) {
+    auto context = getThreadContext();
+
+    // Since this is the local thread, there should be no difference since the start of the test.
+    ASSERT_EQ(decorationsCreated(), 0);
+    ASSERT_EQ(decorationsDestroyed(), 0);
+    ASSERT_EQ(decorationsDestroyedOffThread(), 0);
+}
+
+TEST_F(ThreadContextTest, HasNewThreadContext) {
+    launchAndJoinThread([&] {
+        auto context = getThreadContext();
+
+        ASSERT_EQ(decorationsCreated(), 1);
+        ASSERT_EQ(decorationsDestroyed(), 0);
+        ASSERT_EQ(decorationsDestroyedOffThread(), 0);
+    });
+
+    ASSERT_EQ(decorationsCreated(), 1);
+    ASSERT_EQ(decorationsDestroyed(), 1);
+    ASSERT_EQ(decorationsDestroyedOffThread(), 0);
+}
+
+TEST_F(ThreadContextTest, CanExtendThreadContextLifetime) {
+    boost::intrusive_ptr<ThreadContext> context;
+
+    launchAndJoinThread([&] {
+        context = getThreadContext();
+
+        ASSERT_EQ(decorationsCreated(), 1);
+        ASSERT_EQ(decorationsDestroyed(), 0);
+        ASSERT_EQ(decorationsDestroyedOffThread(), 0);
+    });
+
+    assertNotAlive(context);
+
+    ASSERT_EQ(decorationsCreated(), 1);
+    ASSERT_EQ(decorationsDestroyed(), 0);
+    ASSERT_EQ(decorationsDestroyedOffThread(), 0);
+
+    context.reset();
+
+    // The context is gone.
+    ASSERT_EQ(decorationsCreated(), 1);
+    ASSERT_EQ(decorationsDestroyed(), 1);
+    ASSERT_EQ(decorationsDestroyedOffThread(), 1);
+}
+
+TEST_F(ThreadContextTest, AreThreadContextsUnique) {
+    boost::intrusive_ptr<ThreadContext> contextA;
+    boost::intrusive_ptr<ThreadContext> contextB;
+
+    launchAndJoinThread([&] {
+        contextA = getThreadContext();
+
+        ASSERT_EQ(decorationsCreated(), 1);
+        ASSERT_EQ(decorationsDestroyed(), 0);
+        ASSERT_EQ(decorationsDestroyedOffThread(), 0);
+    });
+
+    launchAndJoinThread([&] {
+        contextB = getThreadContext();
+
+        ASSERT_EQ(decorationsCreated(), 2);
+        ASSERT_EQ(decorationsDestroyed(), 0);
+        ASSERT_EQ(decorationsDestroyedOffThread(), 0);
+    });
+
+    assertNotAlive(contextA);
+    assertNotAlive(contextB);
+
+    ASSERT_NE(ThreadContext::get()->threadId(), contextA->threadId())
+        << "The context for the local thread should be different than the one for thread A";
+    ASSERT_NE(contextA->threadId(), contextB->threadId())
+        << "The context for thread A should be different than the one for thread B";
+
+    contextA.reset();
+    contextB.reset();
+
+    ASSERT_EQ(decorationsCreated(), 2);
+    ASSERT_EQ(decorationsDestroyed(), 2);
+    ASSERT_EQ(decorationsDestroyedOffThread(), 2);
+}
+
+// This check runs in pre-init and then we check it in a test post-init.
+const bool gHasAThreadContextPreInit = [] { return static_cast<bool>(ThreadContext::get()); }();
+TEST_F(ThreadContextTest, HasNoPreInitThreadContext) {
+    ASSERT(!gHasAThreadContextPreInit);
+}
+
+}  // namespace
+}  // namespace mongo