SERVER-74990 Make query sampling handle legacy WouldChangeOwningShard updates

5 files changed, 43 insertions, 13 deletions

diff --git a/src/mongo/db/operation_context.h b/src/mongo/db/operation_context.h
index 08d53972e56..0027cb683b1 100644
--- a/src/mongo/db/operation_context.h
+++ b/src/mongo/db/operation_context.h
@@ -634,16 +634,17 @@ public:
         _killOpsExempt = true;
     }
 
-    bool explicitlyOptedIntoQuerySampling() const {
-        return _explicitlyOptIntoQuerySampling;
+    // The query sampling options for operations on this opCtx. 'optIn' makes the operations
+    // eligible for query sampling regardless of whether the client is considered as internal by
+    // the sampler. 'optOut' does the opposite.
+    enum QuerySamplingOptions { kOptIn, kOptOut };
+
+    boost::optional<QuerySamplingOptions> getQuerySamplingOptions() {
+        return _querySamplingOpts;
     }
 
-    /**
-     * Makes operations on this opCtx eligible for query sampling regardless of whether the client
-     * is considered as internal by the sampler.
-     */
-    void setExplicitlyOptIntoQuerySampling() {
-        _explicitlyOptIntoQuerySampling = true;
+    void setQuerySamplingOptions(QuerySamplingOptions option) {
+        _querySamplingOpts = option;
     }
 
     /**
@@ -850,9 +851,8 @@ private:
     // to refresh locks for prepared tranasctions or abort in-progress transactions.
     bool _killOpsExempt = false;
 
-    // Make operations on this opCtx eligible for query sampling regardless of whether the client
-    // is considered as internal by the sampler.
-    bool _explicitlyOptIntoQuerySampling = false;
+    // The query sampling options for operations on this opCtx.
+    boost::optional<QuerySamplingOptions> _querySamplingOpts;
 };
 
 // Gets a TimeZoneDatabase pointer from the ServiceContext.
diff --git a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
index c0f9ca86f62..586908ba6b8 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
@@ -74,6 +74,8 @@
 namespace mongo {
 namespace {
 
+using QuerySamplingOptions = OperationContext::QuerySamplingOptions;
+
 const ReadPreferenceSetting kPrimaryOnlyReadPreference(ReadPreference::PrimaryOnly);
 const char kLegacyRuntimeConstantsField[] = "runtimeConstants";
 
@@ -575,6 +577,8 @@ void FindAndModifyCmd::_constructResult(OperationContext* opCtx,
             handleWouldChangeOwningShardError(opCtx, shardId, nss, cmdObj, responseStatus, result);
         } else {
             // TODO SERVER-67429: Remove this branch.
+            opCtx->setQuerySamplingOptions(QuerySamplingOptions::kOptOut);
+
             if (isRetryableWrite) {
                 _handleWouldChangeOwningShardErrorRetryableWriteLegacy(
                     opCtx, shardId, shardVersion, dbVersion, nss, cmdObj, result);
diff --git a/src/mongo/s/commands/cluster_write_cmd.cpp b/src/mongo/s/commands/cluster_write_cmd.cpp
index 3d75dcc9eaf..0a7fd140fd3 100644
--- a/src/mongo/s/commands/cluster_write_cmd.cpp
+++ b/src/mongo/s/commands/cluster_write_cmd.cpp
@@ -63,6 +63,8 @@
 namespace mongo {
 namespace {
 
+using QuerySamplingOptions = OperationContext::QuerySamplingOptions;
+
 MONGO_FAIL_POINT_DEFINE(hangAfterThrowWouldChangeOwningShardRetryableWrite);
 
 void batchErrorToNotPrimaryErrorTracker(const BatchedCommandRequest& request,
@@ -334,6 +336,8 @@ bool ClusterWriteCmd::handleWouldChangeOwningShardError(OperationContext* opCtx,
         }
     } else {
         // TODO SERVER-67429: Delete this branch.
+        opCtx->setQuerySamplingOptions(QuerySamplingOptions::kOptOut);
+
         if (isRetryableWrite) {
             if (MONGO_unlikely(hangAfterThrowWouldChangeOwningShardRetryableWrite.shouldFail())) {
                 LOGV2(22759, "Hit hangAfterThrowWouldChangeOwningShardRetryableWrite failpoint");
diff --git a/src/mongo/s/query_analysis_sampler.cpp b/src/mongo/s/query_analysis_sampler.cpp
index 5fb0bb5e316..23e8727fd77 100644
--- a/src/mongo/s/query_analysis_sampler.cpp
+++ b/src/mongo/s/query_analysis_sampler.cpp
@@ -46,6 +46,8 @@ namespace analyze_shard_key {
 
 namespace {
 
+using QuerySamplingOptions = OperationContext::QuerySamplingOptions;
+
 MONGO_FAIL_POINT_DEFINE(disableQueryAnalysisSampler);
 MONGO_FAIL_POINT_DEFINE(overwriteQueryAnalysisSamplerAvgLastCountToZero);
 
@@ -329,11 +331,17 @@ void QueryAnalysisSampler::_incrementCounters(OperationContext* opCtx,
 boost::optional<UUID> QueryAnalysisSampler::tryGenerateSampleId(OperationContext* opCtx,
                                                                 const NamespaceString& nss,
                                                                 SampledCommandNameEnum cmdName) {
-    if (!opCtx->getClient()->session() && !opCtx->explicitlyOptedIntoQuerySampling()) {
+    auto opts = opCtx->getQuerySamplingOptions();
+
+    if (!opCtx->getClient()->session() && opts != QuerySamplingOptions::kOptIn) {
         // Do not generate a sample id for an internal query unless it has explicitly opted into
         // query sampling.
         return boost::none;
     }
+    if (opts == QuerySamplingOptions::kOptOut) {
+        // Do not generate a sample id for a query that has explicitly opted out of query sampling.
+        return boost::none;
+    }
 
     stdx::lock_guard<Latch> lk(_mutex);
     auto it = _sampleRateLimiters.find(nss);
diff --git a/src/mongo/s/query_analysis_sampler_test.cpp b/src/mongo/s/query_analysis_sampler_test.cpp
index 97b432a8303..3144dde2052 100644
--- a/src/mongo/s/query_analysis_sampler_test.cpp
+++ b/src/mongo/s/query_analysis_sampler_test.cpp
@@ -50,6 +50,8 @@ namespace mongo {
 namespace analyze_shard_key {
 namespace {
 
+using QuerySamplingOptions = OperationContext::QuerySamplingOptions;
+
 const auto smoothingFactor = gQueryAnalysisQueryStatsSmoothingFactor;
 
 TEST(QueryAnalysisSamplerQueryStatsTest, RefreshBasic) {
@@ -914,6 +916,12 @@ TEST_F(QueryAnalysisSamplerTest, TryGenerateSampleIdExternalClient) {
     ASSERT_FALSE(sampler.tryGenerateSampleId(opCtx, nss1, SampledCommandNameEnum::kFind));
     // This collection doesn't have sampling enabled.
     ASSERT_FALSE(sampler.tryGenerateSampleId(opCtx, nss2, SampledCommandNameEnum::kFind));
+
+    advanceTime(Milliseconds(1000));
+    // The number of tokens available in the bucket for rateLimiter0 right after the refill is 0
+    // + 1.0. Cannot sample since the client has explicitly opted out of query sampling.
+    opCtx->setQuerySamplingOptions(QuerySamplingOptions::kOptOut);
+    ASSERT_FALSE(sampler.tryGenerateSampleId(opCtx, nss0, SampledCommandNameEnum::kFind));
 }
 
 TEST_F(QueryAnalysisSamplerTest, TryGenerateSampleIdInternalClient) {
@@ -936,11 +944,17 @@ TEST_F(QueryAnalysisSamplerTest, TryGenerateSampleIdInternalClient) {
     // + 1.0. Cannot sample since the client is internal.
     ASSERT_FALSE(sampler.tryGenerateSampleId(opCtx, nss0, SampledCommandNameEnum::kFind));
 
-    opCtx->setExplicitlyOptIntoQuerySampling();
+    opCtx->setQuerySamplingOptions(QuerySamplingOptions::kOptIn);
     // Can sample now since the client has explicitly opted into query sampling.
     ASSERT(sampler.tryGenerateSampleId(opCtx, nss0, SampledCommandNameEnum::kFind));
     // Cannot sample since there are no tokens left.
     ASSERT_FALSE(sampler.tryGenerateSampleId(opCtx, nss0, SampledCommandNameEnum::kFind));
+
+    advanceTime(Milliseconds(1000));
+    opCtx->setQuerySamplingOptions(QuerySamplingOptions::kOptOut);
+    // The number of tokens available in the bucket for rateLimiter0 right after the refill is 0
+    // + 1.0. Cannot sample since the client has explicitly opted into query sampling.
+    ASSERT_FALSE(sampler.tryGenerateSampleId(opCtx, nss0, SampledCommandNameEnum::kFind));
 }
 
 TEST_F(QueryAnalysisSamplerTest, RefreshConfigurationsNewCollectionUuid) {