SERVER-49935 Add startup method to ReplicaSetAwareService interface

7 files changed, 41 insertions, 2 deletions

diff --git a/src/mongo/db/repl/primary_only_service.h b/src/mongo/db/repl/primary_only_service.h
index d07da959def..7c96c61a604 100644
--- a/src/mongo/db/repl/primary_only_service.h
+++ b/src/mongo/db/repl/primary_only_service.h
@@ -291,8 +291,9 @@ public:
      */
     void shutdown();
 
+    void onStartup(OperationContext*) final {}
     void onStepUpBegin(OperationContext*, long long term) final {}
-    void onBecomeArbiter() final{};
+    void onBecomeArbiter() final {}
     void onStepUpComplete(OperationContext*, long long term) final;
     void onStepDown() final;
 
diff --git a/src/mongo/db/repl/replica_set_aware_service.cpp b/src/mongo/db/repl/replica_set_aware_service.cpp
index db781bb286f..056ec48d1b4 100644
--- a/src/mongo/db/repl/replica_set_aware_service.cpp
+++ b/src/mongo/db/repl/replica_set_aware_service.cpp
@@ -50,6 +50,12 @@ ReplicaSetAwareServiceRegistry& ReplicaSetAwareServiceRegistry::get(
     return registryDecoration(serviceContext);
 }
 
