SERVER-34760 Send setShardVersion on retries of shardCollection

6 files changed, 109 insertions, 21 deletions

diff --git a/buildscripts/resmokeconfig/suites/sharding_continuous_config_stepdown.yml b/buildscripts/resmokeconfig/suites/sharding_continuous_config_stepdown.yml
index 5f8e868fe89..ae5cf3c375b 100644
--- a/buildscripts/resmokeconfig/suites/sharding_continuous_config_stepdown.yml
+++ b/buildscripts/resmokeconfig/suites/sharding_continuous_config_stepdown.yml
@@ -155,6 +155,8 @@ selector:
   # Test expects a specific chunk distribution after shardCollection and it can be broken when
   # a step down occurs.
   - jstests/sharding/regex_targeting.js
+  # Test does not expect the primary config server node to step down.
+  - jstests/sharding/shard_collection_sets_shard_version_on_retry.js
 
 executor:
   config:
diff --git a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
index 3193b27eb50..bcbb4b9f6fd 100644
--- a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
+++ b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
@@ -72,6 +72,7 @@ selector:
   - jstests/sharding/secondary_shard_versioning.js
   - jstests/sharding/session_info_in_oplog.js
   - jstests/sharding/sessions_collection_auto_healing.js
+  - jstests/sharding/shard_collection_sets_shard_version_on_retry.js
   - jstests/sharding/shard_config_db_collections.js
   - jstests/sharding/time_zone_info_mongos.js
   - jstests/sharding/verify_sessions_expiration_sharded.js
