Delete ShardNameStage and move functionality to ShardFilterStage.

Author: krk

src/mongo/db/SConscript

@@ -1027,7 +1027,6 @@ env.Library(
         'exec/return_key.cpp',
         'exec/shard_filter.cpp',
         'exec/shard_filterer_impl.cpp',
-        'exec/shard_name.cpp',
         'exec/skip.cpp',
         'exec/sort.cpp',
         'exec/sort_key_generator.cpp',

src/mongo/db/exec/shard_filter.cpp

@@ -38,6 +38,7 @@
 #include "mongo/db/exec/filter.h"
 #include "mongo/db/exec/scoped_timer.h"
 #include "mongo/db/exec/working_set_common.h"
+#include "mongo/db/s/sharding_state.h"
 #include "mongo/s/shard_key_pattern.h"
 #include "mongo/util/log.h"
 
@@ -53,8 +54,10 @@ const char* ShardFilterStage::kStageType = "SHARDING_FILTER";
 ShardFilterStage::ShardFilterStage(OperationContext* opCtx,
                                    ScopedCollectionMetadata metadata,
                                    WorkingSet* ws,
-                                   std::unique_ptr<PlanStage> child)
-    : PlanStage(kStageType, opCtx), _ws(ws), _shardFilterer(std::move(metadata)) {
+                                   std::unique_ptr<PlanStage> child,
+                                   bool wantShardName)
+    : PlanStage(kStageType, opCtx), _ws(ws), _shardFilterer(std::move(metadata)),
+      _wantShardName(wantShardName) {
     _children.emplace_back(std::move(child));
 }
 
@@ -104,6 +107,13 @@ PlanStage::StageState ShardFilterStage::doWork(WorkingSetID* out) {
                 ++_specificStats.chunkSkips;
                 return PlanStage::NEED_TIME;
             }
+
+            if(wantShardName()) {
+                auto sharding = ShardingState::get(this->getOpCtx());
+
+                // Populate the working set member with the shard name and return it.
+                member->metadata().setShardName(sharding->shardId());
+            }
         }
 
         // If we're here either we have shard state and our doc passed, or we have no shard

src/mongo/db/exec/shard_filter.h

@@ -74,7 +74,8 @@ class ShardFilterStage final : public PlanStage {
     ShardFilterStage(OperationContext* opCtx,
                      ScopedCollectionMetadata metadata,
                      WorkingSet* ws,
-                     std::unique_ptr<PlanStage> child);
+                     std::unique_ptr<PlanStage> child,
+                     bool wantShardName);
     ~ShardFilterStage();
 
     bool isEOF() final;
@@ -91,6 +92,10 @@ class ShardFilterStage final : public PlanStage {
     static const char* kStageType;
 
 private:
+    bool wantShardName() const {
+        return _wantShardName;
+    }
+
     WorkingSet* _ws;
 
     // Stats
@@ -103,6 +108,8 @@ class ShardFilterStage final : public PlanStage {
     // ScopedCollectionMetadata for the entire query, it'd be possible for data which the query
     // needs to read to be deleted while it's still running.
     ShardFiltererImpl _shardFilterer;
+
+    bool _wantShardName;
 };
 
 }  // namespace mongo

src/mongo/db/exec/shard_name.cpp

@@ -167,12 +167,12 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> createRandomCursorEx
             sampleSize / (numRecords * kMaxSampleRatioForRandCursor), kMaxSampleRatioForRandCursor);
         // The trial plan is SHARDING_FILTER-MULTI_ITERATOR.
         auto randomCursorPlan =
-            std::make_unique<ShardFilterStage>(opCtx, shardMetadata, ws.get(), std::move(root));
+            std::make_unique<ShardFilterStage>(opCtx, shardMetadata, ws.get(), std::move(root), false);
         // The backup plan is SHARDING_FILTER-COLLSCAN.
         std::unique_ptr<PlanStage> collScanPlan = std::make_unique<CollectionScan>(
             opCtx, coll, CollectionScanParams{}, ws.get(), nullptr);
         collScanPlan = std::make_unique<ShardFilterStage>(
-            opCtx, shardMetadata, ws.get(), std::move(collScanPlan));
+            opCtx, shardMetadata, ws.get(), std::move(collScanPlan), false);
         // Place a TRIAL stage at the root of the plan tree, and pass it the trial and backup plans.
         root = std::make_unique<TrialStage>(opCtx,
                                             ws.get(),

src/mongo/db/exec/shard_name.h

@@ -382,12 +382,14 @@ StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx,
 
         // Might have to filter out orphaned docs.
         if (plannerParams.options & QueryPlannerParams::INCLUDE_SHARD_FILTER) {
+            auto wantShardName = canonicalQuery->metadataDeps()[DocumentMetadataFields::kShardName];
             root = std::make_unique<ShardFilterStage>(
                 opCtx,
                 CollectionShardingState::get(opCtx, canonicalQuery->nss())
                     ->getOrphansFilter(opCtx, collection),
                 ws,
-                std::move(root));
+                std::move(root),
+                wantShardName);
         }
 
         // Add a SortKeyGeneratorStage if the query requested sortKey metadata.

