SERVER-22659 removing _uncommittedMetadata local variable

6 files changed, 93 insertions, 42 deletions

diff --git a/src/mongo/db/s/collection_metadata.cpp b/src/mongo/db/s/collection_metadata.cpp
index 5bd6d51720b..cc0983eeba5 100644
--- a/src/mongo/db/s/collection_metadata.cpp
+++ b/src/mongo/db/s/collection_metadata.cpp
@@ -370,6 +370,23 @@ bool CollectionMetadata::getNextChunk(const BSONObj& lookupKey, ChunkType* chunk
     return false;
 }
 
+bool CollectionMetadata::getDifferentChunk(const BSONObj& chunkMinKey,
+                                           ChunkType* differentChunk) const {
+    RangeMap::const_iterator upperChunkIt = _chunksMap.end();
+    RangeMap::const_iterator lowerChunkIt = _chunksMap.begin();
+
+    while (lowerChunkIt != upperChunkIt) {
+        if (lowerChunkIt->first.woCompare(chunkMinKey) != 0) {
+            differentChunk->setMin(lowerChunkIt->first);
+            differentChunk->setMax(lowerChunkIt->second);
+            return true;
+        }
+        ++lowerChunkIt;
+    }
+
+    return false;
+}
+
 BSONObj CollectionMetadata::toBSON() const {
     BSONObjBuilder bb;
     toBSON(bb);
diff --git a/src/mongo/db/s/collection_metadata.h b/src/mongo/db/s/collection_metadata.h
index e33207fe753..a84eeb61849 100644
--- a/src/mongo/db/s/collection_metadata.h
+++ b/src/mongo/db/s/collection_metadata.h
@@ -145,6 +145,11 @@ public:
     bool getNextChunk(const BSONObj& lookupKey, ChunkType* chunk) const;
 
     /**
+     * Given a chunk identifying key "chunkMinKey", finds a different chunk if one exists.
+     */
+    bool getDifferentChunk(const BSONObj& chunkMinKey, ChunkType* differentChunk) const;
+
+    /**
      * Given a key in the shard key range, get the next range which overlaps or is greater than
      * this key.
      *
diff --git a/src/mongo/db/s/collection_metadata_test.cpp b/src/mongo/db/s/collection_metadata_test.cpp
index c831bc42581..c1ee18605b7 100644
--- a/src/mongo/db/s/collection_metadata_test.cpp
+++ b/src/mongo/db/s/collection_metadata_test.cpp
@@ -126,6 +126,11 @@ TEST_F(NoChunkFixture, getNextFromEmpty) {
     ASSERT(!getCollMetadata().getNextChunk(getCollMetadata().getMinKey(), &nextChunk));
 }
 
+TEST_F(NoChunkFixture, getDifferentFromEmpty) {
+    ChunkType differentChunk;
+    ASSERT(!getCollMetadata().getDifferentChunk(getCollMetadata().getMinKey(), &differentChunk));
+}
+
 TEST_F(NoChunkFixture, FirstChunkClonePlus) {
     ChunkVersion version(1, 0, getCollMetadata().getCollVersion().epoch());
     unique_ptr<CollectionMetadata> cloned(
@@ -391,6 +396,11 @@ TEST_F(SingleChunkFixture, GetLastChunkIsFalse) {
     ASSERT(!getCollMetadata().getNextChunk(getCollMetadata().getMaxKey(), &nextChunk));
 }
 
+TEST_F(SingleChunkFixture, getDifferentFromOneIsFalse) {
+    ChunkType differentChunk;
+    ASSERT(!getCollMetadata().getDifferentChunk(BSON("a" << 10), &differentChunk));
+}
+
 TEST_F(SingleChunkFixture, DonateLastChunk) {
     ChunkType chunk;
     chunk.setMin(BSON("a" << 10));
@@ -918,6 +928,29 @@ TEST_F(ThreeChunkWithRangeGapFixture, GetNextFromMiddle) {
 TEST_F(ThreeChunkWithRangeGapFixture, GetNextFromLast) {
     ChunkType nextChunk;
     ASSERT(getCollMetadata().getNextChunk(BSON("a" << 30), &nextChunk));
+    ASSERT_EQUALS(0, nextChunk.getMin().woCompare(BSON("a" << 30)));
+    ASSERT_EQUALS(0, nextChunk.getMax().woCompare(BSON("a" << MAXKEY)));
+}
+
+TEST_F(ThreeChunkWithRangeGapFixture, GetDifferentFromBeginning) {
+    ChunkType differentChunk;
+    ASSERT(getCollMetadata().getDifferentChunk(getCollMetadata().getMinKey(), &differentChunk));
+    ASSERT_EQUALS(0, differentChunk.getMin().woCompare(BSON("a" << 10)));
+    ASSERT_EQUALS(0, differentChunk.getMax().woCompare(BSON("a" << 20)));
+}
+
+TEST_F(ThreeChunkWithRangeGapFixture, GetDifferentFromMiddle) {
+    ChunkType differentChunk;
+    ASSERT(getCollMetadata().getDifferentChunk(BSON("a" << 10), &differentChunk));
+    ASSERT_EQUALS(0, differentChunk.getMin().woCompare(BSON("a" << MINKEY)));
+    ASSERT_EQUALS(0, differentChunk.getMax().woCompare(BSON("a" << 10)));
+}
+
+TEST_F(ThreeChunkWithRangeGapFixture, GetDifferentFromLast) {
+    ChunkType differentChunk;
+    ASSERT(getCollMetadata().getDifferentChunk(BSON("a" << 30), &differentChunk));
+    ASSERT_EQUALS(0, differentChunk.getMin().woCompare(BSON("a" << MINKEY)));
+    ASSERT_EQUALS(0, differentChunk.getMax().woCompare(BSON("a" << 10)));
 }
 
 TEST_F(ThreeChunkWithRangeGapFixture, MergeChunkHoleInRange) {
diff --git a/src/mongo/db/s/collection_sharding_state.cpp b/src/mongo/db/s/collection_sharding_state.cpp
index a8743a183a2..b4c924034a0 100644
--- a/src/mongo/db/s/collection_sharding_state.cpp
+++ b/src/mongo/db/s/collection_sharding_state.cpp
@@ -131,12 +131,6 @@ void CollectionShardingState::checkShardVersionOrThrow(OperationContext* txn) co
     ChunkVersion received;
     ChunkVersion wanted;
     if (!_checkShardVersionOk(txn, &errmsg, &received, &wanted)) {
-        // Set migration critical section in case we failed because of migration
-        if (_sourceMgr && _sourceMgr->getMigrationCriticalSection()) {
-            OperationShardingState::get(txn)
-                .setMigrationCriticalSection(_sourceMgr->getMigrationCriticalSection());
-        }
-
         throw SendStaleConfigException(_nss.ns(),
                                        str::stream() << "[" << _nss.ns()
                                                      << "] shard version not ok: " << errmsg,
@@ -218,7 +212,7 @@ bool CollectionShardingState::_checkShardVersionOk(OperationContext* txn,
     }
 
     if (!repl::ReplicationCoordinator::get(txn)->canAcceptWritesForDatabase(_nss.db())) {
-        // right now connections to secondaries aren't versioned at all
+        // Right now connections to secondaries aren't versioned at all.
         return true;
     }
 
@@ -244,8 +238,19 @@ bool CollectionShardingState::_checkShardVersionOk(OperationContext* txn,
         return true;
     }
 
+    // Set this for error messaging purposes before potentially returning false.
     *actualShardVersion = (_metadata ? _metadata->getShardVersion() : ChunkVersion::UNSHARDED());
 
+    if (_sourceMgr && _sourceMgr->getMigrationCriticalSection()) {
+        *errmsg = str::stream() << "migration commit in progress for " << _nss.ns();
+
+        // Set migration critical section on operation sharding state: operation will wait for the
+        // migration to finish before returning failure and retrying.
+        OperationShardingState::get(txn)
+            .setMigrationCriticalSection(_sourceMgr->getMigrationCriticalSection());
+        return false;
+    }
+
     if (expectedShardVersion->isWriteCompatibleWith(*actualShardVersion)) {
         return true;
     }
diff --git a/src/mongo/db/s/migration_source_manager.cpp b/src/mongo/db/s/migration_source_manager.cpp
index 6c097090960..4632479ecd8 100644
--- a/src/mongo/db/s/migration_source_manager.cpp
+++ b/src/mongo/db/s/migration_source_manager.cpp
@@ -148,7 +148,6 @@ MigrationSourceManager::MigrationSourceManager(OperationContext* txn, MoveChunkR
 
 MigrationSourceManager::~MigrationSourceManager() {
     invariant(!_cloneDriver);
-    invariant(!_uncommittedMetadata);
 }
 
 NamespaceString MigrationSourceManager::getNss() const {
@@ -222,7 +221,6 @@ Status MigrationSourceManager::enterCriticalSection(OperationContext* txn) {
         AutoGetCollection autoColl(txn, _args.getNss(), MODE_IX, MODE_X);
 
         auto css = CollectionShardingState::get(txn, _args.getNss().ns());
-
         if (!css->getMetadata() ||
             !css->getMetadata()->getCollVersion().equals(_committedMetadata->getCollVersion())) {
             return {ErrorCodes::IncompatibleShardingMetadata,
@@ -232,26 +230,12 @@ Status MigrationSourceManager::enterCriticalSection(OperationContext* txn) {
                         << ", actual: " << css->getMetadata()->getCollVersion().toString()};
         }
 
-        ChunkVersion uncommittedCollVersion = _committedMetadata->getCollVersion();
-        uncommittedCollVersion.incMajor();
-
-        // Bump the metadata's version up and "forget" about the chunk being moved. This is not the
-        // commit point, but in practice the state in this shard won't change until the commit it
-        // done.
-        ChunkType chunk;
-        chunk.setMin(_args.getMinKey());
-        chunk.setMax(_args.getMaxKey());
-
-        _uncommittedMetadata = _committedMetadata->cloneMigrate(chunk, uncommittedCollVersion);
-
         // IMPORTANT: After this line, the critical section is in place and needs to be rolled back
         // if anything fails, which would prevent commit to the config servers.
         _critSec = std::make_shared<CriticalSectionState>();
-        css->setMetadata(_uncommittedMetadata);
     }
 
-    log() << "Successfully entered critical section. Uncommited version: "
-          << _uncommittedMetadata->getCollVersion();
+    log() << "Successfully entered critical section.";
 
     _state = kCriticalSection;
     scopedGuard.Dismiss();
@@ -276,6 +260,10 @@ Status MigrationSourceManager::commitDonateChunk(OperationContext* txn) {
                 str::stream() << "commit clone failed due to " << commitCloneStatus.toString()};
     }
 
+    // Generate the next collection version.
+    ChunkVersion uncommittedCollVersion = _committedMetadata->getCollVersion();
+    uncommittedCollVersion.incMajor();
+
     // applyOps preparation for reflecting the uncommitted metadata on the config server
 
     // Preconditions
@@ -308,7 +296,7 @@ Status MigrationSourceManager::commitDonateChunk(OperationContext* txn) {
 
         BSONObjBuilder n(op.subobjStart("o"));
         n.append(ChunkType::name(), ChunkType::genID(_args.getNss().ns(), _args.getMinKey()));
-        _uncommittedMetadata->getCollVersion().addToBSON(n, ChunkType::DEPRECATED_lastmod());
+        uncommittedCollVersion.addToBSON(n, ChunkType::DEPRECATED_lastmod());
         n.append(ChunkType::ns(), _args.getNss().ns());
         n.append(ChunkType::min(), _args.getMinKey());
         n.append(ChunkType::max(), _args.getMaxKey());
@@ -327,14 +315,13 @@ Status MigrationSourceManager::commitDonateChunk(OperationContext* txn) {
     // Version at which the next highest lastmod will be set. If the chunk being moved is the last
     // in the shard, nextVersion is that chunk's lastmod otherwise the highest version is from the
     // chunk being bumped on the FROM-shard.
-    ChunkVersion nextVersion = _uncommittedMetadata->getCollVersion();
+    ChunkVersion nextVersion = uncommittedCollVersion;
 
     // If we have chunks left on the FROM shard, update the version of one of them as well. We can
     // figure that out by grabbing the metadata as it has been changed.
-    if (_uncommittedMetadata->getNumChunks()) {
+    if (_committedMetadata->getNumChunks() > 1) {
         ChunkType bumpChunk;
-        invariant(
-            _uncommittedMetadata->getNextChunk(_uncommittedMetadata->getMinKey(), &bumpChunk));
+        invariant(_committedMetadata->getDifferentChunk(_args.getMinKey(), &bumpChunk));
 
         BSONObj bumpMin = bumpChunk.getMin();
         BSONObj bumpMax = bumpChunk.getMax();
@@ -377,12 +364,25 @@ Status MigrationSourceManager::commitDonateChunk(OperationContext* txn) {
     // servers, so crash the server.
     fassertStatusOK(34431, applyOpsStatus);
 
-    // Now that applyOps succeeded, zero out the uncommitted metadata since it is now valid
-    _committedMetadata = std::move(_uncommittedMetadata);
+    // Now that applyOps succeeded and the new collection version is committed, update the
+    // collection metadata to the new collection version and forget the migrated chunk.
+    {
+        ScopedTransaction scopedXact(txn, MODE_IX);
+        AutoGetCollection autoColl(txn, _args.getNss(), MODE_IX, MODE_X);
+
+        ChunkType migratingChunkToForget;
+        migratingChunkToForget.setMin(_args.getMinKey());
+        migratingChunkToForget.setMax(_args.getMaxKey());
+        _committedMetadata =
+            _committedMetadata->cloneMigrate(migratingChunkToForget, uncommittedCollVersion);
+        auto css = CollectionShardingState::get(txn, _args.getNss().ns());
+        css->setMetadata(_committedMetadata);
+    }
 
     MONGO_FAIL_POINT_PAUSE_WHILE_SET(hangBeforeLeavingCriticalSection);
 
-    _critSec->exitCriticalSection();
+    scopedGuard.Dismiss();
+    _cleanup(txn);
 
     grid.catalogManager(txn)
         ->logChange(txn,
@@ -391,9 +391,6 @@ Status MigrationSourceManager::commitDonateChunk(OperationContext* txn) {
                     BSON("min" << _args.getMinKey() << "max" << _args.getMaxKey() << "from"
                                << _args.getFromShardId() << "to" << _args.getToShardId()));
 
-    scopedGuard.Dismiss();
-    _cleanup(txn);
-
     return Status::OK();
 }
 
@@ -426,10 +423,8 @@ void MigrationSourceManager::_cleanup(OperationContext* txn) {
         // collection
         css->clearMigrationSourceManager(txn);
 
-        // Restore the uncommitted metadata on the collection
+        // Leave the critical section.
         if (_state == kCriticalSection) {
-            css->setMetadata(_committedMetadata);
-            _uncommittedMetadata.reset();
             _critSec->exitCriticalSection();
             _critSec.reset();
         }
diff --git a/src/mongo/db/s/migration_source_manager.h b/src/mongo/db/s/migration_source_manager.h
index 00c31694a40..292f3c52045 100644
--- a/src/mongo/db/s/migration_source_manager.h
+++ b/src/mongo/db/s/migration_source_manager.h
@@ -209,10 +209,6 @@ private:
     // completed.
     std::unique_ptr<MigrationChunkClonerSource> _cloneDriver;
 
-    // After the critical section has been entered, contains the uncommitted metadata with a single
-    // chunk bumped by one version. Available after the critical section stage has completed.
-    std::shared_ptr<CollectionMetadata> _uncommittedMetadata;
-
     // Whether the source manager is in a critical section. Tracked as a shared pointer so that
     // callers don't have to hold collection lock in order to wait on it. Available after the
     // critical section stage has completed.