From e2e7dd1895e7a85b089e54c10714c48d2a18ef85 Mon Sep 17 00:00:00 2001 From: Philip Laine Date: Fri, 31 May 2024 12:11:37 +0200 Subject: [PATCH] Cancel Cobra parent context on interrupt --- 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 | 4 +++ src/cmd/tools/crane.go | 4 --- src/internal/agent/start.go | 52 +++++++++++++++++++++---------------- 7 files changed, 39 insertions(+), 61 deletions(-) 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 83fb28e61d..526c3c598b 100644 --- a/src/cmd/root.go +++ b/src/cmd/root.go @@ -8,7 +8,9 @@ import ( "context" "fmt" "os" + "os/signal" "strings" + "syscall" "github.com/defenseunicorns/zarf/src/cmd/common" "github.com/defenseunicorns/zarf/src/cmd/tools" @@ -40,6 +42,8 @@ var rootCmd = &cobra.Command{ // Set the global context for the root command and all child commands ctx := context.Background() + ctx, cancel := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM) + defer cancel() cmd.SetContext(ctx) common.SetupCLI() 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..245c2780b5 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 { +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) { message.Fatal(err, lang.AgentErrStart) } - }() - + 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 { + message.Fatal(err, lang.AgentErrShutdown) + } + 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 }