1
2 /**
3 * Copyright (C) 2018-present MongoDB, Inc.
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the Server Side Public License, version 1,
7 * as published by MongoDB, Inc.
8 *
9 * This program is distributed in the hope that it will be useful,
10 * but WITHOUT ANY WARRANTY; without even the implied warranty of
11 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 * Server Side Public License for more details.
13 *
14 * You should have received a copy of the Server Side Public License
15 * along with this program. If not, see
16 * <http://www.mongodb.com/licensing/server-side-public-license>.
17 *
18 * As a special exception, the copyright holders give permission to link the
19 * code of portions of this program with the OpenSSL library under certain
20 * conditions as described in each individual source file and distribute
21 * linked combinations including the program with the OpenSSL library. You
22 * must comply with the Server Side Public License in all respects for
23 * all of the code used other than as permitted herein. If you modify file(s)
24 * with this exception, you may extend this exception to your version of the
25 * file(s), but you are not obligated to do so. If you do not wish to do so,
26 * delete this exception statement from your version. If you delete this
27 * exception statement from all source files in the program, then also delete
28 * it in the license file.
29 */
30
31 #pragma once
32
33 #include "mongo/bson/bsonobj.h"
34 #include "mongo/bson/bsonobjbuilder.h"
35 #include "mongo/s/catalog/type_chunk.h"
36 #include "mongo/s/ns_targeter.h"
37 #include "mongo/stdx/memory.h"
38 #include "mongo/unittest/unittest.h"
39
40 namespace mongo {
41
42 /**
43 * A MockRange represents a range with endpoint that a MockNSTargeter uses to direct writes to
44 * a particular endpoint.
45 */
46 struct MockRange {
MockRangeMockRange47 MockRange(const ShardEndpoint& endpoint, const BSONObj& minKey, const BSONObj& maxKey)
48 : endpoint(endpoint), range(minKey, maxKey) {}
49
50 const ShardEndpoint endpoint;
51 const ChunkRange range;
52 };
53
54 /**
55 * A MockNSTargeter directs writes to particular endpoints based on a list of MockRanges given
56 * to the mock targeter on initialization.
57 *
58 * No refreshing behavior is currently supported.
59 */
60 class MockNSTargeter : public NSTargeter {
61 public:
init(const NamespaceString & nss,std::vector<MockRange> mockRanges)62 void init(const NamespaceString& nss, std::vector<MockRange> mockRanges) {
63 ASSERT(nss.isValid());
64 _nss = nss;
65
66 ASSERT(!mockRanges.empty());
67 _mockRanges = std::move(mockRanges);
68 }
69
getNS()70 const NamespaceString& getNS() const {
71 return _nss;
72 }
73
74 /**
75 * Returns a ShardEndpoint for the doc from the mock ranges
76 */
targetInsert(OperationContext * opCtx,const BSONObj & doc)77 StatusWith<ShardEndpoint> targetInsert(OperationContext* opCtx,
78 const BSONObj& doc) const override {
79 auto swEndpoints = _targetQuery(doc);
80 if (!swEndpoints.isOK())
81 return swEndpoints.getStatus();
82
83 ASSERT_EQ(1U, swEndpoints.getValue().size());
84 return swEndpoints.getValue().front();
85 }
86
87 /**
88 * Returns the first ShardEndpoint for the query from the mock ranges. Only can handle
89 * queries of the form { field : { $gte : <value>, $lt : <value> } }.
90 */
targetUpdate(OperationContext * opCtx,const write_ops::UpdateOpEntry & updateDoc)91 StatusWith<std::vector<ShardEndpoint>> targetUpdate(
92 OperationContext* opCtx, const write_ops::UpdateOpEntry& updateDoc) const override {
93 return _targetQuery(updateDoc.getQ());
94 }
95
96 /**
97 * Returns the first ShardEndpoint for the query from the mock ranges. Only can handle
98 * queries of the form { field : { $gte : <value>, $lt : <value> } }.
99 */
targetDelete(OperationContext * opCtx,const write_ops::DeleteOpEntry & deleteDoc)100 StatusWith<std::vector<ShardEndpoint>> targetDelete(
101 OperationContext* opCtx, const write_ops::DeleteOpEntry& deleteDoc) const {
102 return _targetQuery(deleteDoc.getQ());
103 }
104
targetCollection()105 StatusWith<std::vector<ShardEndpoint>> targetCollection() const override {
106 // No-op
107 return std::vector<ShardEndpoint>{};
108 }
109
targetAllShards(OperationContext * opCtx)110 StatusWith<std::vector<ShardEndpoint>> targetAllShards(OperationContext* opCtx) const override {
111 std::vector<ShardEndpoint> endpoints;
112 for (const auto& range : _mockRanges) {
113 endpoints.push_back(range.endpoint);
114 }
115
116 return endpoints;
117 }
118
noteCouldNotTarget()119 void noteCouldNotTarget() override {
120 // No-op
121 }
122
noteStaleResponse(const ShardEndpoint & endpoint,const BSONObj & staleInfo)123 void noteStaleResponse(const ShardEndpoint& endpoint, const BSONObj& staleInfo) override {
124 // No-op
125 }
126
refreshIfNeeded(OperationContext * opCtx,bool * wasChanged)127 Status refreshIfNeeded(OperationContext* opCtx, bool* wasChanged) override {
128 // No-op
129 if (wasChanged)
130 *wasChanged = false;
131 return Status::OK();
132 }
133
134 private:
_parseRange(const BSONObj & query)135 static ChunkRange _parseRange(const BSONObj& query) {
136 const StringData fieldName(query.firstElement().fieldName());
137
138 if (query.firstElement().isNumber()) {
139 return {BSON(fieldName << query.firstElement().numberInt()),
140 BSON(fieldName << query.firstElement().numberInt() + 1)};
141 } else if (query.firstElement().type() == Object) {
142 BSONObj queryRange = query.firstElement().Obj();
143
144 ASSERT(!queryRange[GTE.l_].eoo());
145 ASSERT(!queryRange[LT.l_].eoo());
146
147 BSONObjBuilder minKeyB;
148 minKeyB.appendAs(queryRange[GTE.l_], fieldName);
149 BSONObjBuilder maxKeyB;
150 maxKeyB.appendAs(queryRange[LT.l_], fieldName);
151
152 return {minKeyB.obj(), maxKeyB.obj()};
153 }
154
155 FAIL("Invalid query");
156 MONGO_UNREACHABLE;
157 }
158
159 /**
160 * Returns the first ShardEndpoint for the query from the mock ranges. Only handles queries of
161 * the form { field : { $gte : <value>, $lt : <value> } }.
162 */
_targetQuery(const BSONObj & query)163 StatusWith<std::vector<ShardEndpoint>> _targetQuery(const BSONObj& query) const {
164 const ChunkRange queryRange(_parseRange(query));
165
166 std::vector<ShardEndpoint> endpoints;
167
168 for (const auto& range : _mockRanges) {
169 if (queryRange.overlapWith(range.range)) {
170 endpoints.push_back(range.endpoint);
171 }
172 }
173
174 if (endpoints.empty())
175 return {ErrorCodes::UnknownError, "no mock ranges found for query"};
176
177 return endpoints;
178 }
179
180 NamespaceString _nss;
181
182 std::vector<MockRange> _mockRanges;
183 };
184
assertEndpointsEqual(const ShardEndpoint & endpointA,const ShardEndpoint & endpointB)185 inline void assertEndpointsEqual(const ShardEndpoint& endpointA, const ShardEndpoint& endpointB) {
186 ASSERT_EQUALS(endpointA.shardName, endpointB.shardName);
187 ASSERT_EQUALS(endpointA.shardVersion.toLong(), endpointB.shardVersion.toLong());
188 ASSERT_EQUALS(endpointA.shardVersion.epoch(), endpointB.shardVersion.epoch());
189 }
190
191 } // namespace mongo
192