SERVER-48901 Have a client observer destroy ClientOutOfLineExecutor

3 files changed, 38 insertions, 2 deletions

diff --git a/src/mongo/db/client_out_of_line_executor.cpp b/src/mongo/db/client_out_of_line_executor.cpp
index 5aa83a66d18..e2fe618165f 100644
--- a/src/mongo/db/client_out_of_line_executor.cpp
+++ b/src/mongo/db/client_out_of_line_executor.cpp
@@ -33,6 +33,7 @@
 
 #include "mongo/base/error_codes.h"
 #include "mongo/base/status.h"
+#include "mongo/db/service_context.h"
 #include "mongo/logv2/log.h"
 #include "mongo/logv2/log_severity.h"
 #include "mongo/logv2/log_severity_suppressor.h"
@@ -54,6 +55,12 @@ ClientOutOfLineExecutor::ClientOutOfLineExecutor() noexcept
     : _impl{std::make_unique<Impl>()}, _taskQueue{std::make_shared<QueueType>()} {}
 
 ClientOutOfLineExecutor::~ClientOutOfLineExecutor() noexcept {
+    invariant(_isShutdown);
+}
+
+void ClientOutOfLineExecutor::shutdown() {
+    ON_BLOCK_EXIT([this]() mutable { _isShutdown = true; });
+
     // Force producers to consume their tasks beyond this point.
     _taskQueue->closeProducerEnd();
 
@@ -117,4 +124,28 @@ void ClientOutOfLineExecutor::QueueHandle::schedule(Task&& task) {
     }
 }
 
+namespace {
+
+/**
+ * The observer ensures that `ClientOutOfLineExecutor` is always stopped before the client object is
+ * destroyed. This is necessary to guarantee that `ClientOutOfLineExecutor::shutdown()` is executed
+ * before the client decorations are destroyed. See SERVER-48901 for more details.
+ */
+class ClientOutOfLineExecutorClientObserver final : public ServiceContext::ClientObserver {
+    void onCreateClient(Client*) {}
+    void onDestroyClient(Client* client) {
+        ClientOutOfLineExecutor::get(client)->shutdown();
+    }
+    void onCreateOperationContext(OperationContext*) {}
+    void onDestroyOperationContext(OperationContext*) {}
+};
+
+ServiceContext::ConstructorActionRegisterer
+    registerClientOutOfLineExecutorClientObserverConstructor{
+        "ClientOutOfLineExecutorClientObserverConstructor", [](ServiceContext* service) {
+            service->registerClientObserver(
+                std::make_unique<ClientOutOfLineExecutorClientObserver>());
+        }};
+
+}  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/client_out_of_line_executor.h b/src/mongo/db/client_out_of_line_executor.h
index a37a33133b2..08a9a734d82 100644
--- a/src/mongo/db/client_out_of_line_executor.h
+++ b/src/mongo/db/client_out_of_line_executor.h
@@ -55,6 +55,8 @@ public:
 
     using Task = OutOfLineExecutor::Task;
 
+    void shutdown();
+
     void schedule(Task) override;
 
     // Blocks until the executor is done running all scheduled tasks.
@@ -84,6 +86,9 @@ private:
     std::unique_ptr<Impl> _impl;
 
     std::shared_ptr<QueueType> _taskQueue;
+
+    // Provides the means to ensure `shutdown()` always precedes the destructor.
+    bool _isShutdown = false;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/thread_client_test.cpp b/src/mongo/db/thread_client_test.cpp
index fbeb5acac51..6c700a3a741 100644
--- a/src/mongo/db/thread_client_test.cpp
+++ b/src/mongo/db/thread_client_test.cpp
@@ -72,8 +72,8 @@ TEST_F(ThreadClientTest, TestDifferentArgs) {
     }
     {
         ASSERT_FALSE(haveClient());
-        ServiceContext sc;
-        ThreadClient tc("Test", &sc, nullptr);
+        auto sc = ServiceContext::make();
+        ThreadClient tc("Test", sc.get(), nullptr);
         ASSERT_TRUE(haveClient());
     }
 }