fix: viewcontroller gorouting leak in status and get subcommand (#1584)

Signed-off-by: Hui Kang <[email protected]>
argoproj · Nov 29, 2021 · 026e517 · 026e517
1 parent ff3471a
commit 026e517
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 4 deletions.
diff --git a/pkg/kubectl-argo-rollouts/cmd/get/get_rollout.go b/pkg/kubectl-argo-rollouts/cmd/get/get_rollout.go
@@ -11,6 +11,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/argoproj/argo-rollouts/pkg/apiclient/rollout"
+	"github.com/argoproj/argo-rollouts/pkg/kubectl-argo-rollouts/cmd/signals"
 	"github.com/argoproj/argo-rollouts/pkg/kubectl-argo-rollouts/info"
 	"github.com/argoproj/argo-rollouts/pkg/kubectl-argo-rollouts/options"
 	"github.com/argoproj/argo-rollouts/pkg/kubectl-argo-rollouts/viewcontroller"
@@ -44,7 +45,9 @@ func NewCmdGetRollout(o *options.ArgoRolloutsOptions) *cobra.Command {
 			}
 			name := args[0]
 			controller := viewcontroller.NewRolloutViewController(o.Namespace(), name, getOptions.KubeClientset(), getOptions.RolloutsClientset())
-			ctx := context.Background()
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			signals.SetupSignalHandler(cancel)
 			controller.Start(ctx)
 
 			ri, err := controller.GetRolloutInfo()

diff --git a/pkg/kubectl-argo-rollouts/cmd/signals/signal.go b/pkg/kubectl-argo-rollouts/cmd/signals/signal.go
@@ -0,0 +1,19 @@
+package signals
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+func SetupSignalHandler(cancel context.CancelFunc) {
+	c := make(chan os.Signal, 2)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-c
+		cancel()
+		<-c
+		os.Exit(1) // second signal. Exit directly.
+	}()
+}
diff --git a/pkg/kubectl-argo-rollouts/cmd/status/status.go b/pkg/kubectl-argo-rollouts/cmd/status/status.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/argoproj/argo-rollouts/pkg/apiclient/rollout"
+	"github.com/argoproj/argo-rollouts/pkg/kubectl-argo-rollouts/cmd/signals"
 	"github.com/argoproj/argo-rollouts/pkg/kubectl-argo-rollouts/options"
 	"github.com/argoproj/argo-rollouts/pkg/kubectl-argo-rollouts/viewcontroller"
 	"github.com/spf13/cobra"
@@ -50,6 +51,7 @@ func NewCmdStatus(o *options.ArgoRolloutsOptions) *cobra.Command {
 			controller := viewcontroller.NewRolloutViewController(o.Namespace(), name, statusOptions.KubeClientset(), statusOptions.RolloutsClientset())
 			ctx, cancel := context.WithCancel(context.Background())
 			defer cancel()
+			signals.SetupSignalHandler(cancel)
 			controller.Start(ctx)
 
 			ri, err := controller.GetRolloutInfo()
@@ -61,7 +63,6 @@ func NewCmdStatus(o *options.ArgoRolloutsOptions) *cobra.Command {
 				fmt.Fprintln(o.Out, ri.Status)
 			} else {
 				rolloutUpdates := make(chan *rollout.RolloutInfo)
-				defer close(rolloutUpdates)
 				controller.RegisterCallback(func(roInfo *rollout.RolloutInfo) {
 					rolloutUpdates <- roInfo
 				})
@@ -72,7 +73,7 @@ func NewCmdStatus(o *options.ArgoRolloutsOptions) *cobra.Command {
 				if err != nil {
 					return err
 				}
-
+				close(rolloutUpdates)
 				if finalRi.Status == "Degraded" {
 					return fmt.Errorf("The rollout is in a degraded state with message: %s", finalRi.Message)
 				} else if finalRi.Status != "Healthy" {
@@ -88,7 +89,7 @@ func NewCmdStatus(o *options.ArgoRolloutsOptions) *cobra.Command {
 	return cmd
 }
 
-func (o *StatusOptions) WatchStatus(stopCh <-chan struct{}, rolloutUpdates chan *rollout.RolloutInfo) string {
+func (o *StatusOptions) WatchStatus(stopCh <-chan struct{}, rolloutUpdates <-chan *rollout.RolloutInfo) string {
 	timeout := make(chan bool)
 	var roInfo *rollout.RolloutInfo
 	var preventFlicker time.Time