src/mongo/db/pipeline/pipeline_d.cpp

@@ -384,12 +384,6 @@ std::unique_ptr<ProjectionNode> analyzeProjection(const CanonicalQuery& query,
             }
         }
     }
-    auto wantShardName = query.metadataDeps()[DocumentMetadataFields::kShardName];
-    if (wantShardName) {
-        auto shardNameNode = std::make_unique<ShardNameNode>();
-        shardNameNode->children.push_back(solnRoot.release());
-        solnRoot = std::move(shardNameNode);
-    }
     return std::make_unique<ProjectionNodeDefault>(
         addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
         *query.root(),
@@ -767,7 +761,8 @@ std::unique_ptr<QuerySolution> QueryPlannerAnalysis::analyzeDataAccess(
             }
         }
 
-        ShardingFilterNode* sfn = new ShardingFilterNode();
+        auto wantShardName = query.metadataDeps()[DocumentMetadataFields::kShardName];
+        ShardingFilterNode* sfn = new ShardingFilterNode(wantShardName);
         sfn->children.push_back(solnRoot.release());
         solnRoot.reset(sfn);
     }

src/mongo/db/query/get_executor.cpp

@@ -1118,32 +1118,9 @@ QuerySolutionNode* GeoNear2DSphereNode::clone() const {
 void ShardingFilterNode::appendToString(str::stream* ss, int indent) const {
     addIndent(ss, indent);
     *ss << "SHARDING_FILTER\n";
-    if (nullptr != filter) {
-        addIndent(ss, indent + 1);
-        StringBuilder sb;
-        *ss << "filter:\n";
-        filter->debugString(sb, indent + 2);
-        *ss << sb.str();
+    if (wantShardName()) {
+        *ss << "with shard name\n";
     }
-    addCommon(ss, indent);
-    addIndent(ss, indent + 1);
-    *ss << "Child:" << '\n';
-    children[0]->appendToString(ss, indent + 2);
-}
-
-QuerySolutionNode* ShardingFilterNode::clone() const {
-    ShardingFilterNode* copy = new ShardingFilterNode();
-    cloneBaseData(copy);
-    return copy;
-}
-
-//
-// ShardNameNode
-//
-
-void ShardNameNode::appendToString(str::stream* ss, int indent) const {
-    addIndent(ss, indent);
-    *ss << "SHARD_NAME\n";
     if (nullptr != filter) {
         addIndent(ss, indent + 1);
         StringBuilder sb;
@@ -1157,8 +1134,8 @@ void ShardNameNode::appendToString(str::stream* ss, int indent) const {
     children[0]->appendToString(ss, indent + 2);
 }
 
-QuerySolutionNode* ShardNameNode::clone() const {
-    ShardNameNode* copy = new ShardNameNode();
+QuerySolutionNode* ShardingFilterNode::clone() const {
+    ShardingFilterNode* copy = new ShardingFilterNode(this->wantShardName());
     cloneBaseData(copy);
     return copy;
 }

src/mongo/db/query/planner_analysis.cpp

@@ -915,7 +915,7 @@ struct GeoNear2DSphereNode : public QuerySolutionNode {
  * through the pipeline.
  */
 struct ShardingFilterNode : public QuerySolutionNode {
-    ShardingFilterNode() {}
+    ShardingFilterNode(bool wantShardName): _wantShardName(wantShardName) {}
     virtual ~ShardingFilterNode() {}
 
     virtual StageType getType() const {
@@ -935,37 +935,14 @@ struct ShardingFilterNode : public QuerySolutionNode {
     const BSONObjSet& getSort() const {
         return children[0]->getSort();
     }
-
-    QuerySolutionNode* clone() const;
-};
-
-/**
- * If we're answering a query on a that has shardName metadata, docs must be appended with
- * the shard name.
- */
-struct ShardNameNode : public QuerySolutionNode {
-    ShardNameNode() {}
-    virtual ~ShardNameNode() {}
-
-    virtual StageType getType() const {
-        return STAGE_SHARD_NAME;
-    }
-    virtual void appendToString(str::stream* ss, int indent) const;
-
-    bool fetched() const {
-        return children[0]->fetched();
-    }
-    bool hasField(const std::string& field) const {
-        return children[0]->hasField(field);
-    }
-    bool sortedByDiskLoc() const {
-        return children[0]->sortedByDiskLoc();
-    }
-    const BSONObjSet& getSort() const {
-        return children[0]->getSort();
+    bool wantShardName() const {
+        return _wantShardName;
     }
 
     QuerySolutionNode* clone() const;
+
+private:
+    bool _wantShardName;
 };
 
 /**