SERVER-28201 Allow detatching of the current Client

3 files changed, 37 insertions, 12 deletions

diff --git a/src/mongo/db/client.cpp b/src/mongo/db/client.cpp
index 3f26dc039c0..6ecfeb3195b 100644
--- a/src/mongo/db/client.cpp
+++ b/src/mongo/db/client.cpp
@@ -70,7 +70,7 @@ void Client::initThread(StringData desc, transport::SessionHandle session) {
 void Client::initThread(StringData desc,
                         ServiceContext* service,
                         transport::SessionHandle session) {
-    invariant(currentClient.getMake()->get() == nullptr);
+    invariant(!haveClient());
 
     std::string fullDesc;
     if (session) {
@@ -82,12 +82,11 @@ void Client::initThread(StringData desc,
     setThreadName(fullDesc);
 
     // Create the client obj, attach to thread
-    *currentClient.get() = service->makeClient(fullDesc, std::move(session));
+    *currentClient.getMake() = service->makeClient(fullDesc, std::move(session));
 }
 
 void Client::destroy() {
-    invariant(currentClient.get());
-    invariant(currentClient.get()->get());
+    invariant(haveClient());
     currentClient.reset(nullptr);
 }
 
@@ -155,13 +154,22 @@ Client* Client::getCurrent() {
 }
 
 Client& cc() {
-    Client* c = currentClient.getMake()->get();
-    invariant(c);
-    return *c;
+    invariant(haveClient());
+    return *Client::getCurrent();
 }
 
 bool haveClient() {
-    return currentClient.getMake()->get();
+    return currentClient.get() && currentClient.get()->get();
+}
+
+ServiceContext::UniqueClient Client::releaseCurrent() {
+    invariant(haveClient());
+    return ServiceContext::UniqueClient(currentClient.get()->release());
+}
+
+void Client::setCurrent(ServiceContext::UniqueClient client) {
+    invariant(!haveClient());
+    *currentClient.getMake() = std::move(client);
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/client.h b/src/mongo/db/client.h
index 64e2850785f..a4f263f7d92 100644
--- a/src/mongo/db/client.h
+++ b/src/mongo/db/client.h
@@ -79,6 +79,22 @@ public:
                            ServiceContext* serviceContext,
                            transport::SessionHandle session);
 
+    /**
+     * Moves client into the thread_local for this thread. After this call, Client::getCurrent
+     * and cc() will return client.get(). The client will be destroyed with the thread exits
+     * or Client::destroy() is called.
+     */
+    static void setCurrent(ServiceContext::UniqueClient client);
+
+    /**
+     * Releases the client being managed by the thread_local for this thread. After this call
+     * cc() will crash the server and Client::getCurrent() will return nullptr until either
+     * Client::initThread() or Client::setCurrent() is called.
+     *
+     * The client will be released to the caller.
+     */
+    static ServiceContext::UniqueClient releaseCurrent();
+
     static Client* getCurrent();
 
     bool getIsLocalHostConnection() {
diff --git a/src/mongo/transport/service_entry_point_utils.cpp b/src/mongo/transport/service_entry_point_utils.cpp
index a0f3dedb285..61987c37972 100644
--- a/src/mongo/transport/service_entry_point_utils.cpp
+++ b/src/mongo/transport/service_entry_point_utils.cpp
@@ -70,10 +70,13 @@ struct Context {
 void* runFunc(void* ptr) {
     std::unique_ptr<Context> ctx(static_cast<Context*>(ptr));
 
-    auto tl = ctx->session->getTransportLayer();
-    Client::initThread("conn", ctx->session);
+    auto client = getGlobalServiceContext()->makeClient("conn", ctx->session);
     setThreadName(str::stream() << "conn" << ctx->session->id());
 
+    Client::setCurrent(std::move(client));
+
+    auto tl = ctx->session->getTransportLayer();
+
     try {
         ctx->task(ctx->session);
     } catch (const AssertionException& e) {
@@ -97,8 +100,6 @@ void* runFunc(void* ptr) {
               << " now open)";
     }
 
-    Client::destroy();
-
     return nullptr;
 }
 }  // namespace