SERVER-41854 Add meta projection $shardName.

Author: krk

src/mongo/db/SConscript

@@ -985,6 +985,8 @@ env.Library(
         'exec/requires_index_stage.cpp',
         'exec/shard_filter.cpp',
         'exec/shard_filterer_impl.cpp',
+        'exec/shard_name.cpp',
+        'exec/shard_namer_impl.cpp',
         'exec/skip.cpp',
         'exec/sort.cpp',
         'exec/sort_key_generator.cpp',

src/mongo/db/exec/projection.cpp

@@ -80,6 +80,19 @@ double textScore(const WorkingSetMember& member) {
         return 0.0;
 }
 
+StringData shardName(const WorkingSetMember& member) {
+    if (member.hasComputed(WSM_COMPUTED_SHARD_NAME)) {
+        StringData shardName = static_cast<const ShardNameComputedData*>(
+                   member.getComputed(WSM_COMPUTED_SHARD_NAME))
+            ->getShardName();
+        return shardName;
+    }
+    // It is permitted to request a shardName when none has been computed. Empty string is returned
+    // as an empty value in this case.
+    else
+        return StringData("");
+}
+
 void transitionMemberToOwnedObj(const BSONObj& bo, WorkingSetMember* member) {
     member->keyData.clear();
     member->recordId = RecordId();
@@ -104,7 +117,8 @@ StatusWith<BSONObj> provideMetaFieldsAndPerformExec(const ProjectionExec& exec,
                        exec.needsSortKey() ? sortKey(member) : BSONObj(),
                        exec.needsTextScore() ? boost::optional<const double>(textScore(member))
                                              : boost::none,
-                       member.recordId.repr())
+                       member.recordId.repr(),
+                       exec.needsShardName() ? shardName(member) : StringData(""))
         : exec.projectCovered(
               member.keyData,
               exec.needsGeoNearDistance() ? boost::optional<const double>(geoDistance(member))
@@ -113,7 +127,8 @@ StatusWith<BSONObj> provideMetaFieldsAndPerformExec(const ProjectionExec& exec,
               exec.needsSortKey() ? sortKey(member) : BSONObj(),
               exec.needsTextScore() ? boost::optional<const double>(textScore(member))
                                     : boost::none,
-              member.recordId.repr());
+              member.recordId.repr(),
+              exec.needsShardName() ? shardName(member) : StringData(""));
 }
 }  // namespace
 

src/mongo/db/exec/projection_exec.cpp

