Skip to content

Commit 52bba63

Browse files
krkkrk
committed
SERVER-41854 Add meta projection $shardName.
1 parent 8b32d75 commit 52bba63

21 files changed

+512
-18
lines changed

‎src/mongo/db/SConscript‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1027,6 +1027,8 @@ env.Library(
10271027
'exec/return_key.cpp',
10281028
'exec/shard_filter.cpp',
10291029
'exec/shard_filterer_impl.cpp',
1030+
'exec/shard_name.cpp',
1031+
'exec/shard_namer_impl.cpp',
10301032
'exec/skip.cpp',
10311033
'exec/sort.cpp',
10321034
'exec/sort_key_generator.cpp',

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,10 @@ void DocumentMetadataFields::serializeForSorter(BufBuilder& buf) const {
173173
buf.appendNum(static_cast<char>(MetaType::kIndexKey + 1));
174174
getIndexKey().appendSelfToBufBuilder(buf);
175175
}
176+
if (hasShardName()) {
177+
buf.appendNum(static_cast<char>(MetaType::kShardName + 1));
178+
buf.appendStr(getShardName());
179+
}
176180
buf.appendNum(static_cast<char>(0));
177181
}
178182

@@ -201,6 +205,8 @@ void DocumentMetadataFields::deserializeForSorter(BufReader& buf, DocumentMetada
201205
} else if (marker == static_cast<char>(MetaType::kIndexKey) + 1) {
202206
out->setIndexKey(
203207
BSONObj::deserializeForSorter(buf, BSONObj::SorterDeserializeSettings()));
208+
} else if (marker == static_cast<char>(MetaType::kShardName) + 1) {
209+
out->setShardName(buf.readCStr());
204210
} else {
205211
uasserted(28744, "Unrecognized marker, unable to deserialize buffer");
206212
}

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ class DocumentMetadataFields {
6464
kSearchScore,
6565
kSortKey,
6666
kTextScore,
67+
kShardName,
6768

6869
// New fields must be added before the kNumFields sentinel.
6970
kNumFields
@@ -302,6 +303,24 @@ class DocumentMetadataFields {
302303
_holder->recordId = rid;
303304
}
304305

306+
bool hasShardName() const {
307+
return _holder && _holder->metaFields.test(MetaType::kShardName);
308+
}
309+
310+
StringData getShardName() const {
311+
invariant(hasShardName());
312+
return _holder->shardName;
313+
}
314+
315+
void setShardName(StringData shardName) {
316+
if (!_holder) {
317+
_holder = std::make_unique<MetadataHolder>();
318+
}
319+
320+
_holder->metaFields.set(MetaType::kShardName);
321+
_holder->shardName = shardName;
322+
}
323+
305324
void serializeForSorter(BufBuilder& buf) const;
306325

307326
private:
@@ -323,6 +342,7 @@ class DocumentMetadataFields {
323342
Value searchHighlights;
324343
BSONObj indexKey;
325344
RecordId recordId;
345+
StringData shardName;
326346
};
327347

328348
// 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‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ TEST(DocumentMetadataFieldsTest, AllMetadataRoundtripsThroughSerialization) {
4848
metadata.setSearchScore(5.4);
4949
metadata.setSearchHighlights(Value{"foo"_sd});
5050
metadata.setIndexKey(BSON("b" << 1));
51+
metadata.setShardName("bar");
5152

5253
BufBuilder builder;
5354
metadata.serializeForSorter(builder);
@@ -64,6 +65,7 @@ TEST(DocumentMetadataFieldsTest, AllMetadataRoundtripsThroughSerialization) {
6465
ASSERT_EQ(deserialized.getSearchScore(), 5.4);
6566
ASSERT_VALUE_EQ(deserialized.getSearchHighlights(), Value{"foo"_sd});
6667
ASSERT_BSONOBJ_EQ(deserialized.getIndexKey(), BSON("b" << 1));
68+
ASSERT_EQ(deserialized.getShardName(), "bar");
6769
}
6870

6971
TEST(DocumentMetadataFieldsTest, HasMethodsReturnFalseForEmptyMetadata) {
@@ -77,6 +79,7 @@ TEST(DocumentMetadataFieldsTest, HasMethodsReturnFalseForEmptyMetadata) {
7779
ASSERT_FALSE(metadata.hasSearchScore());
7880
ASSERT_FALSE(metadata.hasSearchHighlights());
7981
ASSERT_FALSE(metadata.hasIndexKey());
82+
ASSERT_FALSE(metadata.hasShardName());
8083
}
8184

8285
TEST(DocumentMetadataFieldsTest, HasMethodsReturnTrueForInitializedMetadata) {
@@ -114,6 +117,10 @@ TEST(DocumentMetadataFieldsTest, HasMethodsReturnTrueForInitializedMetadata) {
114117
ASSERT_FALSE(metadata.hasIndexKey());
115118
metadata.setIndexKey(BSON("b" << 1));
116119
ASSERT_TRUE(metadata.hasIndexKey());
120+
121+
ASSERT_FALSE(metadata.hasShardName());
122+
metadata.setShardName("bar");
123+
ASSERT_TRUE(metadata.hasShardName());
117124
}
118125

119126
TEST(DocumentMetadataFieldsTest, MoveConstructor) {
@@ -126,6 +133,7 @@ TEST(DocumentMetadataFieldsTest, MoveConstructor) {
126133
metadata.setSearchScore(5.4);
127134
metadata.setSearchHighlights(Value{"foo"_sd});
128135
metadata.setIndexKey(BSON("b" << 1));
136+
metadata.setShardName("bar");
129137

130138
DocumentMetadataFields moveConstructed(std::move(metadata));
131139
ASSERT_TRUE(moveConstructed);
@@ -138,6 +146,7 @@ TEST(DocumentMetadataFieldsTest, MoveConstructor) {
138146
ASSERT_EQ(moveConstructed.getSearchScore(), 5.4);
139147
ASSERT_VALUE_EQ(moveConstructed.getSearchHighlights(), Value{"foo"_sd});
140148
ASSERT_BSONOBJ_EQ(moveConstructed.getIndexKey(), BSON("b" << 1));
149+
ASSERT_EQ(moveConstructed.getShardName(), "bar");
141150

142151
ASSERT_FALSE(metadata);
143152
}
@@ -152,6 +161,7 @@ TEST(DocumentMetadataFieldsTest, MoveAssignmentOperator) {
152161
metadata.setSearchScore(5.4);
153162
metadata.setSearchHighlights(Value{"foo"_sd});
154163
metadata.setIndexKey(BSON("b" << 1));
164+
metadata.setShardName("bar");
155165

156166
DocumentMetadataFields moveAssigned;
157167
moveAssigned.setTextScore(12.3);
@@ -167,6 +177,7 @@ TEST(DocumentMetadataFieldsTest, MoveAssignmentOperator) {
167177
ASSERT_EQ(moveAssigned.getSearchScore(), 5.4);
168178
ASSERT_VALUE_EQ(moveAssigned.getSearchHighlights(), Value{"foo"_sd});
169179
ASSERT_BSONOBJ_EQ(moveAssigned.getIndexKey(), BSON("b" << 1));
180+
ASSERT_EQ(moveAssigned.getShardName(), "bar");
170181

171182
ASSERT_FALSE(metadata);
172183
}

‎src/mongo/db/exec/projection_exec.cpp‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,9 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx,
116116
} else if (e2.valuestr() == QueryRequest::metaGeoNearDistance) {
117117
_meta[e.fieldName()] = META_GEONEAR_DIST;
118118
_needsGeoNearDistance = true;
119+
} else if (e2.valuestr() == QueryRequest::metaShardName) {
120+
_meta[e.fieldName()] = META_SHARD_NAME;
121+
_needsShardName = true;
119122
} else {
120123
// This shouldn't happen, should be caught by parsing.
121124
MONGO_UNREACHABLE;
@@ -202,7 +205,8 @@ StatusWith<BSONObj> ProjectionExec::project(const BSONObj& in,
202205
Value geoNearPoint,
203206
const BSONObj& sortKey,
204207
const boost::optional<const double> textScore,
205-
const int64_t recordId) const {
208+
const int64_t recordId,
209+
const StringData& shardName) const {
206210
BSONObjBuilder bob;
207211
MatchDetails matchDetails;
208212

@@ -217,15 +221,16 @@ StatusWith<BSONObj> ProjectionExec::project(const BSONObj& in,
217221
if (!projStatus.isOK())
218222
return projStatus;
219223
else
220-
return {addMeta(std::move(bob), geoDistance, geoNearPoint, sortKey, textScore, recordId)};
224+
return {addMeta(std::move(bob), geoDistance, geoNearPoint, sortKey, textScore, recordId, shardName)};
221225
}
222226

223227
StatusWith<BSONObj> ProjectionExec::projectCovered(const std::vector<IndexKeyDatum>& keyData,
224228
const boost::optional<const double> geoDistance,
225229
Value geoNearPoint,
226230
const BSONObj& sortKey,
227231
const boost::optional<const double> textScore,
228-
const int64_t recordId) const {
232+
const int64_t recordId,
233+
const StringData& shardName) const {
229234
invariant(!_include);
230235
BSONObjBuilder bob;
231236
// Go field by field.
@@ -270,15 +275,16 @@ StatusWith<BSONObj> ProjectionExec::projectCovered(const std::vector<IndexKeyDat
270275
}
271276

272277
bob.appendElements(projectedDoc.getObject());
273-
return {addMeta(std::move(bob), geoDistance, geoNearPoint, sortKey, textScore, recordId)};
278+
return {addMeta(std::move(bob), geoDistance, geoNearPoint, sortKey, textScore, recordId, shardName)};
274279
}
275280

276281
BSONObj ProjectionExec::addMeta(BSONObjBuilder bob,
277282
const boost::optional<const double> geoDistance,
278283
Value geoNearPoint,
279284
const BSONObj& sortKey,
280285
const boost::optional<const double> textScore,
281-
const int64_t recordId) const {
286+
const int64_t recordId,
287+
const StringData& shardName) const {
282288
for (MetaMap::const_iterator it = _meta.begin(); it != _meta.end(); ++it) {
283289
switch (it->second) {
284290
case META_GEONEAR_DIST:
@@ -299,6 +305,10 @@ BSONObj ProjectionExec::addMeta(BSONObjBuilder bob,
299305
bob.append(it->first, sortKey);
300306
break;
301307
}
308+
case META_SHARD_NAME: {
309+
bob.append(it->first, shardName);
310+
break;
311+
}
302312
case META_RECORDID:
303313
invariant(recordId != 0);
304314
bob.append(it->first, recordId);

‎src/mongo/db/exec/projection_exec.h‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ class ProjectionExec {
6363
META_RECORDID,
6464
META_SORT_KEY,
6565
META_TEXT_SCORE,
66+
META_SHARD_NAME,
6667
};
6768

6869
/**
@@ -108,6 +109,13 @@ class ProjectionExec {
108109
return _needsTextScore;
109110
}
110111

112+
/**
113+
* Indicates whether 'shardName' is going to be used in 'project()'.
114+
*/
115+
bool needsShardName() const {
116+
return _needsShardName;
117+
}
118+
111119
/**
112120
* Returns false if there are no meta fields to project.
113121
*/
@@ -124,7 +132,8 @@ class ProjectionExec {
124132
Value geoNearPoint = Value{},
125133
const BSONObj& sortKey = BSONObj(),
126134
const boost::optional<const double> textScore = boost::none,
127-
const int64_t recordId = 0) const;
135+
const int64_t recordId = 0,
136+
const StringData& shardName = StringData("")) const;
128137

129138
/**
130139
* Performs a projection given index 'KeyData' to directly retrieve results. This function
@@ -137,7 +146,8 @@ class ProjectionExec {
137146
Value geoNearPoint = Value{},
138147
const BSONObj& sortKey = BSONObj(),
139148
const boost::optional<const double> textScore = boost::none,
140-
const int64_t recordId = 0) const;
149+
const int64_t recordId = 0,
150+
const StringData& shardName = StringData("")) const;
141151

142152
/**
143153
* Determines if calls to the project method require that this object was created with the full
@@ -157,7 +167,8 @@ class ProjectionExec {
157167
Value geoNearPoint,
158168
const BSONObj& sortKey,
159169
const boost::optional<const double> textScore,
160-
const int64_t recordId) const;
170+
const int64_t recordId,
171+
const StringData& shardName) const;
161172

162173
//
163174
// Initialization
@@ -251,6 +262,7 @@ class ProjectionExec {
251262
bool _needsGeoNearDistance = false;
252263
bool _needsGeoNearPoint = false;
253264
bool _needsTextScore = false;
265+
bool _needsShardName = false;
254266

255267
// The collator this projection should use to compare strings. Needed for projection operators
256268
// that perform matching (e.g. elemMatch projection). If null, the collation is a simple binary

‎src/mongo/db/exec/shard_name.cpp‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/**
2+
* Copyright (C) 2019-present MongoDB, Inc.
3+
*
4+
* This program is free software: you can redistribute it and/or modify
5+
* it under the terms of the Server Side Public License, version 1,
6+
* as published by MongoDB, Inc.
7+
*
8+
* This program is distributed in the hope that it will be useful,
9+
* but WITHOUT ANY WARRANTY; without even the implied warranty of
10+
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11+
* Server Side Public License for more details.
12+
*
13+
* You should have received a copy of the Server Side Public License
14+
* along with this program. If not, see
15+
* <http://www.mongodb.com/licensing/server-side-public-license>.
16+
*
17+
* As a special exception, the copyright holders give permission to link the
18+
* code of portions of this program with the OpenSSL library under certain
19+
* conditions as described in each individual source file and distribute
20+
* linked combinations including the program with the OpenSSL library. You
21+
* must comply with the Server Side Public License in all respects for
22+
* all of the code used other than as permitted herein. If you modify file(s)
23+
* with this exception, you may extend this exception to your version of the
24+
* file(s), but you are not obligated to do so. If you do not wish to do so,
25+
* delete this exception statement from your version. If you delete this
26+
* exception statement from all source files in the program, then also delete
27+
* it in the license file.
28+
*/
29+
30+
#include "mongo/platform/basic.h"
31+
32+
#include "mongo/db/exec/shard_name.h"
33+
34+
#include <memory>
35+
36+
#include "mongo/db/exec/filter.h"
37+
#include "mongo/db/exec/scoped_timer.h"
38+
#include "mongo/db/exec/working_set_common.h"
39+
#include "mongo/s/shard_key_pattern.h"
40+
41+
namespace mongo {
42+
43+
using std::shared_ptr;
44+
using std::unique_ptr;
45+
using std::vector;
46+
using std::string;
47+
48+
// static
49+
const char* ShardNameStage::kStageType = "SHARD_NAME";
50+
51+
ShardNameStage::ShardNameStage(OperationContext* opCtx,
52+
ScopedCollectionMetadata metadata,
53+
WorkingSet* ws,
54+
std::unique_ptr<PlanStage> child)
55+
: PlanStage(kStageType, opCtx), _ws(ws), _shardNamer(std::move(metadata)) {
56+
_children.emplace_back(std::move(child));
57+
}
58+
59+
ShardNameStage::~ShardNameStage() {}
60+
61+
bool ShardNameStage::isEOF() {
62+
return child()->isEOF();
63+
}
64+
65+
PlanStage::StageState ShardNameStage::doWork(WorkingSetID* out) {
66+
// If we've returned as many results as we're limited to, isEOF will be true.
67+
if (isEOF()) {
68+
return PlanStage::IS_EOF;
69+
}
70+
71+
StageState status = child()->work(out);
72+
73+
if (PlanStage::ADVANCED == status) {
74+
// If we're sharded make sure to add shardName to the output.
75+
if (_shardNamer.isCollectionSharded()) {
76+
WorkingSetMember* member = _ws->get(*out);
77+
const StringData shardName = _shardNamer.shardName();
78+
79+
// Populate the working set member with the shard name and return it.
80+
member->metadata().setShardName(shardName);
81+
}
82+
83+
// If we're here either we have shard state and added the shardName, or we have no shard
84+
// state. Either way, we advance.
85+
return status;
86+
}
87+
88+
return status;
89+
}
90+
91+
unique_ptr<PlanStageStats> ShardNameStage::getStats() {
92+
_commonStats.isEOF = isEOF();
93+
unique_ptr<PlanStageStats> ret =
94+
std::make_unique<PlanStageStats>(_commonStats, STAGE_SHARD_NAME);
95+
ret->children.emplace_back(child()->getStats());
96+
return ret;
97+
}
98+
99+
const SpecificStats* ShardNameStage::getSpecificStats() const {
100+
// No specific stats are tracked for the shard name stage.
101+
return nullptr;
102+
}
103+
104+
} // namespace mongo

0 commit comments

Comments
 (0)