Commit message: SERVER-60218 add initialize helper function for document_source_group
(cherry picked from commit 867f52afbb79bc00e35c70f8e0681b7d602f97b2)
CVE: None
Vulnerability description: None

Project: mongo
File: document_source_group.cpp

Document DocumentSourceGroup::makeDocument(const Value& id,
                                           const Accumulators& accums,
                                           bool mergeableOutput) {
    const size_t n = _accumulatedFields.size();
    MutableDocument out(1 + n);

    /* add the _id field */
    out.addField("_id", expandId(id));

    /* add the rest of the fields */
    for (size_t i = 0; i < n; ++i) {
        Value val = accums[i]->getValue(mergeableOutput);
        if (val.missing()) {
            // we return null in this case so return objects are predictable
            out.addField(_accumulatedFields[i].fieldName, Value(BSONNULL));
        } else {
            out.addField(_accumulatedFields[i].fieldName, std::move(val));
        }
    }

    _stats.totalOutputDataSizeBytes += out.getApproximateSize();
    return out.freeze();
}

Project: mongo
File: document_source_group.cpp

size_t DocumentSourceGroup::getMaxMemoryUsageBytes() const {
    return _memoryTracker._maxAllowedMemoryUsageBytes;
}

Project: mongo
File: document_source_group.cpp

Value DocumentSourceGroup::serialize(boost::optional<ExplainOptions::Verbosity> explain) const {
    MutableDocument insides;

    // Add the _id.
    if (_idFieldNames.empty()) {
        invariant(_idExpressions.size() == 1);
        insides["_id"] = _idExpressions[0]->serialize(static_cast<bool>(explain));
    } else {
        // Decomposed document case.
        invariant(_idExpressions.size() == _idFieldNames.size());
        MutableDocument md;
        for (size_t i = 0; i < _idExpressions.size(); i++) {
            md[_idFieldNames[i]] = _idExpressions[i]->serialize(static_cast<bool>(explain));
        }
        insides["_id"] = md.freezeToValue();
    }

    // Add the remaining fields.
    for (auto&& accumulatedField : _accumulatedFields) {
        intrusive_ptr<AccumulatorState> accum = accumulatedField.makeAccumulator();
        insides[accumulatedField.fieldName] =
            Value(accum->serialize(accumulatedField.expr.initializer,
                                   accumulatedField.expr.argument,
                                   static_cast<bool>(explain)));
    }

    if (_doingMerge) {
        // This makes the output unparsable (with error) on pre 2.6 shards, but it will never
        // be sent to old shards when this flag is true since they can't do a merge anyway.
        insides["$doingMerge"] = Value(true);
    }

    MutableDocument out;
    out[getSourceName()] = Value(insides.freeze());

    if (explain && *explain >= ExplainOptions::Verbosity::kExecStats) {
        MutableDocument md;

        for (size_t i = 0; i < _accumulatedFields.size(); i++) {
            md[_accumulatedFields[i].fieldName] = Value(static_cast<long long>(
                _memoryTracker[_accumulatedFields[i].fieldName].maxMemoryBytes()));
        }

        out["maxAccumulatorMemoryUsageBytes"] = Value(md.freezeToValue());
        out["totalOutputDataSizeBytes"] =
            Value(static_cast<long long>(_stats.totalOutputDataSizeBytes));
        out["usedDisk"] = Value(_stats.usedDisk);
    }

    return Value(out.freezeToValue());
}

Project: mongo
File: document_source_group.cpp

