C++ WorkingSet::get方法代码示例

        void run() {
            Client::WriteContext ctx(&_txn, ns());

            Database* db = ctx.ctx().db();
            Collection* coll = db->getCollection(&_txn, ns());
            if (!coll) {
                WriteUnitOfWork wuow(&_txn);
                coll = db->createCollection(&_txn, ns());
                wuow.commit();
            }
            WorkingSet ws;

            std::set<WorkingSetID> expectedResultIds;
            std::set<WorkingSetID> resultIds;

            // Create a KeepMutationsStage with an EOF child, and flag 50 objects.  We expect these
            // objects to be returned by the KeepMutationsStage.
            MatchExpression* nullFilter = NULL;
            std::auto_ptr<KeepMutationsStage> keep(new KeepMutationsStage(nullFilter, &ws,
                                                                          new EOFStage()));
            for (size_t i = 0; i < 50; ++i) {
                WorkingSetID id = ws.allocate();
                WorkingSetMember* member = ws.get(id);
                member->state = WorkingSetMember::OWNED_OBJ;
                member->obj = BSON("x" << 1);
                ws.flagForReview(id);
                expectedResultIds.insert(id);
            }

            // Call work() on the KeepMutationsStage.  The stage should start streaming the
            // already-flagged objects.
            WorkingSetID id = getNextResult(keep.get());
            resultIds.insert(id);

            // Flag more objects, then call work() again on the KeepMutationsStage, and expect none
            // of the newly-flagged objects to be returned (the KeepMutationsStage does not
            // incorporate objects flagged since the streaming phase started).
            //
            // This condition triggers SERVER-15580 (the new flagging causes a rehash of the
            // unordered_set "WorkingSet::_flagged", which invalidates all iterators, which were
            // previously being dereferenced in KeepMutationsStage::work()).
            // Note that std::unordered_set<>::insert() triggers a rehash if the new number of
            // elements is greater than or equal to max_load_factor()*bucket_count().
            size_t rehashSize = static_cast<size_t>(ws.getFlagged().max_load_factor() *
                                                    ws.getFlagged().bucket_count());
            while (ws.getFlagged().size() <= rehashSize) {
                WorkingSetID id = ws.allocate();
                WorkingSetMember* member = ws.get(id);
                member->state = WorkingSetMember::OWNED_OBJ;
                member->obj = BSON("x" << 1);
                ws.flagForReview(id);
            }
            while ((id = getNextResult(keep.get())) != WorkingSet::INVALID_ID) {
                resultIds.insert(id);
            }

            // Assert that only the first 50 objects were returned.
            ASSERT(expectedResultIds == resultIds);
        }

    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;

        // Add 10 objects to the collection.
        for (size_t i = 0; i < 10; ++i) {
            insert(BSON("x" << 1));
        }

        // Create 10 objects that are flagged.
        for (size_t i = 0; i < 10; ++i) {
            WorkingSetID id = ws.allocate();
            WorkingSetMember* member = ws.get(id);
            member->state = WorkingSetMember::OWNED_OBJ;
            member->obj = Snapshotted<BSONObj>(SnapshotId(), BSON("x" << 2));
            ws.flagForReview(id);
        }

        // Create a collscan to provide the 10 objects in the collection.
        CollectionScanParams params;
        params.collection = coll;
        params.direction = CollectionScanParams::FORWARD;
        params.tailable = false;
        params.start = RecordId();
        CollectionScan* cs = new CollectionScan(&_txn, params, &ws, NULL);

        // Create a KeepMutations stage to merge in the 10 flagged objects.
        // Takes ownership of 'cs'
        MatchExpression* nullFilter = NULL;
        std::unique_ptr<KeepMutationsStage> keep(new KeepMutationsStage(nullFilter, &ws, cs));

        for (size_t i = 0; i < 10; ++i) {
            WorkingSetID id = getNextResult(keep.get());
            WorkingSetMember* member = ws.get(id);
            ASSERT_FALSE(ws.isFlagged(id));
            ASSERT_EQUALS(member->obj.value()["x"].numberInt(), 1);
        }

        {
            WorkingSetID out;
            ASSERT_EQ(cs->work(&out), PlanStage::IS_EOF);
        }

        // Flagged results *must* be at the end.
        for (size_t i = 0; i < 10; ++i) {
            WorkingSetID id = getNextResult(keep.get());
            WorkingSetMember* member = ws.get(id);
            ASSERT(ws.isFlagged(id));
            ASSERT_EQUALS(member->obj.value()["x"].numberInt(), 2);
        }
    }

    void run() {
        dbtests::WriteContextForTests ctx(&_opCtx, ns());
        Database* db = ctx.db();
        Collection* coll = db->getCollection(&_opCtx, ns());
        if (!coll) {
            WriteUnitOfWork wuow(&_opCtx);
            coll = db->createCollection(&_opCtx, 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>(&_opCtx, &ws);

        // Mock data.
        {
            WorkingSetID id = ws.allocate();
            WorkingSetMember* mockMember = ws.get(id);
            mockMember->recordId = *recordIds.begin();
            mockMember->obj = coll->docFor(&_opCtx, mockMember->recordId);
            ws.transitionToRecordIdAndObj(id);
            // Points into our DB.
            mockStage->pushBack(id);
        }
        {
            WorkingSetID id = ws.allocate();
            WorkingSetMember* mockMember = ws.get(id);
            mockMember->recordId = RecordId();
            mockMember->obj = Snapshotted<BSONObj>(SnapshotId(), BSON("foo" << 6));
            mockMember->transitionToOwnedObj();
            ASSERT_TRUE(mockMember->obj.value().isOwned());
            mockStage->pushBack(id);
        }

        unique_ptr<FetchStage> fetchStage(
            new FetchStage(&_opCtx, &ws, mockStage.release(), NULL, coll));

        WorkingSetID id = WorkingSet::INVALID_ID;
        PlanStage::StageState state;

        // Don't bother doing any fetching if an obj exists already.
        state = fetchStage->work(&id);
        ASSERT_EQUALS(PlanStage::ADVANCED, state);
        state = fetchStage->work(&id);
        ASSERT_EQUALS(PlanStage::ADVANCED, state);

        // No more data to fetch, so, EOF.
        state = fetchStage->work(&id);
        ASSERT_EQUALS(PlanStage::IS_EOF, state);
    }

    void run() {
        OldClientWriteContext ctx(&_txn, ns());

        Collection* coll = ctx.getCollection();

        // Get the RecordIds that would be returned by an in-order scan.
        vector<RecordId> recordIds;
        getRecordIds(coll, CollectionScanParams::FORWARD, &recordIds);

        // Configure the scan.
        CollectionScanParams params;
        params.collection = coll;
        params.direction = CollectionScanParams::FORWARD;
        params.tailable = false;

        WorkingSet ws;
        unique_ptr<CollectionScan> scan(new CollectionScan(&_txn, params, &ws, NULL));

        int count = 0;
        while (count < 10) {
            WorkingSetID id = WorkingSet::INVALID_ID;
            PlanStage::StageState state = scan->work(&id);
            if (PlanStage::ADVANCED == state) {
                WorkingSetMember* member = ws.get(id);
                ASSERT_EQUALS(coll->docFor(&_txn, recordIds[count]).value()["foo"].numberInt(),
                              member->obj.value()["foo"].numberInt());
                ++count;
            }
        }

        // Remove recordIds[count].
        scan->saveState();
        {
            WriteUnitOfWork wunit(&_txn);
            scan->invalidate(&_txn, recordIds[count], INVALIDATION_DELETION);
            wunit.commit();  // to avoid rollback of the invalidate
        }
        remove(coll->docFor(&_txn, recordIds[count]).value());
        scan->restoreState();

        // Skip over recordIds[count].
        ++count;

        // Expect the rest.
        while (!scan->isEOF()) {
            WorkingSetID id = WorkingSet::INVALID_ID;
            PlanStage::StageState state = scan->work(&id);
            if (PlanStage::ADVANCED == state) {
                WorkingSetMember* member = ws.get(id);
                ASSERT_EQUALS(coll->docFor(&_txn, recordIds[count]).value()["foo"].numberInt(),
                              member->obj.value()["foo"].numberInt());
                ++count;
            }
        }

        ASSERT_EQUALS(numObj(), count);
    }

        void run() {
            OperationContextImpl txn;
            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;

            // Add 10 objects to the collection.
            for (size_t i = 0; i < 10; ++i) {
                insert(BSON("x" << 1));
            }

            // Create 10 objects that are flagged.
            for (size_t i = 0; i < 10; ++i) {
                WorkingSetID id = ws.allocate();
                WorkingSetMember* member = ws.get(id);
                member->state = WorkingSetMember::OWNED_OBJ;
                member->obj = BSON("x" << 2);
                ws.flagForReview(id);
            }

            // Create a collscan to provide the 10 objects in the collection.
            CollectionScanParams params;
            params.collection = coll;
            params.direction = CollectionScanParams::FORWARD;
            params.tailable = false;
            params.start = DiskLoc();
            CollectionScan* cs = new CollectionScan(params, &ws, NULL);

            // Create a KeepMutations stage to merge in the 10 flagged objects.
            // Takes ownership of 'cs'
            KeepMutationsStage* keep = new KeepMutationsStage(NULL, &ws, cs);

            for (size_t i = 0; i < 10; ++i) {
                WorkingSetID id = getNextResult(keep);
                WorkingSetMember* member = ws.get(id);
                ASSERT_FALSE(ws.isFlagged(id));
                ASSERT_EQUALS(member->obj["x"].numberInt(), 1);
            }

            ASSERT(cs->isEOF());

            // Flagged results *must* be at the end.
            for (size_t i = 0; i < 10; ++i) {
                WorkingSetID id = getNextResult(keep);
                WorkingSetMember* member = ws.get(id);
                ASSERT(ws.isFlagged(id));
                ASSERT_EQUALS(member->obj["x"].numberInt(), 2);
            }
        }

        void run() {
            Client::WriteContext ctx(&_txn, ns());

            Collection* coll = ctx.ctx().db()->getCollection( &_txn, ns() );

            // Get the DiskLocs that would be returned by an in-order scan.
            vector<DiskLoc> locs;
            getLocs(coll, CollectionScanParams::FORWARD, &locs);

            // Configure the scan.
            CollectionScanParams params;
            params.collection = coll;
            params.direction = CollectionScanParams::FORWARD;
            params.tailable = false;

            WorkingSet ws;
            scoped_ptr<CollectionScan> scan(new CollectionScan(&_txn, params, &ws, NULL));

            int count = 0;
            while (count < 10) {
                WorkingSetID id = WorkingSet::INVALID_ID;
                PlanStage::StageState state = scan->work(&id);
                if (PlanStage::ADVANCED == state) {
                    WorkingSetMember* member = ws.get(id);
                    ASSERT_EQUALS(coll->docFor(&_txn, locs[count])["foo"].numberInt(),
                                  member->obj["foo"].numberInt());
                    ++count;
                }
            }

            // Remove locs[count].
            scan->saveState();
            scan->invalidate(locs[count], INVALIDATION_DELETION);
            remove(coll->docFor(&_txn, locs[count]));
            scan->restoreState(&_txn);

            // Skip over locs[count].
            ++count;

            // Expect the rest.
            while (!scan->isEOF()) {
                WorkingSetID id = WorkingSet::INVALID_ID;
                PlanStage::StageState state = scan->work(&id);
                if (PlanStage::ADVANCED == state) {
                    WorkingSetMember* member = ws.get(id);
                    ASSERT_EQUALS(coll->docFor(&_txn, locs[count])["foo"].numberInt(),
                                  member->obj["foo"].numberInt());
                    ++count;
                }
            }
            ctx.commit();

            ASSERT_EQUALS(numObj(), count);
        }

        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" << 1 << "bar" << i));
            }

            addIndex(BSON("foo" << 1));
            addIndex(BSON("bar" << 1));

            WorkingSet ws;
            scoped_ptr<AndHashStage> ah(new AndHashStage(&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));

            // Intersect with 7 <= bar < 10000
            params.descriptor = getIndex(BSON("bar" << 1), coll);
            params.bounds.startKey = BSON("" << 7);
            params.bounds.endKey = BSON("" << 10000);
            ah->addChild(new IndexScan(params, &ws, NULL));

            WorkingSetID lastId = WorkingSet::INVALID_ID;

            int count = 0;
            while (!ah->isEOF()) {
                WorkingSetID id;
                PlanStage::StageState status = ah->work(&id);
                if (PlanStage::ADVANCED != status) { continue; }
                BSONObj thisObj = ws.get(id)->loc.obj();
                ASSERT_EQUALS(7 + count, thisObj["bar"].numberInt());
                ++count;
                if (WorkingSet::INVALID_ID != lastId) {
                    BSONObj lastObj = ws.get(lastId)->loc.obj();
                    ASSERT_LESS_THAN(lastObj["bar"].woCompare(thisObj["bar"]), 0);
                }
                lastId = id;
            }

            ASSERT_EQUALS(count, 43);
        }

    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() {
        OldClientWriteContext ctx(&_txn, nss.ns());
        addIndex(BSON("b" << 1 << "a" << 1));
        addIndex(BSON("c" << 1 << "a" << 1));

        BSONObj query = fromjson("{a: 1, $or: [{b: 2}, {c: 3}]}");

        // Two of these documents match.
        insert(BSON("_id" << 1 << "a" << 1 << "b" << 2));
        insert(BSON("_id" << 2 << "a" << 2 << "b" << 2));
        insert(BSON("_id" << 3 << "a" << 1 << "c" << 3));
        insert(BSON("_id" << 4 << "a" << 1 << "c" << 4));

        auto qr = stdx::make_unique<QueryRequest>(nss);
        qr->setFilter(query);
        auto cq = unittest::assertGet(CanonicalQuery::canonicalize(
            txn(), std::move(qr), ExtensionsCallbackDisallowExtensions()));

        Collection* collection = ctx.getCollection();

        // Get planner params.
        QueryPlannerParams plannerParams;
        fillOutPlannerParams(&_txn, collection, cq.get(), &plannerParams);

        WorkingSet ws;
        std::unique_ptr<SubplanStage> subplan(
            new SubplanStage(&_txn, collection, &ws, plannerParams, cq.get()));

        // Plan selection should succeed due to falling back on regular planning.
        PlanYieldPolicy yieldPolicy(PlanExecutor::YIELD_MANUAL, _clock);
        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));

        // Work the stage until it produces all results.
        size_t numResults = 0;
        PlanStage::StageState stageState = PlanStage::NEED_TIME;
        while (stageState != PlanStage::IS_EOF) {
            WorkingSetID id = WorkingSet::INVALID_ID;
            stageState = subplan->work(&id);
            ASSERT_NE(stageState, PlanStage::DEAD);
            ASSERT_NE(stageState, PlanStage::FAILURE);

            if (stageState == PlanStage::ADVANCED) {
                ++numResults;
                WorkingSetMember* member = ws.get(id);
                ASSERT(member->hasObj());
                ASSERT(member->obj.value() == BSON("_id" << 1 << "a" << 1 << "b" << 2) ||
                       member->obj.value() == BSON("_id" << 3 << "a" << 1 << "c" << 3));
            }
        }

        ASSERT_EQ(numResults, 2U);
    }

    // static
    void WorkingSetCommon::getStatusMemberObject(const WorkingSet& ws, WorkingSetID wsid,
                                                 BSONObj* objOut) {
        invariant(objOut);

        // Validate ID and working set member.
        if (WorkingSet::INVALID_ID == wsid) {
            return;
        }
        WorkingSetMember* member = ws.get(wsid);
        if (!member->hasOwnedObj()) {
            return;
        }
        BSONObj obj = member->obj.value();
        if (!isValidStatusMemberObject(obj)) {
            return;
        }
        *objOut = obj;
    }

    /**
     * Returns a vector of all of the documents currently in 'collection'.
     *
     * Uses a forward collection scan stage to get the docs, and populates 'out' with
     * the results.
     */
    void getCollContents(Collection* collection, vector<BSONObj>* out) {
        WorkingSet ws;

        CollectionScanParams params;
        params.direction = CollectionScanParams::FORWARD;
        params.tailable = false;

        unique_ptr<CollectionScan> scan(new CollectionScan(&_opCtx, collection, params, &ws, NULL));
        while (!scan->isEOF()) {
            WorkingSetID id = WorkingSet::INVALID_ID;
            PlanStage::StageState state = scan->work(&id);
            if (PlanStage::ADVANCED == state) {
                WorkingSetMember* member = ws.get(id);
                verify(member->hasObj());
                out->push_back(member->obj.value().getOwned());
            }
        }
    }

        /**
         * 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 getRecordIds(Collection* collection,
                      CollectionScanParams::Direction direction,
                      vector<RecordId>* out) {
        WorkingSet ws;

        CollectionScanParams params;
        params.direction = direction;
        params.tailable = false;

        unique_ptr<CollectionScan> scan(new CollectionScan(&_opCtx, collection, params, &ws, NULL));
        while (!scan->isEOF()) {
            WorkingSetID id = WorkingSet::INVALID_ID;
            PlanStage::StageState state = scan->work(&id);
            if (PlanStage::ADVANCED == state) {
                WorkingSetMember* member = ws.get(id);
                verify(member->hasRecordId());
                out->push_back(member->recordId);
            }
        }
    }

        void getLocs() {
            _locs.clear();
            WorkingSet ws;

            CollectionScanParams params;
            params.collection = _coll;
            params.direction = CollectionScanParams::FORWARD;
            params.tailable = false;

            scoped_ptr<CollectionScan> scan(new CollectionScan(&_txn, params, &ws, NULL));
            while (!scan->isEOF()) {
                WorkingSetID id = WorkingSet::INVALID_ID;
                PlanStage::StageState state = scan->work(&id);
                if (PlanStage::ADVANCED == state) {
                    WorkingSetMember* member = ws.get(id);
                    verify(member->hasLoc());
                    _locs.push_back(member->loc);
                }
            }
        }

    static size_t getNumResultsForStage(const WorkingSet& ws,
                                        CachedPlanStage* cachedPlanStage,
                                        CanonicalQuery* cq) {
        size_t numResults = 0;
        PlanStage::StageState state = PlanStage::NEED_TIME;
        while (state != PlanStage::IS_EOF) {
            WorkingSetID id = WorkingSet::INVALID_ID;
            state = cachedPlanStage->work(&id);

            ASSERT_NE(state, PlanStage::FAILURE);
            ASSERT_NE(state, PlanStage::DEAD);

            if (state == PlanStage::ADVANCED) {
                WorkingSetMember* member = ws.get(id);
                ASSERT(cq->root()->matchesBSON(member->obj.value()));
                numResults++;
            }
        }

        return numResults;
    }