SERVER-29637 make shards obtain UUID for sharded collections from config server on setFeatureCompatibilityVersion: 3.6

4 files changed, 70 insertions, 2 deletions

diff --git a/src/mongo/db/catalog/coll_mod.cpp b/src/mongo/db/catalog/coll_mod.cpp
index 4233c92f37a..7fec15938b5 100644
--- a/src/mongo/db/catalog/coll_mod.cpp
+++ b/src/mongo/db/catalog/coll_mod.cpp
@@ -26,6 +26,8 @@
  *    it in the license file.
  */
 
+#define MONGO_LOG_DEFAULT_COMPONENT ::mongo::logger::LogComponent::kCommand
+
 #include "mongo/platform/basic.h"
 
 #include "mongo/db/catalog/coll_mod.h"
@@ -43,9 +45,15 @@
 #include "mongo/db/db_raii.h"
 #include "mongo/db/index/index_descriptor.h"
 #include "mongo/db/repl/replication_coordinator_global.h"
+#include "mongo/db/s/sharding_state.h"
 #include "mongo/db/service_context.h"
 #include "mongo/db/storage/recovery_unit.h"
 #include "mongo/db/views/view_catalog.h"
+#include "mongo/s/catalog/type_collection.h"
+#include "mongo/s/client/shard_registry.h"
+#include "mongo/s/grid.h"
+#include "mongo/s/sharding_initialization.h"
+#include "mongo/util/log.h"
 
 namespace mongo {
 
@@ -415,6 +423,7 @@ Status _collModInternal(OperationContext* opCtx,
 
 void _updateDBSchemaVersion(OperationContext* opCtx,
                             const std::string& dbname,
+                            std::map<std::string, UUID>& collToUUID,
                             bool needUUIDAdded) {
     // Iterate through all collections of database dbname and make necessary UUID changes.
     std::vector<NamespaceString> collNamespaceStrings;
@@ -439,6 +448,7 @@ void _updateDBSchemaVersion(OperationContext* opCtx,
         if (collNSS.db() == "local" || collNSS.coll() == "system.profile") {
             continue;
         }
+
         AutoGetDb autoDb(opCtx, dbname, MODE_X);
         Database* const db = autoDb.getDb();
         Collection* coll = db ? db->getCollection(opCtx, collNSS) : nullptr;
@@ -449,9 +459,16 @@ void _updateDBSchemaVersion(OperationContext* opCtx,
         BSONObjBuilder collModObjBuilder;
         collModObjBuilder.append("collMod", coll->ns().coll());
         BSONObj collModObj = collModObjBuilder.done();
+
         OptionalCollectionUUID uuid = boost::none;
         if (needUUIDAdded) {
-            uuid = UUID::gen();
+            if (collToUUID.find(collNSS.coll().toString()) != collToUUID.end()) {
+                // This is a sharded collection. Use the UUID generated by the config server.
+                uuid = collToUUID[collNSS.coll().toString()];
+            } else {
+                // This is an unsharded collection. Generate a UUID.
+                uuid = UUID::gen();
+            }
         }
         if ((needUUIDAdded && !coll->uuid()) || (!needUUIDAdded && coll->uuid())) {
             uassertStatusOK(collModForUUIDUpgrade(opCtx, coll->ns(), collModObj, uuid));
@@ -549,6 +566,40 @@ void updateUUIDSchemaVersion(OperationContext* opCtx, bool upgrade) {
     if (!enableCollectionUUIDs) {
         return;
     }
+
+    // A map of the form { db1: { collB: UUID, collA: UUID, ... }, db2: { ... } }
+    std::map<std::string, std::map<std::string, UUID>> dbToCollToUUID;
+    if (upgrade && ShardingState::get(opCtx)->enabled()) {
+        log() << "obtaining UUIDs for pre-existing sharded collections from config server";
+
+        // Get UUIDs for all existing sharded collections from the config server. Since the sharded
+        // collections are not stored per-database in config.collections, it's more efficient to
+        // read all the collections at once than to read them by database.
+        auto shardedColls =
+            uassertStatusOK(
+                Grid::get(opCtx)->shardRegistry()->getConfigShard()->exhaustiveFindOnConfig(
+                    opCtx,
+                    ReadPreferenceSetting{ReadPreference::PrimaryOnly},
+                    repl::ReadConcernLevel::kMajorityReadConcern,
+                    NamespaceString(CollectionType::ConfigNS),
+                    BSON("dropped" << false),  // query
+                    BSONObj(),                 // sort
+                    boost::none                // limit
+                    ))
+                .docs;
+
+        for (const auto& coll : shardedColls) {
+            auto collType = uassertStatusOK(CollectionType::fromBSON(coll));
+            uassert(ErrorCodes::InternalError,
+                    str::stream() << "expected entry " << coll << " in config.collections for "
+                                  << collType.getNs().ns()
+                                  << " to have a UUID, but it did not",
+                    collType.getUUID());
+            dbToCollToUUID[collType.getNs().db().toString()][collType.getNs().coll().toString()] =
+                *collType.getUUID();
+        }
+    }
+
     // Update UUIDs on all collections of all databases.
     std::vector<std::string> dbNames;
     StorageEngine* storageEngine = opCtx->getServiceContext()->getGlobalStorageEngine();
@@ -559,7 +610,7 @@ void updateUUIDSchemaVersion(OperationContext* opCtx, bool upgrade) {
 
     for (auto it = dbNames.begin(); it != dbNames.end(); ++it) {
         auto dbName = *it;
-        _updateDBSchemaVersion(opCtx, dbName, upgrade);
+        _updateDBSchemaVersion(opCtx, dbName, dbToCollToUUID[dbName], upgrade);
     }
     const WriteConcernOptions writeConcern(WriteConcernOptions::kMajority,
                                            WriteConcernOptions::SyncMode::UNSET,
diff --git a/src/mongo/db/commands/set_feature_compatibility_version_command.cpp b/src/mongo/db/commands/set_feature_compatibility_version_command.cpp
index f552c87478c..a9d4be0b58a 100644
--- a/src/mongo/db/commands/set_feature_compatibility_version_command.cpp
+++ b/src/mongo/db/commands/set_feature_compatibility_version_command.cpp
@@ -37,8 +37,10 @@
 #include "mongo/db/commands/feature_compatibility_version_command_parser.h"
 #include "mongo/db/concurrency/d_concurrency.h"
 #include "mongo/db/db_raii.h"
+#include "mongo/db/logical_time_validator.h"
 #include "mongo/db/repl/replication_coordinator.h"
 #include "mongo/db/repl/replication_coordinator_global.h"
+#include "mongo/db/s/sharding_state.h"
 #include "mongo/db/server_options.h"
 #include "mongo/rpc/get_status_from_command_result.h"
 #include "mongo/util/scopeguard.h"
@@ -132,6 +134,12 @@ public:
             updateUUIDSchemaVersion(opCtx, /*upgrade*/ false);
         }
 
+        // Ensure we try reading the keys for signing clusterTime immediately on upgrade to 3.6.
+        if (ShardingState::get(opCtx)->enabled() &&
+            version == FeatureCompatibilityVersionCommandParser::kVersion36) {
+            LogicalTimeValidator::get(opCtx)->forceKeyRefreshNow(opCtx);
+        }
+
         return true;
     }
 
diff --git a/src/mongo/db/logical_time_validator.cpp b/src/mongo/db/logical_time_validator.cpp
index 98e4f740ff6..9bbb9021ac7 100644
--- a/src/mongo/db/logical_time_validator.cpp
+++ b/src/mongo/db/logical_time_validator.cpp
@@ -186,4 +186,8 @@ bool LogicalTimeValidator::shouldGossipLogicalTime() {
     return _keyManager->hasSeenKeys();
 }
 
+void LogicalTimeValidator::forceKeyRefreshNow(OperationContext* opCtx) {
+    _keyManager->refreshNow(opCtx);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/logical_time_validator.h b/src/mongo/db/logical_time_validator.h
index 7d11450bf1c..dd171c5de2f 100644
--- a/src/mongo/db/logical_time_validator.h
+++ b/src/mongo/db/logical_time_validator.h
@@ -101,6 +101,11 @@ public:
      */
     bool shouldGossipLogicalTime();
 
+    /**
+     * Makes the KeysCollectionManager refresh synchronously.
+     */
+    void forceKeyRefreshNow(OperationContext* opCtx);
+
 private:
     SignedLogicalTime _getProof(const KeysCollectionDocument& keyDoc, LogicalTime newTime);