From e6c531c7bc01e628052792a5c94c6da3e5779adf Mon Sep 17 00:00:00 2001 From: Misha Ivkov Date: Wed, 3 Jul 2019 18:27:47 -0400 Subject: SERVER-42077 Add 'allowDiskUse' option to find command --- src/mongo/db/exec/sort.cpp | 1 + src/mongo/db/exec/sort.h | 5 ++ src/mongo/db/query/SConscript | 7 ++- src/mongo/db/query/planner_analysis.cpp | 1 + src/mongo/db/query/query_request.cpp | 20 +++++++ src/mongo/db/query/query_request.h | 13 +++++ src/mongo/db/query/query_request_test.cpp | 86 +++++++++++++++++++++++++++++++ src/mongo/db/query/query_solution.cpp | 1 + src/mongo/db/query/query_solution.h | 7 ++- src/mongo/db/query/stage_builder.cpp | 1 + src/mongo/shell/query.js | 13 +++++ 11 files changed, 152 insertions(+), 3 deletions(-) diff --git a/src/mongo/db/exec/sort.cpp b/src/mongo/db/exec/sort.cpp index 89de29e69f9..9bf39c08773 100644 --- a/src/mongo/db/exec/sort.cpp +++ b/src/mongo/db/exec/sort.cpp @@ -75,6 +75,7 @@ SortStage::SortStage(OperationContext* opCtx, _ws(ws), _pattern(params.pattern), _limit(params.limit), + _allowDiskUse(params.allowDiskUse), _sorted(false), _resultIterator(_data.end()), _memUsage(0) { diff --git a/src/mongo/db/exec/sort.h b/src/mongo/db/exec/sort.h index 85c4aa444b9..adadc9ad121 100644 --- a/src/mongo/db/exec/sort.h +++ b/src/mongo/db/exec/sort.h @@ -54,6 +54,9 @@ public: // Equal to 0 for no limit. size_t limit = 0; + + // Whether we allow disk use, disabled by default. + bool allowDiskUse = false; }; /** @@ -99,6 +102,8 @@ private: // Equal to 0 for no limit. size_t _limit; + bool _allowDiskUse; + // // Data storage // diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript index dbba9fe6c18..dafb9ca9d4c 100644 --- a/src/mongo/db/query/SConscript +++ b/src/mongo/db/query/SConscript @@ -135,9 +135,12 @@ env.Library( ], LIBDEPS=[ "$BUILD_DIR/mongo/base", - "$BUILD_DIR/mongo/db/repl/read_concern_args", "$BUILD_DIR/mongo/db/catalog/collection_catalog", - "$BUILD_DIR/mongo/db/pipeline/runtime_constants_idl" + # TODO: This dependency edge can be removed when the 'allowDiskUse' option no longer depends + # on enabling test commands. + "$BUILD_DIR/mongo/db/commands/test_commands_enabled", + "$BUILD_DIR/mongo/db/pipeline/runtime_constants_idl", + "$BUILD_DIR/mongo/db/repl/read_concern_args", ], ) diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp index 2be65fbf8d2..d64f43b4611 100644 --- a/src/mongo/db/query/planner_analysis.cpp +++ b/src/mongo/db/query/planner_analysis.cpp @@ -659,6 +659,7 @@ QuerySolutionNode* QueryPlannerAnalysis::analyzeSort(const CanonicalQuery& query sort->pattern = sortObj; sort->children.push_back(solnRoot); solnRoot = sort; + sort->allowDiskUse = qr.allowDiskUse(); // When setting the limit on the sort, we need to consider both // the limit N and skip count M. The sort should return an ordered list // N + M items so that the skip stage can discard the first M results. diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp index 5072c883430..4fc31cd4965 100644 --- a/src/mongo/db/query/query_request.cpp +++ b/src/mongo/db/query/query_request.cpp @@ -63,6 +63,8 @@ const string QueryRequest::metaRecordId("recordId"); const string QueryRequest::metaSortKey("sortKey"); const string QueryRequest::metaTextScore("textScore"); +const string QueryRequest::kAllowDiskUseField("allowDiskUse"); + const long long QueryRequest::kDefaultBatchSize = 101; namespace { @@ -262,6 +264,17 @@ StatusWith> QueryRequest::parseFromFindCommand(unique_p } qr->_wantMore = !el.boolean(); + } else if (fieldName == kAllowDiskUseField) { + if (!getTestCommandsEnabled()) { + return Status(ErrorCodes::FailedToParse, + "allowDiskUse is not allowed unless test commands are enabled."); + } + Status status = checkFieldType(el, Bool); + if (!status.isOK()) { + return status; + } + + qr->_allowDiskUse = el.boolean(); } else if (fieldName == kCommentField) { Status status = checkFieldType(el, String); if (!status.isOK()) { @@ -495,6 +508,10 @@ void QueryRequest::asFindCommandInternal(BSONObjBuilder* cmdBuilder) const { cmdBuilder->append(kLimitField, *_limit); } + if (_allowDiskUse) { + cmdBuilder->append(kAllowDiskUseField, true); + } + if (_batchSize) { cmdBuilder->append(kBatchSizeField, *_batchSize); } @@ -1132,6 +1149,9 @@ StatusWith QueryRequest::asAggregationCommand() const { if (!_unwrappedReadPref.isEmpty()) { aggregationBuilder.append(QueryRequest::kUnwrappedReadPrefField, _unwrappedReadPref); } + if (_allowDiskUse) { + aggregationBuilder.append(QueryRequest::kAllowDiskUseField, _allowDiskUse); + } if (_runtimeConstants) { BSONObjBuilder rtcBuilder(aggregationBuilder.subobjStart(kRuntimeConstantsField)); _runtimeConstants->serialize(&rtcBuilder); diff --git a/src/mongo/db/query/query_request.h b/src/mongo/db/query/query_request.h index ade62c52ddf..7b0cb74dd80 100644 --- a/src/mongo/db/query/query_request.h +++ b/src/mongo/db/query/query_request.h @@ -140,6 +140,9 @@ public: static const std::string metaSortKey; static const std::string metaTextScore; + // Allow using disk during the find command. + static const std::string kAllowDiskUseField; + const NamespaceString& nss() const { invariant(!_nss.isEmpty()); return _nss; @@ -241,6 +244,14 @@ public: _wantMore = wantMore; } + bool allowDiskUse() const { + return _allowDiskUse; + } + + void setAllowDiskUse(bool allowDiskUse) { + _allowDiskUse = allowDiskUse; + } + bool isExplain() const { return _explain; } @@ -504,6 +515,8 @@ private: // allowed. boost::optional _batchSize; + bool _allowDiskUse = false; + // Set only when parsed from an OP_QUERY find message. The value is computed by driver or shell // and is set to be a min of batchSize and limit provided by user. QR can have set either // ntoreturn or batchSize / limit. diff --git a/src/mongo/db/query/query_request_test.cpp b/src/mongo/db/query/query_request_test.cpp index ed4d369602b..7ee502140f3 100644 --- a/src/mongo/db/query/query_request_test.cpp +++ b/src/mongo/db/query/query_request_test.cpp @@ -35,6 +35,7 @@ #include "mongo/db/catalog/collection_catalog.h" #include "mongo/db/catalog/collection_mock.h" +#include "mongo/db/commands/test_commands_enabled.h" #include "mongo/db/dbmessage.h" #include "mongo/db/json.h" #include "mongo/db/namespace_string.h" @@ -1089,6 +1090,51 @@ TEST(QueryRequestTest, DefaultQueryParametersCorrect) { ASSERT_EQUALS(false, qr->isExhaust()); ASSERT_EQUALS(false, qr->isAllowPartialResults()); ASSERT_EQUALS(false, qr->getRuntimeConstants().has_value()); + ASSERT_EQUALS(false, qr->allowDiskUse()); +} + +TEST(QueryRequestTest, ParseCommandAllowDiskUseTrue) { + const bool oldTestCommandsEnabledVal = getTestCommandsEnabled(); + ON_BLOCK_EXIT([&] { setTestCommandsEnabled(oldTestCommandsEnabledVal); }); + setTestCommandsEnabled(true); + + BSONObj cmdObj = fromjson("{find: 'testns', allowDiskUse: true}"); + const NamespaceString nss("test.testns"); + const bool isExplain = false; + auto result = QueryRequest::makeFromFindCommand(nss, cmdObj, isExplain); + + ASSERT_OK(result.getStatus()); + ASSERT_EQ(true, result.getValue()->allowDiskUse()); +} + +TEST(QueryRequestTest, ParseCommandAllowDiskUseFalse) { + const bool oldTestCommandsEnabledVal = getTestCommandsEnabled(); + ON_BLOCK_EXIT([&] { setTestCommandsEnabled(oldTestCommandsEnabledVal); }); + setTestCommandsEnabled(true); + + BSONObj cmdObj = fromjson("{find: 'testns', allowDiskUse: false}"); + const NamespaceString nss("test.testns"); + const bool isExplain = false; + auto result = QueryRequest::makeFromFindCommand(nss, cmdObj, isExplain); + + ASSERT_OK(result.getStatus()); + ASSERT_EQ(false, result.getValue()->allowDiskUse()); +} + +TEST(QueryRequestTest, ParseCommandAllowDiskUseTestCommandsDisabled) { + const bool oldTestCommandsEnabledVal = getTestCommandsEnabled(); + ON_BLOCK_EXIT([&] { setTestCommandsEnabled(oldTestCommandsEnabledVal); }); + setTestCommandsEnabled(false); + + BSONObj cmdObj = fromjson("{find: 'testns', allowDiskUse: true}"); + const NamespaceString nss("test.testns"); + const bool isExplain = false; + auto result = QueryRequest::makeFromFindCommand(nss, cmdObj, isExplain); + + ASSERT_NOT_OK(result.getStatus()); + ASSERT_EQ(ErrorCodes::FailedToParse, result.getStatus().code()); + ASSERT_STRING_CONTAINS(result.getStatus().toString(), + "allowDiskUse is not allowed unless test commands are enabled."); } // @@ -1384,6 +1430,46 @@ TEST(QueryRequestTest, ConvertToAggregationWithRuntimeConstantsSucceeds) { ASSERT_EQ(ar.getValue().getRuntimeConstants()->getClusterTime(), rtc.getClusterTime()); } +TEST(QueryRequestTest, ConvertToAggregationWithAllowDiskUseTrueSucceeds) { + QueryRequest qr(testns); + qr.setAllowDiskUse(true); + const auto aggCmd = qr.asAggregationCommand(); + ASSERT_OK(aggCmd.getStatus()); + + auto ar = AggregationRequest::parseFromBSON(testns, aggCmd.getValue()); + ASSERT_OK(ar.getStatus()); + ASSERT_EQ(true, ar.getValue().shouldAllowDiskUse()); +} + +TEST(QueryRequestTest, ConvertToAggregationWithAllowDiskUseFalseSucceeds) { + QueryRequest qr(testns); + qr.setAllowDiskUse(false); + const auto aggCmd = qr.asAggregationCommand(); + ASSERT_OK(aggCmd.getStatus()); + + auto ar = AggregationRequest::parseFromBSON(testns, aggCmd.getValue()); + ASSERT_OK(ar.getStatus()); + ASSERT_EQ(false, ar.getValue().shouldAllowDiskUse()); +} + +TEST(QueryRequestTest, ConvertToFindWithAllowDiskUseTrueSucceeds) { + QueryRequest qr(testns); + qr.setAllowDiskUse(true); + const auto findCmd = qr.asFindCommand(); + + BSONElement elem = findCmd[QueryRequest::kAllowDiskUseField]; + ASSERT_EQ(true, elem.isBoolean()); + ASSERT_EQ(true, elem.Bool()); +} + +TEST(QueryRequestTest, ConvertToFindWithAllowDiskUseFalseSucceeds) { + QueryRequest qr(testns); + qr.setAllowDiskUse(false); + const auto findCmd = qr.asFindCommand(); + + ASSERT_FALSE(findCmd[QueryRequest::kAllowDiskUseField]); +} + TEST(QueryRequestTest, ParseFromLegacyObjMetaOpComment) { BSONObj queryObj = fromjson( "{$query: {a: 1}," diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp index 9f1951c14a9..331e94875dd 100644 --- a/src/mongo/db/query/query_solution.cpp +++ b/src/mongo/db/query/query_solution.cpp @@ -974,6 +974,7 @@ QuerySolutionNode* SortNode::clone() const { copy->_sorts = this->_sorts; copy->pattern = this->pattern; copy->limit = this->limit; + copy->allowDiskUse = this->allowDiskUse; return copy; } diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h index a665b823529..16098b0146b 100644 --- a/src/mongo/db/query/query_solution.h +++ b/src/mongo/db/query/query_solution.h @@ -691,7 +691,10 @@ struct SortKeyGeneratorNode : public QuerySolutionNode { }; struct SortNode : public QuerySolutionNode { - SortNode() : _sorts(SimpleBSONObjComparator::kInstance.makeBSONObjSet()), limit(0) {} + SortNode() + : _sorts(SimpleBSONObjComparator::kInstance.makeBSONObjSet()), + limit(0), + allowDiskUse(false) {} virtual ~SortNode() {} @@ -731,6 +734,8 @@ struct SortNode : public QuerySolutionNode { // Sum of both limit and skip count in the parsed query. size_t limit; + + bool allowDiskUse; }; struct LimitNode : public QuerySolutionNode { diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp index 190ed408770..8ab1ca42cfb 100644 --- a/src/mongo/db/query/stage_builder.cpp +++ b/src/mongo/db/query/stage_builder.cpp @@ -129,6 +129,7 @@ PlanStage* buildStages(OperationContext* opCtx, SortStageParams params; params.pattern = sn->pattern; params.limit = sn->limit; + params.allowDiskUse = sn->allowDiskUse; return new SortStage(opCtx, params, ws, childStage); } case STAGE_SORT_KEY_GENERATOR: { diff --git a/src/mongo/shell/query.js b/src/mongo/shell/query.js index 4304903ca36..def25f8f65a 100644 --- a/src/mongo/shell/query.js +++ b/src/mongo/shell/query.js @@ -47,6 +47,7 @@ DBQuery.prototype.help = function() { print("\t.allowPartialResults()"); print("\t.returnKey()"); print("\t.showRecordId() - adds a $recordId field to each returned object"); + print("\t.allowDiskUse() - allow using disk in completing the query"); print("\nCursor methods"); print("\t.toArray() - iterates through docs and returns an array of the results"); @@ -127,6 +128,10 @@ DBQuery.prototype._exec = function() { throw new Error("collation requires use of read commands"); } + if (this._special && this._query._allowDiskUse) { + throw new Error("allowDiskUse option requires use of read commands"); + } + this._cursor = this._mongo.find(this._ns, this._query, this._fields, @@ -225,6 +230,10 @@ DBQuery.prototype._convertToCommand = function(canAttachReadPref) { cmd["collation"] = this._query.collation; } + if ("allowDiskUse" in this._query) { + cmd["allowDiskUse"] = this._query.allowDiskUse; + } + if ((this._options & DBQuery.Option.tailable) != 0) { cmd["tailable"] = true; } @@ -471,6 +480,10 @@ DBQuery.prototype.collation = function(collationSpec) { return this._addSpecial("collation", collationSpec); }; +DBQuery.prototype.allowDiskUse = function() { + return this._addSpecial("allowDiskUse", true); +}; + /** * Sets the read preference for this cursor. * -- cgit v1.2.1