Implement rollback for script run as pre defined stage

docs/rfcs/0011-script-run-stage.md

@@ -90,34 +90,23 @@ spec:
             - "curl -X POST -H 'Content-type: application/json' --data '{"text":"failed to deploy: rollback"}' $SLACK_WEBHOOK_URL"
 ```
 
-**SCRIPT_SYNC stage also rollbacks** when the deployment status is `DeploymentStatus_DEPLOYMENT_CANCELLED` or `DeploymentStatus_DEPLOYMENT_FAILURE` even though other rollback stage is also executed.
+**SCRIPT_RUN stage also rollbacks**. Execute the command to rollback SCRIPT_RUN to the point where the deployment was canceled or failed.
+When there are multiple SCRIPT_RUN stages to be rolled back, they are executed in the same order as SCRIPT_RUN on the pipeline.
 
-For example, here is a deploy pipeline combined with other k8s stages.
-The result status of the pipeline is FAIL or CANCELED, piped rollbacks the stages `K8S_CANARY_ROLLOUT`, `K8S_PRIMARY_ROLLOUT`, and `SCRIPT_RUN`.
+For example, consider when deployment proceeds in the following order from 1 to 7.
+
+1. K8S_CANARY_ROLLOUT
+2. WAIT
+3. SCRIPT_RUN
+4. K8S_PRIMARY_ROLLOUT
+5. SCRIPT_RUN
+6. K8S_CANARY_CLEAN
+7. SCRIPT_RUN
+
+Then
+- If 4 is canceled or fails while running, only SCRIPT_RUN of 3 will be rolled back.
+- If 6 is canceled or fails while running, only SCRIPT_RUNs 3 and 5 will be rolled back.
 
-```yaml
-apiVersion: pipecd.dev/v1beta1
-kind: KubernetesApp
-spec:
-  pipeline:
-    stages:
-      - name: K8S_CANARY_ROLLOUT
-        with:
-          replicas: 10%
-      - name: WAIT_APPROVAL
-        with:
-          timeout: 30m
-      - name: K8S_PRIMARY_ROLLOUT
-      - name: K8S_CANARY_CLEAN
-      - name: SCRIPT_RUN
-        with:
-          env:
-            SLACK_WEBHOOK_URL: ""
-          runs:
-            - "curl -X POST -H 'Content-type: application/json' --data '{"text":"successfully deployed!!"}' $SLACK_WEBHOOK_URL"
-          onRollback:
-            - "curl -X POST -H 'Content-type: application/json' --data '{"text":"failed to deploy: rollback"}' $SLACK_WEBHOOK_URL"
-```
 
 ## prepare environment for execution
 

pkg/app/piped/controller/scheduler.go

@@ -369,34 +369,37 @@ func (s *scheduler) Run(ctx context.Context) error {
 	// we start rollback stage if the auto-rollback option is true.
 	if deploymentStatus == model.DeploymentStatus_DEPLOYMENT_CANCELLED ||
 		deploymentStatus == model.DeploymentStatus_DEPLOYMENT_FAILURE {
-		if stage, ok := s.deployment.FindRollbackStage(); ok {
+
+		if rollbackStages, ok := s.deployment.FindRollbackStages(); ok {
 			// Update to change deployment status to ROLLING_BACK.
 			if err := s.reportDeploymentStatusChanged(ctx, model.DeploymentStatus_DEPLOYMENT_ROLLING_BACK, statusReason); err != nil {
 				return err
 			}
 
-			// Start running rollback stage.
-			var (
-				sig, handler = executor.NewStopSignal()
-				doneCh       = make(chan struct{})
-			)
-			go func() {
-				rbs := *stage
-				rbs.Requires = []string{lastStage.Id}
-				s.executeStage(sig, rbs, func(in executor.Input) (executor.Executor, bool) {
-					return s.executorRegistry.RollbackExecutor(s.deployment.Kind, in)
-				})
-				close(doneCh)
-			}()
-
-			select {
-			case <-ctx.Done():
-				handler.Terminate()
-				<-doneCh
-				return nil
-
-			case <-doneCh:
-				break
+			for _, stage := range rollbackStages {
+				// Start running rollback stage.
+				var (
+					sig, handler = executor.NewStopSignal()
+					doneCh       = make(chan struct{})
+				)
+				go func() {
+					rbs := *stage
+					rbs.Requires = []string{lastStage.Id}
+					s.executeStage(sig, rbs, func(in executor.Input) (executor.Executor, bool) {
+						return s.executorRegistry.RollbackExecutor(s.deployment.Kind, in)
+					})
+					close(doneCh)
+				}()
+
+				select {
+				case <-ctx.Done():
+					handler.Terminate()
+					<-doneCh
+					return nil
+
+				case <-doneCh:
+					break
+				}
 			}
 		}
 	}
@@ -433,6 +436,24 @@ func (s *scheduler) executeStage(sig executor.StopSignal, ps model.PipelineStage
 		lp.Complete(time.Minute)
 	}()
 
+	// Check whether to execute the script rollback stage or not.
+	// If the base stage is executed, the script rollback stage will be executed.
+	if ps.Name == model.StageScriptRunRollback.String() {
+		baseStageID := ps.Metadata["baseStageID"]
+		if baseStageID == "" {
+			return
+		}
+
+		baseStageStatus, ok := s.stageStatuses[baseStageID]
+		if !ok {
+			return
+		}
+
+		if baseStageStatus == model.StageStatus_STAGE_NOT_STARTED_YET || baseStageStatus == model.StageStatus_STAGE_SKIPPED {
+			return
+		}
+	}
+
 	// Update stage status to RUNNING if needed.
 	if model.CanUpdateStageStatus(ps.Status, model.StageStatus_STAGE_RUNNING) {
 		if err := s.reportStageStatus(ctx, ps.Id, model.StageStatus_STAGE_RUNNING, ps.Requires); err != nil {

pkg/app/piped/executor/kubernetes/rollback.go

@@ -16,6 +16,9 @@ package kubernetes
 
 import (
 	"context"
+	"encoding/json"
+	"os"
+	"os/exec"
 	"strings"
 
 	"go.uber.org/zap"
@@ -27,6 +30,8 @@ import (
 
 type rollbackExecutor struct {
 	executor.Input
+
+	appDir string
 }
 
 func (e *rollbackExecutor) Execute(sig executor.StopSignal) model.StageStatus {
@@ -39,7 +44,8 @@ func (e *rollbackExecutor) Execute(sig executor.StopSignal) model.StageStatus {
 	switch model.Stage(e.Stage.Name) {
 	case model.StageRollback:
 		status = e.ensureRollback(ctx)
-
+	case model.StageScriptRunRollback:
+		status = e.ensureScriptRunRollback(ctx)
 	default:
 		e.LogPersister.Errorf("Unsupported stage %s for kubernetes application", e.Stage.Name)
 		return model.StageStatus_STAGE_FAILURE
@@ -74,6 +80,8 @@ func (e *rollbackExecutor) ensureRollback(ctx context.Context) model.StageStatus
 		}
 	}
 
+	e.appDir = ds.AppDir
+
 	loader := provider.NewLoader(e.Deployment.ApplicationName, ds.AppDir, ds.RepoDir, e.Deployment.GitPath.ConfigFilename, appCfg.Input, e.GitClient, e.Logger)
 	e.Logger.Info("start executing kubernetes stage",
 		zap.String("stage-name", e.Stage.Name),
@@ -171,3 +179,45 @@ func (e *rollbackExecutor) ensureRollback(ctx context.Context) model.StageStatus
 	}
 	return model.StageStatus_STAGE_SUCCESS
 }
+
+func (e *rollbackExecutor) ensureScriptRunRollback(ctx context.Context) model.StageStatus {
+	e.LogPersister.Info("Runnnig commands for rollback...")
+
+	onRollback, ok := e.Stage.Metadata["onRollback"]
+	if !ok {
+		e.LogPersister.Error("onRollback metadata is missing")
+		return model.StageStatus_STAGE_FAILURE
+	}
+
+	if onRollback == "" {
+		e.LogPersister.Info("No commands to run")
+		return model.StageStatus_STAGE_SUCCESS
+	}
+
+	envStr, ok := e.Stage.Metadata["env"]
+	env := make(map[string]string, 0)
+	if ok {
+		_ = json.Unmarshal([]byte(envStr), &env)
+	}
+
+	for _, v := range strings.Split(onRollback, "\n") {
+		if v != "" {
+			e.LogPersister.Infof("   %s", v)
+		}
+	}
+
+	envs := make([]string, 0, len(env))
+	for key, value := range env {
+		envs = append(envs, key+"="+value)
+	}
+
+	cmd := exec.Command("/bin/sh", "-l", "-c", onRollback)
+	cmd.Dir = e.appDir
+	cmd.Env = append(os.Environ(), envs...)
+	cmd.Stdout = e.LogPersister
+	cmd.Stderr = e.LogPersister
+	if err := cmd.Run(); err != nil {
+		return model.StageStatus_STAGE_FAILURE
+	}
+	return model.StageStatus_STAGE_SUCCESS
+}

pkg/app/piped/planner/kubernetes/pipeline.go

@@ -15,6 +15,7 @@
 package kubernetes
 
 import (
+	"encoding/json"
 	"fmt"
 	"time"
 
@@ -114,6 +115,31 @@ func buildProgressivePipeline(pp *config.DeploymentPipeline, autoRollback bool,
 			CreatedAt:  now.Unix(),
 			UpdatedAt:  now.Unix(),
 		})
+
+		// Add a stage for rolling back script run stages.
+		for i, s := range pp.Stages {
+			if s.Name == model.StageScriptRun {
+				// Use metadata as a way to pass parameters to the stage.
+				envStr, _ := json.Marshal(s.ScriptRunStageOptions.Env)
+				metadata := map[string]string{
+					"baseStageID": out[i].Id,
+					"onRollback":  s.ScriptRunStageOptions.OnRollback,
+					"env":         string(envStr),
+				}
+				ss, _ := planner.GetPredefinedStage(planner.PredefinedStageScriptRunRollback)
+				out = append(out, &model.PipelineStage{
+					Id:         ss.ID,
+					Name:       ss.Name.String(),
+					Desc:       ss.Desc,
+					Predefined: true,
+					Visible:    false,
+					Status:     model.StageStatus_STAGE_NOT_STARTED_YET,
+					Metadata:   metadata,
+					CreatedAt:  now.Unix(),
+					UpdatedAt:  now.Unix(),
+				})
+			}
+		}
 	}
 
 	return out

pkg/app/piped/planner/predefined_stages.go

@@ -27,6 +27,7 @@ const (
 	PredefinedStageECSSync            = "ECSSync"
 	PredefinedStageRollback           = "Rollback"
 	PredefinedStageCustomSyncRollback = "CustomSyncRollback"
+	PredefinedStageScriptRunRollback  = "ScriptRunRollback"
 )
 
 var predefinedStages = map[string]config.PipelineStage{
@@ -65,6 +66,11 @@ var predefinedStages = map[string]config.PipelineStage{
 		Name: model.StageCustomSyncRollback,
 		Desc: "Rollback the custom stages",
 	},
+	PredefinedStageScriptRunRollback: {
+		ID:   PredefinedStageScriptRunRollback,
+		Name: model.StageScriptRunRollback,
+		Desc: "Rollback the script run stage",
+	},
 }
 
 // GetPredefinedStage finds and returns the predefined stage for the given id.

pkg/model/deployment.go

@@ -154,6 +154,16 @@ func (d *Deployment) FindRollbackStage() (*PipelineStage, bool) {
 	return nil, false
 }
 
+func (d *Deployment) FindRollbackStages() ([]*PipelineStage, bool) {
+	rollbackStages := make([]*PipelineStage, 0, len(d.Stages))
+	for i, stage := range d.Stages {
+		if d.Stages[i].Name == StageRollback.String() || d.Stages[i].Name == StageScriptRunRollback.String() {
+			rollbackStages = append(rollbackStages, stage)
+		}
+	}
+	return rollbackStages, len(rollbackStages) > 0
+}
+
 // DeploymentStatusesFromStrings converts a list of strings to list of DeploymentStatus.
 func DeploymentStatusesFromStrings(statuses []string) ([]DeploymentStatus, error) {
 	out := make([]DeploymentStatus, 0, len(statuses))

pkg/model/deployment_test.go

@@ -587,6 +587,7 @@ func TestFindRollbackStage(t *testing.T) {
 	}
 
 	for _, tt := range tests {
+		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
 			d := &Deployment{
 				Stages: tt.stages,
@@ -597,3 +598,46 @@ func TestFindRollbackStage(t *testing.T) {
 		})
 	}
 }
+
+func TestFindRollbackStags(t *testing.T) {
+	tests := []struct {
+		name           string
+		stages         []*PipelineStage
+		wantStages     []*PipelineStage
+		wantStageFound bool
+	}{
+		{
+			name: "found",
+			stages: []*PipelineStage{
+				{Name: StageK8sSync.String()},
+				{Name: StageRollback.String()},
+				{Name: StageScriptRunRollback.String()},
+			},
+			wantStages: []*PipelineStage{
+				{Name: StageRollback.String()},
+				{Name: StageScriptRunRollback.String()},
+			},
+			wantStageFound: true,
+		},
+		{
+			name: "not found",
+			stages: []*PipelineStage{
+				{Name: StageK8sSync.String()},
+			},
+			wantStages:     []*PipelineStage{},
+			wantStageFound: false,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			d := &Deployment{
+				Stages: tt.stages,
+			}
+			stages, found := d.FindRollbackStages()
+			assert.Equal(t, tt.wantStages, stages)
+			assert.Equal(t, tt.wantStageFound, found)
+		})
+	}
+}

pkg/model/stage.go

@@ -108,6 +108,10 @@ const (
 	// all changes made by the CUSTOM_SYNC stage will be reverted to
 	// bring back the pre-deploy stage.
 	StageCustomSyncRollback Stage = "CUSTOM_SYNC_ROLLBACK"
+	// StageScriptRunRollback represents a state where
+	// all changes made by the SCRIPT_RUN_ROLLBACK stage will be reverted to
+	// bring back the pre-deploy stage.
+	StageScriptRunRollback Stage = "SCRIPT_RUN_ROLLBACK"
 )
 
 func (s Stage) String() string {