Fix scale when force update (#188)

Signed-off-by: haorenfsa <[email protected]>
zilliztech · Sep 19, 2024 · 7e4f54c · 7e4f54c
1 parent af463b9
commit 7e4f54c
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 4 deletions.
diff --git a/pkg/controllers/deploy_ctrl_util.go b/pkg/controllers/deploy_ctrl_util.go
@@ -281,14 +281,40 @@ func (c *DeployControllerBizUtilImpl) ScaleDeployments(ctx context.Context, mc v
 	if err != nil {
 		return err
 	}
+	err = c.checkCanScaleNow(ctx, mc, currentDeployment, lastDeployment)
+	if err != nil {
+		return err
+	}
+	action := c.planNextScaleAction(mc, currentDeployment, lastDeployment)
+	return c.doScaleAction(ctx, action)
+}
+
+func (c *DeployControllerBizUtilImpl) checkCanScaleNow(ctx context.Context, mc v1beta1.Milvus, currentDeployment, lastDeployment *appsv1.Deployment) error {
+	isForceUpdate := mc.Spec.Com.ImageUpdateMode == v1beta1.ImageUpdateModeForce
+	if isForceUpdate {
+		return nil
+	}
 	if v1beta1.Labels().IsComponentRolling(mc, c.component.Name) {
-		err = c.checkDeploymentsStable(ctx, currentDeployment, lastDeployment)
+		err := c.checkDeploymentsStable(ctx, currentDeployment, lastDeployment)
 		if err != nil {
-			return err
+			return errors.Wrap(err, "check deployments stable")
 		}
 	}
-	action := c.planNextScaleAction(mc, currentDeployment, lastDeployment)
-	return c.doScaleAction(ctx, action)
+	return nil
+}
+
+func (c *DeployControllerBizUtilImpl) planScaleForForceUpgrade(mc v1beta1.Milvus, currentDeployment, lastDeployment *appsv1.Deployment) scaleAction {
+	currentDeployReplicas := getDeployReplicas(currentDeployment)
+	expectedReplicas := int(ReplicasValue(c.component.GetReplicas(mc.Spec)))
+	if currentDeployReplicas != expectedReplicas {
+		return scaleAction{deploy: currentDeployment, replicaChange: expectedReplicas - currentDeployReplicas}
+	}
+
+	lastDeployReplicas := getDeployReplicas(lastDeployment)
+	if lastDeployReplicas != 0 {
+		return scaleAction{deploy: lastDeployment, replicaChange: -lastDeployReplicas}
+	}
+	return scaleAction{}
 }
 
 type scaleAction struct {
@@ -307,6 +333,8 @@ func (c *DeployControllerBizUtilImpl) planNextScaleAction(mc v1beta1.Milvus, cur
 		return c.planScaleForHPA(currentDeployment)
 	case scaleKindRollout:
 		return c.planScaleForRollout(mc, currentDeployment, lastDeployment)
+	case scaleKindForce:
+		return c.planScaleForForceUpgrade(mc, currentDeployment, lastDeployment)
 	default:
 		return c.planScaleForNormalState(mc, currentDeployment)
 	}
@@ -318,9 +346,13 @@ const (
 	scaleKindNormal scaleKind = iota
 	scaleKindRollout
 	scaleKindHPA
+	scaleKindForce
 )
 
 func (c *DeployControllerBizUtilImpl) checkScaleKind(mc v1beta1.Milvus) scaleKind {
+	if mc.Spec.Com.ImageUpdateMode == v1beta1.ImageUpdateModeForce {
+		return scaleKindForce
+	}
 	expectedReplicas := int(ReplicasValue(c.component.GetReplicas(mc.Spec)))
 	isHpa := expectedReplicas < 0
 	if isHpa {

diff --git a/pkg/controllers/deploy_ctrl_util_test.go b/pkg/controllers/deploy_ctrl_util_test.go
@@ -1346,6 +1346,51 @@ func TestDeployControllerBizUtilImpl_ScaleDeployements(t *testing.T) {
 	})
 }
 
+func TestDeployControllerBizUtilImpl_planNextScaleAction(t *testing.T) {
+	bizUtil := DeployControllerBizUtilImpl{
+		component: MilvusStandalone,
+	}
+	mc := v1beta1.Milvus{}
+	mc.Spec.Com.Standalone = &v1beta1.MilvusStandalone{}
+
+	currentDeploy := new(appsv1.Deployment)
+	lastDeploy := new(appsv1.Deployment)
+	currentReplicas := &currentDeploy.Spec.Replicas
+	lastReplicas := &lastDeploy.Spec.Replicas
+	expected := &mc.Spec.Com.Standalone.Replicas
+	t.Run("force upgrade/current<expected", func(t *testing.T) {
+		mc.Spec.Com.ImageUpdateMode = v1beta1.ImageUpdateModeForce
+		*expected = int32Ptr(3)
+		*currentReplicas = int32Ptr(0)
+		action := bizUtil.planNextScaleAction(mc, currentDeploy, lastDeploy)
+		assert.Equal(t, scaleAction{
+			deploy:        currentDeploy,
+			replicaChange: 3,
+		}, action)
+	})
+
+	t.Run("force upgrade/current=expected/last!=0", func(t *testing.T) {
+		mc.Spec.Com.ImageUpdateMode = v1beta1.ImageUpdateModeForce
+		*expected = int32Ptr(3)
+		*currentReplicas = int32Ptr(3)
+		*lastReplicas = int32Ptr(3)
+		action := bizUtil.planNextScaleAction(mc, currentDeploy, lastDeploy)
+		assert.Equal(t, scaleAction{
+			deploy:        lastDeploy,
+			replicaChange: -3,
+		}, action)
+	})
+
+	t.Run("force upgrade/current=expected/last=0", func(t *testing.T) {
+		mc.Spec.Com.ImageUpdateMode = v1beta1.ImageUpdateModeForce
+		*expected = int32Ptr(3)
+		*currentReplicas = int32Ptr(3)
+		*lastReplicas = int32Ptr(0)
+		action := bizUtil.planNextScaleAction(mc, currentDeploy, lastDeploy)
+		assert.Equal(t, noScaleAction, action)
+	})
+}
+
 func TestDeployControllerBizUtilImpl_PrepareNewRollout(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	defer mockCtrl.Finish()