CLOUDP-290619 Refactor and unit test logic to mark hosts as non-voting when scaling down (#4146)

Julien-Ben · fealebenpae · commit 5ea1f14120e3 · 2025-03-05T10:08:31.000+01:00
# Summary This is a follow up from one of the MC Sharded features branches. https://jira.mongodb.org/browse/CLOUDP-290619 This PR refactors `prepareScaleDownShardedCluster` to make it easier to unit test. - extract logic to compute members to scale down from `prepareScaleDownShardedCluster` - unit test the new function `computeMembersToScaleDown ` - fix old tests that relied on the resource status, now migrated to a config map ## Proof of Work Unit tests
diff --git a/controllers/operator/mongodbshardedcluster_controller.go b/controllers/operator/mongodbshardedcluster_controller.go
@@ -162,8 +162,23 @@ func (r *ShardedClusterReconcileHelper) initializeMemberClusters(globalMemberClu
 		r.shardsMemberClustersMap, r.allShardsMemberClusters = r.createShardsMemberClusterLists(shardsMap, globalMemberClustersMap, log, r.deploymentState, false)
 	} else {
 		r.shardsMemberClustersMap, r.allShardsMemberClusters = r.createShardsMemberClusterLists(shardsMap, globalMemberClustersMap, log, r.deploymentState, true)
-		r.configSrvMemberClusters = []multicluster.MemberCluster{multicluster.GetLegacyCentralMemberCluster(r.deploymentState.Status.MongodbShardedClusterSizeConfig.ConfigServerCount, 0, r.commonController.client, r.commonController.SecretClient)}
-		r.mongosMemberClusters = []multicluster.MemberCluster{multicluster.GetLegacyCentralMemberCluster(r.deploymentState.Status.MongodbShardedClusterSizeConfig.MongosCount, 0, r.commonController.client, r.commonController.SecretClient)}
+
+		// SizeStatusInClusters is the primary struct for storing state, designed with multi-cluster support in mind.
+		// For a single-cluster setup, we first attempt to read from the fields in SizeStatusInClusters,
+		// falling back to the legacy structure (MongodbShardedClusterSizeConfig) if they are unavailable. The fallback
+		// is to be defensive but with the migration performed at the beginning of the reconcile (if necessary), there
+		// should be no case of having only the legacy fields populated in the state.
+		configSrvCount, configSrvCountExists := r.deploymentState.Status.SizeStatusInClusters.ConfigServerMongodsInClusters[multicluster.LegacyCentralClusterName]
+		if !configSrvCountExists {
+			configSrvCount = r.deploymentState.Status.MongodbShardedClusterSizeConfig.ConfigServerCount
+		}
+		r.configSrvMemberClusters = []multicluster.MemberCluster{multicluster.GetLegacyCentralMemberCluster(configSrvCount, 0, r.commonController.client, r.commonController.SecretClient)}
+
+		mongosCount, mongosCountExists := r.deploymentState.Status.SizeStatusInClusters.MongosCountInClusters[multicluster.LegacyCentralClusterName]
+		if !mongosCountExists {
+			mongosCount = r.deploymentState.Status.MongodbShardedClusterSizeConfig.MongosCount
+		}
+		r.mongosMemberClusters = []multicluster.MemberCluster{multicluster.GetLegacyCentralMemberCluster(mongosCount, 0, r.commonController.client, r.commonController.SecretClient)}
 	}
 
 	r.allMemberClusters = r.createAllMemberClustersList()
@@ -1665,12 +1680,9 @@ func (r *ShardedClusterReconcileHelper) getMongosHostnames(memberCluster multicl
 	}
 }
 
