SERVER-36958 Parts of the command do not retry on network errors

1 files changed, 76 insertions, 25 deletions

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 ef3eda30ff1..5e11bd42bf4 100644
--- a/src/mongo/db/s/config/configsvr_shard_collection_command.cpp
+++ b/src/mongo/db/s/config/configsvr_shard_collection_command.cpp
@@ -34,7 +34,6 @@
 
 #include "mongo/bson/simple_bsonelement_comparator.h"
 #include "mongo/bson/util/bson_extract.h"
-#include "mongo/client/connpool.h"
 #include "mongo/db/audit.h"
 #include "mongo/db/auth/action_type.h"
 #include "mongo/db/auth/authorization_session.h"
@@ -137,7 +136,7 @@ void validateAndDeduceFullRequestOptions(OperationContext* opCtx,
                                          const NamespaceString& nss,
                                          const ShardKeyPattern& shardKeyPattern,
                                          int numShards,
-                                         ScopedDbConnection& conn,
+                                         const std::shared_ptr<Shard>& primaryShard,
                                          ConfigsvrShardCollectionRequest* request) {
     uassert(
         ErrorCodes::InvalidOptions, "cannot have empty shard key", !request->getKey().isEmpty());
@@ -186,8 +185,15 @@ void validateAndDeduceFullRequestOptions(OperationContext* opCtx,
     // collection.
     BSONObj res;
     {
-        std::list<BSONObj> all =
-            conn->getCollectionInfos(nss.db().toString(), BSON("name" << nss.coll()));
+        auto listCollectionsCmd =
+            BSON("listCollections" << 1 << "filter" << BSON("name" << nss.coll()));
+        auto allRes = uassertStatusOK(primaryShard->runExhaustiveCursorCommand(
+            opCtx,
+            ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+            nss.db().toString(),
+            listCollectionsCmd,
+            Milliseconds(-1)));
+        const auto& all = allRes.docs;
         if (!all.empty()) {
             res = all.front().getOwned();
         }
@@ -252,8 +258,7 @@ void validateShardKeyAgainstExistingIndexes(OperationContext* opCtx,
                                             const NamespaceString& nss,
                                             const BSONObj& proposedKey,
                                             const ShardKeyPattern& shardKeyPattern,
-                                            const std::shared_ptr<Shard> primaryShard,
-                                            ScopedDbConnection& conn,
+                                            const std::shared_ptr<Shard>& primaryShard,
                                             const ConfigsvrShardCollectionRequest& request) {
     // The proposed shard key must be validated against the set of existing indexes.
     // In particular, we must ensure the following constraints
@@ -282,7 +287,17 @@ void validateShardKeyAgainstExistingIndexes(OperationContext* opCtx,
     // 5. If the collection is empty, and it's still possible to create an index
     //    on the proposed key, we go ahead and do so.
 
-    std::list<BSONObj> indexes = conn->getIndexSpecs(nss.ns());
+    auto listIndexesCmd = BSON("listIndexes" << nss.coll());
+    auto indexesRes =
+        primaryShard->runExhaustiveCursorCommand(opCtx,
+                                                 ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+                                                 nss.db().toString(),
+                                                 listIndexesCmd,
+                                                 Milliseconds(-1));
+    std::vector<BSONObj> indexes;
+    if (indexesRes.getStatus().code() != ErrorCodes::NamespaceNotFound) {
+        indexes = uassertStatusOK(indexesRes).docs;
+    }
 
     // 1.  Verify consistency with existing unique indexes
     for (const auto& idx : indexes) {
@@ -348,16 +363,31 @@ void validateShardKeyAgainstExistingIndexes(OperationContext* opCtx,
         }
     }
 
+    auto countCmd = BSON("count" << nss.coll());
+    auto countRes =
+        uassertStatusOK(primaryShard->runCommand(opCtx,
+                                                 ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+                                                 nss.db().toString(),
+                                                 countCmd,
+                                                 Shard::RetryPolicy::kIdempotent));
+    const bool isEmpty = (countRes.response["n"].Int() == 0);
+
     if (hasUsefulIndexForKey) {
         // Check 2.iii and 2.iv. Make sure no null entries in the sharding index
         // and that there is a useful, non-multikey index available
         BSONObjBuilder checkShardingIndexCmd;
         checkShardingIndexCmd.append("checkShardingIndex", nss.ns());
         checkShardingIndexCmd.append("keyPattern", proposedKey);
-        BSONObj res;
-        auto success = conn.get()->runCommand("admin", checkShardingIndexCmd.obj(), res);
-        uassert(ErrorCodes::OperationFailed, res["errmsg"].str(), success);
-    } else if (conn->count(nss.ns()) != 0) {
+        auto checkShardingIndexRes = uassertStatusOK(
+            primaryShard->runCommand(opCtx,
+                                     ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+                                     "admin",
+                                     checkShardingIndexCmd.obj(),
+                                     Shard::RetryPolicy::kIdempotent));
+        uassert(ErrorCodes::OperationFailed,
+                checkShardingIndexRes.response["errmsg"].str(),
+                checkShardingIndexRes.commandStatus == Status::OK());
+    } else if (!isEmpty) {
         // 4. if no useful index, and collection is non-empty, fail
         uasserted(ErrorCodes::InvalidOptions,
                   "Please create an index that starts with the proposed shard key before "
@@ -503,20 +533,29 @@ void migrateAndFurtherSplitInitialChunks(OperationContext* opCtx,
         }
     }
 }
-boost::optional<UUID> getUUIDFromPrimaryShard(const NamespaceString& nss,
-                                              ScopedDbConnection& conn) {
+
+boost::optional<UUID> getUUIDFromPrimaryShard(OperationContext* opCtx,
+                                              const NamespaceString& nss,
+                                              const std::shared_ptr<Shard>& primaryShard) {
     // Obtain the collection's UUID from the primary shard's listCollections response.
     BSONObj res;
     {
-        std::list<BSONObj> all =
-            conn->getCollectionInfos(nss.db().toString(), BSON("name" << nss.coll()));
+        auto listCollectionsCmd =
+            BSON("listCollections" << 1 << "filter" << BSON("name" << nss.coll()));
+        auto allRes = uassertStatusOK(primaryShard->runExhaustiveCursorCommand(
+            opCtx,
+            ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+            nss.db().toString(),
+            listCollectionsCmd,
+            Milliseconds(-1)));
+        const auto& all = allRes.docs;
         if (!all.empty()) {
             res = all.front().getOwned();
         }
     }
 
     uassert(ErrorCodes::InternalError,
-            str::stream() << "expected the primary shard host " << conn.getHost()
+            str::stream() << "expected the primary shard host " << primaryShard->getConnString()
                           << " for database "
                           << nss.db()
                           << " to return an entry for "
@@ -653,14 +692,20 @@ public:
                         getTestCommandsEnabled());
 
             auto configShard = uassertStatusOK(shardRegistry->getShard(opCtx, dbType.getPrimary()));
-            ScopedDbConnection configConn(configShard->getConnString());
-            ON_BLOCK_EXIT([&configConn] { configConn.done(); });
+            auto countCmd = BSON("count" << nss.coll());
+            auto countRes = uassertStatusOK(
+                configShard->runCommand(opCtx,
+                                        ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+                                        nss.db().toString(),
+                                        countCmd,
+                                        Shard::RetryPolicy::kIdempotent));
+            auto numDocs = countRes.response["n"].Int();
 
             // If this is a collection on the config db, it must be empty to be sharded,
             // otherwise we might end up with chunks on the config servers.
             uassert(ErrorCodes::IllegalOperation,
                     "collections in the config db must be empty to be sharded",
-                    configConn->count(nss.ns()) == 0);
+                    numDocs == 0);
         }
 
         // For the config db, pick a new host shard for this collection, otherwise
@@ -674,12 +719,10 @@ public:
         }();
 
         auto primaryShard = uassertStatusOK(shardRegistry->getShard(opCtx, primaryShardId));
-        ScopedDbConnection conn(primaryShard->getConnString());
-        ON_BLOCK_EXIT([&conn] { conn.done(); });
 
         // Step 1.
         validateAndDeduceFullRequestOptions(
-            opCtx, nss, shardKeyPattern, shardIds.size(), conn, &request);
+            opCtx, nss, shardKeyPattern, shardIds.size(), primaryShard, &request);
 
         // The collation option should have been set to the collection default collation after being
         // validated.
@@ -750,7 +793,15 @@ public:
             // will not show up in the collection right after the count below has executed. It is
             // left here for backwards compatiblity with pre-4.0.4 clusters, which do not support
             // sharding being performed by the primary shard.
-            const bool isEmpty = (conn->count(nss.ns()) == 0);
+            auto countCmd = BSON("count" << nss.coll());
+            auto countRes = uassertStatusOK(
+                primaryShard->runCommand(opCtx,
+                                         ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+                                         nss.db().toString(),
+                                         countCmd,
+                                         Shard::RetryPolicy::kIdempotent));
+
+            const bool isEmpty = (countRes.response["n"].Int() == 0);
 
             // Map/reduce with output to an empty collection assumes it has full control of the
             // output collection and it would be an unsupported operation if the collection is being
@@ -766,11 +817,11 @@ public:
 
             // Step 3.
             validateShardKeyAgainstExistingIndexes(
-                opCtx, nss, proposedKey, shardKeyPattern, primaryShard, conn, request);
+                opCtx, nss, proposedKey, shardKeyPattern, primaryShard, request);
 
             // Step 4.
             if (request.getGetUUIDfromPrimaryShard()) {
-                uuid = getUUIDFromPrimaryShard(nss, conn);
+                uuid = getUUIDFromPrimaryShard(opCtx, nss, primaryShard);
             } else {
                 uuid = UUID::gen();
             }