SERVER-43174 Make migration destination threads killable

(cherry picked from commit c4467548514ff07721f4de215b3b74d1111ee9d1)

4 files changed, 17 insertions, 5 deletions

diff --git a/src/mongo/db/client.cpp b/src/mongo/db/client.cpp
index 358de1437de..8ed19bb98c9 100644
--- a/src/mongo/db/client.cpp
+++ b/src/mongo/db/client.cpp
@@ -81,6 +81,13 @@ void Client::initThread(StringData desc,
     currentClient = service->makeClient(fullDesc, std::move(session));
 }
 
+void Client::initKillableThread(StringData desc, ServiceContext* service) {
+    initThread(desc, service, nullptr);
+
+    stdx::lock_guard lk(*currentClient);
+    currentClient->setSystemOperationKillable(lk);
+}
+
 namespace {
 int64_t generateSeed(const std::string& desc) {
     size_t seed = 0;
diff --git a/src/mongo/db/client.h b/src/mongo/db/client.h
index e7bed7b706f..75e828c63c3 100644
--- a/src/mongo/db/client.h
+++ b/src/mongo/db/client.h
@@ -76,6 +76,11 @@ public:
                            transport::SessionHandle session);
 
     /**
+     * Same as initThread, but also explicitly sets the client for this thread to be killable.
+     */
+    static void initKillableThread(StringData desc, ServiceContext* serviceContext);
+
+    /**
      * 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 when the thread exits
      * or the ThreadClient RAII helper exits its scope.
diff --git a/src/mongo/db/s/migration_destination_manager.cpp b/src/mongo/db/s/migration_destination_manager.cpp
index 059dfe85723..b526b23816b 100644
--- a/src/mongo/db/s/migration_destination_manager.cpp
+++ b/src/mongo/db/s/migration_destination_manager.cpp
@@ -387,7 +387,8 @@ repl::OpTime MigrationDestinationManager::cloneDocumentsFromDonor(
     repl::OpTime lastOpApplied;
 
     stdx::thread inserterThread{[&] {
-        ThreadClient tc("chunkInserter", opCtx->getServiceContext());
+        Client::initKillableThread("chunkInserter", opCtx->getServiceContext());
+
         auto inserterOpCtx = Client::getCurrent()->makeOperationContext();
         auto consumerGuard = makeGuard([&] {
             batches.closeConsumerEnd();
@@ -692,10 +693,9 @@ void MigrationDestinationManager::cloneCollectionIndexesAndOptions(OperationCont
 }
 
 void MigrationDestinationManager::_migrateThread() {
-    Client::initThread("migrateThread");
+    Client::initKillableThread("migrateThread", getGlobalServiceContext());
     auto opCtx = Client::getCurrent()->makeOperationContext();
 
-
     if (AuthorizationManager::get(opCtx->getServiceContext())->isAuthEnabled()) {
         AuthorizationSession::get(opCtx->getClient())->grantInternalAuthorization(opCtx.get());
     }
diff --git a/src/mongo/db/s/session_catalog_migration_destination.cpp b/src/mongo/db/s/session_catalog_migration_destination.cpp
index 0ff9dcaa737..9c5576c21be 100644
--- a/src/mongo/db/s/session_catalog_migration_destination.cpp
+++ b/src/mongo/db/s/session_catalog_migration_destination.cpp
@@ -384,8 +384,8 @@ void SessionCatalogMigrationDestination::join() {
  * 6. Wait for writes to be committed to majority of the replica set.
  */
 void SessionCatalogMigrationDestination::_retrieveSessionStateFromSource(ServiceContext* service) {
-    Client::initThread(
-        "sessionCatalogMigrationProducer-" + _migrationSessionId.toString(), service, nullptr);
+    Client::initKillableThread("sessionCatalogMigrationProducer-" + _migrationSessionId.toString(),
+                               service);
 
     bool oplogDrainedAfterCommiting = false;
     ProcessOplogResult lastResult;