// Is lhs less than rhs? Note that priority_queue is a max heap by default so we invert
// the return from the expected value.
bool MergeSortStage::StageWithValueComparison::operator()(const MergingRef& lhs,
const MergingRef& rhs) {
WorkingSetMember* lhsMember = _ws->get(lhs->id);
WorkingSetMember* rhsMember = _ws->get(rhs->id);
BSONObjIterator it(_pattern);
while (it.more()) {
BSONElement patternElt = it.next();
string fn = patternElt.fieldName();
BSONElement lhsElt;
verify(lhsMember->getFieldDotted(fn, &lhsElt));
BSONElement rhsElt;
verify(rhsMember->getFieldDotted(fn, &rhsElt));
// false means don't compare field name.
int x = lhsElt.woCompare(rhsElt, false);
if (-1 == patternElt.number()) {
x = -x;
}
if (x != 0) {
return x > 0;
}
}
// A comparator for use with sort is required to model a strict weak ordering, so
// to satisfy irreflexivity we must return 'false' for elements that we consider
// equivalent under the pattern.
return false;
}
void run() {
ScopedTransaction transaction(&_txn, MODE_IX);
Lock::DBLock lk(_txn.lockState(), nsToDatabaseSubstring(ns()), MODE_X);
OldClientContext ctx(&_txn, ns());
Database* db = ctx.db();
Collection* coll = db->getCollection(ns());
if (!coll) {
WriteUnitOfWork wuow(&_txn);
coll = db->createCollection(&_txn, ns());
wuow.commit();
}
WorkingSet ws;
// Add an object to the DB.
insert(BSON("foo" << 5));
set<RecordId> recordIds;
getRecordIds(&recordIds, coll);
ASSERT_EQUALS(size_t(1), recordIds.size());
// Create a mock stage that returns the WSM.
auto mockStage = make_unique<QueuedDataStage>(&_txn, &ws);
// Mock data.
{
WorkingSetID id = ws.allocate();
WorkingSetMember* mockMember = ws.get(id);
mockMember->recordId = *recordIds.begin();
ws.transitionToRecordIdAndIdx(id);
// State is RecordId and index, shouldn't be able to get the foo data inside.
BSONElement elt;
ASSERT_FALSE(mockMember->getFieldDotted("foo", &elt));
mockStage->pushBack(id);
}
// Make the filter.
BSONObj filterObj = BSON("foo" << 6);
const CollatorInterface* collator = nullptr;
StatusWithMatchExpression statusWithMatcher = MatchExpressionParser::parse(
filterObj, ExtensionsCallbackDisallowExtensions(), collator);
verify(statusWithMatcher.isOK());
unique_ptr<MatchExpression> filterExpr = std::move(statusWithMatcher.getValue());
// Matcher requires that foo==6 but we only have data with foo==5.
unique_ptr<FetchStage> fetchStage(
new FetchStage(&_txn, &ws, mockStage.release(), filterExpr.get(), coll));
// First call should return a fetch request as it's not in memory.
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState state;
// Normally we'd return the object but we have a filter that prevents it.
state = fetchStage->work(&id);
ASSERT_EQUALS(PlanStage::NEED_TIME, state);
// No more data to fetch, so, EOF.
state = fetchStage->work(&id);
ASSERT_EQUALS(PlanStage::IS_EOF, state);
}
void run() {
ScopedTransaction transaction(&_txn, MODE_IX);
Lock::DBLock lk(_txn.lockState(), nsToDatabaseSubstring(ns()), MODE_X);
Client::Context ctx(&_txn, ns());
Database* db = ctx.db();
Collection* coll = db->getCollection(&_txn, ns());
if (!coll) {
WriteUnitOfWork wuow(&_txn);
coll = db->createCollection(&_txn, ns());
wuow.commit();
}
WorkingSet ws;
// Add an object to the DB.
insert(BSON("foo" << 5));
set<DiskLoc> locs;
getLocs(&locs, coll);
ASSERT_EQUALS(size_t(1), locs.size());
// Create a mock stage that returns the WSM.
auto_ptr<MockStage> mockStage(new MockStage(&ws));
// Mock data.
{
WorkingSetMember mockMember;
mockMember.state = WorkingSetMember::LOC_AND_IDX;
mockMember.loc = *locs.begin();
// State is loc and index, shouldn't be able to get the foo data inside.
BSONElement elt;
ASSERT_FALSE(mockMember.getFieldDotted("foo", &elt));
mockStage->pushBack(mockMember);
}
// Make the filter.
BSONObj filterObj = BSON("foo" << 6);
StatusWithMatchExpression swme = MatchExpressionParser::parse(filterObj);
verify(swme.isOK());
auto_ptr<MatchExpression> filterExpr(swme.getValue());
// Matcher requires that foo==6 but we only have data with foo==5.
auto_ptr<FetchStage> fetchStage(
new FetchStage(&_txn, &ws, mockStage.release(), filterExpr.get(), coll));
// First call should return a fetch request as it's not in memory.
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState state;
// Normally we'd return the object but we have a filter that prevents it.
state = fetchStage->work(&id);
ASSERT_EQUALS(PlanStage::NEED_TIME, state);
// No more data to fetch, so, EOF.
state = fetchStage->work(&id);
ASSERT_EQUALS(PlanStage::IS_EOF, state);
}
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();
}
/**
* Returns the projected value from the working set that would
* be returned in the 'values' field of the distinct command result.
* Limited to NumberInt BSON types because this is the only
* BSON type used in this suite of tests.
*/
static int getIntFieldDotted(const WorkingSet& ws, WorkingSetID wsid,
const std::string& field) {
// For some reason (at least under OS X clang), we cannot refer to INVALID_ID
// inside the test assertion macro.
WorkingSetID invalid = WorkingSet::INVALID_ID;
ASSERT_NOT_EQUALS(invalid, wsid);
WorkingSetMember* member = ws.get(wsid);
// Distinct hack execution is always covered.
// Key value is retrieved from working set key data
// instead of RecordId.
ASSERT_FALSE(member->hasObj());
BSONElement keyElt;
ASSERT_TRUE(member->getFieldDotted(field, &keyElt));
ASSERT_TRUE(keyElt.isNumber());
return keyElt.numberInt();
}
void run() {
Client::WriteContext ctx(&_txn, ns());
Database* db = ctx.ctx().db();
Collection* coll = db->getCollection(&_txn, ns());
if (!coll) {
coll = db->createCollection(&_txn, ns());
}
WorkingSet ws;
// Sort by foo:1
MergeSortStageParams msparams;
msparams.pattern = BSON("foo" << 1);
auto_ptr<MergeSortStage> ms(new MergeSortStage(msparams, &ws, coll));
IndexScanParams params;
params.bounds.isSimpleRange = true;
params.bounds.startKey = objWithMinKey(1);
params.bounds.endKey = objWithMaxKey(1);
params.bounds.endKeyInclusive = true;
params.direction = 1;
// Index 'a'+i has foo equal to 'i'.
int numIndices = 20;
for (int i = 0; i < numIndices; ++i) {
// 'a', 'b', ...
string index(1, 'a' + i);
insert(BSON(index << 1 << "foo" << i));
BSONObj indexSpec = BSON(index << 1 << "foo" << 1);
addIndex(indexSpec);
params.descriptor = getIndex(indexSpec, coll);
ms->addChild(new IndexScan(&_txn, params, &ws, NULL));
}
set<DiskLoc> locs;
getLocs(&locs, coll);
set<DiskLoc>::iterator it = locs.begin();
ctx.commit();
// Get 10 results. Should be getting results in order of 'locs'.
int count = 0;
while (!ms->isEOF() && count < 10) {
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState status = ms->work(&id);
if (PlanStage::ADVANCED != status) { continue; }
WorkingSetMember* member = ws.get(id);
ASSERT_EQUALS(member->loc, *it);
BSONElement elt;
string index(1, 'a' + count);
ASSERT(member->getFieldDotted(index, &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(count, elt.numberInt());
++count;
++it;
}
// Invalidate locs[11]. Should force a fetch. We don't get it back.
ms->prepareToYield();
ms->invalidate(*it, INVALIDATION_DELETION);
ms->recoverFromYield(&_txn);
// Make sure locs[11] was fetched for us.
{
// TODO: If we have "return upon invalidation" ever triggerable, do the following test.
/*
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState status;
do {
status = ms->work(&id);
} while (PlanStage::ADVANCED != status);
WorkingSetMember* member = ws.get(id);
ASSERT(!member->hasLoc());
ASSERT(member->hasObj());
string index(1, 'a' + count);
BSONElement elt;
ASSERT_TRUE(member->getFieldDotted(index, &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(count, elt.numberInt());
*/
++it;
++count;
}
// And get the rest.
while (!ms->isEOF()) {
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState status = ms->work(&id);
if (PlanStage::ADVANCED != status) { continue; }
WorkingSetMember* member = ws.get(id);
ASSERT_EQUALS(member->loc, *it);
BSONElement elt;
string index(1, 'a' + count);
ASSERT_TRUE(member->getFieldDotted(index, &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(count, elt.numberInt());
++count;
++it;
}
}
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();
}
WorkingSet ws;
// Sort by foo:1
MergeSortStageParams msparams;
msparams.pattern = BSON("foo" << 1);
auto ms = make_unique<MergeSortStage>(&_txn, msparams, &ws, coll);
IndexScanParams params;
params.bounds.isSimpleRange = true;
params.bounds.startKey = objWithMinKey(1);
params.bounds.endKey = objWithMaxKey(1);
params.bounds.endKeyInclusive = true;
params.direction = 1;
// Index 'a'+i has foo equal to 'i'.
int numIndices = 20;
for (int i = 0; i < numIndices; ++i) {
// 'a', 'b', ...
string index(1, 'a' + i);
insert(BSON(index << 1 << "foo" << i));
BSONObj indexSpec = BSON(index << 1 << "foo" << 1);
addIndex(indexSpec);
params.descriptor = getIndex(indexSpec, coll);
ms->addChild(new IndexScan(&_txn, params, &ws, NULL));
}
set<RecordId> recordIds;
getRecordIds(&recordIds, coll);
set<RecordId>::iterator it = recordIds.begin();
// Get 10 results. Should be getting results in order of 'recordIds'.
int count = 0;
while (!ms->isEOF() && count < 10) {
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState status = ms->work(&id);
if (PlanStage::ADVANCED != status) {
continue;
}
WorkingSetMember* member = ws.get(id);
ASSERT_EQUALS(member->recordId, *it);
BSONElement elt;
string index(1, 'a' + count);
ASSERT(member->getFieldDotted(index, &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(count, elt.numberInt());
++count;
++it;
}
// Invalidate recordIds[11]. Should force a fetch and return the deleted document.
ms->saveState();
ms->invalidate(&_txn, *it, INVALIDATION_DELETION);
ms->restoreState();
// Make sure recordIds[11] was fetched for us.
{
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState status;
do {
status = ms->work(&id);
} while (PlanStage::ADVANCED != status);
WorkingSetMember* member = ws.get(id);
ASSERT(!member->hasRecordId());
ASSERT(member->hasObj());
string index(1, 'a' + count);
BSONElement elt;
ASSERT_TRUE(member->getFieldDotted(index, &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(count, elt.numberInt());
++it;
++count;
}
// And get the rest.
while (!ms->isEOF()) {
WorkingSetID id = WorkingSet::INVALID_ID;
PlanStage::StageState status = ms->work(&id);
if (PlanStage::ADVANCED != status) {
continue;
}
WorkingSetMember* member = ws.get(id);
ASSERT_EQUALS(member->recordId, *it);
BSONElement elt;
string index(1, 'a' + count);
ASSERT_TRUE(member->getFieldDotted(index, &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(count, elt.numberInt());
++count;
++it;
}
}
void run() {
Client::WriteContext ctx(ns());
Database* db = ctx.ctx().db();
Collection* coll = db->getCollection(ns());
if (!coll) {
coll = db->createCollection(ns());
}
// Insert a bunch of data
for (int i = 0; i < 50; ++i) {
insert(BSON("foo" << 1 << "bar" << 1));
}
addIndex(BSON("foo" << 1));
addIndex(BSON("bar" << 1));
WorkingSet ws;
scoped_ptr<AndSortedStage> ah(new AndSortedStage(&ws, NULL));
// Scan over foo == 1
IndexScanParams params;
params.descriptor = getIndex(BSON("foo" << 1), coll);
params.bounds.isSimpleRange = true;
params.bounds.startKey = BSON("" << 1);
params.bounds.endKey = BSON("" << 1);
params.bounds.endKeyInclusive = true;
params.direction = 1;
ah->addChild(new IndexScan(params, &ws, NULL));
// Scan over bar == 1
params.descriptor = getIndex(BSON("bar" << 1), coll);
ah->addChild(new IndexScan(params, &ws, NULL));
// Get the set of disklocs in our collection to use later.
set<DiskLoc> data;
getLocs(&data, coll);
// We're making an assumption here that happens to be true because we clear out the
// collection before running this: increasing inserts have increasing DiskLocs.
// This isn't true in general if the collection is not dropped beforehand.
WorkingSetID id;
// Sorted AND looks at the first child, which is an index scan over foo==1.
ah->work(&id);
// The first thing that the index scan returns (due to increasing DiskLoc trick) is the
// very first insert, which should be the very first thing in data. Let's invalidate it
// and make sure it shows up in the flagged results.
ah->prepareToYield();
ah->invalidate(*data.begin());
remove(data.begin()->obj());
ah->recoverFromYield();
// Make sure the nuked obj is actually in the flagged data.
ASSERT_EQUALS(ws.getFlagged().size(), size_t(1));
WorkingSetMember* member = ws.get(*ws.getFlagged().begin());
ASSERT_EQUALS(WorkingSetMember::OWNED_OBJ, member->state);
BSONElement elt;
ASSERT_TRUE(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT_TRUE(member->getFieldDotted("bar", &elt));
ASSERT_EQUALS(1, elt.numberInt());
set<DiskLoc>::iterator it = data.begin();
// Proceed along, AND-ing results.
int count = 0;
while (!ah->isEOF() && count < 10) {
WorkingSetID id;
PlanStage::StageState status = ah->work(&id);
if (PlanStage::ADVANCED != status) {
continue;
}
++count;
++it;
member = ws.get(id);
ASSERT_TRUE(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT_TRUE(member->getFieldDotted("bar", &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT_EQUALS(member->loc, *it);
}
// Move 'it' to a result that's yet to show up.
for (int i = 0; i < count + 10; ++i) {
++it;
}
// Remove a result that's coming up. It's not the 'target' result of the AND so it's
// not flagged.
ah->prepareToYield();
ah->invalidate(*it);
remove(it->obj());
ah->recoverFromYield();
// Get all results aside from the two we killed.
while (!ah->isEOF()) {
WorkingSetID id;
PlanStage::StageState status = ah->work(&id);
if (PlanStage::ADVANCED != status) {
continue;
}
++count;
member = ws.get(id);
ASSERT_TRUE(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(1, elt.numberInt());
ASSERT_TRUE(member->getFieldDotted("bar", &elt));
ASSERT_EQUALS(1, elt.numberInt());
}
ASSERT_EQUALS(count, 48);
ASSERT_EQUALS(size_t(1), ws.getFlagged().size());
}
void run() {
Client::WriteContext ctx(ns());
Database* db = ctx.ctx().db();
Collection* coll = db->getCollection(ns());
if (!coll) {
coll = db->createCollection(ns());
}
for (int i = 0; i < 50; ++i) {
insert(BSON("foo" << i << "bar" << i));
}
addIndex(BSON("foo" << 1));
addIndex(BSON("bar" << 1));
WorkingSet ws;
scoped_ptr<AndHashStage> ah(new AndHashStage(&ws, NULL));
// Foo <= 20
IndexScanParams params;
params.descriptor = getIndex(BSON("foo" << 1), coll);
params.bounds.isSimpleRange = true;
params.bounds.startKey = BSON("" << 20);
params.bounds.endKey = BSONObj();
params.bounds.endKeyInclusive = true;
params.direction = -1;
ah->addChild(new IndexScan(params, &ws, NULL));
// Bar >= 10
params.descriptor = getIndex(BSON("bar" << 1), coll);
params.bounds.startKey = BSON("" << 10);
params.bounds.endKey = BSONObj();
params.bounds.endKeyInclusive = true;
params.direction = 1;
ah->addChild(new IndexScan(params, &ws, NULL));
// ah reads the first child into its hash table.
// ah should read foo=20, foo=19, ..., foo=0 in that order.
// Read half of them...
for (int i = 0; i < 10; ++i) {
WorkingSetID out;
PlanStage::StageState status = ah->work(&out);
ASSERT_EQUALS(PlanStage::NEED_TIME, status);
}
// ...yield
ah->prepareToYield();
// ...invalidate one of the read objects
set<DiskLoc> data;
getLocs(&data, coll);
for (set<DiskLoc>::const_iterator it = data.begin(); it != data.end(); ++it) {
if (it->obj()["foo"].numberInt() == 15) {
ah->invalidate(*it);
remove(it->obj());
break;
}
}
ah->recoverFromYield();
// And expect to find foo==15 it flagged for review.
const unordered_set<WorkingSetID>& flagged = ws.getFlagged();
ASSERT_EQUALS(size_t(1), flagged.size());
// Expect to find the right value of foo in the flagged item.
WorkingSetMember* member = ws.get(*flagged.begin());
ASSERT_TRUE(NULL != member);
ASSERT_EQUALS(WorkingSetMember::OWNED_OBJ, member->state);
BSONElement elt;
ASSERT_TRUE(member->getFieldDotted("foo", &elt));
ASSERT_EQUALS(15, elt.numberInt());
// Now, finish up the AND. Since foo == bar, we would have 11 results, but we subtract
// one because of a mid-plan invalidation, so 10.
int count = 0;
while (!ah->isEOF()) {
WorkingSetID id;
PlanStage::StageState status = ah->work(&id);
if (PlanStage::ADVANCED != status) {
continue;
}
++count;
member = ws.get(id);
ASSERT_TRUE(member->getFieldDotted("foo", &elt));
ASSERT_LESS_THAN_OR_EQUALS(elt.numberInt(), 20);
ASSERT_NOT_EQUALS(15, elt.numberInt());
ASSERT_TRUE(member->getFieldDotted("bar", &elt));
ASSERT_GREATER_THAN_OR_EQUALS(elt.numberInt(), 10);
}
ASSERT_EQUALS(10, count);
}
