Fix script run rollback

pkg/app/piped/controller/scheduler.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"io"
 	"path/filepath"
+	"strings"
 	"time"
 
 	"go.opentelemetry.io/otel"
@@ -285,7 +286,7 @@ func (s *scheduler) Run(ctx context.Context) error {
 		case model.DeploymentStatus_DEPLOYMENT_FAILURE, model.DeploymentStatus_DEPLOYMENT_CANCELLED:
 			span.SetStatus(codes.Error, statusReason)
 		}
-		
+
 		span.End()
 	}()
 
@@ -497,19 +498,28 @@ func (s *scheduler) executeStage(sig executor.StopSignal, ps model.PipelineStage
 	// 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 == "" {
+		baseStageIDs := strings.Split(ps.Metadata["scriptRunBaseStageIDs"], ",")
+		if len(baseStageIDs) == 0 {
 			return
 		}
 
-		baseStageStatus, ok := s.stageStatuses[baseStageID]
-		if !ok {
-			return
-		}
+		targetStageIDs := make([]string, 0, len(baseStageIDs))
+		for _, baseStageID := range baseStageIDs {
+			baseStageStatus, ok := s.stageStatuses[baseStageID]
+			if !ok {
+				continue
+			}
 
-		if baseStageStatus == model.StageStatus_STAGE_NOT_STARTED_YET || baseStageStatus == model.StageStatus_STAGE_SKIPPED {
-			return
+			if baseStageStatus == model.StageStatus_STAGE_NOT_STARTED_YET || baseStageStatus == model.StageStatus_STAGE_SKIPPED {
+				continue
+			}
+
+			targetStageIDs = append(targetStageIDs, baseStageID)
 		}
+
+		fmt.Println("ffjlabo-dev: baseStageIDs", baseStageIDs, len(baseStageIDs))
+		fmt.Println("ffjlabo-dev: targetStageIDs", targetStageIDs, len(targetStageIDs))
+		ps.Metadata["scriptRunTargetStageIDs"] = strings.Join(targetStageIDs, ",")
 	}
 
 	// Update stage status to RUNNING if needed.

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

@@ -17,6 +17,7 @@ package kubernetes
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"os"
 	"os/exec"
 	"strings"
@@ -26,6 +27,7 @@ import (
 	"github.com/pipe-cd/pipecd/pkg/app/piped/executor"
 	"github.com/pipe-cd/pipecd/pkg/app/piped/executor/scriptrun"
 	provider "github.com/pipe-cd/pipecd/pkg/app/piped/platformprovider/kubernetes"
+	"github.com/pipe-cd/pipecd/pkg/config"
 	"github.com/pipe-cd/pipecd/pkg/model"
 )
 
@@ -184,52 +186,69 @@ func (e *rollbackExecutor) ensureRollback(ctx context.Context) model.StageStatus
 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")
+	targetScriptRunStageIDs := strings.Split(e.Stage.Metadata["scriptRunTargetStageIDs"], ",")
+
+	ds, err := e.TargetDSP.Get(ctx, e.LogPersister)
+	if err != nil {
+		e.LogPersister.Errorf("Failed to prepare target deploy source data (%v)", err)
 		return model.StageStatus_STAGE_FAILURE
 	}
 
-	if onRollback == "" {
-		e.LogPersister.Info("No commands to run")
-		return model.StageStatus_STAGE_SUCCESS
-	}
+	appDir := ds.AppDir
 
-	envStr, ok := e.Stage.Metadata["env"]
-	env := make(map[string]string, 0)
-	if ok {
-		_ = json.Unmarshal([]byte(envStr), &env)
-	}
+	for i, baseStageID := range targetScriptRunStageIDs {
+		e.LogPersister.Infof("Start executing rollback command for the %dth SCRIPT_RUN stage", i+1)
+
+		optionKey := fmt.Sprintf("scriptRun.%s.option", baseStageID)
 
-	for _, v := range strings.Split(onRollback, "\n") {
-		if v != "" {
-			e.LogPersister.Infof("   %s", v)
+		scriptRunOpt, ok := e.Stage.Metadata[optionKey]
+		if !ok {
+			e.LogPersister.Error("ScriptRun option metadata is missing")
+			return model.StageStatus_STAGE_FAILURE
 		}
-	}
 
-	ci := scriptrun.NewContextInfo(e.Deployment)
-	ciEnv, err := ci.BuildEnv()
-	if err != nil {
-		e.LogPersister.Errorf("failed to build srcipt run context info: %w", err)
-		return model.StageStatus_STAGE_FAILURE
-	}
+		var opts config.ScriptRunStageOptions
+		if err := json.Unmarshal([]byte(scriptRunOpt), &opts); err != nil {
+			e.LogPersister.Error("Failed to parse ScriptRun option")
+			return model.StageStatus_STAGE_FAILURE
+		}
 
-	envs := make([]string, 0, len(ciEnv)+len(env))
-	for key, value := range ciEnv {
-		envs = append(envs, key+"="+value)
-	}
+		if opts.OnRollback == "" {
+			e.LogPersister.Info("No commands to run")
+			continue
+		}
 
-	for key, value := range env {
-		envs = append(envs, key+"="+value)
-	}
+		for _, v := range strings.Split(opts.OnRollback, "\n") {
+			if v != "" {
+				e.LogPersister.Infof("   %s", v)
+			}
+		}
 
-	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
+		ci := scriptrun.NewContextInfo(e.Deployment)
+		ciEnv, err := ci.BuildEnv()
+		if err != nil {
+			e.LogPersister.Errorf("failed to build srcipt run context info: %w", err)
+			return model.StageStatus_STAGE_FAILURE
+		}
+
+		envs := make([]string, 0, len(ciEnv)+len(opts.Env))
+		for key, value := range ciEnv {
+			envs = append(envs, key+"="+value)
+		}
+
+		for key, value := range opts.Env {
+			envs = append(envs, key+"="+value)
+		}
+
+		cmd := exec.Command("/bin/sh", "-l", "-c", opts.OnRollback)
+		cmd.Dir = 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

@@ -17,6 +17,7 @@ package kubernetes
 import (
 	"encoding/json"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/pipe-cd/pipecd/pkg/app/piped/planner"
@@ -117,29 +118,41 @@ func buildProgressivePipeline(pp *config.DeploymentPipeline, autoRollback bool,
 		})
 
 		// Add a stage for rolling back script run stages.
-		for i, s := range pp.Stages {
-			if s.Name == model.StageScriptRun {
+		metadata := make(map[string]string, 0)
+		baseStageIDs := make([]string, 0, len(pp.Stages))
+		for _, s := range out {
+			// TODO: check the base stage env and command
+			// maybe scriptRunBaseStageIDs => stage-1,stage-2
+			// stage-1.env json Marshaled value
+			// stage-1.onRollback => string
+
+			// scriptRunBaseStageIDs => stage-1,stage-2
+			// scriptRun.stage-1.option => Marshaled value
+
+			if s.Name == model.StageScriptRun.String() {
 				// 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(),
-				})
+				opts := pp.Stages[s.Index].ScriptRunStageOptions
+				optionStr, _ := json.Marshal(opts)
+
+				metadata[fmt.Sprintf("scriptRun.%s.option", s.Id)] = string(optionStr)
+				baseStageIDs = append(baseStageIDs, s.Id)
 			}
 		}
+
+		metadata["scriptRunBaseStageIDs"] = strings.Join(baseStageIDs, ",")
+
+		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