Added graceful shutdown for ctrl-C

cmd/arcaflow/main.go

@@ -6,6 +6,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
 
@@ -164,13 +165,21 @@ Options:
 
 func runWorkflow(flow engine.WorkflowEngine, dirContext map[string][]byte, workflowFile string, logger log.Logger, inputData []byte) int {
 	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	ctrlC := make(chan os.Signal, 1)
+	signal.Notify(ctrlC, os.Interrupt)
+
+	go handleOSInterrupt(ctrlC, cancel, logger)
+	defer func() {
+		close(ctrlC) // Ensure that the goroutine exits
+		cancel()
+	}()
 
 	workflow, err := flow.Parse(dirContext, workflowFile)
 	if err != nil {
 		logger.Errorf("Invalid workflow (%v)", err)
 		return ExitCodeInvalidData
 	}
+
 	outputID, outputData, outputError, err := workflow.Run(ctx, inputData)
 	if err != nil {
 		logger.Errorf("Workflow execution failed (%v)", err)
@@ -193,6 +202,27 @@ func runWorkflow(flow engine.WorkflowEngine, dirContext map[string][]byte, workf
 	return ExitCodeOK
 }
 
+func handleOSInterrupt(ctrlC chan os.Signal, cancel context.CancelFunc, logger log.Logger) {
+	for i := 1; true; i++ {
+		sysSignal, ok := <-ctrlC
+		if !ok || sysSignal != os.Interrupt {
+			return
+		}
+		switch {
+		case i <= 1:
+			logger.Infof("Requesting graceful shutdown.")
+			// Request graceful shutdown
+			cancel()
+		case i == 2:
+			logger.Warningf("Hit CTRL-C again to forcefully exit workflow without cleanup. You may need to manually delete pods or containers.")
+		default:
+			logger.Warningf("Force exiting. You may need to manually delete pods or containers.")
+			// Second request. Exit now.
+			os.Exit(1)
+		}
+	}
+}
+
 func loadContext(dir string) (map[string][]byte, error) {
 	absDir, err := filepath.Abs(dir)
 	if err != nil {

engine.go

@@ -62,7 +62,12 @@ type workflowEngine struct {
 	config       *config.Config
 }
 
-func (w workflowEngine) RunWorkflow(ctx context.Context, input []byte, workflowContext map[string][]byte, workflowFileName string) (outputID string, outputData any, outputError bool, err error) {
+func (w workflowEngine) RunWorkflow(
+	ctx context.Context,
+	input []byte,
+	workflowContext map[string][]byte,
+	workflowFileName string,
+) (outputID string, outputData any, outputError bool, err error) {
 	wf, err := w.Parse(workflowContext, workflowFileName)
 	if err != nil {
 		return "", nil, true, err
@@ -126,7 +131,10 @@ type engineWorkflow struct {
 	workflow workflow.ExecutableWorkflow
 }
 
-func (e engineWorkflow) Run(ctx context.Context, input []byte) (outputID string, outputData any, outputIsError bool, err error) {
+func (e engineWorkflow) Run(
+	ctx context.Context,
+	input []byte,
+) (outputID string, outputData any, outputIsError bool, err error) {
 	decodedInput, err := yaml.New().Parse(input)
 	if err != nil {
 		return "", nil, true, fmt.Errorf("failed to YAML decode input (%w)", err)