void run() {
setup();
insert(fromjson("{_id: 1, x: 5}"));
insert(fromjson("{_id: 2, x: 6}"));
insert(fromjson("{_id: 3, x: 10}"));
std::unique_ptr<IndexScan> ixscan(
createIndexScan(BSON("x" << 5), BSON("x" << 10), false, false));
// Expect to get key {'': 6}.
WorkingSetMember* member = getNext(ixscan.get());
ASSERT_EQ(WorkingSetMember::RID_AND_IDX, member->getState());
ASSERT_BSONOBJ_EQ(member->keyData[0].keyData, BSON("" << 6));
// Save state and insert an indexed doc.
ixscan->saveState();
insert(fromjson("{_id: 4, x: 7}"));
ixscan->restoreState();
member = getNext(ixscan.get());
ASSERT_EQ(WorkingSetMember::RID_AND_IDX, member->getState());
ASSERT_BSONOBJ_EQ(member->keyData[0].keyData, BSON("" << 7));
WorkingSetID id;
ASSERT_EQ(PlanStage::IS_EOF, ixscan->work(&id));
ASSERT(ixscan->isEOF());
}
void run() {
setup();
insert(fromjson("{_id: 1, x: 10}"));
insert(fromjson("{_id: 2, x: 8}"));
insert(fromjson("{_id: 3, x: 3}"));
std::unique_ptr<IndexScan> ixscan(
createIndexScan(BSON("x" << 10), BSON("x" << 5), true, true, -1 /* reverse scan */));
// Expect to get key {'': 10} and then {'': 8}.
WorkingSetMember* member = getNext(ixscan.get());
ASSERT_EQ(WorkingSetMember::RID_AND_IDX, member->getState());
ASSERT_BSONOBJ_EQ(member->keyData[0].keyData, BSON("" << 10));
member = getNext(ixscan.get());
ASSERT_EQ(WorkingSetMember::RID_AND_IDX, member->getState());
ASSERT_BSONOBJ_EQ(member->keyData[0].keyData, BSON("" << 8));
// Save state and insert an indexed doc.
ixscan->saveState();
insert(fromjson("{_id: 4, x: 6}"));
insert(fromjson("{_id: 5, x: 9}"));
ixscan->restoreState();
// Ensure that we don't erroneously return {'': 9} or {'':3}.
member = getNext(ixscan.get());
ASSERT_EQ(WorkingSetMember::RID_AND_IDX, member->getState());
ASSERT_BSONOBJ_EQ(member->keyData[0].keyData, BSON("" << 6));
WorkingSetID id;
ASSERT_EQ(PlanStage::IS_EOF, ixscan->work(&id));
ASSERT(ixscan->isEOF());
}
void run() {
// Various variables we'll need.
dbtests::WriteContextForTests ctx(&_opCtx, nss.ns());
Collection* coll = ctx.getCollection();
ASSERT(coll);
const int targetDocIndex = 0;
const BSONObj query = BSON("foo" << BSON("$gte" << targetDocIndex));
const auto ws = make_unique<WorkingSet>();
const unique_ptr<CanonicalQuery> cq(canonicalize(query));
// Get the RecordIds that would be returned by an in-order scan.
vector<RecordId> recordIds;
getRecordIds(coll, CollectionScanParams::FORWARD, &recordIds);
// Configure a QueuedDataStage to pass the first object in the collection back in a
// RID_AND_OBJ state.
auto qds = make_unique<QueuedDataStage>(&_opCtx, ws.get());
WorkingSetID id = ws->allocate();
WorkingSetMember* member = ws->get(id);
member->recordId = recordIds[targetDocIndex];
const BSONObj oldDoc = BSON("_id" << targetDocIndex << "foo" << targetDocIndex);
member->obj = Snapshotted<BSONObj>(SnapshotId(), oldDoc);
ws->transitionToRecordIdAndObj(id);
qds->pushBack(id);
// Configure the delete.
auto deleteParams = std::make_unique<DeleteStageParams>();
deleteParams->returnDeleted = true;
deleteParams->canonicalQuery = cq.get();
const auto deleteStage = make_unique<DeleteStage>(
&_opCtx, std::move(deleteParams), ws.get(), coll, qds.release());
const DeleteStats* stats = static_cast<const DeleteStats*>(deleteStage->getSpecificStats());
// Should return advanced.
id = WorkingSet::INVALID_ID;
PlanStage::StageState state = deleteStage->work(&id);
ASSERT_EQUALS(PlanStage::ADVANCED, state);
// Make sure the returned value is what we expect it to be.
// Should give us back a valid id.
ASSERT_TRUE(WorkingSet::INVALID_ID != id);
WorkingSetMember* resultMember = ws->get(id);
// With an owned copy of the object, with no RecordId.
ASSERT_TRUE(resultMember->hasOwnedObj());
ASSERT_FALSE(resultMember->hasRecordId());
ASSERT_EQUALS(resultMember->getState(), WorkingSetMember::OWNED_OBJ);
ASSERT_TRUE(resultMember->obj.value().isOwned());
// Should be the old value.
ASSERT_BSONOBJ_EQ(resultMember->obj.value(), oldDoc);
// Should have done the delete.
ASSERT_EQUALS(stats->docsDeleted, 1U);
// That should be it.
id = WorkingSet::INVALID_ID;
ASSERT_EQUALS(PlanStage::IS_EOF, deleteStage->work(&id));
}
void WorkingSetCommon::prepareForSnapshotChange(WorkingSet* workingSet) {
dassert(supportsDocLocking());
for (auto id : workingSet->getAndClearYieldSensitiveIds()) {
if (workingSet->isFree(id)) {
continue;
}
// We may see the same member twice, so anything we do here should be idempotent.
WorkingSetMember* member = workingSet->get(id);
if (member->getState() == WorkingSetMember::LOC_AND_IDX) {
member->isSuspicious = true;
} else if (member->getState() == WorkingSetMember::LOC_AND_OBJ) {
// Need to make sure that the data is owned, as underlying storage can change during a
// yield.
member->makeObjOwned();
}
}
}
StatusWith<BSONObj> SortKeyGenerator::getSortKeyFromIndexKey(const WorkingSetMember& member) const {
invariant(member.getState() == WorkingSetMember::RID_AND_IDX);
invariant(!_sortHasMeta);
BSONObjBuilder sortKeyObj;
for (BSONElement specElt : _rawSortSpec) {
invariant(specElt.isNumber());
BSONElement sortKeyElt;
invariant(member.getFieldDotted(specElt.fieldName(), &sortKeyElt));
sortKeyObj.appendAs(sortKeyElt, "");
}
return sortKeyObj.obj();
}
void WorkingSetCommon::prepareForSnapshotChange(WorkingSet* workingSet) {
if (!supportsDocLocking()) {
// Non doc-locking storage engines use invalidations, so we don't need to examine the
// buffered working set ids. But we do need to clear the set of ids in order to keep our
// memory utilization in check.
workingSet->getAndClearYieldSensitiveIds();
return;
}
for (auto id : workingSet->getAndClearYieldSensitiveIds()) {
if (workingSet->isFree(id)) {
continue;
}
// We may see the same member twice, so anything we do here should be idempotent.
WorkingSetMember* member = workingSet->get(id);
if (member->getState() == WorkingSetMember::RID_AND_IDX) {
member->isSuspicious = true;
}
}
}
PlanStage::StageState DeleteStage::doWork(WorkingSetID* out) {
if (isEOF()) {
return PlanStage::IS_EOF;
}
invariant(_collection); // If isEOF() returns false, we must have a collection.
// It is possible that after a delete was executed, a WriteConflictException occurred
// and prevented us from returning ADVANCED with the old version of the document.
if (_idReturning != WorkingSet::INVALID_ID) {
// We should only get here if we were trying to return something before.
invariant(_params.returnDeleted);
WorkingSetMember* member = _ws->get(_idReturning);
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
*out = _idReturning;
_idReturning = WorkingSet::INVALID_ID;
return PlanStage::ADVANCED;
}
// Either retry the last WSM we worked on or get a new one from our child.
WorkingSetID id;
if (_idRetrying != WorkingSet::INVALID_ID) {
id = _idRetrying;
_idRetrying = WorkingSet::INVALID_ID;
} else {
auto status = child()->work(&id);
switch (status) {
case PlanStage::ADVANCED:
break;
case PlanStage::FAILURE:
case PlanStage::DEAD:
*out = id;
// If a stage fails, it may create a status WSM to indicate why it failed, in which
// case 'id' is valid. If ID is invalid, we create our own error message.
if (WorkingSet::INVALID_ID == id) {
const std::string errmsg = "delete stage failed to read in results from child";
*out = WorkingSetCommon::allocateStatusMember(
_ws, Status(ErrorCodes::InternalError, errmsg));
}
return status;
case PlanStage::NEED_TIME:
return status;
case PlanStage::NEED_YIELD:
*out = id;
return status;
case PlanStage::IS_EOF:
return status;
default:
MONGO_UNREACHABLE;
}
}
// We advanced, or are retrying, and id is set to the WSM to work on.
WorkingSetMember* member = _ws->get(id);
// We want to free this member when we return, unless we need to retry it.
ScopeGuard memberFreer = MakeGuard(&WorkingSet::free, _ws, id);
if (!member->hasRecordId()) {
// We expect to be here because of an invalidation causing a force-fetch.
++_specificStats.nInvalidateSkips;
return PlanStage::NEED_TIME;
}
RecordId recordId = member->recordId;
// Deletes can't have projections. This means that covering analysis will always add
// a fetch. We should always get fetched data, and never just key data.
invariant(member->hasObj());
try {
// If the snapshot changed, then we have to make sure we have the latest copy of the
// doc and that it still matches.
std::unique_ptr<SeekableRecordCursor> cursor;
if (getOpCtx()->recoveryUnit()->getSnapshotId() != member->obj.snapshotId()) {
cursor = _collection->getCursor(getOpCtx());
if (!WorkingSetCommon::fetch(getOpCtx(), _ws, id, cursor)) {
// Doc is already deleted. Nothing more to do.
return PlanStage::NEED_TIME;
}
// Make sure the re-fetched doc still matches the predicate.
if (_params.canonicalQuery &&
!_params.canonicalQuery->root()->matchesBSON(member->obj.value(), NULL)) {
// Doesn't match.
return PlanStage::NEED_TIME;
}
}
// Ensure that the BSONObj underlying the WorkingSetMember is owned because saveState()
// is allowed to free the memory.
if (_params.returnDeleted) {
// Save a copy of the document that is about to get deleted, but keep it in the
// RID_AND_OBJ state in case we need to retry deleting it.
BSONObj deletedDoc = member->obj.value();
member->obj.setValue(deletedDoc.getOwned());
}
// TODO: Do we want to buffer docs and delete them in a group rather than
// saving/restoring state repeatedly?
try {
WorkingSetCommon::prepareForSnapshotChange(_ws);
child()->saveState();
} catch (const WriteConflictException& wce) {
std::terminate();
}
// Do the write, unless this is an explain.
if (!_params.isExplain) {
WriteUnitOfWork wunit(getOpCtx());
_collection->deleteDocument(getOpCtx(), recordId, _params.fromMigrate);
wunit.commit();
}
++_specificStats.docsDeleted;
} catch (const WriteConflictException& wce) {
// When we're doing a findAndModify with a sort, the sort will have a limit of 1, so will
// not produce any more results even if there is another matching document. Re-throw the WCE
// here so that these operations get another chance to find a matching document. The
// findAndModify command should automatically retry if it gets a WCE.
// TODO: this is not necessary if there was no sort specified.
if (_params.returnDeleted) {
throw;
}
_idRetrying = id;
memberFreer.Dismiss(); // Keep this member around so we can retry deleting it.
*out = WorkingSet::INVALID_ID;
return NEED_YIELD;
}
if (_params.returnDeleted) {
// After deleting the document, the RecordId associated with this member is invalid.
// Remove the 'recordId' from the WorkingSetMember before returning it.
member->recordId = RecordId();
member->transitionToOwnedObj();
}
// As restoreState may restore (recreate) cursors, cursors are tied to the
// transaction in which they are created, and a WriteUnitOfWork is a
// transaction, make sure to restore the state outside of the WritUnitOfWork.
try {
child()->restoreState();
} catch (const WriteConflictException& wce) {
// Note we don't need to retry anything in this case since the delete already
// was committed. However, we still need to return the deleted document
// (if it was requested).
if (_params.returnDeleted) {
// member->obj should refer to the deleted document.
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
_idReturning = id;
// Keep this member around so that we can return it on the next work() call.
memberFreer.Dismiss();
}
*out = WorkingSet::INVALID_ID;
return NEED_YIELD;
}
if (_params.returnDeleted) {
// member->obj should refer to the deleted document.
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
memberFreer.Dismiss(); // Keep this member around so we can return it.
*out = id;
return PlanStage::ADVANCED;
}
return PlanStage::NEED_TIME;
}
void run() {
// Populate the collection.
for (int i = 0; i < 50; ++i) {
insert(BSON("_id" << i << "foo" << i));
}
ASSERT_EQUALS(50U, count(BSONObj()));
// Various variables we'll need.
dbtests::WriteContextForTests ctx(&_opCtx, nss.ns());
OpDebug* opDebug = &CurOp::get(_opCtx)->debug();
Collection* coll = ctx.getCollection();
ASSERT(coll);
UpdateRequest request(nss);
const CollatorInterface* collator = nullptr;
UpdateDriver driver(new ExpressionContext(&_opCtx, collator));
const int targetDocIndex = 10;
const BSONObj query = BSON("foo" << BSON("$gte" << targetDocIndex));
const auto ws = make_unique<WorkingSet>();
const unique_ptr<CanonicalQuery> cq(canonicalize(query));
// Get the RecordIds that would be returned by an in-order scan.
vector<RecordId> recordIds;
getRecordIds(coll, CollectionScanParams::FORWARD, &recordIds);
// Populate the request.
request.setQuery(query);
request.setUpdates(fromjson("{$set: {x: 0}}"));
request.setSort(BSONObj());
request.setMulti(false);
request.setReturnDocs(UpdateRequest::RETURN_NEW);
const std::map<StringData, std::unique_ptr<ExpressionWithPlaceholder>> arrayFilters;
ASSERT_DOES_NOT_THROW(driver.parse(request.getUpdates(), arrayFilters, request.isMulti()));
// Configure a QueuedDataStage to pass the first object in the collection back in a
// RID_AND_OBJ state.
auto qds = make_unique<QueuedDataStage>(&_opCtx, ws.get());
WorkingSetID id = ws->allocate();
WorkingSetMember* member = ws->get(id);
member->recordId = recordIds[targetDocIndex];
const BSONObj oldDoc = BSON("_id" << targetDocIndex << "foo" << targetDocIndex);
member->obj = Snapshotted<BSONObj>(SnapshotId(), oldDoc);
ws->transitionToRecordIdAndObj(id);
qds->pushBack(id);
// Configure the update.
UpdateStageParams updateParams(&request, &driver, opDebug);
updateParams.canonicalQuery = cq.get();
auto updateStage =
make_unique<UpdateStage>(&_opCtx, updateParams, ws.get(), coll, qds.release());
// Should return advanced.
id = WorkingSet::INVALID_ID;
PlanStage::StageState state = updateStage->work(&id);
ASSERT_EQUALS(PlanStage::ADVANCED, state);
// Make sure the returned value is what we expect it to be.
// Should give us back a valid id.
ASSERT_TRUE(WorkingSet::INVALID_ID != id);
WorkingSetMember* resultMember = ws->get(id);
// With an owned copy of the object, with no RecordId.
ASSERT_TRUE(resultMember->hasOwnedObj());
ASSERT_FALSE(resultMember->hasRecordId());
ASSERT_EQUALS(resultMember->getState(), WorkingSetMember::OWNED_OBJ);
ASSERT_TRUE(resultMember->obj.value().isOwned());
// Should be the new value.
BSONObj newDoc = BSON("_id" << targetDocIndex << "foo" << targetDocIndex << "x" << 0);
ASSERT_BSONOBJ_EQ(resultMember->obj.value(), newDoc);
// Should have done the update.
vector<BSONObj> objs;
getCollContents(coll, &objs);
ASSERT_BSONOBJ_EQ(objs[targetDocIndex], newDoc);
// That should be it.
id = WorkingSet::INVALID_ID;
ASSERT_EQUALS(PlanStage::IS_EOF, updateStage->work(&id));
}
PlanStage::StageState FetchStage::work(WorkingSetID* out) {
++_commonStats.works;
// Adds the amount of time taken by work() to executionTimeMillis.
ScopedTimer timer(&_commonStats.executionTimeMillis);
if (isEOF()) {
return PlanStage::IS_EOF;
}
// Either retry the last WSM we worked on or get a new one from our child.
WorkingSetID id;
StageState status;
if (_idRetrying == WorkingSet::INVALID_ID) {
status = child()->work(&id);
} else {
status = ADVANCED;
id = _idRetrying;
_idRetrying = WorkingSet::INVALID_ID;
}
if (PlanStage::ADVANCED == status) {
WorkingSetMember* member = _ws->get(id);
// If there's an obj there, there is no fetching to perform.
if (member->hasObj()) {
++_specificStats.alreadyHasObj;
} else {
// We need a valid loc to fetch from and this is the only state that has one.
verify(WorkingSetMember::LOC_AND_IDX == member->getState());
verify(member->hasLoc());
try {
if (!_cursor)
_cursor = _collection->getCursor(getOpCtx());
if (auto fetcher = _cursor->fetcherForId(member->loc)) {
// There's something to fetch. Hand the fetcher off to the WSM, and pass up
// a fetch request.
_idRetrying = id;
member->setFetcher(fetcher.release());
*out = id;
_commonStats.needYield++;
return NEED_YIELD;
}
// The doc is already in memory, so go ahead and grab it. Now we have a RecordId
// as well as an unowned object
if (!WorkingSetCommon::fetch(getOpCtx(), _ws, id, _cursor)) {
_ws->free(id);
_commonStats.needTime++;
return NEED_TIME;
}
} catch (const WriteConflictException& wce) {
_idRetrying = id;
*out = WorkingSet::INVALID_ID;
_commonStats.needYield++;
return NEED_YIELD;
}
}
return returnIfMatches(member, id, out);
} else if (PlanStage::FAILURE == status || PlanStage::DEAD == status) {
*out = id;
// If a stage fails, it may create a status WSM to indicate why it
// failed, in which case 'id' is valid. If ID is invalid, we
// create our own error message.
if (WorkingSet::INVALID_ID == id) {
mongoutils::str::stream ss;
ss << "fetch stage failed to read in results from child";
Status status(ErrorCodes::InternalError, ss);
*out = WorkingSetCommon::allocateStatusMember(_ws, status);
}
return status;
} else if (PlanStage::NEED_TIME == status) {
++_commonStats.needTime;
} else if (PlanStage::NEED_YIELD == status) {
++_commonStats.needYield;
*out = id;
}
return status;
}
PlanStage::StageState DeleteStage::work(WorkingSetID* out) {
++_commonStats.works;
// Adds the amount of time taken by work() to executionTimeMillis.
ScopedTimer timer(&_commonStats.executionTimeMillis);
if (isEOF()) {
return PlanStage::IS_EOF;
}
invariant(_collection); // If isEOF() returns false, we must have a collection.
// It is possible that after a delete was executed, a WriteConflictException occurred
// and prevented us from returning ADVANCED with the old version of the document.
if (_idReturning != WorkingSet::INVALID_ID) {
// We should only get here if we were trying to return something before.
invariant(_params.returnDeleted);
WorkingSetMember* member = _ws->get(_idReturning);
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
*out = _idReturning;
_idReturning = WorkingSet::INVALID_ID;
++_commonStats.advanced;
return PlanStage::ADVANCED;
}
// Either retry the last WSM we worked on or get a new one from our child.
WorkingSetID id;
StageState status;
if (_idRetrying == WorkingSet::INVALID_ID) {
status = child()->work(&id);
} else {
status = ADVANCED;
id = _idRetrying;
_idRetrying = WorkingSet::INVALID_ID;
}
if (PlanStage::ADVANCED == status) {
WorkingSetMember* member = _ws->get(id);
// We want to free this member when we return, unless we need to retry it.
ScopeGuard memberFreer = MakeGuard(&WorkingSet::free, _ws, id);
if (!member->hasLoc()) {
// We expect to be here because of an invalidation causing a force-fetch, and
// doc-locking storage engines do not issue invalidations.
++_specificStats.nInvalidateSkips;
++_commonStats.needTime;
return PlanStage::NEED_TIME;
}
RecordId rloc = member->loc;
// Deletes can't have projections. This means that covering analysis will always add
// a fetch. We should always get fetched data, and never just key data.
invariant(member->hasObj());
try {
// If the snapshot changed, then we have to make sure we have the latest copy of the
// doc and that it still matches.
std::unique_ptr<RecordCursor> cursor;
if (getOpCtx()->recoveryUnit()->getSnapshotId() != member->obj.snapshotId()) {
cursor = _collection->getCursor(getOpCtx());
if (!WorkingSetCommon::fetch(getOpCtx(), _ws, id, cursor)) {
// Doc is already deleted. Nothing more to do.
++_commonStats.needTime;
return PlanStage::NEED_TIME;
}
// Make sure the re-fetched doc still matches the predicate.
if (_params.canonicalQuery &&
!_params.canonicalQuery->root()->matchesBSON(member->obj.value(), NULL)) {
// Doesn't match.
++_commonStats.needTime;
return PlanStage::NEED_TIME;
}
}
// Ensure that the BSONObj underlying the WorkingSetMember is owned because saveState()
// is allowed to free the memory.
if (_params.returnDeleted) {
member->makeObjOwnedIfNeeded();
}
// TODO: Do we want to buffer docs and delete them in a group rather than
// saving/restoring state repeatedly?
try {
if (supportsDocLocking()) {
// Doc-locking engines require this before saveState() since they don't use
// invalidations.
WorkingSetCommon::prepareForSnapshotChange(_ws);
}
child()->saveState();
} catch (const WriteConflictException& wce) {
std::terminate();
}
if (_params.returnDeleted) {
// Save a copy of the document that is about to get deleted.
BSONObj deletedDoc = member->obj.value();
member->obj.setValue(deletedDoc.getOwned());
member->loc = RecordId();
member->transitionToOwnedObj();
}
// Do the write, unless this is an explain.
if (!_params.isExplain) {
WriteUnitOfWork wunit(getOpCtx());
_collection->deleteDocument(getOpCtx(), rloc);
wunit.commit();
}
++_specificStats.docsDeleted;
} catch (const WriteConflictException& wce) {
// Ensure that the BSONObj underlying the WorkingSetMember is owned because it may be
// freed when we yield.
member->makeObjOwnedIfNeeded();
_idRetrying = id;
memberFreer.Dismiss(); // Keep this member around so we can retry deleting it.
*out = WorkingSet::INVALID_ID;
_commonStats.needYield++;
return NEED_YIELD;
}
// As restoreState may restore (recreate) cursors, cursors are tied to the
// transaction in which they are created, and a WriteUnitOfWork is a
// transaction, make sure to restore the state outside of the WritUnitOfWork.
try {
child()->restoreState();
} catch (const WriteConflictException& wce) {
// Note we don't need to retry anything in this case since the delete already
// was committed. However, we still need to return the deleted document
// (if it was requested).
if (_params.returnDeleted) {
// member->obj should refer to the deleted document.
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
_idReturning = id;
// Keep this member around so that we can return it on the next work() call.
memberFreer.Dismiss();
}
*out = WorkingSet::INVALID_ID;
_commonStats.needYield++;
return NEED_YIELD;
}
if (_params.returnDeleted) {
// member->obj should refer to the deleted document.
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
memberFreer.Dismiss(); // Keep this member around so we can return it.
*out = id;
++_commonStats.advanced;
return PlanStage::ADVANCED;
}
++_commonStats.needTime;
return PlanStage::NEED_TIME;
} else if (PlanStage::FAILURE == status || PlanStage::DEAD == status) {
*out = id;
// If a stage fails, it may create a status WSM to indicate why it failed, in which case
// 'id' is valid. If ID is invalid, we create our own error message.
if (WorkingSet::INVALID_ID == id) {
const std::string errmsg = "delete stage failed to read in results from child";
*out = WorkingSetCommon::allocateStatusMember(
_ws, Status(ErrorCodes::InternalError, errmsg));
}
return status;
} else if (PlanStage::NEED_TIME == status) {
++_commonStats.needTime;
} else if (PlanStage::NEED_YIELD == status) {
*out = id;
++_commonStats.needYield;
}
return status;
}
void run() {
OldClientWriteContext ctx(&_txn, ns());
Database* db = ctx.db();
Collection* coll = db->getCollection(ns());
if (!coll) {
WriteUnitOfWork wuow(&_txn);
coll = db->createCollection(&_txn, ns());
wuow.commit();
}
// Insert data.
insert(BSON("_id" << 4 << "a" << 4));
insert(BSON("_id" << 5 << "a" << 5));
insert(BSON("_id" << 6 << "a" << 6));
addIndex(BSON("a" << 1));
std::set<RecordId> rids;
getRecordIds(&rids, coll);
set<RecordId>::iterator it = rids.begin();
WorkingSet ws;
WorkingSetMember* member;
MergeSortStageParams msparams;
msparams.pattern = BSON("a" << 1);
auto ms = stdx::make_unique<MergeSortStage>(&_txn, msparams, &ws, coll);
// First child scans [5, 10].
{
IndexScanParams params;
params.descriptor = getIndex(BSON("a" << 1), coll);
params.bounds.isSimpleRange = true;
params.bounds.startKey = BSON("" << 5);
params.bounds.endKey = BSON("" << 10);
params.bounds.endKeyInclusive = true;
params.direction = 1;
auto fetchStage = stdx::make_unique<FetchStage>(
&_txn, &ws, new IndexScan(&_txn, params, &ws, nullptr), nullptr, coll);
ms->addChild(fetchStage.release());
}
// Second child scans [4, 10].
{
IndexScanParams params;
params.descriptor = getIndex(BSON("a" << 1), coll);
params.bounds.isSimpleRange = true;
params.bounds.startKey = BSON("" << 4);
params.bounds.endKey = BSON("" << 10);
params.bounds.endKeyInclusive = true;
params.direction = 1;
auto fetchStage = stdx::make_unique<FetchStage>(
&_txn, &ws, new IndexScan(&_txn, params, &ws, nullptr), nullptr, coll);
ms->addChild(fetchStage.release());
}
// First doc should be {a: 4}.
member = getNextResult(&ws, ms.get());
ASSERT_EQ(member->getState(), WorkingSetMember::RID_AND_OBJ);
ASSERT_EQ(member->recordId, *it);
ASSERT_EQ(member->obj.value(), BSON("_id" << 4 << "a" << 4));
++it;
// Doc {a: 5} gets invalidated by an update.
ms->invalidate(&_txn, *it, INVALIDATION_MUTATION);
// Invalidated doc {a: 5} should still get returned.
member = getNextResult(&ws, ms.get());
ASSERT_EQ(member->getState(), WorkingSetMember::OWNED_OBJ);
ASSERT_EQ(member->obj.value(), BSON("_id" << 5 << "a" << 5));
++it;
// We correctly dedup the invalidated doc and return {a: 6} next.
member = getNextResult(&ws, ms.get());
ASSERT_EQ(member->getState(), WorkingSetMember::RID_AND_OBJ);
ASSERT_EQ(member->recordId, *it);
ASSERT_EQ(member->obj.value(), BSON("_id" << 6 << "a" << 6));
}
PlanStage::StageState FetchStage::doWork(WorkingSetID* out) {
if (isEOF()) {
return PlanStage::IS_EOF;
}
// Either retry the last WSM we worked on or get a new one from our child.
WorkingSetID id;
StageState status;
if (_idRetrying == WorkingSet::INVALID_ID) {
status = child()->work(&id);
} else {
status = ADVANCED;
id = _idRetrying;
_idRetrying = WorkingSet::INVALID_ID;
}
if (PlanStage::ADVANCED == status) {
WorkingSetMember* member = _ws->get(id);
// If there's an obj there, there is no fetching to perform.
if (member->hasObj()) {
++_specificStats.alreadyHasObj;
} else {
// We need a valid RecordId to fetch from and this is the only state that has one.
verify(WorkingSetMember::RID_AND_IDX == member->getState());
verify(member->hasRecordId());
try {
if (!_cursor)
_cursor = _collection->getCursor(getOpCtx());
if (auto fetcher = _cursor->fetcherForId(member->recordId)) {
// There's something to fetch. Hand the fetcher off to the WSM, and pass up
// a fetch request.
_idRetrying = id;
member->setFetcher(fetcher.release());
*out = id;
return NEED_YIELD;
}
// The doc is already in memory, so go ahead and grab it. Now we have a RecordId
// as well as an unowned object
if (!WorkingSetCommon::fetch(getOpCtx(), _ws, id, _cursor)) {
_ws->free(id);
return NEED_TIME;
}
} catch (const WriteConflictException&) {
// Ensure that the BSONObj underlying the WorkingSetMember is owned because it may
// be freed when we yield.
member->makeObjOwnedIfNeeded();
_idRetrying = id;
*out = WorkingSet::INVALID_ID;
return NEED_YIELD;
}
}
return returnIfMatches(member, id, out);
} else if (PlanStage::FAILURE == status || PlanStage::DEAD == status) {
// The stage which produces a failure is responsible for allocating a working set member
// with error details.
invariant(WorkingSet::INVALID_ID != id);
*out = id;
return status;
} else if (PlanStage::NEED_YIELD == status) {
*out = id;
}
return status;
}
PlanStage::StageState TextOrStage::addTerm(WorkingSetID wsid, WorkingSetID* out) {
WorkingSetMember* wsm = _ws->get(wsid);
invariant(wsm->getState() == WorkingSetMember::LOC_AND_IDX);
invariant(1 == wsm->keyData.size());
const IndexKeyDatum newKeyData = wsm->keyData.back(); // copy to keep it around.
TextRecordData* textRecordData = &_scores[wsm->loc];
double* documentAggregateScore = &textRecordData->score;
if (WorkingSet::INVALID_ID == textRecordData->wsid) {
// We haven't seen this RecordId before. Keep the working set member around (it may be
// force-fetched on saveState()).
textRecordData->wsid = wsid;
if (_filter) {
// We have not seen this document before and need to apply a filter.
bool shouldKeep;
bool wasDeleted = false;
try {
TextMatchableDocument tdoc(
_txn, newKeyData.indexKeyPattern, newKeyData.keyData, _ws, wsid, _recordCursor);
shouldKeep = _filter->matches(&tdoc);
} catch (const WriteConflictException& wce) {
_idRetrying = wsid;
*out = WorkingSet::INVALID_ID;
return NEED_YIELD;
} catch (const TextMatchableDocument::DocumentDeletedException&) {
// We attempted to fetch the document but decided it should be excluded from the
// result set.
shouldKeep = false;
wasDeleted = true;
}
if (!shouldKeep) {
if (wasDeleted || wsm->hasObj()) {
// We had to fetch but we're not going to return it.
++_specificStats.fetches;
}
_ws->free(textRecordData->wsid);
textRecordData->wsid = WorkingSet::INVALID_ID;
*documentAggregateScore = -1;
return NEED_TIME;
}
}
} else {
// We already have a working set member for this RecordId. Free the new WSM and retrieve the
// old one. Note that since we don't keep all index keys, we could get a score that doesn't
// match the document, but this has always been a problem.
// TODO something to improve the situation.
invariant(wsid != textRecordData->wsid);
_ws->free(wsid);
wsm = _ws->get(textRecordData->wsid);
}
if (*documentAggregateScore < 0) {
// We have already rejected this document for not matching the filter.
return NEED_TIME;
}
// Locate score within possibly compound key: {prefix,term,score,suffix}.
BSONObjIterator keyIt(newKeyData.keyData);
for (unsigned i = 0; i < _ftsSpec.numExtraBefore(); i++) {
keyIt.next();
}
keyIt.next(); // Skip past 'term'.
BSONElement scoreElement = keyIt.next();
double documentTermScore = scoreElement.number();
// Aggregate relevance score, term keys.
*documentAggregateScore += documentTermScore;
return NEED_TIME;
}
PlanExecutor::ExecState PlanExecutor::getNextImpl(Snapshotted<BSONObj>* objOut, RecordId* dlOut) {
if (MONGO_FAIL_POINT(planExecutorAlwaysFails)) {
Status status(ErrorCodes::OperationFailed,
str::stream() << "PlanExecutor hit planExecutorAlwaysFails fail point");
*objOut =
Snapshotted<BSONObj>(SnapshotId(), WorkingSetCommon::buildMemberStatusObject(status));
return PlanExecutor::FAILURE;
}
invariant(_currentState == kUsable);
if (isMarkedAsKilled()) {
if (NULL != objOut) {
Status status(ErrorCodes::OperationFailed,
str::stream() << "Operation aborted because: " << *_killReason);
*objOut = Snapshotted<BSONObj>(SnapshotId(),
WorkingSetCommon::buildMemberStatusObject(status));
}
return PlanExecutor::DEAD;
}
if (!_stash.empty()) {
invariant(objOut && !dlOut);
*objOut = {SnapshotId(), _stash.front()};
_stash.pop();
return PlanExecutor::ADVANCED;
}
// When a stage requests a yield for document fetch, it gives us back a RecordFetcher*
// to use to pull the record into memory. We take ownership of the RecordFetcher here,
// deleting it after we've had a chance to do the fetch. For timing-based yields, we
// just pass a NULL fetcher.
unique_ptr<RecordFetcher> fetcher;
// Incremented on every writeConflict, reset to 0 on any successful call to _root->work.
size_t writeConflictsInARow = 0;
for (;;) {
// These are the conditions which can cause us to yield:
// 1) The yield policy's timer elapsed, or
// 2) some stage requested a yield due to a document fetch, or
// 3) we need to yield and retry due to a WriteConflictException.
// In all cases, the actual yielding happens here.
if (_yieldPolicy->shouldYield()) {
if (!_yieldPolicy->yield(fetcher.get())) {
// A return of false from a yield should only happen if we've been killed during the
// yield.
invariant(isMarkedAsKilled());
if (NULL != objOut) {
Status status(ErrorCodes::OperationFailed,
str::stream() << "Operation aborted because: " << *_killReason);
*objOut = Snapshotted<BSONObj>(
SnapshotId(), WorkingSetCommon::buildMemberStatusObject(status));
}
return PlanExecutor::DEAD;
}
}
// We're done using the fetcher, so it should be freed. We don't want to
// use the same RecordFetcher twice.
fetcher.reset();
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState code = _root->work(&id);
if (code != PlanStage::NEED_YIELD)
writeConflictsInARow = 0;
if (PlanStage::ADVANCED == code) {
WorkingSetMember* member = _workingSet->get(id);
bool hasRequestedData = true;
if (NULL != objOut) {
if (WorkingSetMember::RID_AND_IDX == member->getState()) {
if (1 != member->keyData.size()) {
_workingSet->free(id);
hasRequestedData = false;
} else {
// TODO: currently snapshot ids are only associated with documents, and
// not with index keys.
*objOut = Snapshotted<BSONObj>(SnapshotId(), member->keyData[0].keyData);
}
} else if (member->hasObj()) {
*objOut = member->obj;
} else {
_workingSet->free(id);
hasRequestedData = false;
}
}
if (NULL != dlOut) {
if (member->hasRecordId()) {
*dlOut = member->recordId;
} else {
_workingSet->free(id);
hasRequestedData = false;
}
}
if (hasRequestedData) {
_workingSet->free(id);
return PlanExecutor::ADVANCED;
}
// This result didn't have the data the caller wanted, try again.
} else if (PlanStage::NEED_YIELD == code) {
if (id == WorkingSet::INVALID_ID) {
if (!_yieldPolicy->canAutoYield())
throw WriteConflictException();
CurOp::get(_opCtx)->debug().writeConflicts++;
writeConflictsInARow++;
WriteConflictException::logAndBackoff(
writeConflictsInARow, "plan execution", _nss.ns());
} else {
WorkingSetMember* member = _workingSet->get(id);
invariant(member->hasFetcher());
// Transfer ownership of the fetcher. Next time around the loop a yield will
// happen.
fetcher.reset(member->releaseFetcher());
}
// If we're allowed to, we will yield next time through the loop.
if (_yieldPolicy->canAutoYield())
_yieldPolicy->forceYield();
} else if (PlanStage::NEED_TIME == code) {
// Fall through to yield check at end of large conditional.
} else if (PlanStage::IS_EOF == code) {
if (shouldWaitForInserts()) {
const bool locksReacquiredAfterYield = waitForInserts();
if (locksReacquiredAfterYield) {
// There may be more results, try to get more data.
continue;
}
invariant(isMarkedAsKilled());
if (objOut) {
Status status(ErrorCodes::OperationFailed,
str::stream() << "Operation aborted because: " << *_killReason);
*objOut = Snapshotted<BSONObj>(
SnapshotId(), WorkingSetCommon::buildMemberStatusObject(status));
}
return PlanExecutor::DEAD;
} else {
return PlanExecutor::IS_EOF;
}
} else {
invariant(PlanStage::DEAD == code || PlanStage::FAILURE == code);
if (NULL != objOut) {
BSONObj statusObj;
WorkingSetCommon::getStatusMemberObject(*_workingSet, id, &statusObj);
*objOut = Snapshotted<BSONObj>(SnapshotId(), statusObj);
}
return (PlanStage::DEAD == code) ? PlanExecutor::DEAD : PlanExecutor::FAILURE;
}
}
}
PlanStage::StageState DeleteStage::doWork(WorkingSetID* out) {
if (isEOF()) {
return PlanStage::IS_EOF;
}
invariant(_collection); // If isEOF() returns false, we must have a collection.
// It is possible that after a delete was executed, a WriteConflictException occurred
// and prevented us from returning ADVANCED with the old version of the document.
if (_idReturning != WorkingSet::INVALID_ID) {
// We should only get here if we were trying to return something before.
invariant(_params.returnDeleted);
WorkingSetMember* member = _ws->get(_idReturning);
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
*out = _idReturning;
_idReturning = WorkingSet::INVALID_ID;
return PlanStage::ADVANCED;
}
// Either retry the last WSM we worked on or get a new one from our child.
WorkingSetID id;
if (_idRetrying != WorkingSet::INVALID_ID) {
id = _idRetrying;
_idRetrying = WorkingSet::INVALID_ID;
} else {
auto status = child()->work(&id);
switch (status) {
case PlanStage::ADVANCED:
break;
case PlanStage::FAILURE:
case PlanStage::DEAD:
// The stage which produces a failure is responsible for allocating a working set
// member with error details.
invariant(WorkingSet::INVALID_ID != id);
*out = id;
return status;
case PlanStage::NEED_TIME:
return status;
case PlanStage::NEED_YIELD:
*out = id;
return status;
case PlanStage::IS_EOF:
return status;
default:
MONGO_UNREACHABLE;
}
}
// We advanced, or are retrying, and id is set to the WSM to work on.
WorkingSetMember* member = _ws->get(id);
// We want to free this member when we return, unless we need to retry deleting or returning it.
ScopeGuard memberFreer = MakeGuard(&WorkingSet::free, _ws, id);
invariant(member->hasRecordId());
RecordId recordId = member->recordId;
// Deletes can't have projections. This means that covering analysis will always add
// a fetch. We should always get fetched data, and never just key data.
invariant(member->hasObj());
// Ensure the document still exists and matches the predicate.
bool docStillMatches;
try {
docStillMatches = write_stage_common::ensureStillMatches(
_collection, getOpCtx(), _ws, id, _params.canonicalQuery);
} catch (const WriteConflictException&) {
// There was a problem trying to detect if the document still exists, so retry.
memberFreer.Dismiss();
return prepareToRetryWSM(id, out);
}
if (!docStillMatches) {
// Either the document has already been deleted, or it has been updated such that it no
// longer matches the predicate.
if (shouldRestartDeleteIfNoLongerMatches(_params)) {
throw WriteConflictException();
}
return PlanStage::NEED_TIME;
}
// Ensure that the BSONObj underlying the WorkingSetMember is owned because saveState() is
// allowed to free the memory.
if (_params.returnDeleted) {
// Save a copy of the document that is about to get deleted, but keep it in the RID_AND_OBJ
// state in case we need to retry deleting it.
BSONObj deletedDoc = member->obj.value();
member->obj.setValue(deletedDoc.getOwned());
}
// TODO: Do we want to buffer docs and delete them in a group rather than saving/restoring state
// repeatedly?
WorkingSetCommon::prepareForSnapshotChange(_ws);
try {
child()->saveState();
} catch (const WriteConflictException&) {
std::terminate();
}
// Do the write, unless this is an explain.
if (!_params.isExplain) {
try {
WriteUnitOfWork wunit(getOpCtx());
_collection->deleteDocument(getOpCtx(),
_params.stmtId,
recordId,
_params.opDebug,
_params.fromMigrate,
false,
_params.returnDeleted ? Collection::StoreDeletedDoc::On
: Collection::StoreDeletedDoc::Off);
wunit.commit();
} catch (const WriteConflictException&) {
memberFreer.Dismiss(); // Keep this member around so we can retry deleting it.
return prepareToRetryWSM(id, out);
}
}
++_specificStats.docsDeleted;
if (_params.returnDeleted) {
// After deleting the document, the RecordId associated with this member is invalid.
// Remove the 'recordId' from the WorkingSetMember before returning it.
member->recordId = RecordId();
member->transitionToOwnedObj();
}
// As restoreState may restore (recreate) cursors, cursors are tied to the transaction in which
// they are created, and a WriteUnitOfWork is a transaction, make sure to restore the state
// outside of the WriteUnitOfWork.
try {
child()->restoreState();
} catch (const WriteConflictException&) {
// Note we don't need to retry anything in this case since the delete already was committed.
// However, we still need to return the deleted document (if it was requested).
if (_params.returnDeleted) {
// member->obj should refer to the deleted document.
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
_idReturning = id;
// Keep this member around so that we can return it on the next work() call.
memberFreer.Dismiss();
}
*out = WorkingSet::INVALID_ID;
return NEED_YIELD;
}
if (_params.returnDeleted) {
// member->obj should refer to the deleted document.
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
memberFreer.Dismiss(); // Keep this member around so we can return it.
*out = id;
return PlanStage::ADVANCED;
}
return PlanStage::NEED_TIME;
}
PlanStage::StageState TextOrStage::addTerm(WorkingSetID wsid, WorkingSetID* out) {
WorkingSetMember* wsm = _ws->get(wsid);
invariant(wsm->getState() == WorkingSetMember::RID_AND_IDX);
invariant(1 == wsm->keyData.size());
const IndexKeyDatum newKeyData = wsm->keyData.back(); // copy to keep it around.
TextRecordData* textRecordData = &_scores[wsm->recordId];
if (textRecordData->score < 0) {
// We have already rejected this document for not matching the filter.
invariant(WorkingSet::INVALID_ID == textRecordData->wsid);
_ws->free(wsid);
return NEED_TIME;
}
if (WorkingSet::INVALID_ID == textRecordData->wsid) {
// We haven't seen this RecordId before.
invariant(textRecordData->score == 0);
bool shouldKeep = true;
if (_filter) {
// We have not seen this document before and need to apply a filter.
bool wasDeleted = false;
try {
TextMatchableDocument tdoc(getOpCtx(),
newKeyData.indexKeyPattern,
newKeyData.keyData,
_ws,
wsid,
_recordCursor);
shouldKeep = _filter->matches(&tdoc);
} catch (const WriteConflictException& wce) {
// Ensure that the BSONObj underlying the WorkingSetMember is owned because it may
// be freed when we yield.
wsm->makeObjOwnedIfNeeded();
_idRetrying = wsid;
*out = WorkingSet::INVALID_ID;
return NEED_YIELD;
} catch (const TextMatchableDocument::DocumentDeletedException&) {
// We attempted to fetch the document but decided it should be excluded from the
// result set.
shouldKeep = false;
wasDeleted = true;
}
if (wasDeleted || wsm->hasObj()) {
++_specificStats.fetches;
}
}
if (shouldKeep && !wsm->hasObj()) {
// Our parent expects RID_AND_OBJ members, so we fetch the document here if we haven't
// already.
try {
shouldKeep = WorkingSetCommon::fetch(getOpCtx(), _ws, wsid, _recordCursor);
++_specificStats.fetches;
} catch (const WriteConflictException& wce) {
wsm->makeObjOwnedIfNeeded();
_idRetrying = wsid;
*out = WorkingSet::INVALID_ID;
return NEED_YIELD;
}
}
if (!shouldKeep) {
_ws->free(wsid);
textRecordData->score = -1;
return NEED_TIME;
}
textRecordData->wsid = wsid;
// Ensure that the BSONObj underlying the WorkingSetMember is owned in case we yield.
wsm->makeObjOwnedIfNeeded();
} else {
// We already have a working set member for this RecordId. Free the new WSM and retrieve the
// old one. Note that since we don't keep all index keys, we could get a score that doesn't
// match the document, but this has always been a problem.
// TODO something to improve the situation.
invariant(wsid != textRecordData->wsid);
_ws->free(wsid);
wsm = _ws->get(textRecordData->wsid);
}
// Locate score within possibly compound key: {prefix,term,score,suffix}.
BSONObjIterator keyIt(newKeyData.keyData);
for (unsigned i = 0; i < _ftsSpec.numExtraBefore(); i++) {
keyIt.next();
}
keyIt.next(); // Skip past 'term'.
BSONElement scoreElement = keyIt.next();
double documentTermScore = scoreElement.number();
// Aggregate relevance score, term keys.
textRecordData->score += documentTermScore;
return NEED_TIME;
}
void run() {
// Populate the collection.
for (int i = 0; i < 50; ++i) {
insert(BSON("_id" << i << "foo" << i));
}
ASSERT_EQUALS(50U, count(BSONObj()));
// Various variables we'll need.
OldClientWriteContext ctx(&_txn, nss.ns());
OpDebug* opDebug = &CurOp::get(_txn)->debug();
Collection* coll = ctx.getCollection();
UpdateLifecycleImpl updateLifecycle(false, nss);
UpdateRequest request(nss);
UpdateDriver driver((UpdateDriver::Options()));
const int targetDocIndex = 10;
const BSONObj query = BSON("foo" << BSON("$gte" << targetDocIndex));
const unique_ptr<WorkingSet> ws(stdx::make_unique<WorkingSet>());
const unique_ptr<CanonicalQuery> cq(canonicalize(query));
// Get the RecordIds that would be returned by an in-order scan.
vector<RecordId> locs;
getLocs(coll, CollectionScanParams::FORWARD, &locs);
// Populate the request.
request.setQuery(query);
request.setUpdates(fromjson("{$set: {x: 0}}"));
request.setSort(BSONObj());
request.setMulti(false);
request.setReturnDocs(UpdateRequest::RETURN_NEW);
request.setLifecycle(&updateLifecycle);
ASSERT_OK(driver.parse(request.getUpdates(), request.isMulti()));
// Configure a QueuedDataStage to pass the first object in the collection back in a
// LOC_AND_OBJ state.
std::unique_ptr<QueuedDataStage> qds(stdx::make_unique<QueuedDataStage>(ws.get()));
WorkingSetID id = ws->allocate();
WorkingSetMember* member = ws->get(id);
member->loc = locs[targetDocIndex];
const BSONObj oldDoc = BSON("_id" << targetDocIndex << "foo" << targetDocIndex);
member->obj = Snapshotted<BSONObj>(SnapshotId(), oldDoc);
ws->transitionToLocAndObj(id);
qds->pushBack(id);
// Configure the update.
UpdateStageParams updateParams(&request, &driver, opDebug);
updateParams.canonicalQuery = cq.get();
unique_ptr<UpdateStage> updateStage(
stdx::make_unique<UpdateStage>(&_txn, updateParams, ws.get(), coll, qds.release()));
// Should return advanced.
id = WorkingSet::INVALID_ID;
PlanStage::StageState state = updateStage->work(&id);
ASSERT_EQUALS(PlanStage::ADVANCED, state);
// Make sure the returned value is what we expect it to be.
// Should give us back a valid id.
ASSERT_TRUE(WorkingSet::INVALID_ID != id);
WorkingSetMember* resultMember = ws->get(id);
// With an owned copy of the object, with no RecordId.
ASSERT_TRUE(resultMember->hasOwnedObj());
ASSERT_FALSE(resultMember->hasLoc());
ASSERT_EQUALS(resultMember->getState(), WorkingSetMember::OWNED_OBJ);
ASSERT_TRUE(resultMember->obj.value().isOwned());
// Should be the new value.
BSONObj newDoc = BSON("_id" << targetDocIndex << "foo" << targetDocIndex << "x" << 0);
ASSERT_EQUALS(resultMember->obj.value(), newDoc);
// Should have done the update.
vector<BSONObj> objs;
getCollContents(coll, &objs);
ASSERT_EQUALS(objs[targetDocIndex], newDoc);
// That should be it.
id = WorkingSet::INVALID_ID;
ASSERT_EQUALS(PlanStage::IS_EOF, updateStage->work(&id));
}
PlanStage::StageState UpdateStage::doWork(WorkingSetID* out) {
if (isEOF()) {
return PlanStage::IS_EOF;
}
if (doneUpdating()) {
// Even if we're done updating, we may have some inserting left to do.
if (needInsert()) {
// TODO we may want to handle WriteConflictException here. Currently we bounce it
// out to a higher level since if this WCEs it is likely that we raced with another
// upsert that may have matched our query, and therefore this may need to perform an
// update rather than an insert. Bouncing to the higher level allows restarting the
// query in this case.
doInsert();
invariant(isEOF());
if (_params.request->shouldReturnNewDocs()) {
// Want to return the document we just inserted, create it as a WorkingSetMember
// so that we can return it.
BSONObj newObj = _specificStats.objInserted;
*out = _ws->allocate();
WorkingSetMember* member = _ws->get(*out);
member->obj = Snapshotted<BSONObj>(getOpCtx()->recoveryUnit()->getSnapshotId(),
newObj.getOwned());
member->transitionToOwnedObj();
return PlanStage::ADVANCED;
}
}
// At this point either we're done updating and there was no insert to do,
// or we're done updating and we're done inserting. Either way, we're EOF.
invariant(isEOF());
return PlanStage::IS_EOF;
}
// If we're here, then we still have to ask for results from the child and apply
// updates to them. We should only get here if the collection exists.
invariant(_collection);
// It is possible that after an update was applied, a WriteConflictException
// occurred and prevented us from returning ADVANCED with the requested version
// of the document.
if (_idReturning != WorkingSet::INVALID_ID) {
// We should only get here if we were trying to return something before.
invariant(_params.request->shouldReturnAnyDocs());
WorkingSetMember* member = _ws->get(_idReturning);
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
*out = _idReturning;
_idReturning = WorkingSet::INVALID_ID;
return PlanStage::ADVANCED;
}
// Either retry the last WSM we worked on or get a new one from our child.
WorkingSetID id;
StageState status;
if (_idRetrying == WorkingSet::INVALID_ID) {
status = child()->work(&id);
} else {
status = ADVANCED;
id = _idRetrying;
_idRetrying = WorkingSet::INVALID_ID;
}
if (PlanStage::ADVANCED == status) {
// Need to get these things from the result returned by the child.
RecordId recordId;
WorkingSetMember* member = _ws->get(id);
// We want to free this member when we return, unless we need to retry updating or returning
// it.
ScopeGuard memberFreer = MakeGuard(&WorkingSet::free, _ws, id);
if (!member->hasRecordId()) {
// We expect to be here because of an invalidation causing a force-fetch.
++_specificStats.nInvalidateSkips;
return PlanStage::NEED_TIME;
}
recordId = member->recordId;
// Updates can't have projections. This means that covering analysis will always add
// a fetch. We should always get fetched data, and never just key data.
invariant(member->hasObj());
// We fill this with the new RecordIds of moved doc so we don't double-update.
if (_updatedRecordIds && _updatedRecordIds->count(recordId) > 0) {
// Found a RecordId that refers to a document we had already updated. Note that
// we can never remove from _updatedRecordIds because updates by other clients
// could cause us to encounter a document again later.
return PlanStage::NEED_TIME;
}
bool docStillMatches;
try {
docStillMatches = write_stage_common::ensureStillMatches(
_collection, getOpCtx(), _ws, id, _params.canonicalQuery);
} catch (const WriteConflictException&) {
// There was a problem trying to detect if the document still exists, so retry.
memberFreer.Dismiss();
return prepareToRetryWSM(id, out);
}
if (!docStillMatches) {
// Either the document has been deleted, or it has been updated such that it no longer
// matches the predicate.
if (shouldRestartUpdateIfNoLongerMatches(_params)) {
throw WriteConflictException();
}
return PlanStage::NEED_TIME;
}
// Ensure that the BSONObj underlying the WorkingSetMember is owned because saveState()
// is allowed to free the memory.
member->makeObjOwnedIfNeeded();
// Save state before making changes
WorkingSetCommon::prepareForSnapshotChange(_ws);
try {
child()->saveState();
} catch (const WriteConflictException&) {
std::terminate();
}
// If we care about the pre-updated version of the doc, save it out here.
BSONObj oldObj;
if (_params.request->shouldReturnOldDocs()) {
oldObj = member->obj.value().getOwned();
}
BSONObj newObj;
try {
// Do the update, get us the new version of the doc.
newObj = transformAndUpdate(member->obj, recordId);
} catch (const WriteConflictException&) {
memberFreer.Dismiss(); // Keep this member around so we can retry updating it.
return prepareToRetryWSM(id, out);
}
// Set member's obj to be the doc we want to return.
if (_params.request->shouldReturnAnyDocs()) {
if (_params.request->shouldReturnNewDocs()) {
member->obj = Snapshotted<BSONObj>(getOpCtx()->recoveryUnit()->getSnapshotId(),
newObj.getOwned());
} else {
invariant(_params.request->shouldReturnOldDocs());
member->obj.setValue(oldObj);
}
member->recordId = RecordId();
member->transitionToOwnedObj();
}
// This should be after transformAndUpdate to make sure we actually updated this doc.
++_specificStats.nMatched;
// Restore state after modification
// As restoreState may restore (recreate) cursors, make sure to restore the
// state outside of the WritUnitOfWork.
try {
child()->restoreState();
} catch (const WriteConflictException&) {
// Note we don't need to retry updating anything in this case since the update
// already was committed. However, we still need to return the updated document
// (if it was requested).
if (_params.request->shouldReturnAnyDocs()) {
// member->obj should refer to the document we want to return.
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
_idReturning = id;
// Keep this member around so that we can return it on the next work() call.
memberFreer.Dismiss();
}
*out = WorkingSet::INVALID_ID;
return NEED_YIELD;
}
if (_params.request->shouldReturnAnyDocs()) {
// member->obj should refer to the document we want to return.
invariant(member->getState() == WorkingSetMember::OWNED_OBJ);
memberFreer.Dismiss(); // Keep this member around so we can return it.
*out = id;
return PlanStage::ADVANCED;
}
return PlanStage::NEED_TIME;
} else if (PlanStage::IS_EOF == status) {
// The child is out of results, but we might not be done yet because we still might
// have to do an insert.
return PlanStage::NEED_TIME;
} else if (PlanStage::FAILURE == status) {
*out = id;
// If a stage fails, it may create a status WSM to indicate why it failed, in which case
// 'id' is valid. If ID is invalid, we create our own error message.
if (WorkingSet::INVALID_ID == id) {
const std::string errmsg = "update stage failed to read in results from child";
*out = WorkingSetCommon::allocateStatusMember(
_ws, Status(ErrorCodes::InternalError, errmsg));
return PlanStage::FAILURE;
}
return status;
} else if (PlanStage::NEED_YIELD == status) {
*out = id;
}
return status;
}