+void ReplicaSetAwareServiceRegistry::onStartup(OperationContext* opCtx) {
+    std::for_each(_services.begin(), _services.end(), [&](ReplicaSetAwareInterface* service) {
+        service->onStartup(opCtx);
+    });
+}
+
 void ReplicaSetAwareServiceRegistry::onStepUpBegin(OperationContext* opCtx, long long term) {
     std::for_each(_services.begin(), _services.end(), [&](ReplicaSetAwareInterface* service) {
         service->onStepUpBegin(opCtx, term);
diff --git a/src/mongo/db/repl/replica_set_aware_service.h b/src/mongo/db/repl/replica_set_aware_service.h
index 20c9e8a8b8c..df0ad747296 100644
--- a/src/mongo/db/repl/replica_set_aware_service.h
+++ b/src/mongo/db/repl/replica_set_aware_service.h
@@ -72,6 +72,9 @@ namespace mongo {
  *     }
  *
  *     // Mandatory:
+ *     void onStartup(OperationContext* opCtx) final {
+ *         // ...
+ *     }
  *     void onStepUpBegin(OperationContext* opCtx) final {
  *         // ...
  *     }
@@ -105,7 +108,14 @@ namespace mongo {
 class ReplicaSetAwareInterface {
 public:
     /**
-     * Called prior to stepping up as PRIMARY, ie. after drain mode has completed.
+     * Called once during ReplicationCoordinator startup. A place to put startup logic such as
+     * initializing thread pools. Cannot depend on the ReplicaSetConfig being loaded yet. Database
+     * reads and writes to unreplicated collections are permitted.
+     */
+    virtual void onStartup(OperationContext* opCtx) = 0;
+
+    /**
+     * Called prior to stepping up as PRIMARY, i.e. after drain mode has completed.
      */
     virtual void onStepUpBegin(OperationContext* opCtx, long long term) = 0;
 
@@ -171,6 +181,7 @@ public:
 
     static ReplicaSetAwareServiceRegistry& get(ServiceContext* serviceContext);
 
+    void onStartup(OperationContext* opCtx) final;
     void onStepUpBegin(OperationContext* opCtx, long long term) final;
     void onStepUpComplete(OperationContext* opCtx, long long term) final;
     void onStepDown() final;
diff --git a/src/mongo/db/repl/replica_set_aware_service_test.cpp b/src/mongo/db/repl/replica_set_aware_service_test.cpp
index 31bd683d4e8..fd702c0cd55 100644
--- a/src/mongo/db/repl/replica_set_aware_service_test.cpp
+++ b/src/mongo/db/repl/replica_set_aware_service_test.cpp
@@ -39,12 +39,17 @@ namespace {
 template <class ActualService>
 class TestService : public ReplicaSetAwareService<ActualService> {
 public:
+    int numCallsOnStartup{0};
     int numCallsOnStepUpBegin{0};
     int numCallsOnStepUpComplete{0};
     int numCallsOnStepDown{0};
     int numCallsOnBecomeArbiter{0};
 
 protected:
+    void onStartup(OperationContext* opCtx) override {
+        numCallsOnStartup++;
+    }
+
     void onStepUpBegin(OperationContext* opCtx, long long term) override {
         numCallsOnStepUpBegin++;
     }
@@ -120,6 +125,11 @@ private:
         return true;
     }
 
+    void onStartup(OperationContext* opCtx) final {
+        ASSERT_EQ(numCallsOnStartup, ServiceB::get(getServiceContext())->numCallsOnStartup - 1);
+        TestService::onStartup(opCtx);
+    }
+
     void onStepUpBegin(OperationContext* opCtx, long long term) final {
         ASSERT_EQ(numCallsOnStepUpBegin,
                   ServiceB::get(getServiceContext())->numCallsOnStepUpBegin - 1);
@@ -168,21 +178,25 @@ TEST_F(ReplicaSetAwareServiceTest, ReplicaSetAwareService) {
     auto b = ServiceB::get(sc);
     auto c = ServiceC::get(sc);
 
+    ASSERT_EQ(0, a->numCallsOnStartup);
     ASSERT_EQ(0, a->numCallsOnStepUpBegin);
     ASSERT_EQ(0, a->numCallsOnStepUpComplete);
     ASSERT_EQ(0, a->numCallsOnStepDown);
     ASSERT_EQ(0, a->numCallsOnBecomeArbiter);
 
+    ASSERT_EQ(0, b->numCallsOnStartup);
     ASSERT_EQ(0, b->numCallsOnStepUpBegin);
     ASSERT_EQ(0, b->numCallsOnStepUpComplete);
     ASSERT_EQ(0, b->numCallsOnStepDown);
     ASSERT_EQ(0, b->numCallsOnBecomeArbiter);
 
+    ASSERT_EQ(0, c->numCallsOnStartup);
     ASSERT_EQ(0, c->numCallsOnStepUpBegin);
     ASSERT_EQ(0, c->numCallsOnStepUpComplete);
     ASSERT_EQ(0, c->numCallsOnStepDown);
     ASSERT_EQ(0, c->numCallsOnBecomeArbiter);
 
+    ReplicaSetAwareServiceRegistry::get(sc).onStartup(opCtx);
     ReplicaSetAwareServiceRegistry::get(sc).onStepUpBegin(opCtx, 0);
     ReplicaSetAwareServiceRegistry::get(sc).onStepUpBegin(opCtx, 0);
     ReplicaSetAwareServiceRegistry::get(sc).onStepUpBegin(opCtx, 0);
@@ -191,16 +205,19 @@ TEST_F(ReplicaSetAwareServiceTest, ReplicaSetAwareService) {
     ReplicaSetAwareServiceRegistry::get(sc).onStepDown();
     ReplicaSetAwareServiceRegistry::get(sc).onBecomeArbiter();
 
+    ASSERT_EQ(0, a->numCallsOnStartup);
     ASSERT_EQ(0, a->numCallsOnStepUpBegin);
     ASSERT_EQ(0, a->numCallsOnStepUpComplete);
     ASSERT_EQ(0, a->numCallsOnStepDown);
     ASSERT_EQ(0, a->numCallsOnBecomeArbiter);
 
+    ASSERT_EQ(1, b->numCallsOnStartup);
     ASSERT_EQ(3, b->numCallsOnStepUpBegin);
     ASSERT_EQ(2, b->numCallsOnStepUpComplete);
     ASSERT_EQ(1, b->numCallsOnStepDown);
     ASSERT_EQ(1, b->numCallsOnBecomeArbiter);
 
+    ASSERT_EQ(1, c->numCallsOnStartup);
     ASSERT_EQ(3, c->numCallsOnStepUpBegin);
     ASSERT_EQ(2, c->numCallsOnStepUpComplete);
     ASSERT_EQ(1, c->numCallsOnStepDown);
diff --git a/src/mongo/db/repl/replication_coordinator_impl.cpp b/src/mongo/db/repl/replication_coordinator_impl.cpp
index b2bc1f9de76..0140ba94d06 100644
--- a/src/mongo/db/repl/replication_coordinator_impl.cpp
+++ b/src/mongo/db/repl/replication_coordinator_impl.cpp
@@ -882,6 +882,8 @@ void ReplicationCoordinatorImpl::startup(OperationContext* opCtx) {
 
     _replExecutor->startup();
 
+    ReplicaSetAwareServiceRegistry::get(_service).onStartup(opCtx);
+
     bool doneLoadingConfig = _startLoadLocalConfig(opCtx);
     if (doneLoadingConfig) {
         // If we're not done loading the config, then the config state will be set by
diff --git a/src/mongo/db/s/balancer/balancer.h b/src/mongo/db/s/balancer/balancer.h
index 5982f436637..d781ae661bd 100644
--- a/src/mongo/db/s/balancer/balancer.h
+++ b/src/mongo/db/s/balancer/balancer.h
@@ -167,6 +167,7 @@ private:
     /**
      * ReplicaSetAwareService entry points.
      */
+    void onStartup(OperationContext* opCtx) final {}
     void onStepUpBegin(OperationContext* opCtx, long long term) final;
     void onStepUpComplete(OperationContext* opCtx, long long term) final;
     void onStepDown() final;
diff --git a/src/mongo/db/vector_clock_mongod.cpp b/src/mongo/db/vector_clock_mongod.cpp
index 1b34ca5133d..90d4f47f739 100644
--- a/src/mongo/db/vector_clock_mongod.cpp
+++ b/src/mongo/db/vector_clock_mongod.cpp
@@ -85,6 +85,7 @@ protected:
     LogicalTime _tick(Component component, uint64_t nTicks) override;
 
 private:
+    void onStartup(OperationContext* opCtx) override {}
     void onStepUpBegin(OperationContext* opCtx, long long term) override {}
     void onStepUpComplete(OperationContext* opCtx, long long term) override {}
     void onStepDown() override {}