SERVER-64872 Ensure `CollectionUUIDMismatch` from sharded `aggregate` does not incorrectly omit the actual collection

6 files changed, 200 insertions, 51 deletions

diff --git a/jstests/sharding/collection_uuid_aggregate.js b/jstests/sharding/collection_uuid_aggregate.js
new file mode 100644
index 00000000000..11a864cc80c
--- /dev/null
+++ b/jstests/sharding/collection_uuid_aggregate.js
@@ -0,0 +1,52 @@
+/**
+ * Tests the collectionUUID parameter of the aggregate command when not all shards own chunks for
+ * the collection.
+ *
+ * @tags: [
+ *   requires_fcv_60,
+ * ]
+ */
+(function() {
+'use strict';
+
+const st = new ShardingTest({shards: 2});
+
+const db = st.s.getDB(jsTestName());
+assert.commandWorked(st.s.adminCommand({enableSharding: db.getName()}));
+st.ensurePrimaryShard(db.getName(), st.shard0.shardName);
+
+const shardedColl = db.sharded;
+const unshardedColl = db.unsharded;
+
+assert.commandWorked(shardedColl.insert({_id: 0}));
+assert.commandWorked(unshardedColl.insert({_id: 2}));
+
+const uuid = function(coll) {
+    return assert.commandWorked(db.runCommand({listCollections: 1}))
+        .cursor.firstBatch.find(c => c.name === coll.getName())
+        .info.uuid;
+};
+
+assert.commandWorked(
+    st.s.adminCommand({shardCollection: shardedColl.getFullName(), key: {_id: 1}}));
+
+// Move the chunk to shard1.
+assert.commandWorked(st.s.adminCommand(
+    {moveChunk: shardedColl.getFullName(), find: {_id: 0}, to: st.shard1.shardName}));
+
+// Run an aggregate which will only target shard1, since shard0 does not own any chunks.
+let res = assert.commandFailedWithCode(db.runCommand({
+    aggregate: shardedColl.getName(),
+    pipeline: [{$indexStats: {}}],
+    cursor: {},
+    collectionUUID: uuid(unshardedColl),
+}),
+                                       ErrorCodes.CollectionUUIDMismatch);
+jsTestLog('$indexStats result: ' + tojson(res));
+assert.eq(res.db, db.getName());
+assert.eq(res.collectionUUID, uuid(unshardedColl));
+assert.eq(res.expectedCollection, shardedColl.getName());
+assert.eq(res.actualCollection, unshardedColl.getName());
+
+st.stop();
+})();
diff --git a/src/mongo/db/pipeline/sharded_agg_helpers.cpp b/src/mongo/db/pipeline/sharded_agg_helpers.cpp
index e47e5414975..3029d1ee440 100644
--- a/src/mongo/db/pipeline/sharded_agg_helpers.cpp
+++ b/src/mongo/db/pipeline/sharded_agg_helpers.cpp
@@ -54,6 +54,7 @@
 #include "mongo/logv2/log.h"
 #include "mongo/rpc/get_status_from_command_result.h"
 #include "mongo/s/cluster_commands_helpers.h"
+#include "mongo/s/collection_uuid_mismatch.h"
 #include "mongo/s/query/cluster_query_knobs_gen.h"
 #include "mongo/s/query/document_source_merge_cursors.h"
 #include "mongo/s/query/establish_cursors.h"
@@ -1116,6 +1117,9 @@ DispatchShardPipelineResults dispatchShardPipeline(
                 uassert(5858100, *failOnShardedCollection, executionNsRoutingInfo->isSharded());
             }
             throw;
+        } catch (const ExceptionFor<ErrorCodes::CollectionUUIDMismatch>& ex) {
+            uassertStatusOK(populateCollectionUUIDMismatch(opCtx, ex.toStatus()));
+            MONGO_UNREACHABLE_TASSERT(6487201);
         }
         // If we thought the collection was sharded and the shard confirmed this, fail if the query
         // isn't valid on a sharded collection.
