SERVER-42141 Implement sharded metadata updates for refineCollectionShardKey without using a transaction

8 files changed, 179 insertions, 11 deletions

diff --git a/src/mongo/db/s/config/configsvr_refine_collection_shard_key_command.cpp b/src/mongo/db/s/config/configsvr_refine_collection_shard_key_command.cpp
index 910fb096689..21901105103 100644
--- a/src/mongo/db/s/config/configsvr_refine_collection_shard_key_command.cpp
+++ b/src/mongo/db/s/config/configsvr_refine_collection_shard_key_command.cpp
@@ -34,6 +34,7 @@
 #include "mongo/db/auth/authorization_session.h"
 #include "mongo/db/commands.h"
 #include "mongo/db/repl/repl_client_info.h"
+#include "mongo/db/s/config/sharding_catalog_manager.h"
 #include "mongo/db/s/shard_key_util.h"
 #include "mongo/s/catalog/dist_lock_manager.h"
 #include "mongo/s/grid.h"
@@ -140,6 +141,9 @@ public:
                                                                  boost::none,
                                                                  collType.getUnique(),
                                                                  false);  // createIndexIfPossible
+
+            ShardingCatalogManager::get(opCtx)->refineCollectionShardKey(
+                opCtx, nss, newShardKeyPattern);
         }
 
     private:
