SERVER-42416 Move returnKey execution code to separate PlanStage

21 files changed, 360 insertions, 145 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index fef38386595..2112af67b86 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -991,6 +991,7 @@ env.Library(
         'exec/requires_all_indices_stage.cpp',
         'exec/requires_collection_stage.cpp',
         'exec/requires_index_stage.cpp',
+        'exec/return_key.cpp',
         'exec/shard_filter.cpp',
         'exec/shard_filterer_impl.cpp',
         'exec/skip.cpp',
diff --git a/src/mongo/db/exec/idhack.cpp b/src/mongo/db/exec/idhack.cpp
index 4ff011d024e..02e347dda4d 100644
--- a/src/mongo/db/exec/idhack.cpp
+++ b/src/mongo/db/exec/idhack.cpp
@@ -57,9 +57,7 @@ IDHackStage::IDHackStage(OperationContext* opCtx,
       _workingSet(ws),
       _key(query->getQueryObj()["_id"].wrap()) {
     _specificStats.indexName = descriptor->indexName();
-    if (nullptr != query->getProj()) {
-        _addKeyMetadata = query->getProj()->wantIndexKey();
-    }
+    _addKeyMetadata = query->getQueryRequest().returnKey();
 }
 
 IDHackStage::IDHackStage(OperationContext* opCtx,
diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h
index 9ec4d0b69a5..cff5085e47a 100644
--- a/src/mongo/db/exec/plan_stats.h
+++ b/src/mongo/db/exec/plan_stats.h
@@ -447,6 +447,16 @@ struct IDHackStats : public SpecificStats {
     size_t docsExamined;
 };
 
+struct ReturnKeyStats : public SpecificStats {
+    SpecificStats* clone() const final {
+        return new ReturnKeyStats(*this);
+    }
+
+    uint64_t estimateObjectSizeInBytes() const {
+        return sizeof(*this);
+    }
+};
+
 struct IndexScanStats : public SpecificStats {
     IndexScanStats()
         : indexVersion(0),
diff --git a/src/mongo/db/exec/projection.cpp b/src/mongo/db/exec/projection.cpp
index cc7ba51c095..8faa479e9d9 100644
--- a/src/mongo/db/exec/projection.cpp
+++ b/src/mongo/db/exec/projection.cpp
@@ -205,19 +205,7 @@ Status ProjectionStageDefault::transform(WorkingSetMember* member) const {
         return Status(ErrorCodes::InternalError,
                       "sortKey meta-projection requested but no data available");
 
-    if (_exec.returnKey()) {
-        auto keys = _exec.computeReturnKeyProjection(
-            member->metadata().hasIndexKey() ? indexKey(*member) : BSONObj(),
-            _exec.needsSortKey() ? sortKey(*member) : BSONObj());
-        if (!keys.isOK())
-            return keys.getStatus();
-
-        transitionMemberToOwnedObj(keys.getValue(), member);
-        return Status::OK();
-    }
-
     auto projected = provideMetaFieldsAndPerformExec(_exec, *member);
-
     if (!projected.isOK())
         return projected.getStatus();
 
diff --git a/src/mongo/db/exec/projection_exec.cpp b/src/mongo/db/exec/projection_exec.cpp
index 736dd900545..bf212a0b44b 100644
--- a/src/mongo/db/exec/projection_exec.cpp
+++ b/src/mongo/db/exec/projection_exec.cpp
@@ -106,8 +106,7 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx,
                     _meta[e.fieldName()] = META_TEXT_SCORE;
                     _needsTextScore = true;
                 } else if (e2.valuestr() == QueryRequest::metaSortKey) {
-                    _sortKeyMetaFields.push_back(e.fieldName());
-                    _meta[_sortKeyMetaFields.back()] = META_SORT_KEY;
+                    _meta[e.fieldName()] = META_SORT_KEY;
                     _needsSortKey = true;
                 } else if (e2.valuestr() == QueryRequest::metaRecordId) {
                     _meta[e.fieldName()] = META_RECORDID;
@@ -117,8 +116,6 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx,
                 } else if (e2.valuestr() == QueryRequest::metaGeoNearDistance) {
                     _meta[e.fieldName()] = META_GEONEAR_DIST;
                     _needsGeoNearDistance = true;
-                } else if (e2.valuestr() == QueryRequest::metaIndexKey) {
-                    _hasReturnKey = true;
                 } else {
                     // This shouldn't happen, should be caught by parsing.
                     MONGO_UNREACHABLE;
@@ -200,22 +197,6 @@ void ProjectionExec::add(const string& field, int skip, int limit) {
 // Execution
 //
 
-StatusWith<BSONObj> ProjectionExec::computeReturnKeyProjection(const BSONObj& indexKey,
-                                                               const BSONObj& sortKey) const {
-    BSONObjBuilder bob;
-
-    if (!indexKey.isEmpty()) {
-        bob.appendElements(indexKey);
-    }
-
-    // Must be possible to do both returnKey meta-projection and sortKey meta-projection so that
-    // mongos can support returnKey.
-    for (auto fieldName : _sortKeyMetaFields)
-        bob.append(fieldName, sortKey);
-
-    return bob.obj();
-}
-
 StatusWith<BSONObj> ProjectionExec::project(const BSONObj& in,
                                             const boost::optional<const double> geoDistance,
                                             Value geoNearPoint,
diff --git a/src/mongo/db/exec/projection_exec.h b/src/mongo/db/exec/projection_exec.h
index a6ce7b1317c..c1b878aa48c 100644
--- a/src/mongo/db/exec/projection_exec.h
+++ b/src/mongo/db/exec/projection_exec.h
@@ -80,14 +80,6 @@ public:
     ~ProjectionExec();
 
     /**
-     * Indicates whether this is a returnKey projection which should be performed via
-     * 'computeReturnKeyProjection()'.
-     */
-    bool returnKey() const {
-        return _hasReturnKey;
-    }
-
-    /**
      * Indicates whether 'sortKey' must be provided for 'computeReturnKeyProjection()' or
      * 'project()'.
      */
@@ -124,12 +116,6 @@ public:
     }
 
     /**
-     * Performs a returnKey projection and provides index keys rather than projection results.
-     */
-    StatusWith<BSONObj> computeReturnKeyProjection(const BSONObj& indexKey,
-                                                   const BSONObj& sortKey) const;
-
-    /**
      * Performs a projection given a BSONObj source. Meta fields must be provided if necessary.
      * Their necessity can be queried via the 'needs*' functions.
      */
@@ -259,10 +245,6 @@ private:
     // Projections that aren't sourced from the document or index keys.
     MetaMap _meta;
 
-    // Do we have a returnKey projection?  If so we *only* output the index key metadata, and
-    // possibly the sort key for mongos to use.  If it's not found we output nothing.
-    bool _hasReturnKey = false;
-
     // After parsing in the constructor, these fields will indicate the neccesity of metadata
     // for $meta projection.
     bool _needsSortKey = false;
@@ -270,10 +252,6 @@ private:
     bool _needsGeoNearPoint = false;
     bool _needsTextScore = false;
 
-    // The field names associated with any sortKey meta-projection(s). Empty if there is no sortKey
-    // meta-projection.
-    std::vector<StringData> _sortKeyMetaFields;
-
     // The collator this projection should use to compare strings. Needed for projection operators
     // that perform matching (e.g. elemMatch projection). If null, the collation is a simple binary
     // compare.
diff --git a/src/mongo/db/exec/return_key.cpp b/src/mongo/db/exec/return_key.cpp
new file mode 100644
index 00000000000..8f0e185b554
--- /dev/null
+++ b/src/mongo/db/exec/return_key.cpp
@@ -0,0 +1,103 @@
+/**
+ *    Copyright (C) 2019-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.
+ */
+
+#define MONGO_LOG_DEFAULT_COMPONENT ::mongo::logger::LogComponent::kQuery
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/exec/return_key.h"
+
+#include "mongo/db/exec/working_set_common.h"
+#include "mongo/util/log.h"
+
+namespace mongo {
+using namespace fmt::literals;
+
+PlanStage::StageState ReturnKeyStage::doWork(WorkingSetID* out) {
+    WorkingSetID id = WorkingSet::INVALID_ID;
+    StageState status = child()->work(&id);
+
+    if (PlanStage::ADVANCED == status) {
+        WorkingSetMember* member = _ws.get(id);
+        Status indexKeyStatus = _extractIndexKey(member);
+
+        if (!indexKeyStatus.isOK()) {
+            warning() << "Couldn't execute {}, status = {}"_format(kStageName,
+                                                                   redact(indexKeyStatus));
+            *out = WorkingSetCommon::allocateStatusMember(&_ws, indexKeyStatus);
+            return PlanStage::FAILURE;
+        }
+
+        *out = id;
+    } else if (PlanStage::FAILURE == status) {
+        // The stage which produces a failure is responsible for allocating a working set member
+        // with error details.
+        invariant(WorkingSet::INVALID_ID != id);
+        *out = id;
+    } else if (PlanStage::NEED_YIELD == status) {
+        *out = id;
+    }
+
+    return status;
+}
+
+std::unique_ptr<PlanStageStats> ReturnKeyStage::getStats() {
+    _commonStats.isEOF = isEOF();
+
+    auto ret = std::make_unique<PlanStageStats>(_commonStats, stageType());
+    ret->specific = std::make_unique<ReturnKeyStats>(_specificStats);
+    ret->children.emplace_back(child()->getStats());
+    return ret;
+}
+
+Status ReturnKeyStage::_extractIndexKey(WorkingSetMember* member) {
+    if (!_sortKeyMetaFields.empty()) {
+        invariant(member->metadata().hasSortKey());
+    }
+
+    auto indexKey = member->metadata().hasIndexKey() ? member->metadata().getIndexKey() : BSONObj();
+    auto sortKey = member->metadata().hasSortKey() ? member->metadata().getSortKey() : BSONObj();
+    BSONObjBuilder bob;
+
+    if (!indexKey.isEmpty()) {
+        bob.appendElements(indexKey);
+    }
+
+    for (const auto& fieldName : _sortKeyMetaFields) {
+        bob.append(fieldName, sortKey);
+    }
+
+    member->keyData.clear();
+    member->recordId = {};
+    member->resetDocument({}, bob.obj());
+    member->transitionToOwnedObj();
+
+    return Status::OK();
+}
+}  // namespace mongo
diff --git a/src/mongo/db/exec/return_key.h b/src/mongo/db/exec/return_key.h
new file mode 100644
index 00000000000..7212de12843
--- /dev/null
+++ b/src/mongo/db/exec/return_key.h
@@ -0,0 +1,83 @@
+/**
+ *    Copyright (C) 2019-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/exec/plan_stage.h"
+
+namespace mongo {
+/**
+ * This stage returns the index key or keys, used for executing the query, for each document in
+ * the result. If the query does not use an index to perform the read operation, the stage returns
+ * empty documents.
+ *
+ * If 'sortKeyMetaFields' vector is specified while constructing this stage, each element in this
+ * array will be treated as a field name which will be added to the output document and will
+ * hold a sort key for the document in the result set, if the sort key exists. The elements in this
+ * array are the values specified in the 'sortKey' meta-projection.
+ */
+class ReturnKeyStage : public PlanStage {
+public:
+    static constexpr StringData kStageName = "RETURN_KEY"_sd;
+
+    ReturnKeyStage(OperationContext* opCtx,
+                   std::vector<std::string> sortKeyMetaFields,
+                   WorkingSet* ws,
+                   std::unique_ptr<PlanStage> child)
+        : PlanStage(opCtx, std::move(child), kStageName.rawData()),
+          _ws(*ws),
+          _sortKeyMetaFields(std::move(sortKeyMetaFields)) {}
+
+    StageType stageType() const final {
+        return STAGE_RETURN_KEY;
+    }
+
+    bool isEOF() final {
+        return child()->isEOF();
+    }
+
+    StageState doWork(WorkingSetID* out) final;
+
+    std::unique_ptr<PlanStageStats> getStats() final;
+
+    const SpecificStats* getSpecificStats() const final {
+        return &_specificStats;
+    }
+
+private:
+    Status _extractIndexKey(WorkingSetMember* member);
+
+    WorkingSet& _ws;
+    ReturnKeyStats _specificStats;
+
+    // The field names associated with any sortKey meta-projection(s). Empty if there is no sortKey
+    // meta-projection.
+    std::vector<std::string> _sortKeyMetaFields;
+};
+}  // namespace mongo
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index 4b039fa7aaf..4d809846186 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -49,6 +49,7 @@
 #include "mongo/db/exec/multi_plan.h"
 #include "mongo/db/exec/projection.h"
 #include "mongo/db/exec/record_store_fast_count.h"
+#include "mongo/db/exec/return_key.h"
 #include "mongo/db/exec/shard_filter.h"
 #include "mongo/db/exec/sort_key_generator.h"
 #include "mongo/db/exec/subplan.h"
@@ -388,23 +389,32 @@ StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx,
                 std::move(root));
         }
 
-        // There might be a projection. The idhack stage will always fetch the full
-        // document, so we don't support covered projections. However, we might use the
-        // simple inclusion fast path.
-        if (nullptr != canonicalQuery->getProj()) {
-
-            // Add a SortKeyGeneratorStage if there is a $meta sortKey projection.
-            if (canonicalQuery->getProj()->wantSortKey()) {
-                root = std::make_unique<SortKeyGeneratorStage>(
-                    canonicalQuery->getExpCtx(),
-                    std::move(root),
-                    ws,
-                    canonicalQuery->getQueryRequest().getSort());
-            }
+        // Add a SortKeyGeneratorStage if there is a $meta sortKey projection.
+        if (canonicalQuery->getProj() && canonicalQuery->getProj()->wantSortKey()) {
+            root = std::make_unique<SortKeyGeneratorStage>(
+                canonicalQuery->getExpCtx(),
+                std::move(root),
+                ws,
+                canonicalQuery->getQueryRequest().getSort());
+        }
 
+        if (canonicalQuery->getQueryRequest().returnKey()) {
+            // If returnKey was requested, add ReturnKeyStage to return only the index keys in the
+            // resulting documents. If a projection was also specified, it will be ignored, with
+            // the exception the $meta sortKey projection, which can be used along with the
+            // returnKey.
+            root = std::make_unique<ReturnKeyStage>(
+                opCtx,
+                QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(
+                    canonicalQuery->getQueryRequest().getProj()),
+                ws,
+                std::move(root));
+        } else if (canonicalQuery->getProj()) {
+            // There might be a projection. The idhack stage will always fetch the full
+            // document, so we don't support covered projections. However, we might use the
+            // simple inclusion fast path.
             // Stuff the right data into the params depending on what proj impl we use.
             if (canonicalQuery->getProj()->requiresDocument() ||
-                canonicalQuery->getProj()->wantIndexKey() ||
                 canonicalQuery->getProj()->wantSortKey() ||
                 canonicalQuery->getProj()->hasDottedFieldPath()) {
                 root = std::make_unique<ProjectionStageDefault>(
diff --git a/src/mongo/db/query/parsed_projection.cpp b/src/mongo/db/query/parsed_projection.cpp
index 359ad5c23d8..165a27e66bb 100644
--- a/src/mongo/db/query/parsed_projection.cpp
+++ b/src/mongo/db/query/parsed_projection.cpp
@@ -56,8 +56,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
     IncludeExclude includeExclude = IncludeExclude::kUninitialized;
 
     bool requiresDocument = false;
-    bool hasIndexKeyProjection = false;
-
     bool wantTextScore = false;
     bool wantGeoNearPoint = false;
     bool wantGeoNearDistance = false;
@@ -161,7 +159,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
 
                 if (e2.valuestr() != QueryRequest::metaTextScore &&
                     e2.valuestr() != QueryRequest::metaRecordId &&
-                    e2.valuestr() != QueryRequest::metaIndexKey &&
                     e2.valuestr() != QueryRequest::metaGeoNearDistance &&
                     e2.valuestr() != QueryRequest::metaGeoNearPoint &&
                     e2.valuestr() != QueryRequest::metaSortKey) {
@@ -171,8 +168,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
                 // This clobbers everything else.
                 if (e2.valuestr() == QueryRequest::metaTextScore) {
                     wantTextScore = true;
-                } else if (e2.valuestr() == QueryRequest::metaIndexKey) {
-                    hasIndexKeyProjection = true;
                 } else if (e2.valuestr() == QueryRequest::metaGeoNearDistance) {
                     wantGeoNearDistance = true;
                 } else if (e2.valuestr() == QueryRequest::metaGeoNearPoint) {
@@ -268,7 +263,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
     // Save the raw spec.  It should be owned by the QueryRequest.
     verify(spec.isOwned());
     pp->_source = spec;
-    pp->_returnKey = hasIndexKeyProjection;
     pp->_requiresDocument = requiresDocument;
 
     // Add meta-projections.
@@ -308,11 +302,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
         }
     }
 
-    // returnKey clobbers everything except for sortKey meta-projection.
-    if (hasIndexKeyProjection && !wantSortKey) {
-        pp->_requiresDocument = false;
-    }
-
     *out = pp.release();
     return Status::OK();
 }
diff --git a/src/mongo/db/query/parsed_projection.h b/src/mongo/db/query/parsed_projection.h
index 0a5ee6de3c2..1d9063acf05 100644
--- a/src/mongo/db/query/parsed_projection.h
+++ b/src/mongo/db/query/parsed_projection.h
@@ -100,10 +100,6 @@ public:
         return _wantGeoNearPoint;
     }
 
-    bool wantIndexKey() const {
-        return _returnKey;
-    }
-
     bool wantSortKey() const {
         return _wantSortKey;
     }
@@ -193,8 +189,6 @@ private:
 
     bool _wantGeoNearPoint = false;
 
-    bool _returnKey = false;
-
     // Whether this projection includes a sortKey meta-projection.
     bool _wantSortKey = false;
 
diff --git a/src/mongo/db/query/parsed_projection_test.cpp b/src/mongo/db/query/parsed_projection_test.cpp
index 990b665d6ed..8994723a62d 100644
--- a/src/mongo/db/query/parsed_projection_test.cpp
+++ b/src/mongo/db/query/parsed_projection_test.cpp
@@ -234,7 +234,6 @@ TEST(ParsedProjectionTest, ParsedProjectionDefaults) {
     ASSERT_FALSE(parsedProjection->requiresMatchDetails());
     ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
     ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
-    ASSERT_FALSE(parsedProjection->wantIndexKey());
 }
 
 TEST(ParsedProjectionTest, SortKeyMetaProjection) {
@@ -247,7 +246,6 @@ TEST(ParsedProjectionTest, SortKeyMetaProjection) {
     ASSERT_FALSE(parsedProjection->requiresMatchDetails());
     ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
     ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
-    ASSERT_FALSE(parsedProjection->wantIndexKey());
 }
 
 TEST(ParsedProjectionTest, SortKeyMetaProjectionCovered) {
@@ -261,7 +259,6 @@ TEST(ParsedProjectionTest, SortKeyMetaProjectionCovered) {
     ASSERT_FALSE(parsedProjection->requiresMatchDetails());
     ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
     ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
-    ASSERT_FALSE(parsedProjection->wantIndexKey());
 }
 
 TEST(ParsedProjectionTest, SortKeyMetaAndSlice) {
@@ -276,7 +273,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndSlice) {
     ASSERT_FALSE(parsedProjection->requiresMatchDetails());
     ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
     ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
-    ASSERT_FALSE(parsedProjection->wantIndexKey());
 }
 
 TEST(ParsedProjectionTest, SortKeyMetaAndElemMatch) {
@@ -291,7 +287,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndElemMatch) {
     ASSERT_FALSE(parsedProjection->requiresMatchDetails());
     ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
     ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
-    ASSERT_FALSE(parsedProjection->wantIndexKey());
 }
 
 TEST(ParsedProjectionTest, SortKeyMetaAndExclusion) {
@@ -305,7 +300,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndExclusion) {
     ASSERT_FALSE(parsedProjection->requiresMatchDetails());
     ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
     ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
-    ASSERT_FALSE(parsedProjection->wantIndexKey());
 }
 
 //
diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 3487e955675..0d09172ad37 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -306,8 +306,8 @@ auto isCoveredOrAlreadyFetched(const vector<StringData>& fields,
  * Checks all properties that exclude a projection from being simple.
  */
 auto isSimpleProjection(const CanonicalQuery& query) {
-    return !query.getProj()->wantIndexKey() && !query.getProj()->wantSortKey() &&
-        !query.getProj()->hasDottedFieldPath() && !query.getProj()->requiresDocument();
+    return !query.getProj()->wantSortKey() && !query.getProj()->hasDottedFieldPath() &&
+        !query.getProj()->requiresDocument();
 }
 
 /**
@@ -332,6 +332,21 @@ auto produceCoveredKeyObj(QuerySolutionNode* solnRoot) {
 }
 
 /**
+ * Adds a stage to generate the sort key metadata if there's no sort stage but we have a sortKey
+ * meta-projection.
+ */
+std::unique_ptr<QuerySolutionNode> addSortKeyGeneratorStageIfNeeded(
+    const CanonicalQuery& query, bool hasSortStage, std::unique_ptr<QuerySolutionNode> solnRoot) {
+    if (!hasSortStage && query.getProj() && query.getProj()->wantSortKey()) {
+        auto keyGenNode = std::make_unique<SortKeyGeneratorNode>();
+        keyGenNode->sortSpec = query.getQueryRequest().getSort();
+        keyGenNode->children.push_back(solnRoot.release());
+        return keyGenNode;
+    }
+    return solnRoot;
+}
+
+/**
  * When projection needs to be added to the solution tree, this function chooses between the default
  * implementation and one of the fast paths.
  */
@@ -340,24 +355,12 @@ std::unique_ptr<ProjectionNode> analyzeProjection(const CanonicalQuery& query,
                                                   const bool hasSortStage) {
     const QueryRequest& qr = query.getQueryRequest();
 
-    // If there's no sort stage but we have a sortKey meta-projection, we need to add a stage to
-    // generate the sort key metadata.
-    auto addSortKeyGeneratorStageIfNeeded = [&]() {
-        if (!hasSortStage && query.getProj()->wantSortKey()) {
-            auto keyGenNode = std::make_unique<SortKeyGeneratorNode>();
-            keyGenNode->sortSpec = qr.getSort();
-            keyGenNode->children.push_back(solnRoot.release());
-            solnRoot = std::move(keyGenNode);
-        }
-    };
-
     LOG(5) << "PROJECTION: Current plan is:\n" << redact(solnRoot->toString());
 
     // If the projection requires the entire document we add a fetch stage if not present. Otherwise
-    // we add a fetch stage if we are not covered and not returnKey.
+    // we add a fetch stage if we are not covered.
     if ((query.getProj()->requiresDocument() && !solnRoot->fetched()) ||
-        (!isCoveredOrAlreadyFetched(query.getProj()->getRequiredFields(), *solnRoot) &&
-         !query.getProj()->wantIndexKey())) {
+        (!isCoveredOrAlreadyFetched(query.getProj()->getRequiredFields(), *solnRoot))) {
         auto fetch = std::make_unique<FetchNode>();
         fetch->children.push_back(solnRoot.release());
         solnRoot = std::move(fetch);
@@ -372,30 +375,33 @@ std::unique_ptr<ProjectionNode> analyzeProjection(const CanonicalQuery& query,
     if (isSimpleProjection(query)) {
         // If the projection is simple, but not covered, use 'ProjectionNodeSimple'.
         if (solnRoot->fetched()) {
-            addSortKeyGeneratorStageIfNeeded();
             return std::make_unique<ProjectionNodeSimple>(
-                std::move(solnRoot), *query.root(), qr.getProj(), *query.getProj());
+                addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
+                *query.root(),
+                qr.getProj(),
+                *query.getProj());
         } else {
             // If we're here we're not fetched so we're covered. Let's see if we can get out of
             // using the default projType. If 'solnRoot' is an index scan we can use the faster
             // covered impl.
             BSONObj coveredKeyObj = produceCoveredKeyObj(solnRoot.get());
             if (!coveredKeyObj.isEmpty()) {
-                addSortKeyGeneratorStageIfNeeded();
-                return std::make_unique<ProjectionNodeCovered>(std::move(solnRoot),
-                                                               *query.root(),
-                                                               qr.getProj(),
-                                                               *query.getProj(),
-                                                               std::move(coveredKeyObj));
+                return std::make_unique<ProjectionNodeCovered>(
+                    addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
+                    *query.root(),
+                    qr.getProj(),
+                    *query.getProj(),
+                    std::move(coveredKeyObj));
             }
         }
     }
 
-    addSortKeyGeneratorStageIfNeeded();
     return std::make_unique<ProjectionNodeDefault>(
-        std::move(solnRoot), *query.root(), qr.getProj(), *query.getProj());
+        addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
+        *query.root(),
+        qr.getProj(),
+        *query.getProj());
 }
-
 }  // namespace
 
 // static
@@ -797,7 +803,13 @@ std::unique_ptr<QuerySolution> QueryPlannerAnalysis::analyzeDataAccess(
     }
 
     // Project the results.
-    if (query.getProj()) {
+    if (qr.returnKey()) {
+        // We don't need a projection stage if returnKey was requested since the intended behavior
+        // is that the projection is ignored when returnKey is specified.
+        solnRoot = std::make_unique<ReturnKeyNode>(
+            addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
+            QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(qr.getProj()));
+    } else if (query.getProj()) {
         solnRoot = analyzeProjection(query, std::move(solnRoot), hasSortStage);
         // If we don't have a covered project, and we're not allowed to put an uncovered one in,
         // bail out.
diff --git a/src/mongo/db/query/query_planner_common.cpp b/src/mongo/db/query/query_planner_common.cpp
index 35cb1ac0293..2b2117676a8 100644
--- a/src/mongo/db/query/query_planner_common.cpp
+++ b/src/mongo/db/query/query_planner_common.cpp
@@ -75,4 +75,32 @@ void QueryPlannerCommon::reverseScans(QuerySolutionNode* node) {
     }
 }
 
+// TODO SERVER-42422: reimplement by walking a projection AST rather than raw bson.
+std::vector<std::string> QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(
+    const BSONObj& proj) {
+    std::vector<std::string> sortKeyMetaFields;
+
+    for (auto&& elem : proj) {
+        if (elem.type() == BSONType::Object) {
+            BSONObj obj = elem.embeddedObject();
+            // The caller must have already validated the projection, so we expect
+            // to see only one element.
+            invariant(1 == obj.nFields());
+
+            BSONElement firstSubElem = obj.firstElement();
+            if (firstSubElem.fieldNameStringData() == "$meta") {
+                invariant(firstSubElem.type() == BSONType::String);
+                if (firstSubElem.valueStringData() == QueryRequest::metaSortKey) {
+                    invariant(std::find(sortKeyMetaFields.begin(),
+                                        sortKeyMetaFields.end(),
+                                        elem.fieldName()) == sortKeyMetaFields.end());
+                    sortKeyMetaFields.push_back(elem.fieldName());
+                }
+            }
+        }
+    }
+
+    return sortKeyMetaFields;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_planner_common.h b/src/mongo/db/query/query_planner_common.h
index c232291e0e3..76c70742a3c 100644
--- a/src/mongo/db/query/query_planner_common.h
+++ b/src/mongo/db/query/query_planner_common.h
@@ -82,6 +82,14 @@ public:
      * the scan direction and index bounds.
      */
     static void reverseScans(QuerySolutionNode* node);
+
+    /**
+     * Extracts all field names for the sortKey meta-projection and stores them in the returned
+     * array. Returns an empty array if there were no sortKey meta-projection specified in the
+     * given projection 'proj'. For example, given a projection {a:1, b: {$meta: "sortKey"},
+     * c: {$meta: "sortKey"}}, the returned vector will contain two elements ["b", "c"].
+     */
+    static std::vector<std::string> extractSortKeyMetaFieldsFromProjection(const BSONObj& proj);
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp
index b3c87b40ab8..fde6441559b 100644
--- a/src/mongo/db/query/query_request.cpp
+++ b/src/mongo/db/query/query_request.cpp
@@ -58,7 +58,6 @@ const char QueryRequest::queryOptionMaxTimeMS[] = "$maxTimeMS";
 
 const string QueryRequest::metaGeoNearDistance("geoNearDistance");
 const string QueryRequest::metaGeoNearPoint("geoNearPoint");
-const string QueryRequest::metaIndexKey("indexKey");
 const string QueryRequest::metaRecordId("recordId");
 const string QueryRequest::metaSortKey("sortKey");
 const string QueryRequest::metaTextScore("textScore");
@@ -592,16 +591,6 @@ void QueryRequest::asFindCommandInternal(BSONObjBuilder* cmdBuilder) const {
     }
 }
 
-void QueryRequest::addReturnKeyMetaProj() {
-    BSONObjBuilder projBob;
-    projBob.appendElements(_proj);
-    // We use $$ because it's never going to show up in a user's projection.
-    // The exact text doesn't matter.
-    BSONObj indexKey = BSON("$$" << BSON("$meta" << QueryRequest::metaIndexKey));
-    projBob.append(indexKey.firstElement());
-    _proj = projBob.obj();
-}
-
 void QueryRequest::addShowRecordIdMetaProj() {
     BSONObjBuilder projBob;
     projBob.appendElements(_proj);
@@ -937,7 +926,6 @@ Status QueryRequest::initFullQuery(const BSONObj& top) {
                 // Won't throw.
                 if (e.trueValue()) {
                     _returnKey = true;
-                    addReturnKeyMetaProj();
                 }
             } else if (name == "showDiskLoc") {
                 // Won't throw.
@@ -1004,11 +992,6 @@ void QueryRequest::initFromInt(int options) {
 }
 
 void QueryRequest::addMetaProjection() {
-    // We might need to update the projection object with a $meta projection.
-    if (returnKey()) {
-        addReturnKeyMetaProj();
-    }
-
     if (showRecordId()) {
         addShowRecordIdMetaProj();
     }
diff --git a/src/mongo/db/query/query_request.h b/src/mongo/db/query/query_request.h
index 7b0cb74dd80..17e915ace5c 100644
--- a/src/mongo/db/query/query_request.h
+++ b/src/mongo/db/query/query_request.h
@@ -135,7 +135,6 @@ public:
     // Names of the $meta projection values.
     static const std::string metaGeoNearDistance;
     static const std::string metaGeoNearPoint;
-    static const std::string metaIndexKey;
     static const std::string metaRecordId;
     static const std::string metaSortKey;
     static const std::string metaTextScore;
@@ -455,11 +454,6 @@ private:
     Status initFullQuery(const BSONObj& top);
 
     /**
-     * Updates the projection object with a $meta projection for the returnKey option.
-     */
-    void addReturnKeyMetaProj();
-
-    /**
      * Updates the projection object with a $meta projection for the showRecordId option.
      */
     void addShowRecordIdMetaProj();
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index ddbe87074a8..6b076d80e96 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -31,6 +31,8 @@
 
 #include "mongo/db/query/query_solution.h"
 
+#include <boost/algorithm/string/join.hpp>
+
 #include "mongo/bson/bsontypes.h"
 #include "mongo/bson/mutable/document.h"
 #include "mongo/bson/simple_bsonelement_comparator.h"
@@ -848,6 +850,27 @@ bool IndexScanNode::operator==(const IndexScanNode& other) const {
 }
 
 //
+// ReturnKeyNode
+//
+
+void ReturnKeyNode::appendToString(str::stream* ss, int indent) const {
+    addIndent(ss, indent);
+    *ss << "RETURN_KEY\n";
+    addIndent(ss, indent + 1);
+    *ss << "sortKeyMetaFields = [" << boost::algorithm::join(sortKeyMetaFields, ", ") << "]\n";
+    addCommon(ss, indent);
+    addIndent(ss, indent + 1);
+    *ss << "Child:" << '\n';
+    children[0]->appendToString(ss, indent + 2);
+}
+
+QuerySolutionNode* ReturnKeyNode::clone() const {
+    auto copy = std::make_unique<ReturnKeyNode>(
+        std::unique_ptr<QuerySolutionNode>(children[0]->clone()), std::vector(sortKeyMetaFields));
+    return copy.release();
+}
+
+//
 // ProjectionNode
 //
 
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index cba57fd1263..48a0288a317 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -540,6 +540,35 @@ struct IndexScanNode : public QuerySolutionNode {
     std::set<StringData> multikeyFields;
 };
 
+struct ReturnKeyNode : public QuerySolutionNode {
+    ReturnKeyNode(std::unique_ptr<QuerySolutionNode> child,
+                  std::vector<std::string> sortKeyMetaFields)
+        : QuerySolutionNode(std::move(child)), sortKeyMetaFields(std::move(sortKeyMetaFields)) {}
+
+    StageType getType() const final {
+        return STAGE_RETURN_KEY;
+    }
+
+    void appendToString(str::stream* ss, int indent) const final;
+
+    bool fetched() const final {
+        return children[0]->fetched();
+    }
+    bool hasField(const std::string& field) const final {
+        return false;
+    }
+    bool sortedByDiskLoc() const final {
+        return children[0]->sortedByDiskLoc();
+    }
+    const BSONObjSet& getSort() const final {
+        return children[0]->getSort();
+    }
+
+    QuerySolutionNode* clone() const final;
+
+    std::vector<std::string> sortKeyMetaFields;
+};
+
 /**
  * We have a few implementations of the projection functionality. They are chosen by constructing
  * a type derived from this abstract struct. The most general implementation 'ProjectionNodeDefault'
diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index 48cc9e39027..c8a92a75963 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -52,6 +52,7 @@
 #include "mongo/db/exec/merge_sort.h"
 #include "mongo/db/exec/or.h"
 #include "mongo/db/exec/projection.h"
+#include "mongo/db/exec/return_key.h"
 #include "mongo/db/exec/shard_filter.h"
 #include "mongo/db/exec/skip.h"
 #include "mongo/db/exec/sort.h"
@@ -133,6 +134,13 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx,
             return std::make_unique<SortKeyGeneratorStage>(
                 cq.getExpCtx(), std::move(childStage), ws, keyGenNode->sortSpec);
         }
+        case STAGE_RETURN_KEY: {
+            auto returnKeyNode = static_cast<const ReturnKeyNode*>(root);
+            auto childStage =
+                buildStages(opCtx, collection, cq, qsol, returnKeyNode->children[0], ws);
+            return std::make_unique<ReturnKeyStage>(
+                opCtx, std::move(returnKeyNode->sortKeyMetaFields), ws, std::move(childStage));
+        }
         case STAGE_PROJECTION_DEFAULT: {
             auto pn = static_cast<const ProjectionNodeDefault*>(root);
             auto childStage = buildStages(opCtx, collection, cq, qsol, pn->children[0], ws);
diff --git a/src/mongo/db/query/stage_types.h b/src/mongo/db/query/stage_types.h
index 8b34d4dfecc..17a67f1169f 100644
--- a/src/mongo/db/query/stage_types.h
+++ b/src/mongo/db/query/stage_types.h
@@ -87,6 +87,7 @@ enum StageType {
 
     STAGE_QUEUED_DATA,
     STAGE_RECORD_STORE_FAST_COUNT,
+    STAGE_RETURN_KEY,
     STAGE_SHARDING_FILTER,
     STAGE_SKIP,
     STAGE_SORT,