boost::intrusive_ptr<DocumentSource> DocumentSourceMergeCursors::createFromBson(
BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& expCtx) {
if (elem.type() == BSONType::Object) {
// This is the modern serialization format. We de-serialize using the IDL.
auto ownedObj = elem.embeddedObject().getOwned();
auto armParams =
AsyncResultsMergerParams::parse(IDLParserErrorContext(kStageName), ownedObj);
return new DocumentSourceMergeCursors(
Grid::get(expCtx->opCtx)->getExecutorPool()->getArbitraryExecutor(),
std::move(armParams),
expCtx,
std::move(ownedObj));
}
// This is the old serialization format which can still be generated by mongos processes
// older than 4.0.
// TODO SERVER-34009 Remove support for this format.
uassert(17026,
"$mergeCursors stage expected either an array or an object as argument",
elem.type() == BSONType::Array);
const auto serializedRemotes = elem.Array();
uassert(50729,
"$mergeCursors stage expected array with at least one entry",
serializedRemotes.size() > 0);
boost::optional<NamespaceString> nss;
std::vector<RemoteCursor> remotes;
for (auto&& cursor : serializedRemotes) {
BSONElement nsElem;
BSONElement hostElem;
BSONElement idElem;
uassert(17027,
"$mergeCursors stage requires each cursor in array to be an object",
cursor.type() == BSONType::Object);
for (auto&& cursorElem : cursor.Obj()) {
const auto fieldName = cursorElem.fieldNameStringData();
if (fieldName == "ns"_sd) {
nsElem = cursorElem;
} else if (fieldName == "host"_sd) {
hostElem = cursorElem;
} else if (fieldName == "id"_sd) {
idElem = cursorElem;
} else {
uasserted(50730,
str::stream() << "Unrecognized option " << fieldName
<< " within cursor provided to $mergeCursors: "
<< cursor);
}
}
uassert(
50731,
"$mergeCursors stage requires \'ns\' field with type string for each cursor in array",
nsElem.type() == BSONType::String);
// We require each cursor to have the same namespace. This isn't a fundamental limit of the
// system, but needs to be true due to the implementation of AsyncResultsMerger, which
// tracks one namespace for all cursors.
uassert(50720,
"$mergeCursors requires each cursor to have the same namespace",
!nss || nss->ns() == nsElem.valueStringData());
nss = NamespaceString(nsElem.String());
uassert(
50721,
"$mergeCursors stage requires \'host\' field with type string for each cursor in array",
hostElem.type() == BSONType::String);
auto host = uassertStatusOK(HostAndPort::parse(hostElem.valueStringData()));
uassert(50722,
"$mergeCursors stage requires \'id\' field with type long for each cursor in array",
idElem.type() == BSONType::NumberLong);
auto cursorId = idElem.Long();
// We are assuming that none of the cursors have been iterated at all, and so will not have
// any data in the initial batch.
// TODO SERVER-33323 We use a fake shard id because the AsyncResultsMerger won't use it for
// anything, and finding the real one is non-trivial.
RemoteCursor remoteCursor;
remoteCursor.setShardId(ShardId("fakeShardIdForMergeCursors"));
remoteCursor.setHostAndPort(std::move(host));
std::vector<BSONObj> emptyBatch;
remoteCursor.setCursorResponse(CursorResponse{*nss, cursorId, emptyBatch});
remotes.push_back(std::move(remoteCursor));
}
invariant(nss); // We know there is at least one cursor in 'serializedRemotes', and we require
// each cursor to have a 'ns' field.
AsyncResultsMergerParams params;
params.setRemotes(std::move(remotes));
params.setNss(*nss);
return new DocumentSourceMergeCursors(
Grid::get(expCtx->opCtx)->getExecutorPool()->getArbitraryExecutor(),
std::move(params),
expCtx,
elem.embeddedObject().getOwned());
}
std::vector<RemoteCursor> establishCursors(OperationContext* opCtx,
executor::TaskExecutor* executor,
const NamespaceString& nss,
const ReadPreferenceSetting readPref,
const std::vector<std::pair<ShardId, BSONObj>>& remotes,
bool allowPartialResults,
Shard::RetryPolicy retryPolicy) {
// Construct the requests
std::vector<AsyncRequestsSender::Request> requests;
for (const auto& remote : remotes) {
requests.emplace_back(remote.first, remote.second);
}
// Send the requests
MultiStatementTransactionRequestsSender ars(
opCtx, executor, nss.db().toString(), std::move(requests), readPref, retryPolicy);
std::vector<RemoteCursor> remoteCursors;
try {
// Get the responses
while (!ars.done()) {
try {
auto response = ars.next();
// Note the shardHostAndPort may not be populated if there was an error, so be sure
// to do this after parsing the cursor response to ensure the response was ok.
// Additionally, be careful not to push into 'remoteCursors' until we are sure we
// have a valid cursor, since the error handling path will attempt to clean up
// anything in 'remoteCursors'
auto cursors = CursorResponse::parseFromBSONMany(
uassertStatusOK(std::move(response.swResponse)).data);
for (auto& cursor : cursors) {
if (cursor.isOK()) {
RemoteCursor remoteCursor;
remoteCursor.setCursorResponse(std::move(cursor.getValue()));
remoteCursor.setShardId(std::move(response.shardId));
remoteCursor.setHostAndPort(*response.shardHostAndPort);
remoteCursors.push_back(std::move(remoteCursor));
}
}
// Throw if there is any error and then the catch block below will do the cleanup.
for (auto& cursor : cursors) {
uassertStatusOK(cursor.getStatus());
}
} catch (const DBException& ex) {
// Retriable errors are swallowed if 'allowPartialResults' is true.
if (allowPartialResults &&
std::find(RemoteCommandRetryScheduler::kAllRetriableErrors.begin(),
RemoteCommandRetryScheduler::kAllRetriableErrors.end(),
ex.code()) !=
RemoteCommandRetryScheduler::kAllRetriableErrors.end()) {
continue;
}
throw; // Fail this loop.
}
}
return remoteCursors;
} catch (const DBException&) {
// If one of the remotes had an error, we make a best effort to finish retrieving responses
// for other requests that were already sent, so that we can send killCursors to any cursors
// that we know were established.
try {
// Do not schedule any new requests.
ars.stopRetrying();
// Collect responses from all requests that were already sent.
while (!ars.done()) {
auto response = ars.next();
// Check if the response contains an established cursor, and if so, store it.
StatusWith<CursorResponse> swCursorResponse(
response.swResponse.isOK()
? CursorResponse::parseFromBSON(response.swResponse.getValue().data)
: response.swResponse.getStatus());
if (swCursorResponse.isOK()) {
RemoteCursor cursor;
cursor.setShardId(std::move(response.shardId));
cursor.setHostAndPort(*response.shardHostAndPort);
cursor.setCursorResponse(std::move(swCursorResponse.getValue()));
remoteCursors.push_back(std::move(cursor));
}
}
// Schedule killCursors against all cursors that were established.
killRemoteCursors(opCtx, executor, std::move(remoteCursors), nss);
} catch (const DBException&) {
// Ignore the new error and rethrow the original one.
}
throw;
}
}