feat: enhance rollback functionality to clean up CANARY and BASELINE variants (#6660)

Signed-off-by: Mohammed Firdous <124298708+mohammedfirdouss@users.noreply.github.com>

Author: mohammedfirdouss

pkg/app/pipedv1/plugin/kubernetes_multicluster/deployment/rollback.go

@@ -147,8 +147,10 @@ func (p *Plugin) rollback(ctx context.Context, input *sdk.ExecuteStageInput[kube
 	// When addVariantLabelToSelector is true, ensure that all workloads
 	// have the variant label in their selector.
 	var (
-		variantLabel   = cfg.Spec.VariantLabel.Key
-		primaryVariant = cfg.Spec.VariantLabel.PrimaryValue
+		variantLabel    = cfg.Spec.VariantLabel.Key
+		primaryVariant  = cfg.Spec.VariantLabel.PrimaryValue
+		canaryVariant   = cfg.Spec.VariantLabel.CanaryValue
+		baselineVariant = cfg.Spec.VariantLabel.BaselineValue
 	)
 	// TODO: Consider other fields to configure whether to add a variant label to the selector
 	// because the rollback stage is executed in both quick sync and pipeline sync strategies.
@@ -186,17 +188,34 @@ func (p *Plugin) rollback(ctx context.Context, input *sdk.ExecuteStageInput[kube
 	}
 
 	// Create the applier for the target cluster.
-	applier := provider.NewApplier(provider.NewKubectl(kubectlPath), cfg.Spec.Input, deployTargetConfig, input.Logger)
+	kubectl := provider.NewKubectl(kubectlPath)
+	applier := provider.NewApplier(kubectl, cfg.Spec.Input, deployTargetConfig, input.Logger)
 
 	// Start applying all manifests to add or update running resources.
 	if err := applyManifests(ctx, applier, manifests, cfg.Spec.Input.Namespace, lp); err != nil {
 		lp.Errorf("Failed while applying manifests (%v)", err)
 		return sdk.StageStatusFailure
 	}
 
-	// TODO: implement prune resources
-	// TODO: delete all resources of CANARY variant
-	// TODO: delete all resources of BASELINE variant
+	var failed bool
 
+	lp.Info("Start removing CANARY variant resources if exists")
+	if err := deleteVariantResources(ctx, lp, kubectl, deployTargetConfig.KubeConfigPath, applier, input.Request.Deployment.ApplicationID, variantLabel, canaryVariant); err != nil {
+		lp.Errorf("Failed while deleting CANARY variant resources (%v)", err)
+		failed = true
+	}
+
+	lp.Info("Start removing BASELINE variant resources if exists")
+	if err := deleteVariantResources(ctx, lp, kubectl, deployTargetConfig.KubeConfigPath, applier, input.Request.Deployment.ApplicationID, variantLabel, baselineVariant); err != nil {
+		lp.Errorf("Failed while deleting BASELINE variant resources (%v)", err)
+		failed = true
+	}
+
+	// TODO: prune resources which don't exist in the running manifests but exist in the target manifests.
+	// This occurs when the user adds a new resource and the deployment pipeline fails.
+
+	if failed {
+		return sdk.StageStatusFailure
+	}
 	return sdk.StageStatusSuccess
 }

pkg/app/pipedv1/plugin/kubernetes_multicluster/deployment/rollback_test.go

@@ -617,3 +617,167 @@ func TestPlugin_executeK8sMultiRollbackStage_FailureRollback_when_at_least_one_o
 
 	assert.Equal(t, sdk.StageStatusSuccess, status)
 }
+
+func TestPlugin_executeK8sMultiRollbackStage_CleansUpCanaryVariant(t *testing.T) {
+	t.Parallel()
+
+	ctx := t.Context()
+	testRegistry := toolregistrytest.NewTestToolRegistry(t)
+	dtConfig, dynamicClient := setupTestDeployTargetConfigAndDynamicClient(t)
+
+	canaryDir := filepath.Join("testdata", "canary_rollout_without_create_service")
+	simpleDir := filepath.Join("testdata", "simple")
+
+	canaryCfg := sdk.LoadApplicationConfigForTest[kubeconfig.KubernetesApplicationSpec](t, filepath.Join(canaryDir, "app.pipecd.yaml"), "kubernetes_multicluster")
+	simpleCfg := sdk.LoadApplicationConfigForTest[kubeconfig.KubernetesApplicationSpec](t, filepath.Join(simpleDir, "app.pipecd.yaml"), "kubernetes_multicluster")
+
+	// Step 1: deploy canary variant so it exists in the cluster.
+	canaryInput := &sdk.ExecuteStageInput[kubeconfig.KubernetesApplicationSpec]{
+		Request: sdk.ExecuteStageRequest[kubeconfig.KubernetesApplicationSpec]{
+			StageName:   StageK8sMultiCanaryRollout,
+			StageConfig: []byte(`{"replicas": 1}`),
+			TargetDeploymentSource: sdk.DeploymentSource[kubeconfig.KubernetesApplicationSpec]{
+				ApplicationDirectory:      canaryDir,
+				CommitHash:                "target-hash",
+				ApplicationConfig:         canaryCfg,
+				ApplicationConfigFilename: "app.pipecd.yaml",
+			},
+			Deployment: sdk.Deployment{
+				PipedID:       "piped-id",
+				ApplicationID: "app-id",
+			},
+		},
+		Client: sdk.NewClient(nil, "kubernetes_multicluster", "app-id", "stage-id", logpersistertest.NewTestLogPersister(t), testRegistry),
+		Logger: zaptest.NewLogger(t),
+	}
+
+	plugin := &Plugin{}
+	status := plugin.executeK8sMultiCanaryRolloutStage(ctx, canaryInput, []*sdk.DeployTarget[kubeconfig.KubernetesDeployTargetConfig]{
+		{Name: "default", Config: *dtConfig},
+	})
+	require.Equal(t, sdk.StageStatusSuccess, status)
+
+	// Verify canary deployment exists.
+	_, err := dynamicClient.Resource(schema.GroupVersionResource{Group: "apps", Version: "v1", Resource: "deployments"}).Namespace("default").Get(ctx, "simple-canary", metav1.GetOptions{})
+	require.NoError(t, err, "canary deployment should exist before rollback")
+
+	// Step 2: run rollback — should clean up the canary variant.
+	rollbackInput := &sdk.ExecuteStageInput[kubeconfig.KubernetesApplicationSpec]{
+		Request: sdk.ExecuteStageRequest[kubeconfig.KubernetesApplicationSpec]{
+			StageName:   "K8S_MULTI_ROLLBACK",
+			StageConfig: []byte(``),
+			RunningDeploymentSource: sdk.DeploymentSource[kubeconfig.KubernetesApplicationSpec]{
+				ApplicationDirectory:      simpleDir,
+				CommitHash:                "previous-hash",
+				ApplicationConfig:         simpleCfg,
+				ApplicationConfigFilename: "app.pipecd.yaml",
+			},
+			TargetDeploymentSource: sdk.DeploymentSource[kubeconfig.KubernetesApplicationSpec]{
+				ApplicationDirectory:      simpleDir,
+				CommitHash:                "target-hash",
+				ApplicationConfig:         simpleCfg,
+				ApplicationConfigFilename: "app.pipecd.yaml",
+			},
+			Deployment: sdk.Deployment{
+				PipedID:       "piped-id",
+				ApplicationID: "app-id",
+			},
+		},
+		Client: sdk.NewClient(nil, "kubernetes_multicluster", "app-id", "stage-id", logpersistertest.NewTestLogPersister(t), testRegistry),
+		Logger: zaptest.NewLogger(t),
+	}
+
+	status = plugin.executeK8sMultiRollbackStage(ctx, rollbackInput, []*sdk.DeployTarget[kubeconfig.KubernetesDeployTargetConfig]{
+		{Name: "default", Config: *dtConfig},
+	})
+	assert.Equal(t, sdk.StageStatusSuccess, status)
+
+	// Verify canary deployment was deleted.
+	_, err = dynamicClient.Resource(schema.GroupVersionResource{Group: "apps", Version: "v1", Resource: "deployments"}).Namespace("default").Get(ctx, "simple-canary", metav1.GetOptions{})
+	assert.True(t, apierrors.IsNotFound(err), "canary deployment should be deleted after rollback")
+}
+
+func TestPlugin_executeK8sMultiRollbackStage_CleansUpBaselineVariant(t *testing.T) {
+	t.Parallel()
+
+	ctx := t.Context()
+	testRegistry := toolregistrytest.NewTestToolRegistry(t)
+	dtConfig, dynamicClient := setupTestDeployTargetConfigAndDynamicClient(t)
+
+	baselineDir := filepath.Join("testdata", "baseline_rollout_without_create_service")
+	simpleDir := filepath.Join("testdata", "simple")
+
+	baselineCfg := sdk.LoadApplicationConfigForTest[kubeconfig.KubernetesApplicationSpec](t, filepath.Join(baselineDir, "app.pipecd.yaml"), "kubernetes_multicluster")
+	simpleCfg := sdk.LoadApplicationConfigForTest[kubeconfig.KubernetesApplicationSpec](t, filepath.Join(simpleDir, "app.pipecd.yaml"), "kubernetes_multicluster")
+
+	// Step 1: deploy baseline variant so it exists in the cluster.
+	baselineInput := &sdk.ExecuteStageInput[kubeconfig.KubernetesApplicationSpec]{
+		Request: sdk.ExecuteStageRequest[kubeconfig.KubernetesApplicationSpec]{
+			StageName:   StageK8sMultiBaselineRollout,
+			StageConfig: []byte(`{"replicas": 1}`),
+			RunningDeploymentSource: sdk.DeploymentSource[kubeconfig.KubernetesApplicationSpec]{
+				ApplicationDirectory:      baselineDir,
+				CommitHash:                "previous-hash",
+				ApplicationConfig:         baselineCfg,
+				ApplicationConfigFilename: "app.pipecd.yaml",
+			},
+			TargetDeploymentSource: sdk.DeploymentSource[kubeconfig.KubernetesApplicationSpec]{
+				ApplicationDirectory:      baselineDir,
+				CommitHash:                "target-hash",
+				ApplicationConfig:         baselineCfg,
+				ApplicationConfigFilename: "app.pipecd.yaml",
+			},
+			Deployment: sdk.Deployment{
+				PipedID:       "piped-id",
+				ApplicationID: "app-id",
+			},
+		},
+		Client: sdk.NewClient(nil, "kubernetes_multicluster", "app-id", "stage-id", logpersistertest.NewTestLogPersister(t), testRegistry),
+		Logger: zaptest.NewLogger(t),
+	}
+
+	plugin := &Plugin{}
+	status := plugin.executeK8sMultiBaselineRolloutStage(ctx, baselineInput, []*sdk.DeployTarget[kubeconfig.KubernetesDeployTargetConfig]{
+		{Name: "default", Config: *dtConfig},
+	})
+	require.Equal(t, sdk.StageStatusSuccess, status)
+
+	// Verify baseline deployment exists.
+	_, err := dynamicClient.Resource(schema.GroupVersionResource{Group: "apps", Version: "v1", Resource: "deployments"}).Namespace("default").Get(ctx, "simple-baseline", metav1.GetOptions{})
+	require.NoError(t, err, "baseline deployment should exist before rollback")
+
+	// Step 2: run rollback — should clean up the baseline variant.
+	rollbackInput := &sdk.ExecuteStageInput[kubeconfig.KubernetesApplicationSpec]{
+		Request: sdk.ExecuteStageRequest[kubeconfig.KubernetesApplicationSpec]{
+			StageName:   "K8S_MULTI_ROLLBACK",
+			StageConfig: []byte(``),
+			RunningDeploymentSource: sdk.DeploymentSource[kubeconfig.KubernetesApplicationSpec]{
+				ApplicationDirectory:      simpleDir,
+				CommitHash:                "previous-hash",
+				ApplicationConfig:         simpleCfg,
+				ApplicationConfigFilename: "app.pipecd.yaml",
+			},
+			TargetDeploymentSource: sdk.DeploymentSource[kubeconfig.KubernetesApplicationSpec]{
+				ApplicationDirectory:      simpleDir,
+				CommitHash:                "target-hash",
+				ApplicationConfig:         simpleCfg,
+				ApplicationConfigFilename: "app.pipecd.yaml",
+			},
+			Deployment: sdk.Deployment{
+				PipedID:       "piped-id",
+				ApplicationID: "app-id",
+			},
+		},
+		Client: sdk.NewClient(nil, "kubernetes_multicluster", "app-id", "stage-id", logpersistertest.NewTestLogPersister(t), testRegistry),
+		Logger: zaptest.NewLogger(t),
+	}
+
+	status = plugin.executeK8sMultiRollbackStage(ctx, rollbackInput, []*sdk.DeployTarget[kubeconfig.KubernetesDeployTargetConfig]{
+		{Name: "default", Config: *dtConfig},
+	})
+	assert.Equal(t, sdk.StageStatusSuccess, status)
+
+	// Verify baseline deployment was deleted.
+	_, err = dynamicClient.Resource(schema.GroupVersionResource{Group: "apps", Version: "v1", Resource: "deployments"}).Namespace("default").Get(ctx, "simple-baseline", metav1.GetOptions{})
+	assert.True(t, apierrors.IsNotFound(err), "baseline deployment should be deleted after rollback")
+}