SERVER-31447 Use correct collation for update lookup

10 files changed, 111 insertions, 46 deletions

diff --git a/src/mongo/db/db_raii.cpp b/src/mongo/db/db_raii.cpp
index 07e65df9aa1..3ca2b1e5752 100644
--- a/src/mongo/db/db_raii.cpp
+++ b/src/mongo/db/db_raii.cpp
@@ -55,6 +55,21 @@ AutoGetDb::AutoGetDb(OperationContext* opCtx, StringData ns, Lock::DBLock lock)
 
 AutoGetCollection::AutoGetCollection(OperationContext* opCtx,
                                      const NamespaceString& nss,
+                                     const UUID& uuid,
+                                     LockMode modeAll)
+    : _viewMode(ViewMode::kViewsForbidden),
+      _autoDb(opCtx, nss.db(), Lock::DBLock(opCtx, nss.db(), modeAll)),
+      _collLock(opCtx->lockState(), nss.ns(), modeAll),
+      _coll(UUIDCatalog::get(opCtx).lookupCollectionByUUID(uuid)) {
+    // Wait for a configured amount of time after acquiring locks if the failpoint is enabled.
+    MONGO_FAIL_POINT_BLOCK(setAutoGetCollectionWait, customWait) {
+        const BSONObj& data = customWait.getData();
+        sleepFor(Milliseconds(data["waitForMillis"].numberInt()));
+    }
+}
+
+AutoGetCollection::AutoGetCollection(OperationContext* opCtx,
+                                     const NamespaceString& nss,
                                      LockMode modeDB,
                                      LockMode modeColl,
                                      ViewMode viewMode)
diff --git a/src/mongo/db/db_raii.h b/src/mongo/db/db_raii.h
index e251d1fe05d..f764d671d31 100644
--- a/src/mongo/db/db_raii.h
+++ b/src/mongo/db/db_raii.h
@@ -86,6 +86,8 @@ class AutoGetCollection {
     enum class ViewMode;
 
 public:
+    AutoGetCollection(OperationContext*, const NamespaceString&, const UUID&, LockMode modeAll);
+
     AutoGetCollection(OperationContext* opCtx, const NamespaceString& nss, LockMode modeAll)
         : AutoGetCollection(opCtx, nss, modeAll, modeAll, ViewMode::kViewsForbidden) {}
 
diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index 8e4278f040b..c576fd5bda8 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -840,13 +840,15 @@ public:
         virtual std::vector<FieldPath> collectDocumentKeyFields(UUID) const = 0;
 
         /**
-         * Returns zero or one documents matching the input filter, or throws if more than one match
-         * was found. The passed ExpressionContext may use a different namespace than the
-         * ExpressionContext used to construct the MongoProcessInterface. Returns boost::none if no
-         * matching documents were found, including cases where the given namespace does not exist.
+         * Returns zero or one documents with the document key 'documentKey'. 'documentKey' is
+         * treated as a unique identifier of a document, and may include an _id or all fields from
+         * the shard key and an _id. Throws if more than one match was found. Returns boost::none if
+         * no matching documents were found, including cases where the given namespace does not
+         * exist.
          */
-        virtual boost::optional<Document> lookupSingleDocument(
-            const boost::intrusive_ptr<ExpressionContext>& expCtx, const Document& filter) = 0;
+        virtual boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                               UUID collectionUUID,
+                                                               const Document& documentKey) = 0;
 
         // Add new methods as needed.
     };
diff --git a/src/mongo/db/pipeline/document_source_lookup_change_post_image.cpp b/src/mongo/db/pipeline/document_source_lookup_change_post_image.cpp
index 058a959f588..7752fbfd751 100644
--- a/src/mongo/db/pipeline/document_source_lookup_change_post_image.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup_change_post_image.cpp
@@ -105,10 +105,9 @@ Value DocumentSourceLookupChangePostImage::lookupPostImage(const Document& updat
     auto resumeToken =
         ResumeToken::parse(updateOp[DocumentSourceChangeStream::kIdField].getDocument());
 
-    // TODO SERVER-29134 we need to extract the namespace from the document and set them on the new
-    // ExpressionContext if we're getting notifications from an entire database.
-    auto foreignExpCtx = pExpCtx->copyWith(nss, resumeToken.getData().uuid);
-    auto lookedUpDoc = _mongoProcessInterface->lookupSingleDocument(foreignExpCtx, documentKey);
+    invariant(resumeToken.getData().uuid);
+    auto lookedUpDoc =
+        _mongoProcessInterface->lookupSingleDocument(nss, *resumeToken.getData().uuid, documentKey);
 
     // Check whether the lookup returned any documents. Even if the lookup itself succeeded, it may
     // not have returned any results if the document was deleted in the time since the update op.
diff --git a/src/mongo/db/pipeline/document_source_lookup_change_post_image_test.cpp b/src/mongo/db/pipeline/document_source_lookup_change_post_image_test.cpp
index 05b92297359..3d35f875619 100644
--- a/src/mongo/db/pipeline/document_source_lookup_change_post_image_test.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup_change_post_image_test.cpp
@@ -113,9 +113,11 @@ public:
         return Status::OK();
     }
 
