SERVER-41854 Add meta projection $shardName.

src/mongo/db/exec/document_value/document_metadata_fields.cpp

@@ -33,6 +33,8 @@
 
 #include "mongo/bson/bsonobjbuilder.h"
 
+#include "mongo/s/shard_id.h"
+
 namespace mongo {
 
 namespace {
@@ -173,6 +175,10 @@ void DocumentMetadataFields::serializeForSorter(BufBuilder& buf) const {
         buf.appendNum(static_cast<char>(MetaType::kIndexKey + 1));
         getIndexKey().appendSelfToBufBuilder(buf);
     }
+    if (hasShardName()) {
+        buf.appendNum(static_cast<char>(MetaType::kShardName + 1));
+        buf.appendStr(getShardName());
+    }
     buf.appendNum(static_cast<char>(0));
 }
 
@@ -201,6 +207,8 @@ void DocumentMetadataFields::deserializeForSorter(BufReader& buf, DocumentMetada
         } else if (marker == static_cast<char>(MetaType::kIndexKey) + 1) {
             out->setIndexKey(
                 BSONObj::deserializeForSorter(buf, BSONObj::SorterDeserializeSettings()));
+        } else if (marker == static_cast<char>(MetaType::kShardName) + 1) {
+            out->setShardName(ShardId(std::string(buf.readCStr())));
         } else {
             uasserted(28744, "Unrecognized marker, unable to deserialize buffer");
         }

src/mongo/db/exec/document_value/document_metadata_fields.h

@@ -34,6 +34,7 @@
 #include "mongo/bson/bsonobj.h"
 #include "mongo/db/exec/document_value/value.h"
 #include "mongo/db/record_id.h"
+#include "mongo/s/shard_id.h"
 
 namespace mongo {
 /**
@@ -64,6 +65,7 @@ class DocumentMetadataFields {
         kSearchScore,
         kSortKey,
         kTextScore,
+        kShardName,
 
         // New fields must be added before the kNumFields sentinel.
         kNumFields
@@ -302,6 +304,24 @@ class DocumentMetadataFields {
         _holder->recordId = rid;
     }
 
+    bool hasShardName() const {
+        return _holder && _holder->metaFields.test(MetaType::kShardName);
+    }
+
+    ShardId getShardName() const {
+        invariant(hasShardName());
+        return _holder->shardName;
+    }
+
+    void setShardName(ShardId shardName) {
+        if (!_holder) {
+            _holder = std::make_unique<MetadataHolder>();
+        }
+
+        _holder->metaFields.set(MetaType::kShardName);
+        _holder->shardName = shardName;
+    }
+
     void serializeForSorter(BufBuilder& buf) const;
 
 private:
@@ -323,6 +343,7 @@ class DocumentMetadataFields {
         Value searchHighlights;
         BSONObj indexKey;
         RecordId recordId;
+        ShardId shardName;
     };
 
     // Null until the first setter is called, at which point a MetadataHolder struct is allocated.

src/mongo/db/exec/document_value/document_metadata_fields_test.cpp

@@ -48,6 +48,7 @@ TEST(DocumentMetadataFieldsTest, AllMetadataRoundtripsThroughSerialization) {
     metadata.setSearchScore(5.4);
     metadata.setSearchHighlights(Value{"foo"_sd});
     metadata.setIndexKey(BSON("b" << 1));
+    metadata.setShardName("bar");
 
     BufBuilder builder;
     metadata.serializeForSorter(builder);
@@ -64,6 +65,7 @@ TEST(DocumentMetadataFieldsTest, AllMetadataRoundtripsThroughSerialization) {
     ASSERT_EQ(deserialized.getSearchScore(), 5.4);
     ASSERT_VALUE_EQ(deserialized.getSearchHighlights(), Value{"foo"_sd});
     ASSERT_BSONOBJ_EQ(deserialized.getIndexKey(), BSON("b" << 1));
+    ASSERT_EQ(deserialized.getShardName(), "bar");
 }
 
 TEST(DocumentMetadataFieldsTest, HasMethodsReturnFalseForEmptyMetadata) {
@@ -77,6 +79,7 @@ TEST(DocumentMetadataFieldsTest, HasMethodsReturnFalseForEmptyMetadata) {
     ASSERT_FALSE(metadata.hasSearchScore());
     ASSERT_FALSE(metadata.hasSearchHighlights());
     ASSERT_FALSE(metadata.hasIndexKey());
+    ASSERT_FALSE(metadata.hasShardName());
 }
 
 TEST(DocumentMetadataFieldsTest, HasMethodsReturnTrueForInitializedMetadata) {
@@ -114,6 +117,10 @@ TEST(DocumentMetadataFieldsTest, HasMethodsReturnTrueForInitializedMetadata) {
     ASSERT_FALSE(metadata.hasIndexKey());
     metadata.setIndexKey(BSON("b" << 1));
     ASSERT_TRUE(metadata.hasIndexKey());
+
+    ASSERT_FALSE(metadata.hasShardName());
+    metadata.setShardName("bar");
+    ASSERT_TRUE(metadata.hasShardName());
 }
 
 TEST(DocumentMetadataFieldsTest, MoveConstructor) {
@@ -126,6 +133,7 @@ TEST(DocumentMetadataFieldsTest, MoveConstructor) {
     metadata.setSearchScore(5.4);
     metadata.setSearchHighlights(Value{"foo"_sd});
     metadata.setIndexKey(BSON("b" << 1));
+    metadata.setShardName("bar");
 
     DocumentMetadataFields moveConstructed(std::move(metadata));
     ASSERT_TRUE(moveConstructed);
@@ -138,6 +146,7 @@ TEST(DocumentMetadataFieldsTest, MoveConstructor) {
     ASSERT_EQ(moveConstructed.getSearchScore(), 5.4);
     ASSERT_VALUE_EQ(moveConstructed.getSearchHighlights(), Value{"foo"_sd});
     ASSERT_BSONOBJ_EQ(moveConstructed.getIndexKey(), BSON("b" << 1));
+    ASSERT_EQ(moveConstructed.getShardName(), "bar");
 
     ASSERT_FALSE(metadata);
 }
@@ -152,6 +161,7 @@ TEST(DocumentMetadataFieldsTest, MoveAssignmentOperator) {
     metadata.setSearchScore(5.4);
     metadata.setSearchHighlights(Value{"foo"_sd});
     metadata.setIndexKey(BSON("b" << 1));
+    metadata.setShardName("bar");
 
     DocumentMetadataFields moveAssigned;
     moveAssigned.setTextScore(12.3);
@@ -167,6 +177,7 @@ TEST(DocumentMetadataFieldsTest, MoveAssignmentOperator) {
     ASSERT_EQ(moveAssigned.getSearchScore(), 5.4);
     ASSERT_VALUE_EQ(moveAssigned.getSearchHighlights(), Value{"foo"_sd});
     ASSERT_BSONOBJ_EQ(moveAssigned.getIndexKey(), BSON("b" << 1));
+    ASSERT_EQ(moveAssigned.getShardName(), "bar");
 
     ASSERT_FALSE(metadata);
 }

src/mongo/db/exec/projection_exec.cpp

@@ -116,6 +116,9 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx,
                 } else if (e2.valuestr() == QueryRequest::metaGeoNearDistance) {
                     _meta[e.fieldName()] = META_GEONEAR_DIST;
                     _needsGeoNearDistance = true;
+                } else if (e2.valuestr() == QueryRequest::metaShardName) {
+                    _meta[e.fieldName()] = META_SHARD_NAME;
+                    _needsShardName = true;
                 } else {
                     // This shouldn't happen, should be caught by parsing.
                     MONGO_UNREACHABLE;
@@ -202,7 +205,8 @@ StatusWith<BSONObj> ProjectionExec::project(const BSONObj& in,
                                             Value geoNearPoint,
                                             const BSONObj& sortKey,
                                             const boost::optional<const double> textScore,
-                                            const int64_t recordId) const {
+                                            const int64_t recordId,
+                                            const StringData& shardName) const {
     BSONObjBuilder bob;
     MatchDetails matchDetails;
 
@@ -217,15 +221,16 @@ StatusWith<BSONObj> ProjectionExec::project(const BSONObj& in,
     if (!projStatus.isOK())
         return projStatus;
     else
-        return {addMeta(std::move(bob), geoDistance, geoNearPoint, sortKey, textScore, recordId)};
+        return {addMeta(std::move(bob), geoDistance, geoNearPoint, sortKey, textScore, recordId, shardName)};
 }
 
 StatusWith<BSONObj> ProjectionExec::projectCovered(const std::vector<IndexKeyDatum>& keyData,
                                                    const boost::optional<const double> geoDistance,
                                                    Value geoNearPoint,
                                                    const BSONObj& sortKey,
                                                    const boost::optional<const double> textScore,
-                                                   const int64_t recordId) const {
+                                                   const int64_t recordId,
+                                                   const StringData& shardName) const {
     invariant(!_include);
     BSONObjBuilder bob;
     // Go field by field.
@@ -270,15 +275,16 @@ StatusWith<BSONObj> ProjectionExec::projectCovered(const std::vector<IndexKeyDat
     }
 
     bob.appendElements(projectedDoc.getObject());
-    return {addMeta(std::move(bob), geoDistance, geoNearPoint, sortKey, textScore, recordId)};
+    return {addMeta(std::move(bob), geoDistance, geoNearPoint, sortKey, textScore, recordId, shardName)};
 }
 
 BSONObj ProjectionExec::addMeta(BSONObjBuilder bob,
                                 const boost::optional<const double> geoDistance,
                                 Value geoNearPoint,
                                 const BSONObj& sortKey,
                                 const boost::optional<const double> textScore,
-                                const int64_t recordId) const {
+                                const int64_t recordId,
+                                const StringData& shardName) const {
     for (MetaMap::const_iterator it = _meta.begin(); it != _meta.end(); ++it) {
         switch (it->second) {
             case META_GEONEAR_DIST:
@@ -299,6 +305,10 @@ BSONObj ProjectionExec::addMeta(BSONObjBuilder bob,
                 bob.append(it->first, sortKey);
                 break;
             }
+            case META_SHARD_NAME: {
+                bob.append(it->first, shardName);
+                break;
+            }
             case META_RECORDID:
                 invariant(recordId != 0);
                 bob.append(it->first, recordId);

src/mongo/db/exec/projection_exec.h

@@ -63,6 +63,7 @@ class ProjectionExec {
         META_RECORDID,
         META_SORT_KEY,
         META_TEXT_SCORE,
+        META_SHARD_NAME,
     };
 
     /**
@@ -108,6 +109,13 @@ class ProjectionExec {
         return _needsTextScore;
     }
 
+    /**
+     * Indicates whether 'shardName' is going to be used in 'project()'.
+     */
+    bool needsShardName() const {
+        return _needsShardName;
+    }
+
     /**
      * Returns false if there are no meta fields to project.
      */
@@ -124,7 +132,8 @@ class ProjectionExec {
                                 Value geoNearPoint = Value{},
                                 const BSONObj& sortKey = BSONObj(),
                                 const boost::optional<const double> textScore = boost::none,
-                                const int64_t recordId = 0) const;
+                                const int64_t recordId = 0,
+                                const StringData& shardName = StringData("")) const;
 
     /**
      * Performs a projection given index 'KeyData' to directly retrieve results. This function
@@ -137,7 +146,8 @@ class ProjectionExec {
         Value geoNearPoint = Value{},
         const BSONObj& sortKey = BSONObj(),
         const boost::optional<const double> textScore = boost::none,
-        const int64_t recordId = 0) const;
+        const int64_t recordId = 0,
+        const StringData& shardName = StringData("")) const;
 
     /**
      * Determines if calls to the project method require that this object was created with the full
@@ -157,7 +167,8 @@ class ProjectionExec {
                     Value geoNearPoint,
                     const BSONObj& sortKey,
                     const boost::optional<const double> textScore,
-                    const int64_t recordId) const;
+                    const int64_t recordId,
+                    const StringData& shardName) const;
 
     //
     // Initialization
@@ -251,6 +262,7 @@ class ProjectionExec {
     bool _needsGeoNearDistance = false;
     bool _needsGeoNearPoint = false;
     bool _needsTextScore = false;
+    bool _needsShardName = false;
 
     // 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

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().toString());
+            }
         }
 
         // 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/pipeline/expression.cpp

@@ -27,7 +27,6 @@
  *    it in the license file.
  */
 
-
 #include "mongo/platform/basic.h"
 
 #include "mongo/db/pipeline/expression.h"
@@ -2546,6 +2545,7 @@ const std::string geoNearPointName = "geoNearPoint";
 const std::string recordIdName = "recordId";
 const std::string indexKeyName = "indexKey";
 const std::string sortKeyName = "sortKey";
+const std::string shardName = "shardName";
 
 using MetaType = DocumentMetadataFields::MetaType;
 const StringMap<DocumentMetadataFields::MetaType> kMetaNameToMetaType = {
@@ -2558,6 +2558,7 @@ const StringMap<DocumentMetadataFields::MetaType> kMetaNameToMetaType = {
     {searchScoreName, MetaType::kSearchScore},
     {sortKeyName, MetaType::kSortKey},
     {textScoreName, MetaType::kTextScore},
+    {shardName, MetaType::kShardName},
 };
 
 const stdx::unordered_map<DocumentMetadataFields::MetaType, StringData> kMetaTypeToMetaName = {
@@ -2570,6 +2571,7 @@ const stdx::unordered_map<DocumentMetadataFields::MetaType, StringData> kMetaTyp
     {MetaType::kSearchScore, searchScoreName},
     {MetaType::kSortKey, sortKeyName},
     {MetaType::kTextScore, textScoreName},
+    {MetaType::kShardName, shardName},
 };
 
 }  // namespace
@@ -2643,6 +2645,10 @@ Value ExpressionMeta::evaluate(const Document& root, Variables* variables) const
                 ? Value(DocumentMetadataFields::serializeSortKey(metadata.isSingleElementKey(),
                                                                  metadata.getSortKey()))
                 : Value();
+        case MetaType::kShardName:
+            return metadata.hasShardName()
+                ? Value(metadata.getShardName().toString())
+                : Value();
         default:
             MONGO_UNREACHABLE;
     }

src/mongo/db/pipeline/pipeline_d.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(),