SERVER-59618 add LockerNoopClientObserverWithReplacementPolicy and use in bimodal unit test fixtures

6 files changed, 46 insertions, 7 deletions

diff --git a/src/mongo/db/concurrency/locker_noop_client_observer.h b/src/mongo/db/concurrency/locker_noop_client_observer.h
index 2e179601443..6f9159e83d1 100644
--- a/src/mongo/db/concurrency/locker_noop_client_observer.h
+++ b/src/mongo/db/concurrency/locker_noop_client_observer.h
@@ -47,11 +47,31 @@ public:
 
     void onDestroyClient(Client* client) final {}
 
-    void onCreateOperationContext(OperationContext* opCtx) final {
+    void onCreateOperationContext(OperationContext* opCtx) override {
         opCtx->setLockState(std::make_unique<LockerNoop>());
     }
 
     void onDestroyOperationContext(OperationContext* opCtx) final {}
 };
 
+/**
+ * Unlike LockerNoopClientObserver, this ClientObserver will not overwrite any existing Lockers on
+ * the OperationContext.
+ *
+ * This class is suitable for test fixtures that may be used in executables with LockerImpl service
+ * hooks installed.
+ */
+class LockerNoopClientObserverWithReplacementPolicy : public LockerNoopClientObserver {
+public:
+    LockerNoopClientObserverWithReplacementPolicy() = default;
+    ~LockerNoopClientObserverWithReplacementPolicy() = default;
+
+    void onCreateOperationContext(OperationContext* opCtx) final {
+        if (opCtx->lockState()) {
+            return;
+        }
+        LockerNoopClientObserver::onCreateOperationContext(opCtx);
+    }
+};
+
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
index 0782f0bdc42..4c1209f49c0 100644
--- a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
@@ -36,11 +36,17 @@
 
 #include "mongo/db/exec/sbe/sbe_plan_stage_test.h"
 
-
+#include "mongo/db/concurrency/locker_noop_client_observer.h"
 #include "mongo/logv2/log.h"
 
 namespace mongo::sbe {
 
+PlanStageTestFixture::PlanStageTestFixture() {
+    auto service = getServiceContext();
+    service->registerClientObserver(
+        std::make_unique<LockerNoopClientObserverWithReplacementPolicy>());
+}
+
 void PlanStageTestFixture::assertValuesEqual(value::TypeTags lhsTag,
                                              value::Value lhsVal,
                                              value::TypeTags rhsTag,
diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.h b/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
index cf9f6cf9714..05174c73747 100644
--- a/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
+++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
@@ -82,7 +82,7 @@ using MakeStageFn = std::function<std::pair<T, std::unique_ptr<PlanStage>>(
  */
 class PlanStageTestFixture : public ServiceContextTest {
 public:
-    PlanStageTestFixture() = default;
+    PlanStageTestFixture();
 
     void setUp() override {
         ServiceContextTest::setUp();
diff --git a/src/mongo/db/pipeline/aggregation_context_fixture.h b/src/mongo/db/pipeline/aggregation_context_fixture.h
index 287caae977c..f6060260455 100644
--- a/src/mongo/db/pipeline/aggregation_context_fixture.h
+++ b/src/mongo/db/pipeline/aggregation_context_fixture.h
@@ -32,6 +32,7 @@
 #include <boost/intrusive_ptr.hpp>
 #include <memory>
 
+#include "mongo/db/concurrency/locker_noop_client_observer.h"
 #include "mongo/db/pipeline/expression_context_for_test.h"
 #include "mongo/db/service_context_test_fixture.h"
 #include "mongo/unittest/temp_dir.h"
@@ -47,8 +48,12 @@ public:
     AggregationContextFixture()
         : AggregationContextFixture(NamespaceString("unittests.pipeline_test")) {}
 
-    AggregationContextFixture(NamespaceString nss)
-        : _expCtx(new ExpressionContextForTest(_opCtx.get(), nss)) {
+    AggregationContextFixture(NamespaceString nss) {
+        auto service = getServiceContext();
+        service->registerClientObserver(
+            std::make_unique<LockerNoopClientObserverWithReplacementPolicy>());
+        _opCtx = makeOperationContext();
+        _expCtx = make_intrusive<ExpressionContextForTest>(_opCtx.get(), nss);
         unittest::TempDir tempDir("AggregationContextFixture");
         _expCtx->tempDir = tempDir.path();
     }
@@ -66,7 +71,7 @@ public:
     }
 
 private:
-    ServiceContext::UniqueOperationContext _opCtx = makeOperationContext();
+    ServiceContext::UniqueOperationContext _opCtx;
     boost::intrusive_ptr<ExpressionContextForTest> _expCtx;
 };
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_test_service_context.cpp b/src/mongo/db/query/query_test_service_context.cpp
index c44836fa6be..e4ee1b40c49 100644
--- a/src/mongo/db/query/query_test_service_context.cpp
+++ b/src/mongo/db/query/query_test_service_context.cpp
@@ -33,6 +33,7 @@
 
 #include <memory>
 
+#include "mongo/db/concurrency/locker_noop_client_observer.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/query/collation/collator_factory_mock.h"
 
@@ -40,6 +41,8 @@ namespace mongo {
 
 QueryTestServiceContext::QueryTestServiceContext()
     : _service(ServiceContext::make()), _client(_service->makeClient("query_test")) {
+    _service->registerClientObserver(
+        std::make_unique<LockerNoopClientObserverWithReplacementPolicy>());
     CollatorFactoryInterface::set(getServiceContext(), std::make_unique<CollatorFactoryMock>());
 }
 
diff --git a/src/mongo/s/sharding_test_fixture_common.cpp b/src/mongo/s/sharding_test_fixture_common.cpp
index 12db6905098..12ef4fd5006 100644
--- a/src/mongo/s/sharding_test_fixture_common.cpp
+++ b/src/mongo/s/sharding_test_fixture_common.cpp
@@ -33,6 +33,7 @@
 
 #include "mongo/s/sharding_test_fixture_common.h"
 
+#include "mongo/db/concurrency/locker_noop_client_observer.h"
 #include "mongo/logv2/log.h"
 #include "mongo/s/catalog/type_changelog.h"
 #include "mongo/s/write_ops/batched_command_request.h"
@@ -44,7 +45,11 @@ using executor::NetworkTestEnv;
 using executor::RemoteCommandRequest;
 using unittest::assertGet;
 
-ShardingTestFixtureCommon::ShardingTestFixtureCommon() = default;
+ShardingTestFixtureCommon::ShardingTestFixtureCommon() {
+    auto service = getServiceContext();
+    service->registerClientObserver(
+        std::make_unique<LockerNoopClientObserverWithReplacementPolicy>());
+}
 
 ShardingTestFixtureCommon::~ShardingTestFixtureCommon() {
     invariant(!_opCtxHolder,