DocumentSourceGroup::rewriteGroupAsTransformOnFirstDocument() const {
    if (_idExpressions.size() != 1) {
        // This transformation is only intended for $group stages that group on a single field.
        return nullptr;
    }

    auto fieldPathExpr = dynamic_cast<ExpressionFieldPath*>(_idExpressions.front().get());
    if (!fieldPathExpr || fieldPathExpr->isVariableReference()) {
        return nullptr;
    }

    const auto fieldPath = fieldPathExpr->getFieldPath();
    if (fieldPath.getPathLength() == 1) {
        // The path is $$CURRENT or $$ROOT. This isn't really a sensible value to group by (since
        // each document has a unique _id, it will just return the entire collection). We only
        // apply the rewrite when grouping by a single field, so we cannot apply it in this case,
        // where we are grouping by the entire document.
        invariant(fieldPath.getFieldName(0) == "CURRENT" || fieldPath.getFieldName(0) == "ROOT");
        return nullptr;
    }

    const auto groupId = fieldPath.tail().fullPath();

    // We can't do this transformation if there are any non-$first accumulators.
    for (auto&& accumulator : _accumulatedFields) {
        if (AccumulatorDocumentsNeeded::kFirstDocument !=
            accumulator.makeAccumulator()->documentsNeeded()) {
            return nullptr;
        }
    }

    std::vector<std::pair<std::string, boost::intrusive_ptr<Expression>>> fields;

    boost::intrusive_ptr<Expression> idField;
    // The _id field can be specified either as a fieldpath (ex. _id: "$a") or as a singleton
    // object (ex. _id: {v: "$a"}).
    if (_idFieldNames.empty()) {
        idField = ExpressionFieldPath::deprecatedCreate(pExpCtx.get(), groupId);
    } else {
        invariant(_idFieldNames.size() == 1);
        idField = ExpressionObject::create(pExpCtx.get(),
                                           {{_idFieldNames.front(), _idExpressions.front()}});
    }
    fields.push_back(std::make_pair("_id", idField));

    for (auto&& accumulator : _accumulatedFields) {
        fields.push_back(std::make_pair(accumulator.fieldName, accumulator.expr.argument));

        // Since we don't attempt this transformation for non-$first accumulators,
        // the initializer should always be trivial.
    }

    return GroupFromFirstDocumentTransformation::create(pExpCtx, groupId, std::move(fields));
}

Project: mongo
File: document_source_group.cpp

DocumentSource::GetNextResult DocumentSourceGroup::getNextStandard() {
    // Not spilled, and not streaming.
    if (_groups->empty())
        return GetNextResult::makeEOF();

    Document out = makeDocument(groupsIterator->first, groupsIterator->second, pExpCtx->needsMerge);

    if (++groupsIterator == _groups->end())
        dispose();

    return out;
}

Project: mongo
File: document_source_group.cpp

void DocumentSourceGroup::freeMemory() {
    invariant(_groups);
    for (auto&& group : *_groups) {
        for (size_t i = 0; i < group.second.size(); i++) {
            // Subtract the current usage.
            _memoryTracker.update(_accumulatedFields[i].fieldName,
                                  -1 * group.second[i]->getMemUsage());

            group.second[i]->reduceMemoryConsumptionIfAble();

            // Update the memory usage for this AccumulationStatement.
            _memoryTracker.update(_accumulatedFields[i].fieldName, group.second[i]->getMemUsage());
        }
    }
}

Project: mongo
File: document_source_group.cpp

intrusive_ptr<Expression> parseIdExpression(const intrusive_ptr<ExpressionContext>& expCtx,
                                            BSONElement groupField,
                                            const VariablesParseState& vps) {
    if (groupField.type() == Object) {
        // {_id: {}} is treated as grouping on a constant, not an expression
        if (groupField.Obj().isEmpty()) {
            return ExpressionConstant::create(expCtx.get(), Value(groupField));
        }

        const BSONObj idKeyObj = groupField.Obj();
        if (idKeyObj.firstElementFieldName()[0] == '$') {
            // grouping on a $op expression
            return Expression::parseObject(expCtx.get(), idKeyObj, vps);
        } else {
            for (auto&& field : idKeyObj) {
                uassert(17390,
                        "$group does not support inclusion-style expressions",
                        !field.isNumber() && field.type() != Bool);
            }
            return ExpressionObject::parse(expCtx.get(), idKeyObj, vps);
        }
    } else {
        return Expression::parseOperand(expCtx.get(), groupField, vps);
    }
}

Project: mongo
File: document_source_group.cpp

