diff options
Diffstat (limited to 'src/mongo')
-rw-r--r-- | src/mongo/db/commands/find_and_modify.cpp | 2 | ||||
-rw-r--r-- | src/mongo/db/ops/write_ops_exec.cpp | 4 | ||||
-rw-r--r-- | src/mongo/db/s/query_analysis_writer.cpp | 73 | ||||
-rw-r--r-- | src/mongo/db/s/query_analysis_writer.h | 18 | ||||
-rw-r--r-- | src/mongo/db/s/query_analysis_writer_test.cpp | 42 |
5 files changed, 90 insertions, 49 deletions
diff --git a/src/mongo/db/commands/find_and_modify.cpp b/src/mongo/db/commands/find_and_modify.cpp index 7042a844da0..aa6b1a691ad 100644 --- a/src/mongo/db/commands/find_and_modify.cpp +++ b/src/mongo/db/commands/find_and_modify.cpp @@ -474,7 +474,7 @@ write_ops::FindAndModifyCommandReply CmdFindAndModify::Invocation::typedRun( opCtx, ns(), analyze_shard_key::SampledCommandNameEnum::kFindAndModify, req); if (sampleId) { analyze_shard_key::QueryAnalysisWriter::get(opCtx) - ->addFindAndModifyQuery(*sampleId, req) + ->addFindAndModifyQuery(opCtx, *sampleId, req) .getAsync([](auto) {}); } diff --git a/src/mongo/db/ops/write_ops_exec.cpp b/src/mongo/db/ops/write_ops_exec.cpp index 8fc25577e3a..9aec8d49943 100644 --- a/src/mongo/db/ops/write_ops_exec.cpp +++ b/src/mongo/db/ops/write_ops_exec.cpp @@ -1401,7 +1401,7 @@ WriteResult performUpdates(OperationContext* opCtx, opCtx, ns, analyze_shard_key::SampledCommandNameEnum::kUpdate, singleOp); if (sampleId) { analyze_shard_key::QueryAnalysisWriter::get(opCtx) - ->addUpdateQuery(*sampleId, wholeOp, currentOpIndex) + ->addUpdateQuery(opCtx, *sampleId, wholeOp, currentOpIndex) .getAsync([](auto) {}); } @@ -1672,7 +1672,7 @@ WriteResult performDeletes(OperationContext* opCtx, if (auto sampleId = analyze_shard_key::getOrGenerateSampleId( opCtx, ns, analyze_shard_key::SampledCommandNameEnum::kDelete, singleOp)) { analyze_shard_key::QueryAnalysisWriter::get(opCtx) - ->addDeleteQuery(*sampleId, wholeOp, currentOpIndex) + ->addDeleteQuery(opCtx, *sampleId, wholeOp, currentOpIndex) .getAsync([](auto) {}); } diff --git a/src/mongo/db/s/query_analysis_writer.cpp b/src/mongo/db/s/query_analysis_writer.cpp index 0289bec8053..7404066b9c9 100644 --- a/src/mongo/db/s/query_analysis_writer.cpp +++ b/src/mongo/db/s/query_analysis_writer.cpp @@ -94,6 +94,11 @@ BSONObj createIndex(OperationContext* opCtx, const NamespaceString& nss, const B return resObj; } +bool isInternalClient(const OperationContext* opCtx) { + return !opCtx->getClient()->session() || + (opCtx->getClient()->session()->getTags() & transport::Session::kInternalClient); +} + struct SampledCommandRequest { UUID sampleId; NamespaceString nss; @@ -119,7 +124,10 @@ SampledCommandRequest makeSampledReadCommand(const UUID& sampleId, * Returns a sampled update command for the update at 'opIndex' in the given update command. */ SampledCommandRequest makeSampledUpdateCommandRequest( - const UUID& sampleId, const write_ops::UpdateCommandRequest& originalCmd, int opIndex) { + const OperationContext* opCtx, + const UUID& sampleId, + const write_ops::UpdateCommandRequest& originalCmd, + int opIndex) { auto op = originalCmd.getUpdates()[opIndex]; if (op.getSampleId()) { tassert(ErrorCodes::IllegalOperation, @@ -131,10 +139,14 @@ SampledCommandRequest makeSampledUpdateCommandRequest( // If the initial query was a write without shard key, the two phase write protocol modifies the // query in the write phase. In order to get correct metrics, we need to reconstruct the // original query here. - if (originalCmd.getOriginalQuery()) { + if (originalCmd.getOriginalQuery() || originalCmd.getOriginalCollation()) { tassert(7406500, "Found a _clusterWithoutShardKey command with batch size > 1", originalCmd.getUpdates().size() == 1); + uassert(ErrorCodes::InvalidOptions, + "Cannot specify '$_originalQuery' or '$_originalCollation' since they are internal " + "fields", + isInternalClient(opCtx)); op.setQ(*originalCmd.getOriginalQuery()); op.setCollation(originalCmd.getOriginalCollation()); } @@ -151,7 +163,10 @@ SampledCommandRequest makeSampledUpdateCommandRequest( * Returns a sampled delete command for the delete at 'opIndex' in the given delete command. */ SampledCommandRequest makeSampledDeleteCommandRequest( - const UUID& sampleId, const write_ops::DeleteCommandRequest& originalCmd, int opIndex) { + const OperationContext* opCtx, + const UUID& sampleId, + const write_ops::DeleteCommandRequest& originalCmd, + int opIndex) { auto op = originalCmd.getDeletes()[opIndex]; if (op.getSampleId()) { tassert(ErrorCodes::IllegalOperation, @@ -163,10 +178,14 @@ SampledCommandRequest makeSampledDeleteCommandRequest( // If the initial query was a write without shard key, the two phase write protocol modifies the // query in the write phase. In order to get correct metrics, we need to reconstruct the // original query here. - if (originalCmd.getOriginalQuery()) { + if (originalCmd.getOriginalQuery() || originalCmd.getOriginalCollation()) { tassert(7406501, "Found a _clusterWithoutShardKey command with batch size > 1", originalCmd.getDeletes().size() == 1); + uassert(ErrorCodes::InvalidOptions, + "Cannot specify '$_originalQuery' or '$_originalCollation' since they are internal " + "fields", + isInternalClient(opCtx)); op.setQ(*originalCmd.getOriginalQuery()); op.setCollation(originalCmd.getOriginalCollation()); } @@ -183,7 +202,9 @@ SampledCommandRequest makeSampledDeleteCommandRequest( * Returns a sampled findAndModify command for the given findAndModify command. */ SampledCommandRequest makeSampledFindAndModifyCommandRequest( - const UUID& sampleId, const write_ops::FindAndModifyCommandRequest& originalCmd) { + OperationContext* opCtx, + const UUID& sampleId, + const write_ops::FindAndModifyCommandRequest& originalCmd) { write_ops::FindAndModifyCommandRequest sampledCmd(originalCmd.getNamespace()); if (sampledCmd.getSampleId()) { tassert(ErrorCodes::IllegalOperation, @@ -195,7 +216,11 @@ SampledCommandRequest makeSampledFindAndModifyCommandRequest( // If the initial query was a write without shard key, the two phase write protocol modifies the // query in the write phase. In order to get correct metrics, we need to reconstruct the // original query here. - if (originalCmd.getOriginalQuery()) { + if (originalCmd.getOriginalQuery() || originalCmd.getOriginalCollation()) { + uassert(ErrorCodes::InvalidOptions, + "Cannot specify '$_originalQuery' or '$_originalCollation' since they are internal " + "fields", + isInternalClient(opCtx)); sampledCmd.setQuery(*originalCmd.getOriginalQuery()); sampledCmd.setCollation(originalCmd.getOriginalCollation()); } else { @@ -626,12 +651,16 @@ ExecutorFuture<void> QueryAnalysisWriter::_addReadQuery( } ExecutorFuture<void> QueryAnalysisWriter::addUpdateQuery( - const UUID& sampleId, const write_ops::UpdateCommandRequest& updateCmd, int opIndex) { + OperationContext* originalOpCtx, + const UUID& sampleId, + const write_ops::UpdateCommandRequest& updateCmd, + int opIndex) { invariant(_executor); return ExecutorFuture<void>(_executor) .then([this, - sampledUpdateCmd = makeSampledUpdateCommandRequest(sampleId, updateCmd, opIndex)]() { + sampledUpdateCmd = + makeSampledUpdateCommandRequest(originalOpCtx, sampleId, updateCmd, opIndex)]() { auto opCtxHolder = cc().makeOperationContext(); auto opCtx = opCtxHolder.get(); @@ -673,19 +702,23 @@ ExecutorFuture<void> QueryAnalysisWriter::addUpdateQuery( } ExecutorFuture<void> QueryAnalysisWriter::addUpdateQuery( - const write_ops::UpdateCommandRequest& updateCmd, int opIndex) { + OperationContext* opCtx, const write_ops::UpdateCommandRequest& updateCmd, int opIndex) { auto sampleId = updateCmd.getUpdates()[opIndex].getSampleId(); invariant(sampleId); - return addUpdateQuery(*sampleId, updateCmd, opIndex); + return addUpdateQuery(opCtx, *sampleId, updateCmd, opIndex); } ExecutorFuture<void> QueryAnalysisWriter::addDeleteQuery( - const UUID& sampleId, const write_ops::DeleteCommandRequest& deleteCmd, int opIndex) { + OperationContext* originalOpCtx, + const UUID& sampleId, + const write_ops::DeleteCommandRequest& deleteCmd, + int opIndex) { invariant(_executor); return ExecutorFuture<void>(_executor) .then([this, - sampledDeleteCmd = makeSampledDeleteCommandRequest(sampleId, deleteCmd, opIndex)]() { + sampledDeleteCmd = + makeSampledDeleteCommandRequest(originalOpCtx, sampleId, deleteCmd, opIndex)]() { auto opCtxHolder = cc().makeOperationContext(); auto opCtx = opCtxHolder.get(); @@ -727,20 +760,22 @@ ExecutorFuture<void> QueryAnalysisWriter::addDeleteQuery( } ExecutorFuture<void> QueryAnalysisWriter::addDeleteQuery( - const write_ops::DeleteCommandRequest& deleteCmd, int opIndex) { + OperationContext* opCtx, const write_ops::DeleteCommandRequest& deleteCmd, int opIndex) { auto sampleId = deleteCmd.getDeletes()[opIndex].getSampleId(); invariant(sampleId); - return addDeleteQuery(*sampleId, deleteCmd, opIndex); + return addDeleteQuery(opCtx, *sampleId, deleteCmd, opIndex); } ExecutorFuture<void> QueryAnalysisWriter::addFindAndModifyQuery( - const UUID& sampleId, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd) { + OperationContext* originalOpCtx, + const UUID& sampleId, + const write_ops::FindAndModifyCommandRequest& findAndModifyCmd) { invariant(_executor); return ExecutorFuture<void>(_executor) .then([this, - sampledFindAndModifyCmd = - makeSampledFindAndModifyCommandRequest(sampleId, findAndModifyCmd)]() { + sampledFindAndModifyCmd = makeSampledFindAndModifyCommandRequest( + originalOpCtx, sampleId, findAndModifyCmd)]() { auto opCtxHolder = cc().makeOperationContext(); auto opCtx = opCtxHolder.get(); @@ -782,10 +817,10 @@ ExecutorFuture<void> QueryAnalysisWriter::addFindAndModifyQuery( } ExecutorFuture<void> QueryAnalysisWriter::addFindAndModifyQuery( - const write_ops::FindAndModifyCommandRequest& findAndModifyCmd) { + OperationContext* opCtx, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd) { auto sampleId = findAndModifyCmd.getSampleId(); invariant(sampleId); - return addFindAndModifyQuery(*sampleId, findAndModifyCmd); + return addFindAndModifyQuery(opCtx, *sampleId, findAndModifyCmd); } ExecutorFuture<void> QueryAnalysisWriter::addDiff(const UUID& sampleId, diff --git a/src/mongo/db/s/query_analysis_writer.h b/src/mongo/db/s/query_analysis_writer.h index 18f980dd20f..0b5d5f02ac9 100644 --- a/src/mongo/db/s/query_analysis_writer.h +++ b/src/mongo/db/s/query_analysis_writer.h @@ -173,22 +173,28 @@ public: const BSONObj& filter, const BSONObj& collation); - ExecutorFuture<void> addUpdateQuery(const UUID& sampleId, + ExecutorFuture<void> addUpdateQuery(OperationContext* opCtx, + const UUID& sampleId, const write_ops::UpdateCommandRequest& updateCmd, int opIndex); - ExecutorFuture<void> addUpdateQuery(const write_ops::UpdateCommandRequest& updateCmd, + ExecutorFuture<void> addUpdateQuery(OperationContext* opCtx, + const write_ops::UpdateCommandRequest& updateCmd, int opIndex); - ExecutorFuture<void> addDeleteQuery(const UUID& sampleId, + ExecutorFuture<void> addDeleteQuery(OperationContext* opCtx, + const UUID& sampleId, const write_ops::DeleteCommandRequest& deleteCmd, int opIndex); - ExecutorFuture<void> addDeleteQuery(const write_ops::DeleteCommandRequest& deleteCmd, + ExecutorFuture<void> addDeleteQuery(OperationContext* opCtx, + const write_ops::DeleteCommandRequest& deleteCmd, int opIndex); ExecutorFuture<void> addFindAndModifyQuery( - const UUID& sampleId, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd); - ExecutorFuture<void> addFindAndModifyQuery( + OperationContext* opCtx, + const UUID& sampleId, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd); + ExecutorFuture<void> addFindAndModifyQuery( + OperationContext* opCtx, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd); ExecutorFuture<void> addDiff(const UUID& sampleId, const NamespaceString& nss, diff --git a/src/mongo/db/s/query_analysis_writer_test.cpp b/src/mongo/db/s/query_analysis_writer_test.cpp index c21217c8e04..c40e596a7d3 100644 --- a/src/mongo/db/s/query_analysis_writer_test.cpp +++ b/src/mongo/db/s/query_analysis_writer_test.cpp @@ -740,7 +740,7 @@ TEST_F(QueryAnalysisWriterTest, AggregateQuery) { DEATH_TEST_F(QueryAnalysisWriterTest, UpdateQueryNotMarkedForSampling, "invariant") { auto& writer = *QueryAnalysisWriter::get(operationContext()); auto [originalCmd, _] = makeUpdateCommandRequest(nss0, 1, {} /* markForSampling */); - writer.addUpdateQuery(originalCmd, 0).get(); + writer.addUpdateQuery(operationContext(), originalCmd, 0).get(); } TEST_F(QueryAnalysisWriterTest, UpdateQueriesMarkedForSampling) { @@ -750,8 +750,8 @@ TEST_F(QueryAnalysisWriterTest, UpdateQueriesMarkedForSampling) { makeUpdateCommandRequest(nss0, 3, {0, 2} /* markForSampling */); ASSERT_EQ(expectedSampledCmds.size(), 2U); - writer.addUpdateQuery(originalCmd, 0).get(); - writer.addUpdateQuery(originalCmd, 2).get(); + writer.addUpdateQuery(operationContext(), originalCmd, 0).get(); + writer.addUpdateQuery(operationContext(), originalCmd, 2).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 2); writer.flushQueriesForTest(operationContext()); ASSERT_EQ(writer.getQueriesCountForTest(), 0); @@ -768,7 +768,7 @@ TEST_F(QueryAnalysisWriterTest, UpdateQueriesMarkedForSampling) { DEATH_TEST_F(QueryAnalysisWriterTest, DeleteQueryNotMarkedForSampling, "invariant") { auto& writer = *QueryAnalysisWriter::get(operationContext()); auto [originalCmd, _] = makeDeleteCommandRequest(nss0, 1, {} /* markForSampling */); - writer.addDeleteQuery(originalCmd, 0).get(); + writer.addDeleteQuery(operationContext(), originalCmd, 0).get(); } TEST_F(QueryAnalysisWriterTest, DeleteQueriesMarkedForSampling) { @@ -778,8 +778,8 @@ TEST_F(QueryAnalysisWriterTest, DeleteQueriesMarkedForSampling) { makeDeleteCommandRequest(nss0, 3, {1, 2} /* markForSampling */); ASSERT_EQ(expectedSampledCmds.size(), 2U); - writer.addDeleteQuery(originalCmd, 1).get(); - writer.addDeleteQuery(originalCmd, 2).get(); + writer.addDeleteQuery(operationContext(), originalCmd, 1).get(); + writer.addDeleteQuery(operationContext(), originalCmd, 2).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 2); writer.flushQueriesForTest(operationContext()); ASSERT_EQ(writer.getQueriesCountForTest(), 0); @@ -797,7 +797,7 @@ DEATH_TEST_F(QueryAnalysisWriterTest, FindAndModifyQueryNotMarkedForSampling, "i auto& writer = *QueryAnalysisWriter::get(operationContext()); auto [originalCmd, _] = makeFindAndModifyCommandRequest(nss0, true /* isUpdate */, false /* markForSampling */); - writer.addFindAndModifyQuery(originalCmd).get(); + writer.addFindAndModifyQuery(operationContext(), originalCmd).get(); } TEST_F(QueryAnalysisWriterTest, FindAndModifyQueryUpdateMarkedForSampling) { @@ -808,7 +808,7 @@ TEST_F(QueryAnalysisWriterTest, FindAndModifyQueryUpdateMarkedForSampling) { ASSERT_EQ(expectedSampledCmds.size(), 1U); auto [sampleId, expectedSampledCmd] = *expectedSampledCmds.begin(); - writer.addFindAndModifyQuery(originalCmd).get(); + writer.addFindAndModifyQuery(operationContext(), originalCmd).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 1); writer.flushQueriesForTest(operationContext()); ASSERT_EQ(writer.getQueriesCountForTest(), 0); @@ -828,7 +828,7 @@ TEST_F(QueryAnalysisWriterTest, FindAndModifyQueryRemoveMarkedForSampling) { ASSERT_EQ(expectedSampledCmds.size(), 1U); auto [sampleId, expectedSampledCmd] = *expectedSampledCmds.begin(); - writer.addFindAndModifyQuery(originalCmd).get(); + writer.addFindAndModifyQuery(operationContext(), originalCmd).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 1); writer.flushQueriesForTest(operationContext()); ASSERT_EQ(writer.getQueriesCountForTest(), 0); @@ -859,8 +859,8 @@ TEST_F(QueryAnalysisWriterTest, MultipleQueriesAndCollections) { auto originalCountFilter = makeNonEmptyFilter(); auto originalCountCollation = makeNonEmptyCollation(); - writer.addDeleteQuery(originalDeleteCmd, 1).get(); - writer.addUpdateQuery(originalUpdateCmd, 0).get(); + writer.addDeleteQuery(operationContext(), originalDeleteCmd, 1).get(); + writer.addUpdateQuery(operationContext(), originalUpdateCmd, 0).get(); writer.addCountQuery(countSampleId, nss1, originalCountFilter, originalCountCollation).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 3); writer.flushQueriesForTest(operationContext()); @@ -917,7 +917,7 @@ TEST_F(QueryAnalysisWriterTest, DuplicateQueries) { originalFindFilter, originalFindCollation); - writer.addUpdateQuery(originalUpdateCmd, 0).get(); + writer.addUpdateQuery(operationContext(), originalUpdateCmd, 0).get(); writer .addFindQuery(findSampleId, nss0, @@ -1067,10 +1067,10 @@ TEST_F(QueryAnalysisWriterTest, FlushAfterAddUpdateIfExceedsSizeLimit) { std::string(maxMemoryUsageBytes / 2, 'a') /* filterFieldName */); ASSERT_EQ(expectedSampledCmds.size(), 2U); - writer.addUpdateQuery(originalCmd, 0).get(); + writer.addUpdateQuery(operationContext(), originalCmd, 0).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 1); // Adding the next query causes the size to exceed the limit. - writer.addUpdateQuery(originalCmd, 2).get(); + writer.addUpdateQuery(operationContext(), originalCmd, 2).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 0); ASSERT_EQ(getSampledQueryDocumentsCount(nss0), 2); @@ -1095,10 +1095,10 @@ TEST_F(QueryAnalysisWriterTest, FlushAfterAddDeleteIfExceedsSizeLimit) { std::string(maxMemoryUsageBytes / 2, 'a') /* filterFieldName */); ASSERT_EQ(expectedSampledCmds.size(), 2U); - writer.addDeleteQuery(originalCmd, 0).get(); + writer.addDeleteQuery(operationContext(), originalCmd, 0).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 1); // Adding the next query causes the size to exceed the limit. - writer.addDeleteQuery(originalCmd, 1).get(); + writer.addDeleteQuery(operationContext(), originalCmd, 1).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 0); ASSERT_EQ(getSampledQueryDocumentsCount(nss0), 2); @@ -1133,10 +1133,10 @@ TEST_F(QueryAnalysisWriterTest, FlushAfterAddFindAndModifyIfExceedsSizeLimit) { ASSERT_EQ(expectedSampledCmds0.size(), 1U); auto [sampleId1, expectedSampledCmd1] = *expectedSampledCmds1.begin(); - writer.addFindAndModifyQuery(originalCmd0).get(); + writer.addFindAndModifyQuery(operationContext(), originalCmd0).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 1); // Adding the next query causes the size to exceed the limit. - writer.addFindAndModifyQuery(originalCmd1).get(); + writer.addFindAndModifyQuery(operationContext(), originalCmd1).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 0); ASSERT_EQ(getSampledQueryDocumentsCount(nss0), 1); @@ -1539,16 +1539,16 @@ void QueryAnalysisWriterTest::assertNoSampling(const NamespaceString& nss, const ASSERT_EQ(writer.getQueriesCountForTest(), 0); auto originalUpdateCmd = makeUpdateCommandRequest(nss, 1, {0} /* markForSampling */).first; - writer.addUpdateQuery(originalUpdateCmd, 0).get(); + writer.addUpdateQuery(operationContext(), originalUpdateCmd, 0).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 0); auto originalDeleteCmd = makeDeleteCommandRequest(nss, 1, {0} /* markForSampling */).first; - writer.addDeleteQuery(originalDeleteCmd, 0).get(); + writer.addDeleteQuery(operationContext(), originalDeleteCmd, 0).get(); ASSERT_EQ(writer.getQueriesCountForTest(), 0); auto originalFindAndModifyCmd = makeFindAndModifyCommandRequest(nss, true /* isUpdate */, true /* markForSampling */).first; - writer.addFindAndModifyQuery(originalFindAndModifyCmd).get(); + writer.addFindAndModifyQuery(operationContext(), originalFindAndModifyCmd).get(); writer .addDiff(UUID::gen() /* sampleId */, |