diff --git a/cmd/argoexec/commands/wait.go b/cmd/argoexec/commands/wait.go index 170f7e1521ee..9d05aec88174 100644 --- a/cmd/argoexec/commands/wait.go +++ b/cmd/argoexec/commands/wait.go @@ -2,8 +2,6 @@ package commands import ( "context" - "os/signal" - "syscall" "time" "github.com/argoproj/pkg/stats" @@ -16,7 +14,7 @@ func NewWaitCommand() *cobra.Command { Use: "wait", Short: "wait for main container to finish and save artifacts", Run: func(cmd *cobra.Command, args []string) { - ctx := context.Background() + ctx := cmd.Context() err := waitContainer(ctx) if err != nil { log.Fatalf("%+v", err) @@ -32,20 +30,14 @@ func waitContainer(ctx context.Context) error { defer stats.LogStats() stats.StartStatsTicker(5 * time.Minute) - // use a function to constrain the scope of ctx - func() { - // this allows us to gracefully shutdown, capturing artifacts - ctx, cancel := signal.NotifyContext(ctx, syscall.SIGTERM) - defer cancel() - - // Wait for main container to complete - err := wfExecutor.Wait(ctx) - if err != nil { - wfExecutor.AddError(err) - } - }() + // Wait for main container to complete + err := wfExecutor.Wait(ctx) + if err != nil { + wfExecutor.AddError(err) + } + ctx = context.Background() // don't allow cancellation to impact capture of results, parameters,or artifacts // Capture output script result - err := wfExecutor.CaptureScriptResult(ctx) + err = wfExecutor.CaptureScriptResult(ctx) if err != nil { wfExecutor.AddError(err) } diff --git a/cmd/argoexec/main.go b/cmd/argoexec/main.go index 1bd6c581bc19..8ae2e0d512c5 100644 --- a/cmd/argoexec/main.go +++ b/cmd/argoexec/main.go @@ -1,7 +1,10 @@ package main import ( + "context" "os" + "os/signal" + "syscall" "github.com/argoproj/argo-workflows/v3/util/errors" @@ -13,7 +16,9 @@ import ( ) func main() { - err := commands.NewRootCommand().Execute() + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM) + defer stop() + err := commands.NewRootCommand().ExecuteContext(ctx) if err != nil { if exitError, ok := err.(errors.Exited); ok { if exitError.ExitCode() >= 0 {