diff --git a/src/mongo/db/s/config/sharding_catalog_manager.h b/src/mongo/db/s/config/sharding_catalog_manager.h
index 12544482a38..ec084b70cab 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager.h
+++ b/src/mongo/db/s/config/sharding_catalog_manager.h
@@ -304,6 +304,16 @@ public:
                           const CollectionOptions& options);
 
     /**
+     * Refines the shard key of an existing collection with namespace 'nss'. Here, 'shardKey'
+     * denotes the new shard key, which must contain the old shard key as a prefix.
+     *
+     * Throws exception on errors.
+     */
+    void refineCollectionShardKey(OperationContext* opCtx,
+                                  const NamespaceString& nss,
+                                  const ShardKeyPattern& newShardKey);
+
+    /**
      * Creates a ScopedLock on the collection name in _namespaceSerializer. This is to prevent
      * timeouts waiting on the dist lock if multiple threads attempt to create or drop the same
      * collection.
diff --git a/src/mongo/db/s/config/sharding_catalog_manager_collection_operations.cpp b/src/mongo/db/s/config/sharding_catalog_manager_collection_operations.cpp
index c205ce2841a..5993661a884 100644
--- a/src/mongo/db/s/config/sharding_catalog_manager_collection_operations.cpp
+++ b/src/mongo/db/s/config/sharding_catalog_manager_collection_operations.cpp
@@ -574,4 +574,99 @@ void ShardingCatalogManager::createCollection(OperationContext* opCtx,
     repl::ReplClientInfo::forClient(opCtx->getClient()).setLastOpToSystemLastOpTime(opCtx);
 }
 
+void ShardingCatalogManager::refineCollectionShardKey(OperationContext* opCtx,
+                                                      const NamespaceString& nss,
+                                                      const ShardKeyPattern& newShardKeyPattern) {
+    // Take _kChunkOpLock in exclusive mode to prevent concurrent chunk splits, merges, and
+    // migrations. Take _kZoneOpLock in exclusive mode to prevent concurrent zone operations.
+    // TODO(SERVER-25359): Replace with a collection-specific lock map to allow splits/merges/
+    // move chunks on different collections to proceed in parallel.
+    Lock::ExclusiveLock chunkLk(opCtx->lockState(), _kChunkOpLock);
+    Lock::ExclusiveLock zoneLk(opCtx->lockState(), _kZoneOpLock);
+
+    const auto catalogClient = Grid::get(opCtx)->catalogClient();
+    const auto newEpoch = OID::gen();
+
+    auto collType = uassertStatusOK(catalogClient->getCollection(opCtx, nss)).value;
+    const auto oldShardKeyPattern = ShardKeyPattern(collType.getKeyPattern());
+    collType.setEpoch(newEpoch);
+    collType.setKeyPattern(newShardKeyPattern.getKeyPattern());
+
+    uassertStatusOK(ShardingCatalogClientImpl::updateShardingCatalogEntryForCollection(
+        opCtx, nss, collType, false /* upsert */));
+
+    const auto oldFields = oldShardKeyPattern.toBSON();
+    const auto newFields =
+        newShardKeyPattern.toBSON().filterFieldsUndotted(oldFields, false /* inFilter */);
+
+    // Construct query objects for calls to 'updateConfigDocument(s)' below.
+    BSONObjBuilder notGlobalMaxBuilder, isGlobalMaxBuilder;
+    notGlobalMaxBuilder.append(ChunkType::ns.name(), nss.ns());
+    isGlobalMaxBuilder.append(ChunkType::ns.name(), nss.ns());
+    for (const auto& fieldElem : oldFields) {
+        notGlobalMaxBuilder.append("max." + fieldElem.fieldNameStringData(), BSON("$ne" << MAXKEY));
+        isGlobalMaxBuilder.append("max." + fieldElem.fieldNameStringData(), BSON("$eq" << MAXKEY));
+    }
+    const auto notGlobalMaxQuery = notGlobalMaxBuilder.obj();
+    const auto isGlobalMaxQuery = isGlobalMaxBuilder.obj();
+
+    // The defaultBounds object sets the bounds of each new field in the refined key to MinKey. The
+    // globalMaxBounds object corrects the max bounds of the global max chunk/tag to MaxKey.
+    //
+    // Example: oldKeyDoc = {a: 1}
+    //          newKeyDoc = {a: 1, b: 1, c: 1}
+    //          defaultBounds = {min.b: MinKey, min.c: MinKey, max.b: MinKey, max.c: MinKey}
+    //          globalMaxBounds = {min.b: MinKey, min.c: MinKey, max.b: MaxKey, max.c: MaxKey}
+    BSONObjBuilder defaultBoundsBuilder, globalMaxBoundsBuilder;
+    for (const auto& fieldElem : newFields) {
+        defaultBoundsBuilder.appendMinKey("min." + fieldElem.fieldNameStringData());
+        defaultBoundsBuilder.appendMinKey("max." + fieldElem.fieldNameStringData());
+
+        globalMaxBoundsBuilder.appendMinKey("min." + fieldElem.fieldNameStringData());
+        globalMaxBoundsBuilder.appendMaxKey("max." + fieldElem.fieldNameStringData());
+    }
+    const auto defaultBounds = defaultBoundsBuilder.obj();
+    const auto globalMaxBounds = globalMaxBoundsBuilder.obj();
+
+    // Update all config.chunks entries for the given namespace by setting (i) their epoch to the
+    // newly-generated objectid, (ii) their bounds for each new field in the refined key to MinKey
+    // (except for the global max chunk where the max bounds are set to MaxKey), and unsetting (iii)
+    // their jumbo field.
+    uassertStatusOK(catalogClient->updateConfigDocuments(
+        opCtx,
+        ChunkType::ConfigNS,
+        notGlobalMaxQuery,
+        BSON("$set" << BSON(ChunkType::epoch(newEpoch)) << "$max" << defaultBounds << "$unset"
+                    << BSON(ChunkType::jumbo(true))),
+        false,  // upsert
+        ShardingCatalogClient::kLocalWriteConcern));
+
+    uassertStatusOK(catalogClient->updateConfigDocument(
+        opCtx,
+        ChunkType::ConfigNS,
+        isGlobalMaxQuery,
+        BSON("$set" << BSON(ChunkType::epoch(newEpoch)) << "$max" << globalMaxBounds << "$unset"
+                    << BSON(ChunkType::jumbo(true))),
+        false,  // upsert
+        ShardingCatalogClient::kLocalWriteConcern));
+
+    // Update all config.tags entries for the given namespace by setting their bounds for each new
+    // field in the refined key to MinKey (except for the global max tag where the max bounds are
+    // set to MaxKey).
+    uassertStatusOK(
+        catalogClient->updateConfigDocuments(opCtx,
+                                             TagsType::ConfigNS,
+                                             notGlobalMaxQuery,
+                                             BSON("$max" << defaultBounds),
+                                             false,  // upsert
+                                             ShardingCatalogClient::kLocalWriteConcern));
+
+    uassertStatusOK(catalogClient->updateConfigDocument(opCtx,
+                                                        TagsType::ConfigNS,
+                                                        isGlobalMaxQuery,
+                                                        BSON("$max" << globalMaxBounds),
+                                                        false,  // upsert
+                                                        ShardingCatalogClient::kLocalWriteConcern));
+}
+
 }  // namespace mongo