diff --git a/jstests/sharding/shard_collection_sets_shard_version_on_retry.js b/jstests/sharding/shard_collection_sets_shard_version_on_retry.js
new file mode 100644
index 00000000000..63dae26db87
--- /dev/null
+++ b/jstests/sharding/shard_collection_sets_shard_version_on_retry.js
@@ -0,0 +1,40 @@
+/**
+ * Tests that a retry of a successful shardCollection will send setShardVersion to the
+ * primary shard.
+ */
+
+(function() {
+    'use strict';
+
+    let st = new ShardingTest({mongos: 1, shards: {rs0: {nodes: 2}}});
+    const dbName = 'db';
+    const coll = 'foo';
+    const nss = dbName + '.' + coll;
+    const mongos = st.s0;
+    const configPrimary = st.configRS.getPrimary();
+
+    assert.commandWorked(mongos.adminCommand({enableSharding: dbName}));
+
+    assert.commandWorked(configPrimary.getDB('admin').runCommand({
+        configureFailPoint: 'skipSendingSetShardVersionAfterCompletionOfShardCollection',
+        mode: 'alwaysOn'
+    }));
+
+    // Returns true if the shard is aware that the collection is sharded.
+    const isShardAware = (shard, coll) => {
+        const res =
+            assert.commandWorked(shard.adminCommand({getShardVersion: coll, fullMetadata: true}));
+        return res.metadata.hasOwnProperty("collVersion");
+    };
+
+    assert.commandWorked(mongos.adminCommand({shardCollection: nss, key: {aKey: 1}}));
+    assert(!isShardAware(st.rs0.getPrimary(), nss));
+    assert.eq(
+        configPrimary.getDB('config').chunks.find({ns: nss, shard: st.shard0.shardName}).itcount(),
+        1);
+
+    assert.commandWorked(mongos.adminCommand({shardCollection: nss, key: {aKey: 1}}));
+    assert(isShardAware(st.rs0.getPrimary(), nss));
+
+    st.stop();
+})();
diff --git a/src/mongo/db/s/config/configsvr_shard_collection_command.cpp b/src/mongo/db/s/config/configsvr_shard_collection_command.cpp
index a42e0b87659..45c2a714666 100644
--- a/src/mongo/db/s/config/configsvr_shard_collection_command.cpp
+++ b/src/mongo/db/s/config/configsvr_shard_collection_command.cpp
@@ -833,6 +833,28 @@ public:
 
         // Step 2.
         if (auto existingColl = checkIfAlreadyShardedWithSameOptions(opCtx, nss, request)) {
+            auto findResponse = uassertStatusOK(
+                Grid::get(opCtx)->shardRegistry()->getConfigShard()->exhaustiveFindOnConfig(
+                    opCtx,
+                    ReadPreferenceSetting{ReadPreference::PrimaryOnly},
+                    repl::ReadConcernLevel::kLocalReadConcern,
+                    NamespaceString(ChunkType::ConfigNS),
+                    BSON("ns" << nss.ns() << "shard" << primaryShardId),
+                    BSON(ChunkType::lastmod << -1),
+                    1));
+
+            const auto& chunksVector = findResponse.docs;
+
+            // If the vector is empty, it implies that the shard has moved all chunks away. We
+            // accordingly send the "has no chunks" shard version, (0, 0, collection epoch).
+            const auto chunkVersion = chunksVector.empty()
+                ? ChunkVersion(0, 0, existingColl->getEpoch())
+                : uassertStatusOK(ChunkType::fromConfigBSON(chunksVector.front())).getVersion();
+
+            // Send setShardVersion just in case this command invocation is a retry and the primary
+            // shard has not yet received its new shardVersion.
+            catalogManager->trySetShardVersionOnPrimaryShard(
+                opCtx, nss, primaryShardId, chunkVersion);
             result << "collectionsharded" << nss.ns();
             if (existingColl->getUUID()) {
                 result << "collectionUUID" << *existingColl->getUUID();
diff --git a/src/mongo/s/catalog/sharding_catalog_manager.h b/src/mongo/s/catalog/sharding_catalog_manager.h
index a7d764fe2db..ada04fb0720 100644
--- a/src/mongo/s/catalog/sharding_catalog_manager.h
+++ b/src/mongo/s/catalog/sharding_catalog_manager.h
@@ -266,6 +266,15 @@ public:
                          const ShardId& dbPrimaryShardId);
 
     /**
+     * Attempts to sets the shard version to the given chunk version with authoritative: true on
+     * the primary shard. If it fails, emit a warning then move on.
+     */
+    void trySetShardVersionOnPrimaryShard(OperationContext* opCtx,
+                                          const NamespaceString& nss,
+                                          const ShardId& dbPrimaryShardId,
+                                          const ChunkVersion& collVersion);
+
+    /**
      * Iterates through each entry in config.collections that does not have a UUID, generates a UUID
      * for the collection, and updates the entry with the generated UUID.
      *
diff --git a/src/mongo/s/catalog/sharding_catalog_manager_collection_operations.cpp b/src/mongo/s/catalog/sharding_catalog_manager_collection_operations.cpp
index 4952c9a5aa1..7b6488ab9a8 100644
--- a/src/mongo/s/catalog/sharding_catalog_manager_collection_operations.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_manager_collection_operations.cpp
@@ -61,6 +61,7 @@
 #include "mongo/s/set_shard_version_request.h"
 #include "mongo/s/shard_key_pattern.h"
 #include "mongo/s/shard_util.h"
+#include "mongo/util/fail_point_service.h"
 #include "mongo/util/log.h"
 #include "mongo/util/mongoutils/str.h"
 #include "mongo/util/scopeguard.h"
@@ -74,6 +75,8 @@ using std::set;
 
 namespace {
 
+MONGO_FP_DECLARE(skipSendingSetShardVersionAfterCompletionOfShardCollection);
+
 const ReadPreferenceSetting kConfigReadSelector(ReadPreference::Nearest, TagSet{});
 const WriteConcernOptions kNoWaitWriteConcern(1, WriteConcernOptions::SyncMode::UNSET, Seconds(0));
 
@@ -295,37 +298,48 @@ void ShardingCatalogManager::shardCollection(OperationContext* opCtx,
     auto shard = uassertStatusOK(shardRegistry->getShard(opCtx, dbPrimaryShardId));
     invariant(!shard->isConfig());
 
-    // Tell the primary mongod to refresh its data
-    // TODO:  Think the real fix here is for mongos to just
-    //        assume that all collections are sharded, when we get there
+    if (!MONGO_FAIL_POINT(skipSendingSetShardVersionAfterCompletionOfShardCollection)) {
+        // Inform primary shard that the collection has been sharded.
+        trySetShardVersionOnPrimaryShard(opCtx, NamespaceString(ns), dbPrimaryShardId, collVersion);
+    }
+
+    catalogClient
+        ->logChange(opCtx,
+                    "shardCollection.end",
+                    ns,
+                    BSON("version" << collVersion.toString()),
+                    ShardingCatalogClient::kMajorityWriteConcern)
+        .transitional_ignore();
+}
+
+void ShardingCatalogManager::trySetShardVersionOnPrimaryShard(OperationContext* opCtx,
+                                                              const NamespaceString& nss,
+                                                              const ShardId& dbPrimaryShardId,
+                                                              const ChunkVersion& collVersion) {
+    const auto shardRegistry = Grid::get(opCtx)->shardRegistry();
+    auto primaryShard = uassertStatusOK(shardRegistry->getShard(opCtx, dbPrimaryShardId));
     SetShardVersionRequest ssv = SetShardVersionRequest::makeForVersioningNoPersist(
         shardRegistry->getConfigServerConnectionString(),
         dbPrimaryShardId,
         primaryShard->getConnString(),
-        NamespaceString(ns),
+        nss,
         collVersion,
-        true);
-
-    auto ssvResponse =
-        shard->runCommandWithFixedRetryAttempts(opCtx,
-                                                ReadPreferenceSetting{ReadPreference::PrimaryOnly},
-                                                "admin",
-                                                ssv.toBSON(),
-                                                Shard::RetryPolicy::kIdempotent);
+        true  // isAuthoritative
+        );
+
+    auto ssvResponse = primaryShard->runCommandWithFixedRetryAttempts(
+        opCtx,
+        ReadPreferenceSetting{ReadPreference::PrimaryOnly},
+        "admin",
+        ssv.toBSON(),
+        Shard::RetryPolicy::kIdempotent);
+
     auto status = ssvResponse.isOK() ? std::move(ssvResponse.getValue().commandStatus)
                                      : std::move(ssvResponse.getStatus());
     if (!status.isOK()) {
-        warning() << "could not update initial version of " << ns << " on shard primary "
+        warning() << "could not update initial version of " << nss << " on shard primary "
                   << dbPrimaryShardId << causedBy(redact(status));
     }
-
-    catalogClient
-        ->logChange(opCtx,
-                    "shardCollection.end",
-                    ns,
-                    BSON("version" << collVersion.toString()),
-                    ShardingCatalogClient::kMajorityWriteConcern)
-        .transitional_ignore();
 }
 
 void ShardingCatalogManager::generateUUIDsForExistingShardedCollections(OperationContext* opCtx) {