From 4fc906eb48b87223e2800b55546ec80423bce187 Mon Sep 17 00:00:00 2001 From: Philip Laine Date: Mon, 3 Jun 2024 21:42:58 +0200 Subject: [PATCH] fix: cancel Cobra parent context on interrupt (#2567) ## Description Right now different functions handle interrupts in different ways. Some will register their own signal listeners while others will use the global signal handler that exits the program. This means that contexts are never cancelled giving the process time to shut down and clean up. This change adds a signal handler to the Cobra parent context and removes the other signal handlers. Now all commands will use the same signal handler. ## Related Issue Fixes #2505 Relates to #2520 ## Checklist before merging - [x] Test, docs, adr added or updated as needed - [x] [Contributor Guide Steps](https://github.com/defenseunicorns/zarf/blob/main/.github/CONTRIBUTING.md#developer-workflow) followed Co-authored-by: Austin Abro <37223396+AustinAbro321@users.noreply.github.com> Signed-off-by: Austin Abro --- main.go | 8 +++++- src/cmd/common/setup.go | 2 -- src/cmd/common/utils.go | 19 ------------- src/cmd/connect.go | 11 ++------ src/cmd/internal.go | 8 +++--- src/cmd/root.go | 15 ++++------- src/cmd/tools/crane.go | 4 --- src/internal/agent/start.go | 54 ++++++++++++++++++++----------------- 8 files changed, 48 insertions(+), 73 deletions(-) diff --git a/main.go b/main.go index b9f58f6e44..0449ccd119 100644 --- a/main.go +++ b/main.go @@ -5,7 +5,10 @@ package main import ( + "context" "embed" + "os/signal" + "syscall" "github.com/defenseunicorns/zarf/src/cmd" "github.com/defenseunicorns/zarf/src/config" @@ -19,7 +22,10 @@ var cosignPublicKey string var zarfSchema embed.FS func main() { + ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer cancel() + config.CosignPublicKey = cosignPublicKey lint.ZarfSchema = zarfSchema - cmd.Execute() + cmd.Execute(ctx) } diff --git a/src/cmd/common/setup.go b/src/cmd/common/setup.go index 72fd0d6722..6fab6e134b 100644 --- a/src/cmd/common/setup.go +++ b/src/cmd/common/setup.go @@ -21,8 +21,6 @@ var LogLevelCLI string // SetupCLI sets up the CLI logging, interrupt functions, and more func SetupCLI() { - ExitOnInterrupt() - match := map[string]message.LogLevel{ "warn": message.WarnLevel, "info": message.InfoLevel, diff --git a/src/cmd/common/utils.go b/src/cmd/common/utils.go index 4fe6fa5043..26f2ce6707 100644 --- a/src/cmd/common/utils.go +++ b/src/cmd/common/utils.go @@ -6,18 +6,11 @@ package common import ( "context" - "os" - "os/signal" - "syscall" - "github.com/defenseunicorns/zarf/src/config/lang" "github.com/defenseunicorns/zarf/src/pkg/cluster" "github.com/defenseunicorns/zarf/src/pkg/message" ) -// SuppressGlobalInterrupt suppresses the global error on an interrupt -var SuppressGlobalInterrupt = false - // SetBaseDirectory sets the base directory. This is a directory with a zarf.yaml. func SetBaseDirectory(args []string) string { if len(args) > 0 { @@ -26,18 +19,6 @@ func SetBaseDirectory(args []string) string { return "." } -// ExitOnInterrupt catches an interrupt and exits with fatal error -func ExitOnInterrupt() { - c := make(chan os.Signal, 1) - signal.Notify(c, os.Interrupt, syscall.SIGTERM) - go func() { - <-c - if !SuppressGlobalInterrupt { - message.Fatal(lang.ErrInterrupt, lang.ErrInterrupt.Error()) - } - }() -} - // NewClusterOrDie creates a new Cluster instance and waits for the cluster to be ready or throws a fatal error. func NewClusterOrDie(ctx context.Context) *cluster.Cluster { timeoutCtx, cancel := context.WithTimeout(ctx, cluster.DefaultTimeout) diff --git a/src/cmd/connect.go b/src/cmd/connect.go index 19422df967..ba6eed3d2d 100644 --- a/src/cmd/connect.go +++ b/src/cmd/connect.go @@ -7,8 +7,6 @@ package cmd import ( "fmt" "os" - "os/signal" - "syscall" "github.com/defenseunicorns/zarf/src/cmd/common" "github.com/defenseunicorns/zarf/src/config/lang" @@ -72,17 +70,12 @@ var ( } } - // Keep this open until an interrupt signal is received. - interruptChan := make(chan os.Signal, 1) - signal.Notify(interruptChan, os.Interrupt, syscall.SIGTERM) - common.SuppressGlobalInterrupt = true - // Wait for the interrupt signal or an error. select { + case <-ctx.Done(): + spinner.Successf(lang.CmdConnectTunnelClosed, url) case err = <-tunnel.ErrChan(): spinner.Fatalf(err, lang.CmdConnectErrService, err.Error()) - case <-interruptChan: - spinner.Successf(lang.CmdConnectTunnelClosed, url) } os.Exit(0) }, diff --git a/src/cmd/internal.go b/src/cmd/internal.go index 97263d0b08..59c1981b0a 100644 --- a/src/cmd/internal.go +++ b/src/cmd/internal.go @@ -38,8 +38,8 @@ var agentCmd = &cobra.Command{ Use: "agent", Short: lang.CmdInternalAgentShort, Long: lang.CmdInternalAgentLong, - Run: func(_ *cobra.Command, _ []string) { - agent.StartWebhook() + RunE: func(cmd *cobra.Command, _ []string) error { + return agent.StartWebhook(cmd.Context()) }, } @@ -47,8 +47,8 @@ var httpProxyCmd = &cobra.Command{ Use: "http-proxy", Short: lang.CmdInternalProxyShort, Long: lang.CmdInternalProxyLong, - Run: func(_ *cobra.Command, _ []string) { - agent.StartHTTPProxy() + RunE: func(cmd *cobra.Command, _ []string) error { + return agent.StartHTTPProxy(cmd.Context()) }, } diff --git a/src/cmd/root.go b/src/cmd/root.go index 52e10044c1..447a9d76d2 100644 --- a/src/cmd/root.go +++ b/src/cmd/root.go @@ -10,6 +10,9 @@ import ( "os" "strings" + "github.com/pterm/pterm" + "github.com/spf13/cobra" + "github.com/defenseunicorns/zarf/src/cmd/common" "github.com/defenseunicorns/zarf/src/cmd/tools" "github.com/defenseunicorns/zarf/src/config" @@ -17,8 +20,6 @@ import ( "github.com/defenseunicorns/zarf/src/pkg/layout" "github.com/defenseunicorns/zarf/src/pkg/message" "github.com/defenseunicorns/zarf/src/types" - "github.com/pterm/pterm" - "github.com/spf13/cobra" ) var ( @@ -33,16 +34,10 @@ var rootCmd = &cobra.Command{ if common.CheckVendorOnlyFromPath(cmd) { return } - // Don't log the help command if cmd.Parent() == nil { config.SkipLogFile = true } - - // Set the global context for the root command and all child commands - ctx := context.Background() - cmd.SetContext(ctx) - common.SetupCLI() }, Short: lang.RootCmdShort, @@ -67,8 +62,8 @@ var rootCmd = &cobra.Command{ } // Execute is the entrypoint for the CLI. -func Execute() { - cmd, err := rootCmd.ExecuteC() +func Execute(ctx context.Context) { + cmd, err := rootCmd.ExecuteContextC(ctx) if err == nil { return } diff --git a/src/cmd/tools/crane.go b/src/cmd/tools/crane.go index 3b750964d3..a89041bb26 100644 --- a/src/cmd/tools/crane.go +++ b/src/cmd/tools/crane.go @@ -10,7 +10,6 @@ import ( "strings" "github.com/AlecAivazis/survey/v2" - "github.com/defenseunicorns/zarf/src/cmd/common" "github.com/defenseunicorns/zarf/src/config" "github.com/defenseunicorns/zarf/src/config/lang" "github.com/defenseunicorns/zarf/src/internal/packager/images" @@ -39,9 +38,6 @@ func init() { Aliases: []string{"r", "crane"}, Short: lang.CmdToolsRegistryShort, PersistentPreRun: func(cmd *cobra.Command, _ []string) { - - common.ExitOnInterrupt() - // The crane options loading here comes from the rootCmd of crane craneOptions = append(craneOptions, crane.WithContext(cmd.Context())) // TODO(jonjohnsonjr): crane.Verbose option? diff --git a/src/internal/agent/start.go b/src/internal/agent/start.go index 35b57e5fe2..12ee3bcc8b 100644 --- a/src/internal/agent/start.go +++ b/src/internal/agent/start.go @@ -6,10 +6,11 @@ package agent import ( "context" + "errors" "net/http" - "os" - "os/signal" - "syscall" + "time" + + "golang.org/x/sync/errgroup" "github.com/defenseunicorns/zarf/src/config/lang" agentHttp "github.com/defenseunicorns/zarf/src/internal/agent/http" @@ -27,35 +28,40 @@ const ( ) // StartWebhook launches the Zarf agent mutating webhook in the cluster. -func StartWebhook() { +func StartWebhook(ctx context.Context) error { message.Debug("agent.StartWebhook()") - - startServer(agentHttp.NewAdmissionServer(httpPort)) + return startServer(ctx, agentHttp.NewAdmissionServer(httpPort)) } // StartHTTPProxy launches the zarf agent proxy in the cluster. -func StartHTTPProxy() { +func StartHTTPProxy(ctx context.Context) error { message.Debug("agent.StartHttpProxy()") - - startServer(agentHttp.NewProxyServer(httpPort)) + return startServer(ctx, agentHttp.NewProxyServer(httpPort)) } -func startServer(server *http.Server) { - go func() { - if err := server.ListenAndServeTLS(tlsCert, tlsKey); err != nil && err != http.ErrServerClosed { - message.Fatal(err, lang.AgentErrStart) +func startServer(ctx context.Context, srv *http.Server) error { + g, gCtx := errgroup.WithContext(ctx) + g.Go(func() error { + err := srv.ListenAndServeTLS(tlsCert, tlsKey) + if err != nil && !errors.Is(err, http.ErrServerClosed) { + return err } - }() - + return nil + }) + g.Go(func() error { + <-gCtx.Done() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + err := srv.Shutdown(ctx) + if err != nil { + return err + } + return nil + }) message.Infof(lang.AgentInfoPort, httpPort) - - // listen shutdown signal - signalChan := make(chan os.Signal, 1) - signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM) - <-signalChan - - message.Infof(lang.AgentInfoShutdown) - if err := server.Shutdown(context.Background()); err != nil { - message.Fatal(err, lang.AgentErrShutdown) + err := g.Wait() + if err != nil { + return err } + return nil }