add a rety option to the step command

cmd/step/main.go

@@ -71,7 +71,8 @@ func main() {
 	flags.StringSliceVarP(&step.UploadFile, "upload", "u", []string{}, "Upload file as a kubernetes secret")
 	flags.StringVar(&step.WaitFile, "wait-on", "", "The path to a file to indicate this step can be run")
 	flags.StringSliceVarP(&step.Commands, "command", "c", []string{}, "Command to execute")
-
+	flags.IntVar(&step.RetryAttempts, "retry-attempts", 0, "Number of times to retry the commands")
+	flags.DurationVar(&step.RetryBackoff, "retry-backoff", 0, "Duration to wait between retry attempts")
 	if err := cmd.Execute(); err != nil {
 		fmt.Fprintf(os.Stderr, "[Error] %s\n", err)
 
@@ -128,49 +129,78 @@ func Run(ctx context.Context, step Step) error {
 	}
 
 	for i, command := range step.Commands {
-		//nolint:gosec
-		cmd := exec.CommandContext(ctx, step.Shell, "-c", command)
-		cmd.Env = os.Environ()
+		attempt := 0
+		var lastErr error
+
+		for attempt <= step.RetryAttempts {
+			if attempt > 0 {
+				log.WithFields(log.Fields{
+					"attempt": attempt,
+					"command": i,
+					"backoff": step.RetryBackoff,
+				}).Info("retrying command")
+
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				case <-time.After(step.RetryBackoff):
+				}
+			}
 
-		logger := log.WithField("command", i)
+			//nolint:gosec
+			cmd := exec.CommandContext(ctx, step.Shell, "-c", command)
+			cmd.Env = os.Environ()
 
-		stdout, err := cmd.StdoutPipe()
-		if err != nil {
-			logger.WithError(err).Error("failed to acquire stdout pipe on command")
+			logger := log.WithFields(log.Fields{
+				"command": i,
+				"attempt": attempt,
+			})
 
-			return err
-		}
-		stderr, err := cmd.StderrPipe()
-		if err != nil {
-			logger.WithError(err).Error("failed to acquire stderr pipe on command")
+			stdout, err := cmd.StdoutPipe()
+			if err != nil {
+				logger.WithError(err).Error("failed to acquire stdout pipe on command")
+				return err
+			}
+			stderr, err := cmd.StderrPipe()
+			if err != nil {
+				logger.WithError(err).Error("failed to acquire stderr pipe on command")
+				return err
+			}
 
-			return err
-		}
+			//nolint:errcheck
+			go io.Copy(os.Stdout, stdout)
+			//nolint:errcheck
+			go io.Copy(os.Stdout, stderr)
 
-		//nolint:errcheck
-		go io.Copy(os.Stdout, stdout)
-		//nolint:errcheck
-		go io.Copy(os.Stdout, stderr)
+			if err := cmd.Start(); err != nil {
+				logger.WithError(err).Error("failed to execute the command")
+				lastErr = err
+				attempt++
+				continue
+			}
 
-		if err := cmd.Start(); err != nil {
-			logger.WithError(err).Error("failed to execute the command")
+			// @step: wait for the command to finish
+			if err := cmd.Wait(); err != nil {
+				logger.WithError(err).Error("command execution failed")
+				lastErr = err
+				attempt++
+				continue
+			}
 
-			return err
+			// Command succeeded, break the retry loop
+			lastErr = nil
+			break
 		}
 
-		// @step: wait for the command to finish
-		if err := cmd.Wait(); err != nil {
-			logger.WithError(err).Error("failed to execute command successfully")
-
+		// If we exhausted all retries and still have an error
+		if lastErr != nil {
 			if step.ErrorFile != "" {
 				if err := utils.TouchFile(step.ErrorFile); err != nil {
-					logger.WithError(err).WithField("file", step.ErrorFile).Error("failed to create error file")
-
+					log.WithError(err).WithField("file", step.ErrorFile).Error("failed to create error file")
 					return err
 				}
 			}
-
-			return err
+			return fmt.Errorf("command failed after %d attempts: %w", attempt, lastErr)
 		}
 	}
 	log.Info("successfully executed the step")

cmd/step/types.go

@@ -28,6 +28,10 @@ import (
 type Step struct {
 	// Commands is the commands and arguments to run
 	Commands []string
+	// RetryBackoff is the backoff time to retry ANY of the commands
+	RetryBackoff time.Duration
+	// RetryAttempts is the number of times to retry the commands before giving up
+	RetryAttempts int
 	// Comment adds a banner to the stage
 	Comment string
 	// ErrorFile is the path to a file which is created when the command failed
@@ -57,6 +61,12 @@ func (s Step) IsValid() error {
 	case s.Timeout < 0:
 		return errors.New("timeout must be greater than 0")
 
+	case s.RetryAttempts < 0:
+		return errors.New("retry attempts must be greater than or equal to 0")
+
+	case s.RetryBackoff < 0:
+		return errors.New("retry backoff must be greater than or equal to 0")
+
 	case len(s.UploadFile) > 0 && s.Namespace == "":
 		return errors.New("namespace must be specified when uploading files")