diff --git a/src/mongo/s/catalog/sharding_catalog_client.h b/src/mongo/s/catalog/sharding_catalog_client.h
index febbf82fe06..d11bedac634 100644
--- a/src/mongo/s/catalog/sharding_catalog_client.h
+++ b/src/mongo/s/catalog/sharding_catalog_client.h
@@ -335,8 +335,8 @@ public:
                                                        const WriteConcernOptions& writeConcern) = 0;
 
     /**
-     * Updates a single document in the specified namespace on the config server. The document must
-     * have an _id index. Must only be used for updates to the 'config' database.
+     * Updates a single document in the specified namespace on the config server. Must only be used
+     * for updates to the 'config' database.
      *
      * This method retries the operation on NotMaster or network errors, so it should only be used
      * with modifications which are idempotent.
@@ -356,6 +356,19 @@ public:
                                                   const WriteConcernOptions& writeConcern) = 0;
 
     /**
+     * Updates multiple documents in the specified namespace on the config server. Must only be used
+     * for updates to the 'config' database.
+     *
+     * Read the comment for 'updateConfigDocument' for additional information.
+     */
+    virtual StatusWith<bool> updateConfigDocuments(OperationContext* opCtx,
+                                                   const NamespaceString& nss,
+                                                   const BSONObj& query,
+                                                   const BSONObj& update,
+                                                   bool upsert,
+                                                   const WriteConcernOptions& writeConcern) = 0;
+
+    /**
      * Removes documents matching a particular query predicate from the specified namespace on the
      * config server. Must only be used for deletions from the 'config' database.
      *
diff --git a/src/mongo/s/catalog/sharding_catalog_client_impl.cpp b/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
index 65020a8ec1b..23db44c5453 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
@@ -203,7 +203,8 @@ Status ShardingCatalogClientImpl::updateShardingCatalogEntryForCollection(
                                         BSON(CollectionType::fullNs(nss.ns())),
                                         coll.toBSON(),
                                         upsert,
-                                        ShardingCatalogClient::kMajorityWriteConcern);
+                                        ShardingCatalogClient::kMajorityWriteConcern,
+                                        false /* multi */);
     return status.getStatus().withContext(str::stream() << "Collection metadata write failed");
 }
 
@@ -925,16 +926,29 @@ StatusWith<bool> ShardingCatalogClientImpl::updateConfigDocument(
     const BSONObj& update,
     bool upsert,
     const WriteConcernOptions& writeConcern) {
-    return _updateConfigDocument(opCtx, nss, query, update, upsert, writeConcern);
+    return _updateConfigDocument(
+        opCtx, nss, query, update, upsert, writeConcern, false /* useMultiUpdate */);
 }
 