@@ -118,6 +118,9 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx,
                     _needsGeoNearDistance = true;
                 } else if (e2.valuestr() == QueryRequest::metaIndexKey) {
                     _hasReturnKey = 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;
@@ -220,7 +223,8 @@ StatusWith<BSONObj> ProjectionExec::project(const BSONObj& in,
                                             const BSONObj& 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;
 
@@ -235,15 +239,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,
                                                    const BSONObj& 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.
@@ -288,15 +293,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,
                                 const BSONObj& 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:
@@ -322,6 +328,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

@@ -62,6 +62,7 @@ class ProjectionExec {
         META_RECORDID,
         META_SORT_KEY,
         META_TEXT_SCORE,
+        META_SHARD_NAME,
     };
 
     /**
@@ -115,6 +116,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.
      */
@@ -137,7 +145,8 @@ class ProjectionExec {
                                 const BSONObj& geoNearPoint = BSONObj(),
                                 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
@@ -150,7 +159,8 @@ class ProjectionExec {
         const BSONObj& geoNearPoint = BSONObj(),
         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
@@ -170,7 +180,8 @@ class ProjectionExec {
                     const BSONObj& geoNearPoint,
                     const BSONObj& sortKey,
                     const boost::optional<const double> textScore,
-                    const int64_t recordId) const;
+                    const int64_t recordId,
+                    const StringData& shardName) const;
 
     //
     // Initialization
@@ -268,6 +279,7 @@ class ProjectionExec {
     bool _needsGeoNearDistance = false;
     bool _needsGeoNearPoint = false;
     bool _needsTextScore = false;
+    bool _needsShardName = false;
 
     // The field names associated with any sortKey meta-projection(s). Empty if there is no sortKey
     // meta-projection.

src/mongo/db/exec/shard_name.cpp

@@ -0,0 +1,108 @@
+/**
+ *    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/shard_name.h"
+
+#include <memory>
+
+#include "mongo/db/exec/filter.h"
+#include "mongo/db/exec/scoped_timer.h"
+#include "mongo/db/exec/working_set_common.h"
+#include "mongo/db/exec/working_set_computed_data.h"
+#include "mongo/s/shard_key_pattern.h"
+#include "mongo/util/log.h"
+
+namespace mongo {
+
+using std::shared_ptr;
+using std::unique_ptr;
+using std::vector;
+using std::string;
+
+// static
+const char* ShardNameStage::kStageType = "SHARD_NAME";
+
+ShardNameStage::    ShardNameStage(OperationContext* opCtx,
+                               ScopedCollectionMetadata metadata,
+                               WorkingSet* ws,
+                               PlanStage* child)
+    : PlanStage(kStageType, opCtx), _ws(ws), _shardNamer(std::move(metadata)) {
+    _children.emplace_back(child);
+}
+
+ShardNameStage::~ShardNameStage() {}
+
+bool ShardNameStage::isEOF() {
+    return child()->isEOF();
+}
+
+PlanStage::StageState ShardNameStage::doWork(WorkingSetID* out) {
+    // If we've returned as many results as we're limited to, isEOF will be true.
+    if (isEOF()) {
+        return PlanStage::IS_EOF;
+    }
+
+    StageState status = child()->work(out);
+
+    if (PlanStage::ADVANCED == status) {
+        // If we're sharded make sure to add shardName to the output.
+        if (_shardNamer.isCollectionSharded()) {
+            WorkingSetMember* member = _ws->get(*out);
+            const StringData shardName = _shardNamer.shardName();
+
+            // Populate the working set member with the shard name and return it.
+            member->addComputed(new ShardNameComputedData(shardName));
+        }
+
+        // If we're here either we have shard state and added the shardName, or we have no shard
+        // state.  Either way, we advance.
+        return status;
+    }
+
+    return status;
+}
+
+unique_ptr<PlanStageStats> ShardNameStage::getStats() {
+    _commonStats.isEOF = isEOF();
+    unique_ptr<PlanStageStats> ret =
+        std::make_unique<PlanStageStats>(_commonStats, STAGE_SHARD_NAME);
+    ret->children.emplace_back(child()->getStats());
+    return ret;
+}
+
+const SpecificStats* ShardNameStage::getSpecificStats() const {
+    // No specific stats are tracked for the shard name stage.
+    return nullptr;
+}
+
+}  // namespace mongo

src/mongo/db/exec/shard_name.h

@@ -0,0 +1,70 @@
+/**
+ *    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"
+#include "mongo/db/exec/shard_namer_impl.h"
+
+namespace mongo {
+
+class ShardNameStage final : public PlanStage {
+public:
+    ShardNameStage(OperationContext* opCtx,
+                   ScopedCollectionMetadata metadata,
+                   WorkingSet* ws,
+                   PlanStage* child);
+    ~ShardNameStage();
+
+    bool isEOF() final;
+    StageState doWork(WorkingSetID* out) final;
+
+    StageType stageType() const final {
+        return STAGE_SHARD_NAME;
+    }
+
+    std::unique_ptr<PlanStageStats> getStats() final;
+
+    const SpecificStats* getSpecificStats() const final;
+
+    static const char* kStageType;
+
+private:
+    WorkingSet* _ws;
+
+    // Note: it is important that this owns the ScopedCollectionMetadata from the time this stage
+    // is constructed. See ScopedCollectionMetadata class comment and MetadataManager comment for
+    // details. The existence of the ScopedCollectionMetadata prevents data which may have been
+    // migrated from being deleted while the query is still active. If we didn't hold one
+    // 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.
+    ShardNamerImpl _shardNamer;
+};
+
+}  // namespace mongo