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 #include "mongo/platform/basic.h"
32 
33 #include "mongo/base/owned_pointer_map.h"
34 #include "mongo/db/operation_context_noop.h"
35 #include "mongo/s/write_ops/batch_write_op.h"
36 #include "mongo/s/write_ops/batched_command_request.h"
37 #include "mongo/s/write_ops/mock_ns_targeter.h"
38 #include "mongo/s/write_ops/write_error_detail.h"
39 #include "mongo/unittest/unittest.h"
40 
41 namespace mongo {
42 namespace {
43 
initTargeterFullRange(const NamespaceString & nss,const ShardEndpoint & endpoint,MockNSTargeter * targeter)44 void initTargeterFullRange(const NamespaceString& nss,
45                            const ShardEndpoint& endpoint,
46                            MockNSTargeter* targeter) {
47     targeter->init(nss, {MockRange(endpoint, BSON("x" << MINKEY), BSON("x" << MAXKEY))});
48 }
49 
initTargeterSplitRange(const NamespaceString & nss,const ShardEndpoint & endpointA,const ShardEndpoint & endpointB,MockNSTargeter * targeter)50 void initTargeterSplitRange(const NamespaceString& nss,
51                             const ShardEndpoint& endpointA,
52                             const ShardEndpoint& endpointB,
53                             MockNSTargeter* targeter) {
54     targeter->init(nss,
55                    {MockRange(endpointA, BSON("x" << MINKEY), BSON("x" << 0)),
56                     MockRange(endpointB, BSON("x" << 0), BSON("x" << MAXKEY))});
57 }
58 
initTargeterHalfRange(const NamespaceString & nss,const ShardEndpoint & endpoint,MockNSTargeter * targeter)59 void initTargeterHalfRange(const NamespaceString& nss,
60                            const ShardEndpoint& endpoint,
61                            MockNSTargeter* targeter) {
62     // x >= 0 values are untargetable
63     targeter->init(nss, {MockRange(endpoint, BSON("x" << MINKEY), BSON("x" << 0))});
64 }
65 
buildDelete(const BSONObj & query,bool multi)66 write_ops::DeleteOpEntry buildDelete(const BSONObj& query, bool multi) {
67     write_ops::DeleteOpEntry entry;
68     entry.setQ(query);
69     entry.setMulti(multi);
70     return entry;
71 }
72 
buildUpdate(const BSONObj & query,bool multi)73 write_ops::UpdateOpEntry buildUpdate(const BSONObj& query, bool multi) {
74     write_ops::UpdateOpEntry entry;
75     entry.setQ(query);
76     entry.setU(BSONObj());
77     entry.setMulti(multi);
78     return entry;
79 }
80 
buildUpdate(const BSONObj & query,const BSONObj & updateExpr,bool multi)81 write_ops::UpdateOpEntry buildUpdate(const BSONObj& query, const BSONObj& updateExpr, bool multi) {
82     write_ops::UpdateOpEntry entry;
83     entry.setQ(query);
84     entry.setU(updateExpr);
85     entry.setMulti(multi);
86     return entry;
87 }
88 
buildResponse(int n,BatchedCommandResponse * response)89 void buildResponse(int n, BatchedCommandResponse* response) {
90     response->clear();
91     response->setOk(true);
92     response->setN(n);
93     ASSERT(response->isValid(NULL));
94 }
95 
buildErrResponse(int code,const std::string & message,BatchedCommandResponse * response)96 void buildErrResponse(int code, const std::string& message, BatchedCommandResponse* response) {
97     response->clear();
98     response->setOk(false);
99     response->setN(0);
100     response->setErrCode(code);
101     response->setErrMessage(message);
102     ASSERT(response->isValid(NULL));
103 }
104 
addError(int code,const std::string & message,int index,BatchedCommandResponse * response)105 void addError(int code, const std::string& message, int index, BatchedCommandResponse* response) {
106     std::unique_ptr<WriteErrorDetail> error(new WriteErrorDetail);
107     error->setErrCode(code);
108     error->setErrMessage(message);
109     error->setIndex(index);
110 
111     response->addToErrDetails(error.release());
112 }
113 
addWCError(BatchedCommandResponse * response)114 void addWCError(BatchedCommandResponse* response) {
115     std::unique_ptr<WriteConcernErrorDetail> error(new WriteConcernErrorDetail);
116     error->setErrCode(ErrorCodes::WriteConcernFailed);
117     error->setErrMessage("mock wc error");
118 
119     response->setWriteConcernError(error.release());
120 }
121 
122 class WriteOpTestFixture : public unittest::Test {
123 protected:
operationContext()124     OperationContext* operationContext() {
125         return &_opCtx;
126     }
127 
128 private:
129     OperationContextNoop _opCtx;
130 };
131 
132 using BatchWriteOpTest = WriteOpTestFixture;
133 
TEST_F(BatchWriteOpTest,SingleOp)134 TEST_F(BatchWriteOpTest, SingleOp) {
135     NamespaceString nss("foo.bar");
136     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
137     MockNSTargeter targeter;
138     initTargeterFullRange(nss, endpoint, &targeter);
139 
140     // Do single-target, single doc batch write op
141     BatchedCommandRequest request([&] {
142         write_ops::Insert insertOp(nss);
143         insertOp.setDocuments({BSON("x" << 1)});
144         return insertOp;
145     }());
146 
147     BatchWriteOp batchOp(operationContext(), request);
148 
149     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
150     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
151     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
152     ASSERT(!batchOp.isFinished());
153     ASSERT_EQUALS(targeted.size(), 1u);
154     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpoint);
155 
156     BatchedCommandResponse response;
157     buildResponse(1, &response);
158 
159     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
160     ASSERT(batchOp.isFinished());
161 
162     BatchedCommandResponse clientResponse;
163     batchOp.buildClientResponse(&clientResponse);
164     ASSERT(clientResponse.getOk());
165 }
166 
TEST_F(BatchWriteOpTest,SingleError)167 TEST_F(BatchWriteOpTest, SingleError) {
168     NamespaceString nss("foo.bar");
169     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
170     MockNSTargeter targeter;
171     initTargeterFullRange(nss, endpoint, &targeter);
172 
173     // Do single-target, single doc batch write op
174     BatchedCommandRequest request([&] {
175         write_ops::Delete deleteOp(nss);
176         deleteOp.setDeletes({buildDelete(BSON("x" << 1), false)});
177         return deleteOp;
178     }());
179 
180     BatchWriteOp batchOp(operationContext(), request);
181 
182     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
183     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
184     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
185     ASSERT(!batchOp.isFinished());
186     ASSERT_EQUALS(targeted.size(), 1u);
187     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpoint);
188 
189     BatchedCommandResponse response;
190     buildErrResponse(ErrorCodes::UnknownError, "message", &response);
191 
192     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
193     ASSERT(batchOp.isFinished());
194 
195     BatchedCommandResponse clientResponse;
196     batchOp.buildClientResponse(&clientResponse);
197 
198     ASSERT(clientResponse.getOk());
199     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
200     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(), response.getErrCode());
201     ASSERT(clientResponse.getErrDetailsAt(0)->getErrMessage().find(response.getErrMessage()) !=
202            std::string::npos);
203     ASSERT_EQUALS(clientResponse.getN(), 0);
204 }
205 
TEST_F(BatchWriteOpTest,SingleTargetError)206 TEST_F(BatchWriteOpTest, SingleTargetError) {
207     NamespaceString nss("foo.bar");
208     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
209     MockNSTargeter targeter;
210     initTargeterHalfRange(nss, endpoint, &targeter);
211 
212     // Do untargetable delete op
213     BatchedCommandRequest request([&] {
214         write_ops::Delete deleteOp(nss);
215         deleteOp.setDeletes({buildDelete(BSON("x" << 1), false)});
216         return deleteOp;
217     }());
218 
219     BatchWriteOp batchOp(operationContext(), request);
220 
221     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
222     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
223     ASSERT_NOT_OK(batchOp.targetBatch(targeter, false, &targeted));
224     ASSERT(!batchOp.isFinished());
225     ASSERT_EQUALS(targeted.size(), 0u);
226 
227     // Record targeting failures
228     ASSERT_OK(batchOp.targetBatch(targeter, true, &targeted));
229     ASSERT(batchOp.isFinished());
230     ASSERT_EQUALS(targeted.size(), 0u);
231 
232     BatchedCommandResponse clientResponse;
233     batchOp.buildClientResponse(&clientResponse);
234     ASSERT(clientResponse.getOk());
235     ASSERT_EQUALS(clientResponse.getN(), 0);
236     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
237 }
238 
239 // Write concern error test - we should pass write concern to sub-batches, and pass up the write
240 // concern error if one occurs.
TEST_F(BatchWriteOpTest,SingleWriteConcernErrorOrdered)241 TEST_F(BatchWriteOpTest, SingleWriteConcernErrorOrdered) {
242     NamespaceString nss("foo.bar");
243     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
244     MockNSTargeter targeter;
245     initTargeterFullRange(nss, endpoint, &targeter);
246 
247     BatchedCommandRequest request([&] {
248         write_ops::Insert insertOp(nss);
249         insertOp.setDocuments({BSON("x" << 1)});
250         return insertOp;
251     }());
252     request.setWriteConcern(BSON("w" << 3));
253 
254     BatchWriteOp batchOp(operationContext(), request);
255 
256     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
257     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
258     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
259     ASSERT(!batchOp.isFinished());
260     ASSERT_EQUALS(targeted.size(), 1u);
261     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpoint);
262 
263     BatchedCommandRequest targetBatch = batchOp.buildBatchRequest(*targeted.begin()->second);
264     ASSERT(targetBatch.getWriteConcern().woCompare(request.getWriteConcern()) == 0);
265 
266     BatchedCommandResponse response;
267     buildResponse(1, &response);
268     addWCError(&response);
269 
270     // First stale response comes back, we should retry
271     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
272     ASSERT(batchOp.isFinished());
273 
274     BatchedCommandResponse clientResponse;
275     batchOp.buildClientResponse(&clientResponse);
276     ASSERT(clientResponse.getOk());
277     ASSERT_EQUALS(clientResponse.getN(), 1);
278     ASSERT(!clientResponse.isErrDetailsSet());
279     ASSERT(clientResponse.isWriteConcernErrorSet());
280 }
281 
282 // Single-op stale version test. We should retry the same batch until we're not stale.
TEST_F(BatchWriteOpTest,SingleStaleError)283 TEST_F(BatchWriteOpTest, SingleStaleError) {
284     NamespaceString nss("foo.bar");
285     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
286     MockNSTargeter targeter;
287     initTargeterFullRange(nss, endpoint, &targeter);
288 
289     BatchedCommandRequest request([&] {
290         write_ops::Insert insertOp(nss);
291         insertOp.setDocuments({BSON("x" << 1)});
292         return insertOp;
293     }());
294 
295     BatchWriteOp batchOp(operationContext(), request);
296 
297     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
298     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
299     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
300 
301     BatchedCommandResponse response;
302     buildResponse(0, &response);
303     addError(ErrorCodes::StaleShardVersion, "mock stale error", 0, &response);
304 
305     // First stale response comes back, we should retry
306     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
307     ASSERT(!batchOp.isFinished());
308 
309     targetedOwned.clear();
310     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
311 
312     // Respond again with a stale response
313     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
314     ASSERT(!batchOp.isFinished());
315 
316     targetedOwned.clear();
317     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
318 
319     buildResponse(1, &response);
320 
321     // Respond with an 'ok' response
322     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
323     ASSERT(batchOp.isFinished());
324 
325     BatchedCommandResponse clientResponse;
326     batchOp.buildClientResponse(&clientResponse);
327     ASSERT(clientResponse.getOk());
328     ASSERT_EQUALS(clientResponse.getN(), 1);
329     ASSERT(!clientResponse.isErrDetailsSet());
330 }
331 
332 //
333 // Multi-operation batches
334 //
335 
336 // Multi-op targeting test (ordered)
TEST_F(BatchWriteOpTest,MultiOpSameShardOrdered)337 TEST_F(BatchWriteOpTest, MultiOpSameShardOrdered) {
338     NamespaceString nss("foo.bar");
339     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
340     MockNSTargeter targeter;
341     initTargeterFullRange(nss, endpoint, &targeter);
342 
343     // Do single-target, multi-doc batch write op
344     BatchedCommandRequest request([&] {
345         write_ops::Update updateOp(nss);
346         updateOp.setUpdates(
347             {buildUpdate(BSON("x" << 1), false), buildUpdate(BSON("x" << 2), false)});
348         return updateOp;
349     }());
350 
351     BatchWriteOp batchOp(operationContext(), request);
352 
353     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
354     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
355     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
356     ASSERT(!batchOp.isFinished());
357     ASSERT_EQUALS(targeted.size(), 1u);
358     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 2u);
359     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpoint);
360 
361     BatchedCommandResponse response;
362     buildResponse(2, &response);
363 
364     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
365     ASSERT(batchOp.isFinished());
366 
367     BatchedCommandResponse clientResponse;
368     batchOp.buildClientResponse(&clientResponse);
369     ASSERT(clientResponse.getOk());
370     ASSERT_EQUALS(clientResponse.getN(), 2);
371 }
372 
373 // Multi-op targeting test (unordered)
TEST_F(BatchWriteOpTest,MultiOpSameShardUnordered)374 TEST_F(BatchWriteOpTest, MultiOpSameShardUnordered) {
375     NamespaceString nss("foo.bar");
376     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
377     MockNSTargeter targeter;
378     initTargeterFullRange(nss, endpoint, &targeter);
379 
380     // Do single-target, multi-doc batch write op
381     BatchedCommandRequest request([&] {
382         write_ops::Update updateOp(nss);
383         updateOp.setWriteCommandBase([] {
384             write_ops::WriteCommandBase wcb;
385             wcb.setOrdered(false);
386             return wcb;
387         }());
388         updateOp.setUpdates(
389             {buildUpdate(BSON("x" << 1), false), buildUpdate(BSON("x" << 2), false)});
390         return updateOp;
391     }());
392 
393     BatchWriteOp batchOp(operationContext(), request);
394 
395     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
396     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
397     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
398     ASSERT(!batchOp.isFinished());
399     ASSERT_EQUALS(targeted.size(), 1u);
400     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 2u);
401     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpoint);
402 
403     BatchedCommandResponse response;
404     buildResponse(2, &response);
405 
406     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
407     ASSERT(batchOp.isFinished());
408 
409     BatchedCommandResponse clientResponse;
410     batchOp.buildClientResponse(&clientResponse);
411     ASSERT(clientResponse.getOk());
412     ASSERT_EQUALS(clientResponse.getN(), 2);
413 }
414 
415 // Multi-op, multi-endpoing targeting test (ordered). There should be two sets of single batches
416 // (one to each shard, one-by-one)
TEST_F(BatchWriteOpTest,MultiOpTwoShardsOrdered)417 TEST_F(BatchWriteOpTest, MultiOpTwoShardsOrdered) {
418     NamespaceString nss("foo.bar");
419     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
420     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
421     MockNSTargeter targeter;
422     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
423 
424     // Do multi-target, multi-doc batch write op
425     BatchedCommandRequest request([&] {
426         write_ops::Insert insertOp(nss);
427         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 1)});
428         return insertOp;
429     }());
430 
431     BatchWriteOp batchOp(operationContext(), request);
432 
433     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
434     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
435     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
436     ASSERT(!batchOp.isFinished());
437     ASSERT_EQUALS(targeted.size(), 1u);
438     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 1u);
439     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpointA);
440 
441     BatchedCommandResponse response;
442     buildResponse(1, &response);
443 
444     // Respond to first targeted batch
445     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
446     ASSERT(!batchOp.isFinished());
447 
448     targetedOwned.clear();
449 
450     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
451     ASSERT(!batchOp.isFinished());
452     ASSERT_EQUALS(targeted.size(), 1u);
453     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 1u);
454     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpointB);
455 
456     // Respond to second targeted batch
457     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
458     ASSERT(batchOp.isFinished());
459 
460     BatchedCommandResponse clientResponse;
461     batchOp.buildClientResponse(&clientResponse);
462     ASSERT(clientResponse.getOk());
463     ASSERT_EQUALS(clientResponse.getN(), 2);
464 }
465 
verifyTargetedBatches(std::map<ShardId,size_t> expected,const std::map<ShardId,TargetedWriteBatch * > & targeted)466 void verifyTargetedBatches(std::map<ShardId, size_t> expected,
467                            const std::map<ShardId, TargetedWriteBatch*>& targeted) {
468     // 'expected' contains each ShardId that was expected to be targeted and the size of the batch
469     // that was expected to be targeted to it.
470     // We check that each ShardId in 'targeted' corresponds to one in 'expected', in that it
471     // contains a batch of the correct size.
472     // Finally, we ensure that no additional ShardIds are present in 'targeted' than 'expected'.
473     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
474         ASSERT_EQUALS(expected[it->second->getEndpoint().shardName],
475                       it->second->getWrites().size());
476         ASSERT_EQUALS(ChunkVersion::IGNORED(), it->second->getEndpoint().shardVersion);
477         expected.erase(expected.find(it->second->getEndpoint().shardName));
478     }
479     ASSERT(expected.empty());
480 }
481 
482 // Multi-op, multi-endpoint targeting test (unordered). There should be one set of two batches (one
483 // to each shard).
TEST_F(BatchWriteOpTest,MultiOpTwoShardsUnordered)484 TEST_F(BatchWriteOpTest, MultiOpTwoShardsUnordered) {
485     NamespaceString nss("foo.bar");
486     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
487     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
488     MockNSTargeter targeter;
489     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
490 
491     // Do multi-target, multi-doc batch write op
492     BatchedCommandRequest request([&] {
493         write_ops::Insert insertOp(nss);
494         insertOp.setWriteCommandBase([] {
495             write_ops::WriteCommandBase wcb;
496             wcb.setOrdered(false);
497             return wcb;
498         }());
499         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 1)});
500         return insertOp;
501     }());
502 
503     BatchWriteOp batchOp(operationContext(), request);
504 
505     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
506     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
507     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
508     ASSERT(!batchOp.isFinished());
509     ASSERT_EQUALS(targeted.size(), 2u);
510     verifyTargetedBatches({{endpointA.shardName, 1u}, {endpointB.shardName, 1u}}, targeted);
511 
512     BatchedCommandResponse response;
513     buildResponse(1, &response);
514 
515     // Respond to both targeted batches
516     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
517         ASSERT(!batchOp.isFinished());
518         batchOp.noteBatchResponse(*it->second, response, NULL);
519     }
520     ASSERT(batchOp.isFinished());
521 
522     BatchedCommandResponse clientResponse;
523     batchOp.buildClientResponse(&clientResponse);
524     ASSERT(clientResponse.getOk());
525     ASSERT_EQUALS(clientResponse.getN(), 2);
526 }
527 
528 // Multi-op (ordered) targeting test where each op goes to both shards. There should be two sets of
529 // two batches to each shard (two for each delete op).
TEST_F(BatchWriteOpTest,MultiOpTwoShardsEachOrdered)530 TEST_F(BatchWriteOpTest, MultiOpTwoShardsEachOrdered) {
531     NamespaceString nss("foo.bar");
532     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
533     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
534     MockNSTargeter targeter;
535     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
536 
537     // Do multi-target, multi-doc batch write op
538     BatchedCommandRequest request([&] {
539         write_ops::Delete deleteOp(nss);
540         deleteOp.setDeletes({buildDelete(BSON("x" << GTE << -1 << LT << 2), true),
541                              buildDelete(BSON("x" << GTE << -2 << LT << 1), true)});
542         return deleteOp;
543     }());
544 
545     BatchWriteOp batchOp(operationContext(), request);
546 
547     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
548     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
549     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
550     ASSERT(!batchOp.isFinished());
551     ASSERT_EQUALS(targeted.size(), 2u);
552     verifyTargetedBatches({{endpointA.shardName, 1u}, {endpointB.shardName, 1u}}, targeted);
553 
554     BatchedCommandResponse response;
555     buildResponse(1, &response);
556 
557     // Respond to both targeted batches for first multi-delete
558     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
559         ASSERT(!batchOp.isFinished());
560         batchOp.noteBatchResponse(*it->second, response, NULL);
561     }
562     ASSERT(!batchOp.isFinished());
563 
564     targetedOwned.clear();
565 
566     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
567     ASSERT(!batchOp.isFinished());
568     ASSERT_EQUALS(targeted.size(), 2u);
569     verifyTargetedBatches({{endpointA.shardName, 1u}, {endpointB.shardName, 1u}}, targeted);
570 
571     // Respond to second targeted batches for second multi-delete
572     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
573         ASSERT(!batchOp.isFinished());
574         batchOp.noteBatchResponse(*it->second, response, NULL);
575     }
576     ASSERT(batchOp.isFinished());
577 
578     BatchedCommandResponse clientResponse;
579     batchOp.buildClientResponse(&clientResponse);
580     ASSERT(clientResponse.getOk());
581     ASSERT_EQUALS(clientResponse.getN(), 4);
582 }
583 
584 // Multi-op (unaordered) targeting test where each op goes to both shards. There should be one set
585 // of two batches to each shard (containing writes for both ops).
TEST_F(BatchWriteOpTest,MultiOpTwoShardsEachUnordered)586 TEST_F(BatchWriteOpTest, MultiOpTwoShardsEachUnordered) {
587     NamespaceString nss("foo.bar");
588     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
589     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
590     MockNSTargeter targeter;
591     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
592 
593     // Do multi-target, multi-doc batch write op
594     BatchedCommandRequest request([&] {
595         write_ops::Update updateOp(nss);
596         updateOp.setWriteCommandBase([] {
597             write_ops::WriteCommandBase wcb;
598             wcb.setOrdered(false);
599             return wcb;
600         }());
601         updateOp.setUpdates({buildUpdate(BSON("x" << GTE << -1 << LT << 2), true),
602                              buildUpdate(BSON("x" << GTE << -2 << LT << 1), true)});
603         return updateOp;
604     }());
605 
606     BatchWriteOp batchOp(operationContext(), request);
607 
608     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
609     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
610     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
611     ASSERT(!batchOp.isFinished());
612     ASSERT_EQUALS(targeted.size(), 2u);
613     verifyTargetedBatches({{endpointA.shardName, 2u}, {endpointB.shardName, 2u}}, targeted);
614 
615     BatchedCommandResponse response;
616     buildResponse(2, &response);
617 
618     // Respond to both targeted batches, each containing two ops
619     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
620         ASSERT(!batchOp.isFinished());
621         batchOp.noteBatchResponse(*it->second, response, NULL);
622     }
623     ASSERT(batchOp.isFinished());
624 
625     BatchedCommandResponse clientResponse;
626     batchOp.buildClientResponse(&clientResponse);
627     ASSERT(clientResponse.getOk());
628     ASSERT_EQUALS(clientResponse.getN(), 4);
629 }
630 
631 // Multi-op (ordered) targeting test where first two ops go to one shard, second two ops go to two
632 // shards. Should batch the first two ops, then second ops should be batched separately, then last
633 // ops should be batched together.
TEST_F(BatchWriteOpTest,MultiOpOneOrTwoShardsOrdered)634 TEST_F(BatchWriteOpTest, MultiOpOneOrTwoShardsOrdered) {
635     NamespaceString nss("foo.bar");
636     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
637     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
638     MockNSTargeter targeter;
639     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
640 
641     BatchedCommandRequest request([&] {
642         write_ops::Delete deleteOp(nss);
643         deleteOp.setDeletes({
644             // These go to the same shard
645             buildDelete(BSON("x" << -1), false),
646             buildDelete(BSON("x" << -2), false),
647             // These go to both shards
648             buildDelete(BSON("x" << GTE << -1 << LT << 2), true),
649             buildDelete(BSON("x" << GTE << -2 << LT << 1), true),
650             // These go to the same shard
651             buildDelete(BSON("x" << 1), false),
652             buildDelete(BSON("x" << 2), false),
653         });
654         return deleteOp;
655     }());
656 
657     BatchWriteOp batchOp(operationContext(), request);
658 
659     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
660     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
661     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
662     ASSERT(!batchOp.isFinished());
663     ASSERT_EQUALS(targeted.size(), 1u);
664     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 2u);
665     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpointA);
666 
667     BatchedCommandResponse response;
668     // Emulate one-write-per-delete-per-host
669     buildResponse(2, &response);
670 
671     // Respond to first targeted batch containing the two single-host deletes
672     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
673     ASSERT(!batchOp.isFinished());
674 
675     targetedOwned.clear();
676 
677     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
678     ASSERT(!batchOp.isFinished());
679     ASSERT_EQUALS(targeted.size(), 2u);
680     verifyTargetedBatches({{endpointA.shardName, 1u}, {endpointB.shardName, 1u}}, targeted);
681 
682     // Emulate one-write-per-delete-per-host
683     buildResponse(1, &response);
684 
685     // Respond to two targeted batches for first multi-delete
686     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
687         ASSERT(!batchOp.isFinished());
688         batchOp.noteBatchResponse(*it->second, response, NULL);
689     }
690     ASSERT(!batchOp.isFinished());
691 
692     targetedOwned.clear();
693 
694     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
695     ASSERT(!batchOp.isFinished());
696     ASSERT_EQUALS(targeted.size(), 2u);
697     verifyTargetedBatches({{endpointA.shardName, 1u}, {endpointB.shardName, 1u}}, targeted);
698 
699     // Respond to two targeted batches for second multi-delete
700     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
701         ASSERT(!batchOp.isFinished());
702         batchOp.noteBatchResponse(*it->second, response, NULL);
703     }
704     ASSERT(!batchOp.isFinished());
705 
706     targetedOwned.clear();
707 
708     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
709     ASSERT(!batchOp.isFinished());
710     ASSERT_EQUALS(targeted.size(), 1u);
711     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 2u);
712     assertEndpointsEqual(targeted.begin()->second->getEndpoint(), endpointB);
713 
714     // Emulate one-write-per-delete-per-host
715     buildResponse(2, &response);
716 
717     // Respond to final targeted batch containing the last two single-host deletes
718     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
719     ASSERT(batchOp.isFinished());
720 
721     BatchedCommandResponse clientResponse;
722     batchOp.buildClientResponse(&clientResponse);
723     ASSERT(clientResponse.getOk());
724     ASSERT_EQUALS(clientResponse.getN(), 8);
725 }
726 
727 // Multi-op (unordered) targeting test where first two ops go to one shard, second two ops go to two
728 // shards. Should batch all the ops together into two batches of four ops for each shard.
TEST_F(BatchWriteOpTest,MultiOpOneOrTwoShardsUnordered)729 TEST_F(BatchWriteOpTest, MultiOpOneOrTwoShardsUnordered) {
730     NamespaceString nss("foo.bar");
731     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
732     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
733     MockNSTargeter targeter;
734     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
735 
736     BatchedCommandRequest request([&] {
737         write_ops::Update updateOp(nss);
738         updateOp.setWriteCommandBase([] {
739             write_ops::WriteCommandBase wcb;
740             wcb.setOrdered(false);
741             return wcb;
742         }());
743         updateOp.setUpdates({// These go to the same shard
744                              buildUpdate(BSON("x" << -1), false),
745                              buildUpdate(BSON("x" << -2), false),
746                              // These go to both shards
747                              buildUpdate(BSON("x" << GTE << -1 << LT << 2), true),
748                              buildUpdate(BSON("x" << GTE << -2 << LT << 1), true),
749                              // These go to the same shard
750                              buildUpdate(BSON("x" << 1), false),
751                              buildUpdate(BSON("x" << 2), false)});
752         return updateOp;
753     }());
754 
755     BatchWriteOp batchOp(operationContext(), request);
756 
757     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
758     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
759     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
760     ASSERT(!batchOp.isFinished());
761     ASSERT_EQUALS(targeted.size(), 2u);
762     verifyTargetedBatches({{endpointA.shardName, 4u}, {endpointB.shardName, 4u}}, targeted);
763 
764     BatchedCommandResponse response;
765     // Emulate one-write-per-delete-per-host
766     buildResponse(4, &response);
767 
768     // Respond to first targeted batch containing the two single-host deletes
769     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
770         ASSERT(!batchOp.isFinished());
771         batchOp.noteBatchResponse(*it->second, response, NULL);
772     }
773     ASSERT(batchOp.isFinished());
774 
775     BatchedCommandResponse clientResponse;
776     batchOp.buildClientResponse(&clientResponse);
777     ASSERT(clientResponse.getOk());
778     ASSERT_EQUALS(clientResponse.getN(), 8);
779 }
780 
781 // Multi-op targeting test where two ops go to two separate shards and there's an error on one op on
782 // one shard. There should be one set of two batches to each shard and an error reported.
TEST_F(BatchWriteOpTest,MultiOpSingleShardErrorUnordered)783 TEST_F(BatchWriteOpTest, MultiOpSingleShardErrorUnordered) {
784     NamespaceString nss("foo.bar");
785     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
786     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
787     MockNSTargeter targeter;
788     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
789 
790     BatchedCommandRequest request([&] {
791         write_ops::Insert insertOp(nss);
792         insertOp.setWriteCommandBase([] {
793             write_ops::WriteCommandBase wcb;
794             wcb.setOrdered(false);
795             return wcb;
796         }());
797         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 1)});
798         return insertOp;
799     }());
800 
801     BatchWriteOp batchOp(operationContext(), request);
802 
803     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
804     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
805     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
806     ASSERT(!batchOp.isFinished());
807     ASSERT_EQUALS(targeted.size(), 2u);
808     verifyTargetedBatches({{endpointA.shardName, 1u}, {endpointB.shardName, 1u}}, targeted);
809 
810     BatchedCommandResponse response;
811     buildResponse(1, &response);
812 
813     // Respond to batches.
814     auto targetedIt = targeted.begin();
815 
816     // No error on first shard
817     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
818     ASSERT(!batchOp.isFinished());
819 
820     buildResponse(0, &response);
821     addError(ErrorCodes::UnknownError, "mock error", 0, &response);
822 
823     // Error on second write on second shard
824     ++targetedIt;
825     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
826     ASSERT(batchOp.isFinished());
827     ASSERT(++targetedIt == targeted.end());
828 
829     BatchedCommandResponse clientResponse;
830     batchOp.buildClientResponse(&clientResponse);
831     ASSERT(clientResponse.getOk());
832     ASSERT_EQUALS(clientResponse.getN(), 1);
833     ASSERT(clientResponse.isErrDetailsSet());
834     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
835     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(),
836                   response.getErrDetailsAt(0)->getErrCode());
837     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrMessage(),
838                   response.getErrDetailsAt(0)->getErrMessage());
839     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 1);
840 }
841 
842 // Multi-op targeting test where two ops go to two separate shards and there's an error on each op
843 // on each shard. There should be one set of two batches to each shard and and two errors reported.
TEST_F(BatchWriteOpTest,MultiOpTwoShardErrorsUnordered)844 TEST_F(BatchWriteOpTest, MultiOpTwoShardErrorsUnordered) {
845     NamespaceString nss("foo.bar");
846     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
847     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
848     MockNSTargeter targeter;
849     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
850 
851     BatchedCommandRequest request([&] {
852         write_ops::Insert insertOp(nss);
853         insertOp.setWriteCommandBase([] {
854             write_ops::WriteCommandBase wcb;
855             wcb.setOrdered(false);
856             return wcb;
857         }());
858         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 1)});
859         return insertOp;
860     }());
861 
862     BatchWriteOp batchOp(operationContext(), request);
863 
864     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
865     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
866     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
867     ASSERT(!batchOp.isFinished());
868     ASSERT_EQUALS(targeted.size(), 2u);
869     verifyTargetedBatches({{endpointA.shardName, 1u}, {endpointB.shardName, 1u}}, targeted);
870 
871     BatchedCommandResponse response;
872     buildResponse(0, &response);
873     addError(ErrorCodes::UnknownError, "mock error", 0, &response);
874 
875     // Error on first write on first shard and second write on second shard.
876     for (auto it = targeted.begin(); it != targeted.end(); ++it) {
877         ASSERT(!batchOp.isFinished());
878         batchOp.noteBatchResponse(*it->second, response, NULL);
879     }
880     ASSERT(batchOp.isFinished());
881 
882     BatchedCommandResponse clientResponse;
883     batchOp.buildClientResponse(&clientResponse);
884     ASSERT(clientResponse.getOk());
885     ASSERT_EQUALS(clientResponse.getN(), 0);
886     ASSERT(clientResponse.isErrDetailsSet());
887     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 2u);
888     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(),
889                   response.getErrDetailsAt(0)->getErrCode());
890     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrMessage(),
891                   response.getErrDetailsAt(0)->getErrMessage());
892     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 0);
893     ASSERT_EQUALS(clientResponse.getErrDetailsAt(1)->getErrCode(),
894                   response.getErrDetailsAt(0)->getErrCode());
895     ASSERT_EQUALS(clientResponse.getErrDetailsAt(1)->getErrMessage(),
896                   response.getErrDetailsAt(0)->getErrMessage());
897     ASSERT_EQUALS(clientResponse.getErrDetailsAt(1)->getIndex(), 1);
898 }
899 
900 // Multi-op targeting test where each op goes to both shards and there's an error on one op on one
901 // shard. There should be one set of two batches to each shard and an error reported.
TEST_F(BatchWriteOpTest,MultiOpPartialSingleShardErrorUnordered)902 TEST_F(BatchWriteOpTest, MultiOpPartialSingleShardErrorUnordered) {
903     NamespaceString nss("foo.bar");
904     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
905     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
906     MockNSTargeter targeter;
907     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
908 
909     BatchedCommandRequest request([&] {
910         write_ops::Delete deleteOp(nss);
911         deleteOp.setWriteCommandBase([] {
912             write_ops::WriteCommandBase wcb;
913             wcb.setOrdered(false);
914             return wcb;
915         }());
916         deleteOp.setDeletes({buildDelete(BSON("x" << GTE << -1 << LT << 2), true),
917                              buildDelete(BSON("x" << GTE << -2 << LT << 1), true)});
918         return deleteOp;
919     }());
920 
921     BatchWriteOp batchOp(operationContext(), request);
922 
923     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
924     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
925     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
926     ASSERT(!batchOp.isFinished());
927     ASSERT_EQUALS(targeted.size(), 2u);
928     verifyTargetedBatches({{endpointA.shardName, 2u}, {endpointB.shardName, 2u}}, targeted);
929 
930     // Respond to batches.
931     auto targetedIt = targeted.begin();
932 
933     BatchedCommandResponse response;
934     buildResponse(2, &response);
935 
936     // No errors on first shard
937     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
938     ASSERT(!batchOp.isFinished());
939 
940     buildResponse(1, &response);
941     addError(ErrorCodes::UnknownError, "mock error", 1, &response);
942 
943     // Error on second write on second shard
944     ++targetedIt;
945     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
946     ASSERT(batchOp.isFinished());
947     ASSERT(++targetedIt == targeted.end());
948 
949     BatchedCommandResponse clientResponse;
950     batchOp.buildClientResponse(&clientResponse);
951     ASSERT(clientResponse.getOk());
952     ASSERT_EQUALS(clientResponse.getN(), 3);
953     ASSERT(clientResponse.isErrDetailsSet());
954     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
955     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(),
956                   response.getErrDetailsAt(0)->getErrCode());
957     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrMessage(),
958                   response.getErrDetailsAt(0)->getErrMessage());
959     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 1);
960 }
961 
962 // Multi-op targeting test where each op goes to both shards and there's an error on one op on one
963 // shard. There should be one set of two batches to each shard and an error reported, the second op
964 // should not get run.
TEST_F(BatchWriteOpTest,MultiOpPartialSingleShardErrorOrdered)965 TEST_F(BatchWriteOpTest, MultiOpPartialSingleShardErrorOrdered) {
966     NamespaceString nss("foo.bar");
967     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
968     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
969     MockNSTargeter targeter;
970     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
971 
972     BatchedCommandRequest request([&] {
973         write_ops::Delete deleteOp(nss);
974         deleteOp.setDeletes({buildDelete(BSON("x" << GTE << -1 << LT << 2), true),
975                              buildDelete(BSON("x" << GTE << -2 << LT << 1), true)});
976         return deleteOp;
977     }());
978 
979     BatchWriteOp batchOp(operationContext(), request);
980 
981     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
982     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
983     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
984     ASSERT(!batchOp.isFinished());
985     ASSERT_EQUALS(targeted.size(), 2u);
986     verifyTargetedBatches({{endpointA.shardName, 1u}, {endpointB.shardName, 1u}}, targeted);
987 
988     // Respond to batches.
989     auto targetedIt = targeted.begin();
990 
991     BatchedCommandResponse response;
992     buildResponse(1, &response);
993 
994     // No errors on first shard
995     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
996     ASSERT(!batchOp.isFinished());
997 
998     buildResponse(0, &response);
999     addError(ErrorCodes::UnknownError, "mock error", 0, &response);
1000 
1001     // Error on second write on second shard
1002     ++targetedIt;
1003     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
1004     ASSERT(batchOp.isFinished());
1005     ASSERT(++targetedIt == targeted.end());
1006 
1007     BatchedCommandResponse clientResponse;
1008     batchOp.buildClientResponse(&clientResponse);
1009     ASSERT(clientResponse.getOk());
1010     ASSERT_EQUALS(clientResponse.getN(), 1);
1011     ASSERT(clientResponse.isErrDetailsSet());
1012     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
1013     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(),
1014                   response.getErrDetailsAt(0)->getErrCode());
1015     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrMessage(),
1016                   response.getErrDetailsAt(0)->getErrMessage());
1017     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 0);
1018 }
1019 
1020 //
1021 // Tests of edge-case functionality, lifecycle is assumed to be behaving normally
1022 //
1023 
1024 // Multi-op (unordered) error and write concern error test. We never report the write concern error
1025 // for single-doc batches, since the error means there's no write concern applied. Don't suppress
1026 // the error if ordered : false.
TEST_F(BatchWriteOpTest,MultiOpErrorAndWriteConcernErrorUnordered)1027 TEST_F(BatchWriteOpTest, MultiOpErrorAndWriteConcernErrorUnordered) {
1028     NamespaceString nss("foo.bar");
1029     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
1030     MockNSTargeter targeter;
1031     initTargeterFullRange(nss, endpoint, &targeter);
1032 
1033     BatchedCommandRequest request([&] {
1034         write_ops::Insert insertOp(nss);
1035         insertOp.setWriteCommandBase([] {
1036             write_ops::WriteCommandBase wcb;
1037             wcb.setOrdered(false);
1038             return wcb;
1039         }());
1040         insertOp.setDocuments({BSON("x" << 1), BSON("x" << 1)});
1041         return insertOp;
1042     }());
1043     request.setWriteConcern(BSON("w" << 3));
1044 
1045     BatchWriteOp batchOp(operationContext(), request);
1046 
1047     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1048     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1049     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1050 
1051     BatchedCommandResponse response;
1052     buildResponse(1, &response);
1053     addError(ErrorCodes::UnknownError, "mock error", 1, &response);
1054     addWCError(&response);
1055 
1056     // First stale response comes back, we should retry
1057     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1058     ASSERT(batchOp.isFinished());
1059 
1060     // Unordered reports write concern error
1061     BatchedCommandResponse clientResponse;
1062     batchOp.buildClientResponse(&clientResponse);
1063     ASSERT(clientResponse.getOk());
1064     ASSERT_EQUALS(clientResponse.getN(), 1);
1065     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
1066     ASSERT(clientResponse.isWriteConcernErrorSet());
1067 }
1068 
1069 // Single-op (ordered) error and write concern error test. Suppress the write concern error if
1070 // ordered and we also have an error
TEST_F(BatchWriteOpTest,SingleOpErrorAndWriteConcernErrorOrdered)1071 TEST_F(BatchWriteOpTest, SingleOpErrorAndWriteConcernErrorOrdered) {
1072     NamespaceString nss("foo.bar");
1073     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
1074     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
1075     MockNSTargeter targeter;
1076     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
1077 
1078     BatchedCommandRequest request([&] {
1079         write_ops::Update updateOp(nss);
1080         updateOp.setWriteCommandBase([] {
1081             write_ops::WriteCommandBase wcb;
1082             wcb.setOrdered(false);
1083             return wcb;
1084         }());
1085         updateOp.setUpdates({buildUpdate(BSON("x" << GTE << -1 << LT << 2), true)});
1086         return updateOp;
1087     }());
1088     request.setWriteConcern(BSON("w" << 3));
1089 
1090     BatchWriteOp batchOp(operationContext(), request);
1091 
1092     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1093     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1094     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1095 
1096     // Respond to batches.
1097     auto targetedIt = targeted.begin();
1098 
1099     BatchedCommandResponse response;
1100     buildResponse(1, &response);
1101     addWCError(&response);
1102 
1103     // First response comes back with write concern error
1104     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
1105     ASSERT(!batchOp.isFinished());
1106 
1107     buildResponse(0, &response);
1108     addError(ErrorCodes::UnknownError, "mock error", 0, &response);
1109 
1110     // Second response comes back with write error
1111     ++targetedIt;
1112     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
1113     ASSERT(batchOp.isFinished());
1114     ASSERT(++targetedIt == targeted.end());
1115 
1116     // Ordered doesn't report write concern error
1117     BatchedCommandResponse clientResponse;
1118     batchOp.buildClientResponse(&clientResponse);
1119     ASSERT(clientResponse.getOk());
1120     ASSERT_EQUALS(clientResponse.getN(), 1);
1121     ASSERT(clientResponse.isErrDetailsSet());
1122     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
1123     ASSERT(!clientResponse.isWriteConcernErrorSet());
1124 }
1125 
1126 // Targeting failure on second op in batch op (ordered)
TEST_F(BatchWriteOpTest,MultiOpFailedTargetOrdered)1127 TEST_F(BatchWriteOpTest, MultiOpFailedTargetOrdered) {
1128     NamespaceString nss("foo.bar");
1129     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
1130     MockNSTargeter targeter;
1131     initTargeterHalfRange(nss, endpoint, &targeter);
1132 
1133     BatchedCommandRequest request([&] {
1134         write_ops::Insert insertOp(nss);
1135         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 2), BSON("x" << -2)});
1136         return insertOp;
1137     }());
1138 
1139     // Do single-target, multi-doc batch write op
1140 
1141     BatchWriteOp batchOp(operationContext(), request);
1142 
1143     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1144     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1145     ASSERT_NOT_OK(batchOp.targetBatch(targeter, false, &targeted));
1146 
1147     // First targeting round fails since we may be stale
1148     ASSERT(!batchOp.isFinished());
1149 
1150     targetedOwned.clear();
1151     ASSERT_OK(batchOp.targetBatch(targeter, true, &targeted));
1152 
1153     // Second targeting round is ok, but should stop at first write
1154     ASSERT(!batchOp.isFinished());
1155     ASSERT_EQUALS(targeted.size(), 1u);
1156     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 1u);
1157 
1158     BatchedCommandResponse response;
1159     buildResponse(1, &response);
1160 
1161     // First response ok
1162     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1163     ASSERT(!batchOp.isFinished());
1164 
1165     targetedOwned.clear();
1166     ASSERT_OK(batchOp.targetBatch(targeter, true, &targeted));
1167 
1168     // Second targeting round results in an error which finishes the batch
1169     ASSERT(batchOp.isFinished());
1170     ASSERT_EQUALS(targeted.size(), 0u);
1171 
1172     BatchedCommandResponse clientResponse;
1173     batchOp.buildClientResponse(&clientResponse);
1174     ASSERT(clientResponse.getOk());
1175     ASSERT_EQUALS(clientResponse.getN(), 1);
1176     ASSERT(clientResponse.isErrDetailsSet());
1177     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
1178     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 1);
1179 }
1180 
1181 // Targeting failure on second op in batch op (unordered)
TEST_F(BatchWriteOpTest,MultiOpFailedTargetUnordered)1182 TEST_F(BatchWriteOpTest, MultiOpFailedTargetUnordered) {
1183     NamespaceString nss("foo.bar");
1184     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
1185     MockNSTargeter targeter;
1186     initTargeterHalfRange(nss, endpoint, &targeter);
1187 
1188     BatchedCommandRequest request([&] {
1189         write_ops::Insert insertOp(nss);
1190         insertOp.setWriteCommandBase([] {
1191             write_ops::WriteCommandBase wcb;
1192             wcb.setOrdered(false);
1193             return wcb;
1194         }());
1195         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 2), BSON("x" << -2)});
1196         return insertOp;
1197     }());
1198 
1199     // Do single-target, multi-doc batch write op
1200 
1201     BatchWriteOp batchOp(operationContext(), request);
1202 
1203     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1204     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1205     ASSERT_NOT_OK(batchOp.targetBatch(targeter, false, &targeted));
1206 
1207     // First targeting round fails since we may be stale
1208     ASSERT(!batchOp.isFinished());
1209 
1210     targetedOwned.clear();
1211     ASSERT_OK(batchOp.targetBatch(targeter, true, &targeted));
1212 
1213     // Second targeting round is ok, and should record an error
1214     ASSERT(!batchOp.isFinished());
1215     ASSERT_EQUALS(targeted.size(), 1u);
1216     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 2u);
1217 
1218     BatchedCommandResponse response;
1219     buildResponse(2, &response);
1220 
1221     // Response is ok for first and third write
1222     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1223     ASSERT(batchOp.isFinished());
1224 
1225     BatchedCommandResponse clientResponse;
1226     batchOp.buildClientResponse(&clientResponse);
1227     ASSERT(clientResponse.getOk());
1228     ASSERT_EQUALS(clientResponse.getN(), 2);
1229     ASSERT(clientResponse.isErrDetailsSet());
1230     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
1231     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 1);
1232 }
1233 
1234 // Batch failure (ok : 0) reported in a multi-op batch (ordered). Expect this gets translated down
1235 // into write errors for first affected write.
TEST_F(BatchWriteOpTest,MultiOpFailedBatchOrdered)1236 TEST_F(BatchWriteOpTest, MultiOpFailedBatchOrdered) {
1237     NamespaceString nss("foo.bar");
1238     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
1239     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
1240     MockNSTargeter targeter;
1241     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
1242 
1243     BatchedCommandRequest request([&] {
1244         write_ops::Insert insertOp(nss);
1245         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 2), BSON("x" << 3)});
1246         return insertOp;
1247     }());
1248 
1249     BatchWriteOp batchOp(operationContext(), request);
1250 
1251     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1252     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1253     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1254 
1255     BatchedCommandResponse response;
1256     buildResponse(1, &response);
1257 
1258     // First shard batch is ok
1259     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1260     ASSERT(!batchOp.isFinished());
1261 
1262     targetedOwned.clear();
1263     ASSERT_OK(batchOp.targetBatch(targeter, true, &targeted));
1264 
1265     buildErrResponse(ErrorCodes::UnknownError, "mock error", &response);
1266 
1267     // Second shard batch fails
1268     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1269     ASSERT(batchOp.isFinished());
1270 
1271     // We should have recorded an error for the second write
1272     BatchedCommandResponse clientResponse;
1273     batchOp.buildClientResponse(&clientResponse);
1274     ASSERT(clientResponse.getOk());
1275     ASSERT_EQUALS(clientResponse.getN(), 1);
1276     ASSERT(clientResponse.isErrDetailsSet());
1277     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
1278     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 1);
1279     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(), response.getErrCode());
1280 }
1281 
1282 // Batch failure (ok : 0) reported in a multi-op batch (unordered). Expect this gets translated down
1283 // into write errors for all affected writes.
TEST_F(BatchWriteOpTest,MultiOpFailedBatchUnordered)1284 TEST_F(BatchWriteOpTest, MultiOpFailedBatchUnordered) {
1285     NamespaceString nss("foo.bar");
1286     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
1287     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
1288     MockNSTargeter targeter;
1289     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
1290 
1291     BatchedCommandRequest request([&] {
1292         write_ops::Insert insertOp(nss);
1293         insertOp.setWriteCommandBase([] {
1294             write_ops::WriteCommandBase wcb;
1295             wcb.setOrdered(false);
1296             return wcb;
1297         }());
1298         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 2), BSON("x" << 3)});
1299         return insertOp;
1300     }());
1301 
1302     BatchWriteOp batchOp(operationContext(), request);
1303 
1304     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1305     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1306     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1307 
1308     // Respond to batches.
1309     auto targetedIt = targeted.begin();
1310 
1311     BatchedCommandResponse response;
1312     buildResponse(1, &response);
1313 
1314     // First shard batch is ok
1315     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
1316     ASSERT(!batchOp.isFinished());
1317 
1318     buildErrResponse(ErrorCodes::UnknownError, "mock error", &response);
1319 
1320     // Second shard batch fails
1321     ++targetedIt;
1322     batchOp.noteBatchResponse(*targetedIt->second, response, NULL);
1323     ASSERT(batchOp.isFinished());
1324     ASSERT(++targetedIt == targeted.end());
1325 
1326     // We should have recorded an error for the second and third write
1327     BatchedCommandResponse clientResponse;
1328     batchOp.buildClientResponse(&clientResponse);
1329     ASSERT(clientResponse.getOk());
1330     ASSERT_EQUALS(clientResponse.getN(), 1);
1331     ASSERT(clientResponse.isErrDetailsSet());
1332     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 2u);
1333     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 1);
1334     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(), response.getErrCode());
1335     ASSERT_EQUALS(clientResponse.getErrDetailsAt(1)->getIndex(), 2);
1336     ASSERT_EQUALS(clientResponse.getErrDetailsAt(1)->getErrCode(), response.getErrCode());
1337 }
1338 
1339 // Batch aborted (ordered). Expect this gets translated down into write error for first affected
1340 // write.
TEST_F(BatchWriteOpTest,MultiOpAbortOrdered)1341 TEST_F(BatchWriteOpTest, MultiOpAbortOrdered) {
1342     NamespaceString nss("foo.bar");
1343     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
1344     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
1345     MockNSTargeter targeter;
1346     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
1347 
1348     BatchedCommandRequest request([&] {
1349         write_ops::Insert insertOp(nss);
1350         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 2), BSON("x" << 3)});
1351         return insertOp;
1352     }());
1353 
1354     BatchWriteOp batchOp(operationContext(), request);
1355 
1356     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1357     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1358     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1359 
1360     BatchedCommandResponse response;
1361     buildResponse(1, &response);
1362 
1363     // First shard batch is ok
1364     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1365     ASSERT(!batchOp.isFinished());
1366 
1367     WriteErrorDetail abortError;
1368     abortError.setErrCode(ErrorCodes::UnknownError);
1369     abortError.setErrMessage("mock abort");
1370     batchOp.abortBatch(abortError);
1371     ASSERT(batchOp.isFinished());
1372 
1373     // We should have recorded an error for the second write
1374     BatchedCommandResponse clientResponse;
1375     batchOp.buildClientResponse(&clientResponse);
1376     ASSERT(clientResponse.getOk());
1377     ASSERT_EQUALS(clientResponse.getN(), 1);
1378     ASSERT(clientResponse.isErrDetailsSet());
1379     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 1u);
1380     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 1);
1381     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(), abortError.getErrCode());
1382 }
1383 
1384 // Batch aborted (unordered). Expect this gets translated down into write errors for all affected
1385 // writes.
TEST_F(BatchWriteOpTest,MultiOpAbortUnordered)1386 TEST_F(BatchWriteOpTest, MultiOpAbortUnordered) {
1387     NamespaceString nss("foo.bar");
1388     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
1389     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
1390     MockNSTargeter targeter;
1391     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
1392 
1393     BatchedCommandRequest request([&] {
1394         write_ops::Insert insertOp(nss);
1395         insertOp.setWriteCommandBase([] {
1396             write_ops::WriteCommandBase wcb;
1397             wcb.setOrdered(false);
1398             return wcb;
1399         }());
1400         insertOp.setDocuments({BSON("x" << -1), BSON("x" << -2)});
1401         return insertOp;
1402     }());
1403 
1404     BatchWriteOp batchOp(operationContext(), request);
1405 
1406     WriteErrorDetail abortError;
1407     abortError.setErrCode(ErrorCodes::UnknownError);
1408     abortError.setErrMessage("mock abort");
1409     batchOp.abortBatch(abortError);
1410     ASSERT(batchOp.isFinished());
1411 
1412     // We should have recorded an error for the first and second write
1413     BatchedCommandResponse clientResponse;
1414     batchOp.buildClientResponse(&clientResponse);
1415     ASSERT(clientResponse.getOk());
1416     ASSERT_EQUALS(clientResponse.getN(), 0);
1417     ASSERT(clientResponse.isErrDetailsSet());
1418     ASSERT_EQUALS(clientResponse.sizeErrDetails(), 2u);
1419     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getIndex(), 0);
1420     ASSERT_EQUALS(clientResponse.getErrDetailsAt(0)->getErrCode(), abortError.getErrCode());
1421     ASSERT_EQUALS(clientResponse.getErrDetailsAt(1)->getIndex(), 1);
1422     ASSERT_EQUALS(clientResponse.getErrDetailsAt(1)->getErrCode(), abortError.getErrCode());
1423 }
1424 
1425 // Multi-op targeting test where each op goes to both shards and both return a write concern error
TEST_F(BatchWriteOpTest,MultiOpTwoWCErrors)1426 TEST_F(BatchWriteOpTest, MultiOpTwoWCErrors) {
1427     NamespaceString nss("foo.bar");
1428     ShardEndpoint endpointA(ShardId("shardA"), ChunkVersion::IGNORED());
1429     ShardEndpoint endpointB(ShardId("shardB"), ChunkVersion::IGNORED());
1430     MockNSTargeter targeter;
1431     initTargeterSplitRange(nss, endpointA, endpointB, &targeter);
1432 
1433     BatchedCommandRequest request([&] {
1434         write_ops::Insert insertOp(nss);
1435         insertOp.setDocuments({BSON("x" << -1), BSON("x" << 2)});
1436         return insertOp;
1437     }());
1438     request.setWriteConcern(BSON("w" << 3));
1439 
1440     BatchWriteOp batchOp(operationContext(), request);
1441 
1442     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1443     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1444     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1445 
1446     BatchedCommandResponse response;
1447     buildResponse(1, &response);
1448     addWCError(&response);
1449 
1450     // First shard write write concern fails.
1451     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1452     ASSERT(!batchOp.isFinished());
1453 
1454     targetedOwned.clear();
1455     ASSERT_OK(batchOp.targetBatch(targeter, true, &targeted));
1456 
1457     // Second shard write write concern fails.
1458     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1459     ASSERT(batchOp.isFinished());
1460 
1461     BatchedCommandResponse clientResponse;
1462     batchOp.buildClientResponse(&clientResponse);
1463     ASSERT(clientResponse.getOk());
1464     ASSERT_EQUALS(clientResponse.getN(), 2);
1465     ASSERT(!clientResponse.isErrDetailsSet());
1466     ASSERT(clientResponse.isWriteConcernErrorSet());
1467 }
1468 
1469 //
1470 // Tests of batch size limit functionality
1471 //
1472 
1473 using BatchWriteOpLimitTests = WriteOpTestFixture;
1474 
1475 // Big single operation test - should go through
TEST_F(BatchWriteOpLimitTests,OneBigDoc)1476 TEST_F(BatchWriteOpLimitTests, OneBigDoc) {
1477     NamespaceString nss("foo.bar");
1478     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
1479     MockNSTargeter targeter;
1480     initTargeterFullRange(nss, endpoint, &targeter);
1481 
1482     // Create a BSONObj (slightly) bigger than the maximum size by including a max-size string
1483     const std::string bigString(BSONObjMaxUserSize, 'x');
1484 
1485     // Do single-target, single doc batch write op
1486     BatchedCommandRequest request([&] {
1487         write_ops::Insert insertOp(nss);
1488         insertOp.setWriteCommandBase([] {
1489             write_ops::WriteCommandBase wcb;
1490             wcb.setOrdered(false);
1491             return wcb;
1492         }());
1493         insertOp.setDocuments({BSON("x" << 1 << "data" << bigString)});
1494         return insertOp;
1495     }());
1496 
1497     BatchWriteOp batchOp(operationContext(), request);
1498 
1499     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1500     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1501     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1502     ASSERT_EQUALS(targeted.size(), 1u);
1503 
1504     BatchedCommandResponse response;
1505     buildResponse(1, &response);
1506 
1507     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1508     ASSERT(batchOp.isFinished());
1509 }
1510 
1511 // Big doc with smaller additional doc - should go through as two batches
TEST_F(BatchWriteOpLimitTests,OneBigOneSmall)1512 TEST_F(BatchWriteOpLimitTests, OneBigOneSmall) {
1513     NamespaceString nss("foo.bar");
1514     ShardEndpoint endpoint(ShardId("shard"), ChunkVersion::IGNORED());
1515     MockNSTargeter targeter;
1516     initTargeterFullRange(nss, endpoint, &targeter);
1517 
1518     // Create a BSONObj (slightly) bigger than the maximum size by including a max-size string
1519     const std::string bigString(BSONObjMaxUserSize, 'x');
1520 
1521     BatchedCommandRequest request([&] {
1522         write_ops::Update updateOp(nss);
1523         updateOp.setUpdates({buildUpdate(BSON("x" << 1), BSON("data" << bigString), false),
1524                              buildUpdate(BSON("x" << 2), BSONObj(), false)});
1525         return updateOp;
1526     }());
1527 
1528     BatchWriteOp batchOp(operationContext(), request);
1529 
1530     OwnedPointerMap<ShardId, TargetedWriteBatch> targetedOwned;
1531     std::map<ShardId, TargetedWriteBatch*>& targeted = targetedOwned.mutableMap();
1532     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1533     ASSERT_EQUALS(targeted.size(), 1u);
1534     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 1u);
1535 
1536     BatchedCommandResponse response;
1537     buildResponse(1, &response);
1538 
1539     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1540     ASSERT(!batchOp.isFinished());
1541 
1542     targetedOwned.clear();
1543     ASSERT_OK(batchOp.targetBatch(targeter, false, &targeted));
1544     ASSERT_EQUALS(targeted.size(), 1u);
1545     ASSERT_EQUALS(targeted.begin()->second->getWrites().size(), 1u);
1546 
1547     batchOp.noteBatchResponse(*targeted.begin()->second, response, NULL);
1548     ASSERT(batchOp.isFinished());
1549 }
1550 
1551 }  // namespace
1552 }  // namespace mongo
1553