-StatusWith<bool> ShardingCatalogClientImpl::_updateConfigDocument(
+StatusWith<bool> ShardingCatalogClientImpl::updateConfigDocuments(
     OperationContext* opCtx,
     const NamespaceString& nss,
     const BSONObj& query,
     const BSONObj& update,
     bool upsert,
     const WriteConcernOptions& writeConcern) {
+    return _updateConfigDocument(
+        opCtx, nss, query, update, upsert, writeConcern, true /* useMultiUpdate */);
+}
+
+StatusWith<bool> ShardingCatalogClientImpl::_updateConfigDocument(
+    OperationContext* opCtx,
+    const NamespaceString& nss,
+    const BSONObj& query,
+    const BSONObj& update,
+    bool upsert,
+    const WriteConcernOptions& writeConcern,
+    bool useMultiUpdate) {
     invariant(nss.db() == NamespaceString::kConfigDb);
 
     BatchedCommandRequest request([&] {
@@ -944,7 +958,7 @@ StatusWith<bool> ShardingCatalogClientImpl::_updateConfigDocument(
             entry.setQ(query);
             entry.setU(update);
             entry.setUpsert(upsert);
-            entry.setMulti(false);
+            entry.setMulti(useMultiUpdate);
             return entry;
         }()});
         return updateOp;
@@ -961,8 +975,14 @@ StatusWith<bool> ShardingCatalogClientImpl::_updateConfigDocument(
     }
 
     const auto nSelected = response.getN();
-    invariant(nSelected == 0 || nSelected == 1);
-    return (nSelected == 1);
+
+    if (useMultiUpdate) {
+        invariant(nSelected >= 0);
+        return (nSelected > 0);
+    } else {
+        invariant(nSelected == 0 || nSelected == 1);
+        return (nSelected == 1);
+    }
 }
 
 Status ShardingCatalogClientImpl::removeConfigDocuments(OperationContext* opCtx,
diff --git a/src/mongo/s/catalog/sharding_catalog_client_impl.h b/src/mongo/s/catalog/sharding_catalog_client_impl.h
index ab4e9506594..3f70d6bae86 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_impl.h
+++ b/src/mongo/s/catalog/sharding_catalog_client_impl.h
@@ -155,6 +155,13 @@ public:
                                           bool upsert,
                                           const WriteConcernOptions& writeConcern) override;
 
+    StatusWith<bool> updateConfigDocuments(OperationContext* opCtx,
+                                           const NamespaceString& nss,
+                                           const BSONObj& query,
+                                           const BSONObj& update,
+                                           bool upsert,
+                                           const WriteConcernOptions& writeConcern) override;
+
     Status removeConfigDocuments(OperationContext* opCtx,
                                  const NamespaceString& nss,
                                  const BSONObj& query,
@@ -170,8 +177,9 @@ public:
 
 private:
     /**
-     * Updates a single document in the specified namespace on the config server. The document must
-     * have an _id index. Must only be used for updates to the 'config' database.
+     * Updates a single document (if useMultiUpdate is false) or multiple documents (if
+     * useMultiUpdate is true) in the specified namespace on the config server. Must only be used
+     * for updates to the 'config' database.
      *
      * This method retries the operation on NotMaster or network errors, so it should only be used
      * with modifications which are idempotent.
@@ -186,7 +194,8 @@ private:
                                                   const BSONObj& query,
                                                   const BSONObj& update,
                                                   bool upsert,
-                                                  const WriteConcernOptions& writeConcern);
+                                                  const WriteConcernOptions& writeConcern,
+                                                  bool useMultiUpdate);
 
     StatusWith<repl::OpTimeWith<std::vector<BSONObj>>> _exhaustiveFindOnConfig(
         OperationContext* opCtx,
diff --git a/src/mongo/s/catalog/sharding_catalog_client_mock.cpp b/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
index 481b9448033..511b8872000 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
@@ -180,6 +180,16 @@ StatusWith<bool> ShardingCatalogClientMock::updateConfigDocument(
     return {ErrorCodes::InternalError, "Method not implemented"};
 }
 
+StatusWith<bool> ShardingCatalogClientMock::updateConfigDocuments(
+    OperationContext* opCtx,
+    const NamespaceString& nss,
+    const BSONObj& query,
+    const BSONObj& update,
+    bool upsert,
+    const WriteConcernOptions& writeConcern) {
+    return {ErrorCodes::InternalError, "Method not implemented"};
+}
+
 Status ShardingCatalogClientMock::removeConfigDocuments(OperationContext* opCtx,
                                                         const NamespaceString& nss,
                                                         const BSONObj& query,
diff --git a/src/mongo/s/catalog/sharding_catalog_client_mock.h b/src/mongo/s/catalog/sharding_catalog_client_mock.h
index f8c84300cc1..dcb681762f1 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_mock.h
+++ b/src/mongo/s/catalog/sharding_catalog_client_mock.h
@@ -128,6 +128,13 @@ public:
                                           bool upsert,
                                           const WriteConcernOptions& writeConcern) override;
 
+    StatusWith<bool> updateConfigDocuments(OperationContext* opCtx,
+                                           const NamespaceString& nss,
+                                           const BSONObj& query,
+                                           const BSONObj& update,
+                                           bool upsert,
+                                           const WriteConcernOptions& writeConcern) override;
+
     Status removeConfigDocuments(OperationContext* opCtx,
                                  const NamespaceString& nss,
                                  const BSONObj& query,