-    boost::optional<Document> lookupSingleDocument(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx, const Document& filter) {
-        auto swPipeline = makePipeline({BSON("$match" << filter)}, expCtx);
+    boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                   UUID collectionUUID,
+                                                   const Document& documentKey) {
+        boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest(nss));
+        auto swPipeline = makePipeline({BSON("$match" << documentKey)}, expCtx);
         if (swPipeline == ErrorCodes::NamespaceNotFound) {
             return boost::none;
         }
@@ -124,11 +126,12 @@ public:
         auto lookedUpDocument = pipeline->getNext();
         if (auto next = pipeline->getNext()) {
             uasserted(ErrorCodes::TooManyMatchingDocuments,
-                      str::stream() << "found more than one document matching " << filter.toString()
+                      str::stream() << "found more than one document matching "
+                                    << documentKey.toString()
                                     << " ["
-                                    << (*lookedUpDocument).toString()
+                                    << lookedUpDocument->toString()
                                     << ", "
-                                    << (*next).toString()
+                                    << next->toString()
                                     << "]");
         }
         return lookedUpDocument;
diff --git a/src/mongo/db/pipeline/expression_context.cpp b/src/mongo/db/pipeline/expression_context.cpp
index 64c813ed5a4..e2b2e1e758d 100644
--- a/src/mongo/db/pipeline/expression_context.cpp
+++ b/src/mongo/db/pipeline/expression_context.cpp
@@ -29,6 +29,7 @@
 #include "mongo/platform/basic.h"
 
 #include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/query/collation/collation_spec.h"
 #include "mongo/db/query/collation/collator_factory_interface.h"
 
 namespace mongo {
@@ -121,8 +122,10 @@ void ExpressionContext::setCollator(const CollatorInterface* collator) {
     _valueComparator = ValueComparator(_collator);
 }
 
-intrusive_ptr<ExpressionContext> ExpressionContext::copyWith(NamespaceString ns,
-                                                             boost::optional<UUID> uuid) const {
+intrusive_ptr<ExpressionContext> ExpressionContext::copyWith(
+    NamespaceString ns,
+    boost::optional<UUID> uuid,
+    boost::optional<std::unique_ptr<CollatorInterface>> collator) const {
     intrusive_ptr<ExpressionContext> expCtx =
         new ExpressionContext(std::move(ns), timeZoneDatabase);
 
@@ -140,11 +143,17 @@ intrusive_ptr<ExpressionContext> ExpressionContext::copyWith(NamespaceString ns,
 
     expCtx->opCtx = opCtx;
 
-    expCtx->collation = collation;
-    if (_ownedCollator) {
-        expCtx->setCollator(_ownedCollator->clone());
-    } else if (_collator) {
-        expCtx->setCollator(_collator);
+    if (collator) {
+        expCtx->collation =
+            *collator ? (*collator)->getSpec().toBSON() : CollationSpec::kSimpleSpec;
+        expCtx->setCollator(std::move(*collator));
+    } else {
+        expCtx->collation = collation;
+        if (_ownedCollator) {
+            expCtx->setCollator(_ownedCollator->clone());
+        } else if (_collator) {
+            expCtx->setCollator(_collator);
+        }
     }
 
     expCtx->_resolvedNamespaces = _resolvedNamespaces;
diff --git a/src/mongo/db/pipeline/expression_context.h b/src/mongo/db/pipeline/expression_context.h
index 59e85edfe90..a2e16163a73 100644
--- a/src/mongo/db/pipeline/expression_context.h
+++ b/src/mongo/db/pipeline/expression_context.h
@@ -140,7 +140,9 @@ public:
      * separate aggregation pipeline on 'ns' with the optional 'uuid'.
      */
     boost::intrusive_ptr<ExpressionContext> copyWith(
-        NamespaceString ns, boost::optional<UUID> uuid = boost::none) const;
+        NamespaceString ns,
+        boost::optional<UUID> uuid = boost::none,
+        boost::optional<std::unique_ptr<CollatorInterface>> collator = boost::none) const;
 
     /**
      * Returns the ResolvedNamespace corresponding to 'nss'. It is an error to call this method on a
diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp
index ee1ded37c99..65214a154bf 100644
--- a/src/mongo/db/pipeline/pipeline_d.cpp
+++ b/src/mongo/db/pipeline/pipeline_d.cpp
@@ -384,9 +384,14 @@ public:
         return result;
     }
 
-    boost::optional<Document> lookupSingleDocument(const intrusive_ptr<ExpressionContext>& expCtx,
-                                                   const Document& filter) final {
-        auto swPipeline = makePipeline({BSON("$match" << filter)}, expCtx);
+    boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                   UUID collectionUUID,
+                                                   const Document& documentKey) final {
+        // Be sure to do the lookup using the collection default collation.
+        auto foreignExpCtx =
+            _ctx->copyWith(nss, collectionUUID, _getCollectionDefaultCollator(nss, collectionUUID));
+        auto swPipeline = makePipeline({BSON("$match" << documentKey)}, foreignExpCtx);
+
         if (swPipeline == ErrorCodes::NamespaceNotFound) {
             return boost::none;
         }
@@ -395,7 +400,8 @@ public:
         auto lookedUpDocument = pipeline->getNext();
         if (auto next = pipeline->getNext()) {
             uasserted(ErrorCodes::TooManyMatchingDocuments,
-                      str::stream() << "found more than one document matching " << filter.toString()
+                      str::stream() << "found more than one document with document key "
+                                    << documentKey.toString()
                                     << " ["
                                     << lookedUpDocument->toString()
                                     << ", "
@@ -406,8 +412,34 @@ public:
     }
 
 private:
+    /**
+     * Looks up the collection default collator for the collection given by 'collectionUUID'. A
+     * collection's default collation is not allowed to change, so we cache the result to allow for
+     * quick lookups in the future. Looks up the collection by UUID, and returns 'nullptr' if the
+     * collection does not exist or if the collection's default collation is the simple collation.
+     */
+    std::unique_ptr<CollatorInterface> _getCollectionDefaultCollator(const NamespaceString& nss,
+                                                                     UUID collectionUUID) {
+        if (_collatorCache.find(collectionUUID) == _collatorCache.end()) {
+            AutoGetCollection autoColl(_ctx->opCtx, nss, collectionUUID, MODE_IS);
+            if (!autoColl.getCollection()) {
+                // This collection doesn't exist - since we looked up by UUID, it will never exist
+                // in the future, so we cache a null pointer as the default collation.
+                _collatorCache[collectionUUID] = nullptr;
+            } else {
+                auto defaultCollator = autoColl.getCollection()->getDefaultCollator();
+                // Clone the collator so that we can safely use the pointer if the collection
+                // disappears right after we release the lock.
+                _collatorCache[collectionUUID] =
+                    defaultCollator ? defaultCollator->clone() : nullptr;
+            }
+        }
+        return _collatorCache[collectionUUID] ? _collatorCache[collectionUUID]->clone() : nullptr;
+    }
+
     intrusive_ptr<ExpressionContext> _ctx;
     DBDirectClient _client;
+    std::map<UUID, std::unique_ptr<const CollatorInterface>> _collatorCache;
 };
 
 /**
diff --git a/src/mongo/db/pipeline/stub_mongo_process_interface.h b/src/mongo/db/pipeline/stub_mongo_process_interface.h
index 06e22b24bdd..bebf8f02a34 100644
--- a/src/mongo/db/pipeline/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/stub_mongo_process_interface.h
@@ -118,8 +118,9 @@ public:
         MONGO_UNREACHABLE;
     }
 
-    boost::optional<Document> lookupSingleDocument(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx, const Document& filter) {
+    boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                   UUID collectionUUID,
+                                                   const Document& documentKey) {
         MONGO_UNREACHABLE;
     }
 };
diff --git a/src/mongo/s/commands/pipeline_s.cpp b/src/mongo/s/commands/pipeline_s.cpp
index c8c58168045..8503d041d9c 100644
--- a/src/mongo/s/commands/pipeline_s.cpp
+++ b/src/mongo/s/commands/pipeline_s.cpp
@@ -31,6 +31,7 @@
 #include "mongo/s/commands/pipeline_s.h"
 
 #include "mongo/db/pipeline/document_source.h"
+#include "mongo/db/query/collation/collation_spec.h"
 #include "mongo/executor/task_executor_pool.h"
 #include "mongo/s/catalog_cache.h"
 #include "mongo/s/commands/cluster_commands_helpers.h"
@@ -53,7 +54,7 @@ std::pair<ShardId, ChunkVersion> getSingleTargetedShardForQuery(
     OperationContext* opCtx, const CachedCollectionRoutingInfo& routingInfo, BSONObj query) {
     if (auto chunkMgr = routingInfo.cm()) {
         std::set<ShardId> shardIds;
-        chunkMgr->getShardIdsForQuery(opCtx, query, BSONObj(), &shardIds);
+        chunkMgr->getShardIdsForQuery(opCtx, query, CollationSpec::kSimpleSpec, &shardIds);
         uassert(ErrorCodes::InternalError,
                 str::stream() << "Unable to target lookup query to a single shard: "
                               << query.toString(),
@@ -171,21 +172,20 @@ public:
         MONGO_UNREACHABLE;
     }
 
-    boost::optional<Document> lookupSingleDocument(const intrusive_ptr<ExpressionContext>& expCtx,
+    boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                   UUID collectionUUID,
                                                    const Document& filter) final {
-        // The passed ExpressionContext may use a different namespace than the _expCtx used to
-        // construct the MongosProcessInterface, but both should be using the same opCtx.
-        invariant(expCtx->opCtx == _expCtx->opCtx);
+        auto foreignExpCtx = _expCtx->copyWith(nss, collectionUUID);
 
         // Create the find command to be dispatched to the shard in order to return the post-change
         // document.
         auto filterObj = filter.toBson();
         BSONObjBuilder cmdBuilder;
-        bool findCmdIsByUuid(expCtx->uuid);
+        bool findCmdIsByUuid(foreignExpCtx->uuid);
         if (findCmdIsByUuid) {
-            expCtx->uuid->appendToBuilder(&cmdBuilder, "find");
+            foreignExpCtx->uuid->appendToBuilder(&cmdBuilder, "find");
         } else {
-            cmdBuilder.append("find", expCtx->ns.coll());
+            cmdBuilder.append("find", nss.coll());
         }
         cmdBuilder.append("filter", filterObj);
 
@@ -194,8 +194,8 @@ public:
         size_t numAttempts = 0;
         do {
             // Verify that the collection exists, with the UUID passed in the expCtx.
-            auto catalogCache = Grid::get(expCtx->opCtx)->catalogCache();
-            auto swRoutingInfo = getCollectionRoutingInfo(expCtx);
+            auto catalogCache = Grid::get(_expCtx->opCtx)->catalogCache();
+            auto swRoutingInfo = getCollectionRoutingInfo(foreignExpCtx);
             if (swRoutingInfo == ErrorCodes::NamespaceNotFound) {
                 return boost::none;
             }
@@ -207,21 +207,21 @@ public:
                 // the same name created through a different mongos, the shard version will be
                 // detected as stale, as shard versions contain an 'epoch' field unique to the
                 // collection.
-                findCmd = findCmd.addField(BSON("find" << expCtx->ns.coll()).firstElement());
+                findCmd = findCmd.addField(BSON("find" << nss.coll()).firstElement());
                 findCmdIsByUuid = false;
             }
 
             // Get the ID and version of the single shard to which this query will be sent.
-            auto shardInfo = getSingleTargetedShardForQuery(expCtx->opCtx, routingInfo, filterObj);
+            auto shardInfo = getSingleTargetedShardForQuery(_expCtx->opCtx, routingInfo, filterObj);
 
             // Dispatch the request. This will only be sent to a single shard and only a single
             // result will be returned. The 'establishCursors' method conveniently prepares the
             // result into a cursor response for us.
             swShardResult = establishCursors(
-                expCtx->opCtx,
-                Grid::get(expCtx->opCtx)->getExecutorPool()->getArbitraryExecutor(),
-                expCtx->ns,
-                ReadPreferenceSetting::get(expCtx->opCtx),
+                _expCtx->opCtx,
+                Grid::get(_expCtx->opCtx)->getExecutorPool()->getArbitraryExecutor(),
+                nss,
+                ReadPreferenceSetting::get(_expCtx->opCtx),
                 {{shardInfo.first, appendShardVersion(findCmd, shardInfo.second)}},
                 false,
                 nullptr);