intrusive_ptr<DocumentSourceGroup> DocumentSourceGroup::create(
    const intrusive_ptr<ExpressionContext>& expCtx,
    const boost::intrusive_ptr<Expression>& groupByExpression,
    std::vector<AccumulationStatement> accumulationStatements,
    boost::optional<size_t> maxMemoryUsageBytes) {
    size_t memoryBytes = maxMemoryUsageBytes ? *maxMemoryUsageBytes
                                             : internalDocumentSourceGroupMaxMemoryBytes.load();
    intrusive_ptr<DocumentSourceGroup> groupStage(new DocumentSourceGroup(expCtx, memoryBytes));
    groupStage->setIdExpression(groupByExpression);
    for (auto&& statement : accumulationStatements) {
        groupStage->addAccumulator(statement);
        groupStage->_memoryTracker.set(statement.fieldName, 0);
    }

    return groupStage;
}

Project: mongo
File: document_source_group.cpp

intrusive_ptr<DocumentSource> DocumentSourceGroup::createFromBson(
    BSONElement elem, const intrusive_ptr<ExpressionContext>& expCtx) {
    uassert(15947, "a group's fields must be specified in an object", elem.type() == Object);

    intrusive_ptr<DocumentSourceGroup> groupStage(new DocumentSourceGroup(expCtx));

    BSONObj groupObj(elem.Obj());
    BSONObjIterator groupIterator(groupObj);
    VariablesParseState vps = expCtx->variablesParseState;
    while (groupIterator.more()) {
        BSONElement groupField(groupIterator.next());
        StringData pFieldName = groupField.fieldNameStringData();
        if (pFieldName == "_id") {
            uassert(15948,
                    "a group's _id may only be specified once",
                    groupStage->_idExpressions.empty());
            groupStage->setIdExpression(parseIdExpression(expCtx, groupField, vps));
            invariant(!groupStage->_idExpressions.empty());
        } else if (pFieldName == "$doingMerge") {
            massert(17030, "$doingMerge should be true if present", groupField.Bool());

            groupStage->setDoingMerge(true);
        } else {
            // Any other field will be treated as an accumulator specification.
            groupStage->addAccumulator(
                AccumulationStatement::parseAccumulationStatement(expCtx.get(), groupField, vps));
            groupStage->_memoryTracker.set(pFieldName, 0);
        }
    }

    uassert(
        15955, "a group specification must include an _id", !groupStage->_idExpressions.empty());
    return groupStage;
}

Project: mongo
File: document_source_group.cpp

shared_ptr<Sorter<Value, Value>::Iterator> DocumentSourceGroup::spill() {
    _stats.usedDisk = true;
    vector<const GroupsMap::value_type*> ptrs;  // using pointers to speed sorting
    ptrs.reserve(_groups->size());
    for (GroupsMap::const_iterator it = _groups->begin(), end = _groups->end(); it != end; ++it) {
        ptrs.push_back(&*it);
    }

    stable_sort(ptrs.begin(), ptrs.end(), SpillSTLComparator(pExpCtx->getValueComparator()));

    SortedFileWriter<Value, Value> writer(SortOptions().TempDir(pExpCtx->tempDir), _file);
    switch (_accumulatedFields.size()) {  // same as ptrs[i]->second.size() for all i.
        case 0:                           // no values, essentially a distinct
            for (size_t i = 0; i < ptrs.size(); i++) {
                writer.addAlreadySorted(ptrs[i]->first, Value());
            }
            break;

        case 1:  // just one value, use optimized serialization as single Value
            for (size_t i = 0; i < ptrs.size(); i++) {
                writer.addAlreadySorted(ptrs[i]->first,
                                        ptrs[i]->second[0]->getValue(/*toBeMerged=*/true));
            }
            break;

        default:  // multiple values, serialize as array-typed Value
            for (size_t i = 0; i < ptrs.size(); i++) {
                vector<Value> accums;
                for (size_t j = 0; j < ptrs[i]->second.size(); j++) {
                    accums.push_back(ptrs[i]->second[j]->getValue(/*toBeMerged=*/true));
                }
                writer.addAlreadySorted(ptrs[i]->first, Value(std::move(accums)));
            }
            break;
    }

    auto& metricsCollector = ResourceConsumption::MetricsCollector::get(pExpCtx->opCtx);
    metricsCollector.incrementKeysSorted(ptrs.size());
    metricsCollector.incrementSorterSpills(1);

    _groups->clear();
    // Zero out the current per-accumulation statement memory consumption, as the memory has been
    // freed by spilling.
    for (auto accum : _accumulatedFields) {
        _memoryTracker.set(accum.fieldName, 0);
    }

    Sorter<Value, Value>::Iterator* iteratorPtr = writer.done();
    return shared_ptr<Sorter<Value, Value>::Iterator>(iteratorPtr);
}

