SERVER-21366 Periodically yield when applying migration deletions

3 files changed, 71 insertions, 14 deletions

diff --git a/src/mongo/db/client.cpp b/src/mongo/db/client.cpp
index b357c317c33..142dd1f40bb 100644
--- a/src/mongo/db/client.cpp
+++ b/src/mongo/db/client.cpp
@@ -200,6 +200,13 @@ Client::Context::Context(OperationContext* txn, const string& ns, bool doVersion
 AutoGetDb::AutoGetDb(OperationContext* txn, const StringData& ns, LockMode mode)
     : _dbLock(txn->lockState(), ns, mode), _db(dbHolder().get(txn, ns)) {}
 
+AutoGetCollection::AutoGetCollection(OperationContext* txn,
+                                     const NamespaceString& nss,
+                                     LockMode mode)
+    : _autoDb(txn, nss.db(), mode),
+      _collLock(txn->lockState(), nss.ns(), mode),
+      _coll(_autoDb.getDb() ? _autoDb.getDb()->getCollection(nss) : nullptr) {}
+
 AutoGetOrCreateDb::AutoGetOrCreateDb(OperationContext* txn, const StringData& ns, LockMode mode)
     : _transaction(txn, MODE_IX),
       _dbLock(txn->lockState(), ns, mode),
diff --git a/src/mongo/db/client.h b/src/mongo/db/client.h
index 61e82acc01f..0ddbd2df084 100644
--- a/src/mongo/db/client.h
+++ b/src/mongo/db/client.h
@@ -89,6 +89,33 @@ private:
 };
 
 /**
+ * RAII-style class, which acquires a locks on the specified database and collection in the
+ * requested mode and obtains references to both.
+ *
+ * It is guaranteed that locks will be released when this object goes out of scope, therefore
+ * the database and the collection references returned by this class should not be retained.
+ */
+class AutoGetCollection {
+    MONGO_DISALLOW_COPYING(AutoGetCollection);
+
+public:
+    AutoGetCollection(OperationContext* txn, const NamespaceString& nss, LockMode mode);
+
+    Database* getDb() const {
+        return _autoDb.getDb();
+    }
+
+    Collection* getCollection() const {
+        return _coll;
+    }
+
+private:
+    const AutoGetDb _autoDb;
+    const Lock::CollectionLock _collLock;
+    Collection* const _coll;
+};
+
+/**
  * RAII-style class, which acquires a lock on the specified database in the requested mode and
  * obtains a reference to the database, creating it was non-existing. Used as a shortcut for
  * calls to dbHolder().openDb(), taking care of locking details. The requested mode must be
diff --git a/src/mongo/s/d_migrate.cpp b/src/mongo/s/d_migrate.cpp
index 77beca54897..080a877454d 100644
--- a/src/mongo/s/d_migrate.cpp
+++ b/src/mongo/s/d_migrate.cpp
@@ -1897,7 +1897,7 @@ public:
                 _errmsg = "UNKNOWN ERROR";
             }
 
-            error() << "migrate failed with unknown exception" << migrateLog;
+            severe() << "migrate failed with unknown exception" << migrateLog;
         }
 
         if (getState() != DONE) {
@@ -2385,33 +2385,54 @@ public:
         bool didAnything = false;
 
         if (xfer["deleted"].isABSONObj()) {
-            ScopedTransaction transaction(txn, MODE_IX);
-            Lock::DBLock dlk(txn->lockState(), nsToDatabaseSubstring(ns), MODE_IX);
-            Helpers::RemoveSaver rs("moveChunk", ns, "removedDuring");
+            scoped_ptr<Helpers::RemoveSaver> removeSaver;
+            if (serverGlobalParams.moveParanoia) {
+                removeSaver.reset(new Helpers::RemoveSaver("moveChunk", ns, "removedDuring"));
+            }
+
+            scoped_ptr<ScopedTransaction> autoXact;
+            scoped_ptr<AutoGetCollection> autoColl;
+
+            const int maxItersBeforeYield =
+                std::max(static_cast<int>(internalQueryExecYieldIterations), 1);
+            int opCounter = 0;
 
             BSONObjIterator i(xfer["deleted"].Obj());
             while (i.more()) {
-                Lock::CollectionLock clk(txn->lockState(), ns, MODE_X);
-                Client::Context ctx(txn, ns);
+                if (opCounter % maxItersBeforeYield == 0) {
+                    autoXact.reset();
+                    autoColl.reset();
+                    autoXact.reset(new ScopedTransaction(txn, MODE_IX));
+                    autoColl.reset(new AutoGetCollection(txn, NamespaceString(ns), MODE_X));
+                }
+
+                // If the collection does not exist there won't be any documents to remove. This
+                // condition can only happen if deletions are happening for an empty chunk or in the
+                // unlikely event that someone manually deleted the collection.
+                if (!autoColl->getCollection()) {
+                    break;
+                }
 
                 BSONObj id = i.next().Obj();
 
+                // Increment the counter before doing either the read or delete, so we yield the WT
+                // snapshot more aggressively
+                opCounter++;
+
                 // do not apply deletes if they do not belong to the chunk being migrated
                 BSONObj fullObj;
-                if (Helpers::findById(txn, ctx.db(), ns.c_str(), id, fullObj)) {
+                if (Helpers::findById(txn, autoColl->getDb(), ns.c_str(), id, fullObj)) {
                     if (!isInRange(fullObj, min, max, shardKeyPattern)) {
-                        log() << "not applying out of range deletion: " << fullObj << migrateLog;
-
                         continue;
                     }
                 }
 
-                if (serverGlobalParams.moveParanoia) {
-                    rs.goingToDelete(fullObj);
+                if (removeSaver) {
+                    removeSaver->goingToDelete(fullObj);
                 }
 
                 deleteObjects(txn,
-                              ctx.db(),
+                              autoColl->getDb(),
                               ns,
                               id,
                               PlanExecutor::YIELD_MANUAL,
@@ -2420,7 +2441,7 @@ public:
                               false /* god */,
                               true /* fromMigrate */);
 
-                *lastOpApplied = ctx.getClient()->getLastOp().asDate();
+                *lastOpApplied = txn->getClient()->getLastOp().asDate();
                 didAnything = true;
             }
         }
@@ -2636,7 +2657,9 @@ void migrateThread(std::string ns,
                    std::string fromShard,
                    OID epoch,
                    WriteConcernOptions writeConcern) {
-    Client::initThread("migrateThread");
+    const std::string migrateThreadName(str::stream() << "migrateThread-" << ns);
+    Client::initThread(migrateThreadName.c_str());
+
     OperationContextImpl txn;
     if (getGlobalAuthorizationManager()->isAuthEnabled()) {
         ShardedConnectionInfo::addHook();