SERVER-56738 Create template classes for resharding unittest helpers

4 files changed, 278 insertions, 102 deletions

diff --git a/src/mongo/db/s/SConscript b/src/mongo/db/s/SConscript
index 6cc70f920f4..c3010924d8c 100644
--- a/src/mongo/db/s/SConscript
+++ b/src/mongo/db/s/SConscript
@@ -486,6 +486,7 @@ env.CppUnitTest(
         'resharding/resharding_recipient_service_external_state_test.cpp',
         'resharding/resharding_recipient_service_test.cpp',
         'resharding/resharding_donor_service_test.cpp',
+        'resharding/resharding_service_test_helpers.cpp',
         'resharding/resharding_txn_cloner_test.cpp',
         'session_catalog_migration_destination_test.cpp',
         'session_catalog_migration_source_test.cpp',
diff --git a/src/mongo/db/s/resharding/resharding_donor_service_test.cpp b/src/mongo/db/s/resharding/resharding_donor_service_test.cpp
index 220989ad509..79de3f0fecb 100644
--- a/src/mongo/db/s/resharding/resharding_donor_service_test.cpp
+++ b/src/mongo/db/s/resharding/resharding_donor_service_test.cpp
@@ -47,6 +47,7 @@
 #include "mongo/db/s/operation_sharding_state.h"
 #include "mongo/db/s/resharding/resharding_data_copy_util.h"
 #include "mongo/db/s/resharding/resharding_donor_service.h"
+#include "mongo/db/s/resharding/resharding_service_test_helpers.h"
 #include "mongo/db/s/resharding_util.h"
 #include "mongo/logv2/log.h"
 #include "mongo/unittest/unittest.h"
@@ -54,107 +55,12 @@
 namespace mongo {
 namespace {
 
-class OpObserverForTest;
-class PauseDuringStateTransitions;
-
-class DonorStateTransitionController {
-public:
-    DonorStateTransitionController() = default;
-
-    void waitUntilStateIsReached(DonorStateEnum state) {
-        stdx::unique_lock lk(_mutex);
-        _waitUntilUnpausedCond.wait(lk, [this, state] { return _state == state; });
-    }
-
-private:
-    friend OpObserverForTest;
-    friend PauseDuringStateTransitions;
-
-    void setPauseDuringTransition(DonorStateEnum state) {
-        stdx::lock_guard lk(_mutex);
-        _pauseDuringTransition.insert(state);
-    }
-
-    void unsetPauseDuringTransition(DonorStateEnum state) {
-        stdx::lock_guard lk(_mutex);
-        _pauseDuringTransition.erase(state);
-        _pauseDuringTransitionCond.notify_all();
-    }
-
-    void notifyNewStateAndWaitUntilUnpaused(OperationContext* opCtx, DonorStateEnum newState) {
-        stdx::unique_lock lk(_mutex);
-        _state = newState;
-        _waitUntilUnpausedCond.notify_all();
-        opCtx->waitForConditionOrInterrupt(_pauseDuringTransitionCond, lk, [this, newState] {
-            return _pauseDuringTransition.count(newState) == 0;
-        });
-    }
-
-    Mutex _mutex = MONGO_MAKE_LATCH("DonorStateTransitionController::_mutex");
-    stdx::condition_variable _pauseDuringTransitionCond;
-    stdx::condition_variable _waitUntilUnpausedCond;
-
-    std::set<DonorStateEnum> _pauseDuringTransition;
-    DonorStateEnum _state = DonorStateEnum::kUnused;
-};
-
-class PauseDuringStateTransitions {
-public:
-    PauseDuringStateTransitions(DonorStateTransitionController* controller, DonorStateEnum state)
-        : PauseDuringStateTransitions(controller, std::vector<DonorStateEnum>{state}) {}
-
-    PauseDuringStateTransitions(DonorStateTransitionController* controller,
-                                std::vector<DonorStateEnum> states)
-        : _controller{controller}, _states{std::move(states)} {
-        for (auto state : _states) {
-            _controller->setPauseDuringTransition(state);
-        }
-    }
-
-    ~PauseDuringStateTransitions() {
-        for (auto state : _states) {
-            _controller->unsetPauseDuringTransition(state);
-        }
-    }
-
-    PauseDuringStateTransitions(const PauseDuringStateTransitions&) = delete;
-    PauseDuringStateTransitions& operator=(const PauseDuringStateTransitions&) = delete;
-
-    PauseDuringStateTransitions(PauseDuringStateTransitions&&) = delete;
-    PauseDuringStateTransitions& operator=(PauseDuringStateTransitions&&) = delete;
-
-    void wait(DonorStateEnum state) {
-        _controller->waitUntilStateIsReached(state);
-    }
-
-    void unset(DonorStateEnum state) {
-        _controller->unsetPauseDuringTransition(state);
-    }
-
-private:
-    DonorStateTransitionController* const _controller;
-    const std::vector<DonorStateEnum> _states;
-};
-
-class OpObserverForTest : public OpObserverNoop {
-public:
-    OpObserverForTest(std::shared_ptr<DonorStateTransitionController> controller)
-        : _controller{std::move(controller)} {}
-
-    void onUpdate(OperationContext* opCtx, const OplogUpdateEntryArgs& args) override {
-        if (args.nss != NamespaceString::kDonorReshardingOperationsNamespace) {
-            return;
-        }
-
-        auto doc =
-            ReshardingDonorDocument::parse({"OpObserverForTest"}, args.updateArgs.updatedDoc);
-
-        _controller->notifyNewStateAndWaitUntilUnpaused(opCtx, doc.getMutableState().getState());
-    }
-
-private:
-    std::shared_ptr<DonorStateTransitionController> _controller;
-};
+using DonorStateTransitionController =
+    resharding_service_test_helpers::StateTransitionController<DonorStateEnum>;
+using OpObserverForTest =
+    resharding_service_test_helpers::OpObserverForTest<DonorStateEnum, ReshardingDonorDocument>;
+using PauseDuringStateTransitions =
+    resharding_service_test_helpers::PauseDuringStateTransitions<DonorStateEnum>;
 
 class ExternalStateForTest : public ReshardingDonorService::DonorStateMachineExternalState {
 public:
@@ -171,6 +77,17 @@ public:
                                    const BSONObj& update) override {}
 };
 
+class DonorOpObserverForTest : public OpObserverForTest {
+public:
+    DonorOpObserverForTest(std::shared_ptr<DonorStateTransitionController> controller)
+        : OpObserverForTest(std::move(controller),
+                            NamespaceString::kDonorReshardingOperationsNamespace) {}
+
+    DonorStateEnum getState(const ReshardingDonorDocument& donorDoc) override {
+        return donorDoc.getMutableState().getState();
+    }
+};
+
 class ReshardingDonorServiceForTest : public ReshardingDonorService {
 public:
     explicit ReshardingDonorServiceForTest(ServiceContext* serviceContext)
@@ -202,7 +119,7 @@ public:
         repl::StorageInterface::set(serviceContext, std::move(storageMock));
 
         _controller = std::make_shared<DonorStateTransitionController>();
-        _opObserverRegistry->addObserver(std::make_unique<OpObserverForTest>(_controller));
+        _opObserverRegistry->addObserver(std::make_unique<DonorOpObserverForTest>(_controller));
     }
 
     void stepUp() {
diff --git a/src/mongo/db/s/resharding/resharding_service_test_helpers.cpp b/src/mongo/db/s/resharding/resharding_service_test_helpers.cpp
new file mode 100644
index 00000000000..7af58b266fc
--- /dev/null
+++ b/src/mongo/db/s/resharding/resharding_service_test_helpers.cpp
@@ -0,0 +1,140 @@
+/**
+ *    Copyright (C) 2021-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTest
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/s/resharding/resharding_service_test_helpers.h"
+
+namespace mongo {
+namespace resharding_service_test_helpers {
+
+// -----------------------------------------------
+//  StateTransitionController
+// -----------------------------------------------
+
+template <class StateEnum>
+void StateTransitionController<StateEnum>::waitUntilStateIsReached(StateEnum state) {
+    stdx::unique_lock lk(_mutex);
+    _waitUntilUnpausedCond.wait(lk, [this, state] { return _state == state; });
+}
+
+template <class StateEnum>
+void StateTransitionController<StateEnum>::_setPauseDuringTransition(StateEnum state) {
+    stdx::lock_guard lk(_mutex);
+    _pauseDuringTransition.insert(state);
+}
+
+template <class StateEnum>
+void StateTransitionController<StateEnum>::_unsetPauseDuringTransition(StateEnum state) {
+    stdx::lock_guard lk(_mutex);
+    _pauseDuringTransition.erase(state);
+    _pauseDuringTransitionCond.notify_all();
+}
+
+template <class StateEnum>
+void StateTransitionController<StateEnum>::_notifyNewStateAndWaitUntilUnpaused(
+    OperationContext* opCtx, StateEnum newState) {
+    stdx::unique_lock lk(_mutex);
+    _state = newState;
+    _waitUntilUnpausedCond.notify_all();
+    opCtx->waitForConditionOrInterrupt(_pauseDuringTransitionCond, lk, [this, newState] {
+        return _pauseDuringTransition.count(newState) == 0;
+    });
+}
+
+// -----------------------------------------------
+//  PauseDuringStateTransitions
+// -----------------------------------------------
+
+template <class StateEnum>
+PauseDuringStateTransitions<StateEnum>::PauseDuringStateTransitions(
+    StateTransitionController<StateEnum>* controller, StateEnum state)
+    : PauseDuringStateTransitions<StateEnum>(controller, std::vector<StateEnum>{state}) {}
+
+template <class StateEnum>
+PauseDuringStateTransitions<StateEnum>::PauseDuringStateTransitions(
+    StateTransitionController<StateEnum>* controller, std::vector<StateEnum> states)
+    : _controller{controller}, _states{std::move(states)} {
+    for (auto state : _states) {
+        _controller->_setPauseDuringTransition(state);
+    }
+}
+
+template <class StateEnum>
+PauseDuringStateTransitions<StateEnum>::~PauseDuringStateTransitions() {
+    for (auto state : _states) {
+        _controller->_unsetPauseDuringTransition(state);
+    }
+}
+
+template <class StateEnum>
+void PauseDuringStateTransitions<StateEnum>::wait(StateEnum state) {
+    _controller->waitUntilStateIsReached(state);
+}
+
+template <class StateEnum>
+void PauseDuringStateTransitions<StateEnum>::unset(StateEnum state) {
+    _controller->_unsetPauseDuringTransition(state);
+}
+
+// -----------------------------------------------
+// OpObserverForTest
+// -----------------------------------------------
+
+template <class StateEnum, class ReshardingDocument>
+OpObserverForTest<StateEnum, ReshardingDocument>::OpObserverForTest(
+    std::shared_ptr<StateTransitionController<StateEnum>> controller,
+    NamespaceString reshardingDocumentNss)
+    : _controller{std::move(controller)},
+      _reshardingDocumentNss{std::move(reshardingDocumentNss)} {}
+
+template <class StateEnum, class ReshardingDocument>
+void OpObserverForTest<StateEnum, ReshardingDocument>::onUpdate(OperationContext* opCtx,
+                                                                const OplogUpdateEntryArgs& args) {
+    if (args.nss != _reshardingDocumentNss) {
+        return;
+    }
+
+    auto doc = ReshardingDocument::parse({"OpObserverForTest"}, args.updateArgs.updatedDoc);
+    _controller->_notifyNewStateAndWaitUntilUnpaused(opCtx, getState(doc));
+}
+
+template class StateTransitionController<DonorStateEnum>;
+template class StateTransitionController<RecipientStateEnum>;
+
+template class PauseDuringStateTransitions<DonorStateEnum>;
+template class PauseDuringStateTransitions<RecipientStateEnum>;
+
+template class OpObserverForTest<DonorStateEnum, ReshardingDonorDocument>;
+template class OpObserverForTest<RecipientStateEnum, ReshardingRecipientDocument>;
+
+}  // namespace resharding_service_test_helpers
+}  // namespace mongo
diff --git a/src/mongo/db/s/resharding/resharding_service_test_helpers.h b/src/mongo/db/s/resharding/resharding_service_test_helpers.h
new file mode 100644
index 00000000000..d91f1672000
--- /dev/null
+++ b/src/mongo/db/s/resharding/resharding_service_test_helpers.h
@@ -0,0 +1,118 @@
+/**
+ *    Copyright (C) 2021-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/op_observer_noop.h"
+#include "mongo/db/s/resharding/resharding_coordinator_service.h"
+#include "mongo/db/s/resharding/resharding_donor_service.h"
+#include "mongo/db/s/resharding/resharding_recipient_service.h"
+
+namespace mongo {
+namespace resharding_service_test_helpers {
+template <class StateEnum, class ReshardingDocument>
+class OpObserverForTest;
+
+template <class StateEnum>
+class PauseDuringStateTransitions;
+
+template <class StateEnum>
+class StateTransitionController;
+
+template <class StateEnum>
+class StateTransitionController {
+public:
+    StateTransitionController() = default;
+
+    void waitUntilStateIsReached(StateEnum state);
+
+private:
+    template <class Enum, class ReshardingDocument>
+    friend class OpObserverForTest;
+
+    template <class Enum>
+    friend class PauseDuringStateTransitions;
+
+    void _setPauseDuringTransition(StateEnum state);
+
+    void _unsetPauseDuringTransition(StateEnum state);
+
+    void _notifyNewStateAndWaitUntilUnpaused(OperationContext* opCtx, StateEnum newState);
+
+    Mutex _mutex = MONGO_MAKE_LATCH("StateTransitionController::_mutex");
+    stdx::condition_variable _pauseDuringTransitionCond;
+    stdx::condition_variable _waitUntilUnpausedCond;
+
+    std::set<StateEnum> _pauseDuringTransition;
+    StateEnum _state = StateEnum::kUnused;
+};
+
+template <class StateEnum>
+class PauseDuringStateTransitions {
+public:
+    PauseDuringStateTransitions(StateTransitionController<StateEnum>* controller, StateEnum state);
+    PauseDuringStateTransitions(StateTransitionController<StateEnum>* controller,
+                                std::vector<StateEnum> states);
+
+    ~PauseDuringStateTransitions();
+
+    PauseDuringStateTransitions(const PauseDuringStateTransitions&) = delete;
+    PauseDuringStateTransitions& operator=(const PauseDuringStateTransitions&) = delete;
+
+    PauseDuringStateTransitions(PauseDuringStateTransitions&&) = delete;
+    PauseDuringStateTransitions& operator=(PauseDuringStateTransitions&&) = delete;
+
+    void wait(StateEnum state);
+
+    void unset(StateEnum state);
+
+private:
+    StateTransitionController<StateEnum>* const _controller;
+    const std::vector<StateEnum> _states;
+};
+
+template <class StateEnum, class ReshardingDocument>
+class OpObserverForTest : public OpObserverNoop {
+public:
+    OpObserverForTest(std::shared_ptr<StateTransitionController<StateEnum>> controller,
+                      NamespaceString reshardingDocumentNss);
+
+    void onUpdate(OperationContext* opCtx, const OplogUpdateEntryArgs& args) override;
+
+    virtual StateEnum getState(const ReshardingDocument& reshardingDoc) = 0;
+
+private:
+    std::shared_ptr<StateTransitionController<StateEnum>> _controller;
+    const NamespaceString _reshardingDocumentNss;
+};
+
+}  // namespace resharding_service_test_helpers
+}  // namespace mongo