diff --git a/src/mongo/s/SConscript b/src/mongo/s/SConscript
index 3c51d49e125..187e802d613 100644
--- a/src/mongo/s/SConscript
+++ b/src/mongo/s/SConscript
@@ -47,6 +47,7 @@ env.Library(
     target='sharding_router_api',
     source=[
         'cluster_commands_helpers.cpp',
+        'collection_uuid_mismatch.cpp',
         'multi_statement_transaction_requests_sender.cpp',
         'router_transactions_metrics.cpp',
         'router_transactions_stats.idl',
diff --git a/src/mongo/s/collection_uuid_mismatch.cpp b/src/mongo/s/collection_uuid_mismatch.cpp
new file mode 100644
index 00000000000..e9fee9363eb
--- /dev/null
+++ b/src/mongo/s/collection_uuid_mismatch.cpp
@@ -0,0 +1,93 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/s/collection_uuid_mismatch.h"
+
+#include "mongo/db/bson/dotted_path_support.h"
+#include "mongo/db/catalog/collection_uuid_mismatch_info.h"
+#include "mongo/db/list_collections_gen.h"
+#include "mongo/s/cluster_commands_helpers.h"
+#include "mongo/s/grid.h"
+
+namespace mongo {
+Status populateCollectionUUIDMismatch(OperationContext* opCtx,
+                                      const Status& collectionUUIDMismatch) {
+    tassert(6487200,
+            str::stream() << "Expected CollectionUUIDMismatch but got " << collectionUUIDMismatch,
+            collectionUUIDMismatch == ErrorCodes::CollectionUUIDMismatch);
+
+    auto info = collectionUUIDMismatch.extraInfo<CollectionUUIDMismatchInfo>();
+    if (info->actualCollection()) {
+        return collectionUUIDMismatch;
+    }
+
+    // The listCollections command cannot be run in multi-document transactions, so run it using an
+    // alternative client.
+    auto client = opCtx->getServiceContext()->makeClient("populateCollectionUUIDMismatch");
+    auto alternativeOpCtx = client->makeOperationContext();
+    opCtx = alternativeOpCtx.get();
+    AlternativeClientRegion acr{client};
+
+    auto swDbInfo = Grid::get(opCtx)->catalogCache()->getDatabase(opCtx, info->db());
+    if (!swDbInfo.isOK()) {
+        return swDbInfo.getStatus();
+    }
+
+    ListCollections listCollections;
+    listCollections.setDbName(info->db());
+    listCollections.setFilter(BSON("info.uuid" << info->collectionUUID()));
+
+    auto response =
+        executeCommandAgainstDatabasePrimary(opCtx,
+                                             info->db(),
+                                             swDbInfo.getValue(),
+                                             listCollections.toBSON({}),
+                                             ReadPreferenceSetting{ReadPreference::PrimaryOnly},
+                                             Shard::RetryPolicy::kIdempotent);
+    if (!response.swResponse.isOK()) {
+        return response.swResponse.getStatus();
+    }
+
+    if (auto status = getStatusFromCommandResult(response.swResponse.getValue().data);
+        !status.isOK()) {
+        return status;
+    }
+
+    if (auto actualCollectionElem = dotted_path_support::extractElementAtPath(
+            response.swResponse.getValue().data, "cursor.firstBatch.0.name")) {
+        return {CollectionUUIDMismatchInfo{info->db(),
+                                           info->collectionUUID(),
+                                           info->expectedCollection(),
+                                           actualCollectionElem.str()},
+                collectionUUIDMismatch.reason()};
+    }
+
+    return collectionUUIDMismatch;
+}
+}  // namespace mongo
diff --git a/src/mongo/s/collection_uuid_mismatch.h b/src/mongo/s/collection_uuid_mismatch.h
new file mode 100644
index 00000000000..6eaee37841a
--- /dev/null
+++ b/src/mongo/s/collection_uuid_mismatch.h
@@ -0,0 +1,42 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/operation_context.h"
+
+namespace mongo {
+/**
+ * Attempts to populate the actualCollection field of a CollectionUUIDMismatch error (if not already
+ * populated) by contacting the primary shard. Returns the populated CollectionUUIDMismatch or the
+ * error that caused attemping to populate the Status to fail. Will never return an OK Status.
+ */
+Status populateCollectionUUIDMismatch(OperationContext* opCtx,
+                                      const Status& collectionUUIDMismatch);
+}  // namespace mongo
diff --git a/src/mongo/s/write_ops/batch_write_op.cpp b/src/mongo/s/write_ops/batch_write_op.cpp
index fd370663ef9..b034bcd0ee4 100644
--- a/src/mongo/s/write_ops/batch_write_op.cpp
+++ b/src/mongo/s/write_ops/batch_write_op.cpp
@@ -34,14 +34,12 @@
 #include <numeric>
 
 #include "mongo/base/error_codes.h"
-#include "mongo/db/bson/dotted_path_support.h"
 #include "mongo/db/catalog/collection_uuid_mismatch_info.h"
-#include "mongo/db/list_collections_gen.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/ops/write_ops.h"
 #include "mongo/s/client/num_hosts_targeted_metrics.h"
 #include "mongo/s/cluster_commands_helpers.h"
-#include "mongo/s/grid.h"
+#include "mongo/s/collection_uuid_mismatch.h"
 #include "mongo/s/transaction_router.h"
 #include "mongo/util/transitional_tools_do_not_use/vector_spooling.h"
 
@@ -281,55 +279,14 @@ void populateCollectionUUIDMismatch(OperationContext* opCtx,
         return;
     }
 
-    // The listCollections command cannot be run in multi-document transactions, so run it using an
-    // alternative client.
-    auto client = opCtx->getServiceContext()->makeClient("populateCollectionUUIDMismatch");
-    auto alternativeOpCtx = client->makeOperationContext();
-    opCtx = alternativeOpCtx.get();
-    AlternativeClientRegion acr{client};
-
-    auto swDbInfo = Grid::get(opCtx)->catalogCache()->getDatabase(opCtx, info->db());
-    if (!swDbInfo.isOK()) {
-        error->setStatus(swDbInfo.getStatus());
-        return;
-    }
-
-    ListCollections listCollections;
-    listCollections.setDbName(info->db());
-    listCollections.setFilter(BSON("info.uuid" << info->collectionUUID()));
-
-    auto response =
-        executeCommandAgainstDatabasePrimary(opCtx,
-                                             info->db(),
-                                             swDbInfo.getValue(),
-                                             listCollections.toBSON({}),
-                                             ReadPreferenceSetting{ReadPreference::PrimaryOnly},
-                                             Shard::RetryPolicy::kIdempotent);
-    if (!response.swResponse.isOK()) {
-        error->setStatus(response.swResponse.getStatus());
-        return;
-    }
-
-
-    if (auto status = getStatusFromCommandResult(response.swResponse.getValue().data);
-        !status.isOK()) {
-        error->setStatus(status);
-        return;
-    }
-
-    if (auto actualCollectionElem = dotted_path_support::extractElementAtPath(
-            response.swResponse.getValue().data, "cursor.firstBatch.0.name")) {
-        *actualCollection = actualCollectionElem.str();
-        error->setStatus({CollectionUUIDMismatchInfo{info->db(),
-                                                     info->collectionUUID(),
-                                                     info->expectedCollection(),
-                                                     **actualCollection},
-                          error->getStatus().reason()});
-        return;
+    error->setStatus(populateCollectionUUIDMismatch(opCtx, error->getStatus()));
+    if (error->getStatus() == ErrorCodes::CollectionUUIDMismatch) {
+        *hasContactedPrimaryShard = true;
+        if (auto& populatedActualCollection =
+                error->getStatus().extraInfo<CollectionUUIDMismatchInfo>()->actualCollection()) {
+            *actualCollection = populatedActualCollection;
+        }
     }
-
-    *hasContactedPrimaryShard = true;
-    return;
 }
 
 }  // namespace