-// prepareScaleDownShardedCluster collects all replicasets members to scale down, from configservers and shards, accross
-// all clusters, and pass them to PrepareScaleDownFromMap, which sets their votes and priorities to 0
-func (r *ShardedClusterReconcileHelper) prepareScaleDownShardedCluster(omClient om.Connection, log *zap.SugaredLogger) error {
+func (r *ShardedClusterReconcileHelper) computeMembersToScaleDown(configSrvMemberClusters []multicluster.MemberCluster, shardsMemberClustersMap map[int][]multicluster.MemberCluster, log *zap.SugaredLogger) map[string][]string {
 	membersToScaleDown := make(map[string][]string)
-
-	for _, memberCluster := range r.configSrvMemberClusters {
+	for _, memberCluster := range configSrvMemberClusters {
 		currentReplicas := memberCluster.Replicas
 		desiredReplicas := scale.ReplicasThisReconciliation(r.GetConfigSrvScaler(memberCluster))
 		_, currentPodNames := r.getConfigSrvHostnames(memberCluster, currentReplicas)
@@ -1686,7 +1698,7 @@ func (r *ShardedClusterReconcileHelper) prepareScaleDownShardedCluster(omClient
 	}
 
 	// Scaledown size of each shard
-	for shardIdx, memberClusters := range r.shardsMemberClustersMap {
+	for shardIdx, memberClusters := range shardsMemberClustersMap {
 		for _, memberCluster := range memberClusters {
 			currentReplicas := memberCluster.Replicas
 			desiredReplicas := scale.ReplicasThisReconciliation(r.GetShardScaler(shardIdx, memberCluster))
@@ -1703,6 +1715,14 @@ func (r *ShardedClusterReconcileHelper) prepareScaleDownShardedCluster(omClient
 		}
 	}
 
+	return membersToScaleDown
+}
+
+// prepareScaleDownShardedCluster collects all replicasets members to scale down, from configservers and shards, across
+// all clusters, and pass them to PrepareScaleDownFromMap, which sets their votes and priorities to 0
+func (r *ShardedClusterReconcileHelper) prepareScaleDownShardedCluster(omClient om.Connection, log *zap.SugaredLogger) error {
+	membersToScaleDown := r.computeMembersToScaleDown(r.configSrvMemberClusters, r.shardsMemberClustersMap, log)
+
 	if len(membersToScaleDown) > 0 {
 		healthyProcessesToWaitForReadyState := r.getHealthyProcessNamesToWaitForReadyState(omClient, log)
 		if err := replicaset.PrepareScaleDownFromMap(omClient, membersToScaleDown, healthyProcessesToWaitForReadyState, log); err != nil {
diff --git a/controllers/operator/mongodbshardedcluster_controller_multi_test.go b/controllers/operator/mongodbshardedcluster_controller_multi_test.go
@@ -52,9 +52,22 @@ func newShardedClusterReconcilerForMultiCluster(ctx context.Context, sc *mdbv1.M
 // createMockStateConfigMap creates a configMap with the sizeStatusInClusters populated based on the cluster state
 // passed in parameters, to simulate it is the current state of the cluster
 func createMockStateConfigMap(kubeClient client.Client, namespace, scName string, state MultiClusterShardedScalingStep) error {
+	sumMap := func(m map[string]int) int {
+		sum := 0
+		for _, val := range m {
+			sum += val
+		}
+		return sum
+	}
+
+	configServerSum := sumMap(state.configServerDistribution)
+	mongosSum := sumMap(state.mongosDistribution)
+
 	sizeStatus := map[string]interface{}{
 		"status": map[string]interface{}{
-			"shardCount": state.shardCount,
+			"shardCount":        state.shardCount,
+			"configServerCount": configServerSum,
+			"mongosCount":       mongosSum,
 			"sizeStatusInClusters": map[string]interface{}{
 				"shardMongodsInClusters":        state.shardDistribution,
 				"mongosCountInClusters":         state.mongosDistribution,
@@ -2477,6 +2490,125 @@ func checkCorrectShardDistributionInStatus(t *testing.T, sc *mdbv1.MongoDB) {
 	assert.Equal(t, sumSlice(util.Transform(sc.Spec.ShardSpec.ClusterSpecList, clusterSpecItemToMembers)), sc.Status.MongodbShardedClusterSizeConfig.MongodsPerShardCount)
 }
 
+func TestComputeMembersToScaleDown(t *testing.T) {
+	ctx := context.Background()
+	memberCluster1 := "cluster1"
+	memberCluster2 := "cluster2"
+	memberClusterNames := []string{memberCluster1, memberCluster2}
+
+	type testCase struct {
+		name                        string
+		shardCount                  int
+		cfgServerCurrentClusters    []multicluster.MemberCluster
+		shardsCurrentClusters       map[int][]multicluster.MemberCluster
+		targetCfgServerDistribution map[string]int
+		targetShardDistribution     map[string]int
+		expected                    map[string][]string
+	}
+
+	testCases := []testCase{
+		{
+			name:       "Case 1: Downscale config server and shard",
+			shardCount: 1,
+			cfgServerCurrentClusters: []multicluster.MemberCluster{
+				{Name: memberCluster1, Index: 0, Replicas: 5},
+				{Name: memberCluster2, Index: 1, Replicas: 0},
+			},
+			shardsCurrentClusters: map[int][]multicluster.MemberCluster{
+				0: {
+					{Name: memberCluster1, Index: 0, Replicas: 3},
+					{Name: memberCluster2, Index: 1, Replicas: 2},
+				},
+			},
+			targetCfgServerDistribution: map[string]int{
+				memberCluster1: 2,
+				memberCluster2: 1,
+			},
+			targetShardDistribution: map[string]int{
+				memberCluster1: 1,
+				memberCluster2: 2,
+			},
+			expected: map[string][]string{
+				// For the config replica set: downscale from 5 to 2 means remove members with indices 2, 3, 4
+				test.SCBuilderDefaultName + "-config": {
+					test.SCBuilderDefaultName + "-config-0-2",
+					test.SCBuilderDefaultName + "-config-0-3",
+					test.SCBuilderDefaultName + "-config-0-4",
+				},
+				// For the shard replica set (shard 0): downscale from 3 to 1, so remove two members
+				test.SCBuilderDefaultName + "-0": {
+					test.SCBuilderDefaultName + "-0-0-1",
+					test.SCBuilderDefaultName + "-0-0-2",
+				},
+			},
+		},
+		{
+			name:       "Case 2: Scale down and move replicas among clusters",
+			shardCount: 2,
+			cfgServerCurrentClusters: []multicluster.MemberCluster{
+				{Name: memberCluster1, Index: 0, Replicas: 2},
+				{Name: memberCluster2, Index: 1, Replicas: 1},
+			},
+			shardsCurrentClusters: map[int][]multicluster.MemberCluster{
+				0: {
+					{Name: memberCluster1, Index: 0, Replicas: 3},
+					{Name: memberCluster2, Index: 1, Replicas: 2},
+				},
+				1: {
+					{Name: memberCluster1, Index: 0, Replicas: 3},
+					{Name: memberCluster2, Index: 1, Replicas: 2},
+				},
+			},
+			targetCfgServerDistribution: map[string]int{
+				memberCluster1: 1,
+				memberCluster2: 2,
+			},
+			targetShardDistribution: map[string]int{
+				memberCluster1: 3,
+				memberCluster2: 0,
+			},
+			expected: map[string][]string{
+				test.SCBuilderDefaultName + "-config": {
+					test.SCBuilderDefaultName + "-config" + "-0" + "-1",
+				},
+				// For each shard replica set, we remove two members from cluster with index 1
+				test.SCBuilderDefaultName + "-0": {
+					test.SCBuilderDefaultName + "-0-1-0",
+					test.SCBuilderDefaultName + "-0-1-1",
+				},
+				test.SCBuilderDefaultName + "-1": {
+					test.SCBuilderDefaultName + "-1-1-0",
+					test.SCBuilderDefaultName + "-1-1-1",
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			targetSpec := test.DefaultClusterBuilder().
+				SetTopology(mdbv1.ClusterTopologyMultiCluster).
+				SetShardCountSpec(tc.shardCount).
+				SetMongodsPerShardCountSpec(0).
+				SetConfigServerCountSpec(0).
+				SetMongosCountSpec(0).
+				SetConfigSrvClusterSpec(test.CreateClusterSpecList(memberClusterNames, tc.targetCfgServerDistribution)).
+				SetShardClusterSpec(test.CreateClusterSpecList(memberClusterNames, tc.targetShardDistribution)).
+				Build()
+
+			_, omConnectionFactory := mock.NewDefaultFakeClient(targetSpec)
+			memberClusterMap := getFakeMultiClusterMapWithClusters(memberClusterNames, omConnectionFactory)
+
+			_, reconcileHelper, _, _, err := defaultClusterReconciler(ctx, targetSpec, memberClusterMap)
+			assert.NoError(t, err)
+
+			membersToScaleDown := reconcileHelper.computeMembersToScaleDown(tc.cfgServerCurrentClusters, tc.shardsCurrentClusters, zap.S())
+
+			assert.Equal(t, tc.expected, membersToScaleDown)
+		})
+	}
+}
+
 func sumSlice[T constraints.Integer](s []T) int {
 	result := 0
 	for i := range s {
diff --git a/controllers/operator/mongodbshardedcluster_controller_test.go b/controllers/operator/mongodbshardedcluster_controller_test.go
@@ -450,26 +450,34 @@ func getEmptyDeploymentOptions() deploymentOptions {
 // TestPrepareScaleDownShardedCluster tests the scale down operation for config servers and mongods per shard. It checks
 // that all members that will be removed are marked as unvoted
 func TestPrepareScaleDownShardedCluster_ConfigMongodsUp(t *testing.T) {
-	t.Skip("This test is too fragile to be executed; it's based on status and not deployment state and test internal interactions that are no longer true. Either we rewrite it to full Reconcile or remove it.")
 	ctx := context.Background()
+
+	initialState := MultiClusterShardedScalingStep{
+		shardCount: 2,
+		configServerDistribution: map[string]int{
+			multicluster.LegacyCentralClusterName: 3,
+		},
+		shardDistribution: map[string]int{
+			multicluster.LegacyCentralClusterName: 4,
+		},
+	}
+
 	scBeforeScale := test.DefaultClusterBuilder().
-		SetConfigServerCountStatus(3).
 		SetConfigServerCountSpec(3).
-		SetMongodsPerShardCountStatus(4).
 		SetMongodsPerShardCountSpec(4).
 		Build()
 
-	omConnectionFactory := om.NewCachedOMConnectionFactoryWithInitializedConnection(om.NewMockedOmConnection(createDeploymentFromShardedCluster(t, scBeforeScale)))
-	_, reconcileHelper, _, _, _ := defaultClusterReconciler(ctx, scBeforeScale, nil)
-
-	// TODO prepareScaleDownShardedCluster is getting data from deployment state so modify it instead of passing state in MongoDB object
 	scAfterScale := test.DefaultClusterBuilder().
-		SetConfigServerCountStatus(3).
 		SetConfigServerCountSpec(2).
-		SetMongodsPerShardCountStatus(4).
 		SetMongodsPerShardCountSpec(3).
 		Build()
 
+	omConnectionFactory := om.NewCachedOMConnectionFactoryWithInitializedConnection(om.NewMockedOmConnection(createDeploymentFromShardedCluster(t, scBeforeScale)))
+	kubeClient, _ := mock.NewDefaultFakeClient(scAfterScale)
+	// Store the initial scaling status in state configmap
+	assert.NoError(t, createMockStateConfigMap(kubeClient, mock.TestNamespace, scBeforeScale.Name, initialState))
+	_, reconcileHelper, err := newShardedClusterReconcilerFromResource(ctx, scAfterScale, nil, kubeClient, omConnectionFactory)
+	assert.NoError(t, err)
 	assert.NoError(t, reconcileHelper.prepareScaleDownShardedCluster(omConnectionFactory.GetConnection(), zap.S()))
 
 	// create the expected deployment from the sharded cluster that has not yet scaled
@@ -493,16 +501,31 @@ func TestPrepareScaleDownShardedCluster_ConfigMongodsUp(t *testing.T) {
 // TestPrepareScaleDownShardedCluster_ShardsUpMongodsDown checks the situation when shards count increases and mongods
 // count per shard is decreased - scale down operation is expected to be called only for existing shards
 func TestPrepareScaleDownShardedCluster_ShardsUpMongodsDown(t *testing.T) {
-	t.Skip("This test is too fragile to be executed; it's based on status and not deployment state and test internal interactions that are no longer true. Either we rewrite it to full Reconcile or remove it.")
 	ctx := context.Background()
+
+	initialState := MultiClusterShardedScalingStep{
+		shardCount: 4,
+		shardDistribution: map[string]int{
+			multicluster.LegacyCentralClusterName: 4,
+		},
+	}
+
 	scBeforeScale := test.DefaultClusterBuilder().
-		SetShardCountStatus(4).
 		SetShardCountSpec(4).
-		SetMongodsPerShardCountStatus(4).
 		SetMongodsPerShardCountSpec(4).
 		Build()
 
-	_, reconcileHelper, _, omConnectionFactory, _ := defaultClusterReconciler(ctx, scBeforeScale, nil)
+	scAfterScale := test.DefaultClusterBuilder().
+		SetShardCountSpec(2).
+		SetMongodsPerShardCountSpec(3).
+		Build()
+
+	omConnectionFactory := om.NewCachedOMConnectionFactoryWithInitializedConnection(om.NewMockedOmConnection(createDeploymentFromShardedCluster(t, scBeforeScale)))
+	kubeClient, _ := mock.NewDefaultFakeClient(scAfterScale)
+	assert.NoError(t, createMockStateConfigMap(kubeClient, mock.TestNamespace, scBeforeScale.Name, initialState))
+	_, reconcileHelper, err := newShardedClusterReconcilerFromResource(ctx, scAfterScale, nil, kubeClient, omConnectionFactory)
+	assert.NoError(t, err)
+
 	omConnectionFactory.SetPostCreateHook(func(connection om.Connection) {
 		deployment := createDeploymentFromShardedCluster(t, scBeforeScale)
 		if _, err := connection.UpdateDeployment(deployment); err != nil {
@@ -512,17 +535,6 @@ func TestPrepareScaleDownShardedCluster_ShardsUpMongodsDown(t *testing.T) {
 		connection.(*om.MockedOmConnection).CleanHistory()
 	})
 
-	// TODO prepareScaleDownShardedCluster is getting data from deployment state so modify it instead of passing state in MongoDB object
-	scAfterScale := test.DefaultClusterBuilder().
-		SetShardCountStatus(4).
-		SetShardCountSpec(2).
-		SetMongodsPerShardCountStatus(4).
-		SetMongodsPerShardCountSpec(3).
-		Build()
-
-	// necessary otherwise next omConnectionFactory.GetConnection() will return nil as the connectionFactoryFunc hasn't been called yet
-	initializeOMConnection(t, ctx, reconcileHelper, scAfterScale, zap.S(), omConnectionFactory)
-
 	assert.NoError(t, reconcileHelper.prepareScaleDownShardedCluster(omConnectionFactory.GetConnection(), zap.S()))
 
 	// expected change of state: rs members are marked unvoted only for two shards (old state)
@@ -589,17 +601,34 @@ func initializeOMConnection(t *testing.T, ctx context.Context, reconcileHelper *
 // TestUpdateOmDeploymentShardedCluster_HostsRemovedFromMonitoring verifies that if scale down operation was performed -
 // hosts are removed
 func TestUpdateOmDeploymentShardedCluster_HostsRemovedFromMonitoring(t *testing.T) {
-	t.Skip("This test is too fragile to be executed; it's based on status and not deployment state and test internal interactions that are no longer true. Either we rewrite it to full Reconcile or remove it.")
 	ctx := context.Background()
-	// TODO use deployment state instead of status
+
+	initialState := MultiClusterShardedScalingStep{
+		mongosDistribution: map[string]int{
+			multicluster.LegacyCentralClusterName: 2,
+		},
+		configServerDistribution: map[string]int{
+			multicluster.LegacyCentralClusterName: 4,
+		},
+	}
+
 	sc := test.DefaultClusterBuilder().
-		SetMongosCountStatus(2).
 		SetMongosCountSpec(2).
-		SetConfigServerCountStatus(4).
 		SetConfigServerCountSpec(4).
 		Build()
 
-	_, reconcileHelper, _, omConnectionFactory, _ := defaultClusterReconciler(ctx, sc, nil)
+	// we need to create a different sharded cluster that is currently in the process of scaling down
+	scScaledDown := test.DefaultClusterBuilder().
+		SetMongosCountSpec(1).
+		SetConfigServerCountSpec(3).
+		Build()
+
+	omConnectionFactory := om.NewCachedOMConnectionFactoryWithInitializedConnection(om.NewMockedOmConnection(createDeploymentFromShardedCluster(t, sc)))
+	kubeClient, _ := mock.NewDefaultFakeClient(sc)
+	assert.NoError(t, createMockStateConfigMap(kubeClient, mock.TestNamespace, sc.Name, initialState))
+	_, reconcileHelper, err := newShardedClusterReconcilerFromResource(ctx, scScaledDown, nil, kubeClient, omConnectionFactory)
+	assert.NoError(t, err)
+
 	omConnectionFactory.SetPostCreateHook(func(connection om.Connection) {
 		// the initial deployment we create should have all processes
 		deployment := createDeploymentFromShardedCluster(t, sc)
@@ -609,20 +638,6 @@ func TestUpdateOmDeploymentShardedCluster_HostsRemovedFromMonitoring(t *testing.
 		connection.(*om.MockedOmConnection).AddHosts(deployment.GetAllHostnames())
 		connection.(*om.MockedOmConnection).CleanHistory()
 	})
-	// necessary otherwise next omConnectionFactory.GetConnection() will return nil as the connectionFactoryFunc hasn't been called yet
-	initializeOMConnection(t, ctx, reconcileHelper, sc, zap.S(), omConnectionFactory)
-
-	// we need to create a different sharded cluster that is currently in the process of scaling down
-	// TODO use deployment state instead of status
-	scScaledDown := test.DefaultClusterBuilder().
-		SetMongosCountStatus(2).
-		SetMongosCountSpec(1).
-		SetConfigServerCountStatus(4).
-		SetConfigServerCountSpec(3).
-		Build()
-
-	// necessary otherwise next omConnectionFactory.GetConnection() will return nil as the connectionFactoryFunc hasn't been called yet
-	initializeOMConnection(t, ctx, reconcileHelper, scScaledDown, zap.S(), omConnectionFactory)
 
 	// updateOmDeploymentShardedCluster checks an element from ac.Auth.DeploymentAuthMechanisms
 	// so we need to ensure it has a non-nil value. An empty list implies no authentication