Project: mongo
File: document_source_group.cpp

DocumentSourceGroup::DocumentSourceGroup(const intrusive_ptr<ExpressionContext>& expCtx,
                                         boost::optional<size_t> maxMemoryUsageBytes)
    : DocumentSource(kStageName, expCtx),
      _doingMerge(false),
      _memoryTracker{expCtx->allowDiskUse && !expCtx->inMongos,
                     maxMemoryUsageBytes ? *maxMemoryUsageBytes
                                         : internalDocumentSourceGroupMaxMemoryBytes.load()},
      // We spill to disk in debug mode, regardless of allowDiskUse, to stress the system.
      _file(
          !expCtx->inMongos && (expCtx->allowDiskUse || kDebugBuild)
              ? std::make_shared<Sorter<Value, Value>::File>(expCtx->tempDir + "/" + nextFileName())
              : nullptr),
      _initialized(false),
      _groups(expCtx->getValueComparator().makeUnorderedValueMap<Accumulators>()),
      _spilled(false) {}

Project: mongo
File: document_source_group.cpp

bool DocumentSourceGroup::shouldSpillWithAttemptToSaveMemory() {
    if (!_memoryTracker._allowDiskUse &&
        (_memoryTracker.currentMemoryBytes() >
         static_cast<long long>(_memoryTracker._maxAllowedMemoryUsageBytes))) {
        freeMemory();
    }

    if (_memoryTracker.currentMemoryBytes() >
        static_cast<long long>(_memoryTracker._maxAllowedMemoryUsageBytes)) {
        uassert(ErrorCodes::QueryExceededMemoryLimitNoDiskUseAllowed,
                "Exceeded memory limit for $group, but didn't allow external sort."
                " Pass allowDiskUse:true to opt in.",
                _memoryTracker._allowDiskUse);
        _memoryTracker.resetCurrent();
        return true;
    }
    return false;
}

Project: mongo
File: document_source_group.cpp

boost::optional<DocumentSource::DistributedPlanLogic> DocumentSourceGroup::distributedPlanLogic() {
    intrusive_ptr<DocumentSourceGroup> mergingGroup(new DocumentSourceGroup(pExpCtx));
    mergingGroup->setDoingMerge(true);

    VariablesParseState vps = pExpCtx->variablesParseState;
    /* the merger will use the same grouping key */
    mergingGroup->setIdExpression(ExpressionFieldPath::parse(pExpCtx.get(), "$$ROOT._id", vps));

    for (auto&& accumulatedField : _accumulatedFields) {
        // The merger's output field names will be the same, as will the accumulator factories.
        // However, for some accumulators, the expression to be accumulated will be different. The
        // original accumulator may be collecting an expression based on a field expression or
        // constant.  Here, we accumulate the output of the same name from the prior group.
        auto copiedAccumulatedField = accumulatedField;
        copiedAccumulatedField.expr.argument = ExpressionFieldPath::parse(
            pExpCtx.get(), "$$ROOT." + copiedAccumulatedField.fieldName, vps);
        mergingGroup->addAccumulator(copiedAccumulatedField);
        mergingGroup->_memoryTracker.set(copiedAccumulatedField.fieldName, 0);
    }

    // {shardsStage, mergingStage, sortPattern}
    return DistributedPlanLogic{this, mergingGroup, boost::none};
}

Project: mongo
File: document_source_group.cpp

