SERVER-46200 implement basic VectorClock service

2 files changed, 18 insertions, 2 deletions

diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index 1e80510985c..74e1a643016 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -156,17 +156,27 @@ void warnOnMultiVersion(const vector<ClusterStatistics::ShardStatistics>& cluste
                   "shardVersions"_attr = shardVersions.done());
 }
 
-ReplicaSetAwareServiceRegistry::Registerer<Balancer> balancerRegisterer("Balancer");
+const auto _balancerDecoration = ServiceContext::declareDecoration<Balancer>();
+
+const ReplicaSetAwareServiceRegistry::Registerer<Balancer> _balancerRegisterer("Balancer");
 
 }  // namespace
 
+Balancer* Balancer::get(ServiceContext* serviceContext) {
+    return &_balancerDecoration(serviceContext);
+}
+
+Balancer* Balancer::get(OperationContext* operationContext) {
+    return get(operationContext->getServiceContext());
+}
+
 Balancer::Balancer()
     : _balancedLastTime(0),
       _random(std::random_device{}()),
       _clusterStats(std::make_unique<ClusterStatisticsImpl>(_random)),
       _chunkSelectionPolicy(
           std::make_unique<BalancerChunkSelectionPolicyImpl>(_clusterStats.get(), _random)),
-      _migrationManager(getServiceContext()) {}
+      _migrationManager(_balancerDecoration.owner(this)) {}
 
 
 Balancer::~Balancer() {
diff --git a/src/mongo/db/s/balancer/balancer.h b/src/mongo/db/s/balancer/balancer.h
index fdbbb79d1d5..eef505a014a 100644
--- a/src/mongo/db/s/balancer/balancer.h
+++ b/src/mongo/db/s/balancer/balancer.h
@@ -59,6 +59,12 @@ class Balancer : public ReplicaSetAwareServiceConfigSvr<Balancer> {
     Balancer& operator=(const Balancer&) = delete;
 
 public:
+    /**
+     * Provide access to the Balancer decoration on ServiceContext.
+     */
+    static Balancer* get(ServiceContext* serviceContext);
+    static Balancer* get(OperationContext* operationContext);
+
     Balancer();
     ~Balancer();