SERVER-74990 Make query sampling handle legacy WouldChangeOwningShard updates

2 files changed, 8 insertions, 0 deletions

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");