MONGO_COMPILER_NOINLINE DocumentSource::GetNextResult DocumentSourceGroup::initializeSelf(
    GetNextResult input) {
    const size_t numAccumulators = _accumulatedFields.size();
    // Barring any pausing, this loop exhausts 'pSource' and populates '_groups'.
    for (; input.isAdvanced(); input = pSource->getNext()) {
        if (shouldSpillWithAttemptToSaveMemory()) {
            _sortedFiles.push_back(spill());
        }

        // We release the result document here so that it does not outlive the end of this loop
        // iteration. Not releasing could lead to an array copy when this group follows an unwind.
        auto rootDocument = input.releaseDocument();
        Value id = computeId(rootDocument);

        // Look for the _id value in the map. If it's not there, add a new entry with a blank
        // accumulator. This is done in a somewhat odd way in order to avoid hashing 'id' and
        // looking it up in '_groups' multiple times.
        const size_t oldSize = _groups->size();
        vector<intrusive_ptr<AccumulatorState>>& group = (*_groups)[id];
        const bool inserted = _groups->size() != oldSize;

        vector<uint64_t> oldAccumMemUsage(numAccumulators, 0);
        if (inserted) {
            _memoryTracker.set(_memoryTracker.currentMemoryBytes() + id.getApproximateSize());

            // Initialize and add the accumulators
            Value expandedId = expandId(id);
            Document idDoc =
                expandedId.getType() == BSONType::Object ? expandedId.getDocument() : Document();
            group.reserve(numAccumulators);
            for (auto&& accumulatedField : _accumulatedFields) {
                auto accum = accumulatedField.makeAccumulator();
                Value initializerValue =
                    accumulatedField.expr.initializer->evaluate(idDoc, &pExpCtx->variables);
                accum->startNewGroup(initializerValue);
                group.push_back(accum);
            }
        } else {
            for (size_t i = 0; i < group.size(); i++) {
                // subtract old mem usage. New usage added back after processing.
                _memoryTracker.update(_accumulatedFields[i].fieldName,
                                      -1 * group[i]->getMemUsage());
            }
        }

        /* tickle all the accumulators for the group we found */
        dassert(numAccumulators == group.size());

        for (size_t i = 0; i < numAccumulators; i++) {
            group[i]->process(
                _accumulatedFields[i].expr.argument->evaluate(rootDocument, &pExpCtx->variables),
                _doingMerge);
            _memoryTracker.update(_accumulatedFields[i].fieldName, group[i]->getMemUsage());
        }

        if (kDebugBuild && !storageGlobalParams.readOnly) {
            // In debug mode, spill every time we have a duplicate id to stress merge logic.
            if (!inserted &&           // is a dup
                !pExpCtx->inMongos &&  // can't spill to disk in mongos
                !_memoryTracker
                     ._allowDiskUse &&       // don't change behavior when testing external sort
                _sortedFiles.size() < 20) {  // don't open too many FDs

                _sortedFiles.push_back(spill());
            }
        }
    }

    switch (input.getStatus()) {
        case DocumentSource::GetNextResult::ReturnStatus::kAdvanced: {
            MONGO_UNREACHABLE;  // We consumed all advances above.
        }
        case DocumentSource::GetNextResult::ReturnStatus::kPauseExecution: {
            return input;  // Propagate pause.
        }
        case DocumentSource::GetNextResult::ReturnStatus::kEOF: {
            // Do any final steps necessary to prepare to output results.
            if (!_sortedFiles.empty()) {
                _spilled = true;
                if (!_groups->empty()) {
                    _sortedFiles.push_back(spill());
                }

                // We won't be using groups again so free its memory.
                _groups = pExpCtx->getValueComparator().makeUnorderedValueMap<Accumulators>();

                _sorterIterator.reset(Sorter<Value, Value>::Iterator::merge(
                    _sortedFiles, SortOptions(), SorterComparator(pExpCtx->getValueComparator())));

                // prepare current to accumulate data
                _currentAccumulators.reserve(numAccumulators);
                for (auto&& accumulatedField : _accumulatedFields) {
                    _currentAccumulators.push_back(accumulatedField.makeAccumulator());
                }

                verify(_sorterIterator->more());  // we put data in, we should get something out.
                _firstPartOfNextGroup = _sorterIterator->next();
            } else {
                // start the group iterator
                groupsIterator = _groups->begin();
            }

            // This must happen last so that, unless control gets here, we will re-enter
            // initialization after getting a GetNextResult::ResultState::kPauseExecution.
            _initialized = true;
            return input;
        }
    }
    MONGO_UNREACHABLE;
}