feat(osctl): handle ^C by aborting context (#693)

This provides a bit better handling for the handing grpc
requests (or just slow requests):

```
$ osctl-linux-amd64 --talosconfig talosconfig version
Client:
	Tag:         ad410fb-dirty
	SHA:         ad410fb-dirty
	Built:
	Go version:  go1.12.5
	OS/Arch:     linux/amd64

^CSignal received, aborting, press Ctrl+C once again to abort immediately...
error getting version: rpc error: code = Canceled desc = context canceled
```

For now we catch `SIGINT` & `SIGTERM`. Second signal kills process
immediately as signal handler is removed.

Signed-off-by: Andrey Smirnov <smirnov.andrey@gmail.com>

cmd/osctl/cmd/df.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 	"math"
 	"os"
@@ -30,7 +29,7 @@ var dfCmd = &cobra.Command{
 		}
 
 		setupClient(func(c *client.Client) {
-			dfRender(c.DF(context.TODO()))
+			dfRender(c.DF(globalCtx))
 		})
 	},
 }

cmd/osctl/cmd/dmesg.go

@@ -5,7 +5,6 @@
 package cmd
 
 import (
-	"context"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -25,7 +24,7 @@ var dmesgCmd = &cobra.Command{
 		}
 
 		setupClient(func(c *client.Client) {
-			msg, err := c.Dmesg(context.TODO())
+			msg, err := c.Dmesg(globalCtx)
 			if err != nil {
 				helpers.Fatalf("error getting dmesg: %s", err)
 			}

cmd/osctl/cmd/kubeconfig.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -26,7 +25,7 @@ var kubeconfigCmd = &cobra.Command{
 		}
 
 		setupClient(func(c *client.Client) {
-			kubeconfig, err := c.Kubeconfig(context.TODO())
+			kubeconfig, err := c.Kubeconfig(globalCtx)
 			if err != nil {
 				helpers.Fatalf("error fetching kubeconfig: %s", err)
 			}

cmd/osctl/cmd/logs.go

@@ -5,7 +5,6 @@
 package cmd
 
 import (
-	"context"
 	"io"
 	"os"
 
@@ -35,7 +34,7 @@ var logsCmd = &cobra.Command{
 				namespace = constants.SystemContainerdNamespace
 			}
 
-			stream, err := c.Logs(context.TODO(), namespace, args[0])
+			stream, err := c.Logs(globalCtx, namespace, args[0])
 			if err != nil {
 				helpers.Fatalf("error fetching logs: %s", err)
 			}

cmd/osctl/cmd/ps.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"text/tabwriter"
@@ -37,7 +36,7 @@ var psCmd = &cobra.Command{
 			} else {
 				namespace = constants.SystemContainerdNamespace
 			}
-			reply, err := c.Processes(context.TODO(), namespace)
+			reply, err := c.Processes(globalCtx, namespace)
 
 			if err != nil {
 				helpers.Fatalf("error getting process list: %s", err)

cmd/osctl/cmd/reboot.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -26,7 +25,7 @@ var rebootCmd = &cobra.Command{
 		}
 
 		setupClient(func(c *client.Client) {
-			if err := c.Reboot(context.TODO()); err != nil {
+			if err := c.Reboot(globalCtx); err != nil {
 				helpers.Fatalf("error executing reboot: %s", err)
 			}
 		})

cmd/osctl/cmd/reset.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -26,7 +25,7 @@ var resetCmd = &cobra.Command{
 		}
 
 		setupClient(func(c *client.Client) {
-			if err := c.Reset(context.TODO()); err != nil {
+			if err := c.Reset(globalCtx); err != nil {
 				helpers.Fatalf("error executing reset: %s", err)
 			}
 		})

cmd/osctl/cmd/restart.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"os"
 
 	criconstants "github.com/containerd/cri/pkg/constants"
@@ -40,7 +39,7 @@ var restartCmd = &cobra.Command{
 			} else {
 				namespace = constants.SystemContainerdNamespace
 			}
-			if err := c.Restart(context.TODO(), namespace, args[0], timeout); err != nil {
+			if err := c.Restart(globalCtx, namespace, args[0], timeout); err != nil {
 				helpers.Fatalf("error restarting process: %s", err)
 			}
 		})

cmd/osctl/cmd/root.go

@@ -5,9 +5,13 @@
 package cmd
 
 import (
+	"context"
+	"fmt"
 	"os"
+	"os/signal"
 	"os/user"
 	"path"
+	"syscall"
 
 	"github.com/spf13/cobra"
 	"github.com/talos-systems/talos/cmd/osctl/pkg/client"
@@ -38,9 +42,33 @@ var rootCmd = &cobra.Command{
 	Long:  ``,
 }
 
+// Global context to be used in the commands.
+//
+// Cobra doesn't have a way to pass it around, so we have to use global variable.
+// Context is initialized in Execute, and initial value is failsafe default.
+var globalCtx = context.Background()
+
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
+	var globalCtxCancel context.CancelFunc
+	globalCtx, globalCtxCancel = context.WithCancel(context.Background())
+	defer globalCtxCancel()
+
+	// listen for ^C and SIGTERM and abort context
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+
+	go func() {
+		select {
+		case <-sigCh:
+			signal.Stop(sigCh)
+			fmt.Fprintln(os.Stderr, "Signal received, aborting, press Ctrl+C once again to abort immediately...")
+			globalCtxCancel()
+		case <-globalCtx.Done():
+		}
+	}()
+
 	var (
 		defaultTalosConfig string
 		ok                 bool

cmd/osctl/cmd/routes.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"text/tabwriter"
@@ -29,7 +28,7 @@ var routesCmd = &cobra.Command{
 		}
 
 		setupClient(func(c *client.Client) {
-			reply, err := c.Routes(context.TODO())
+			reply, err := c.Routes(globalCtx)
 			if err != nil {
 				helpers.Fatalf("error getting routes: %s", err)
 			}

cmd/osctl/cmd/service.go

@@ -5,7 +5,6 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"text/tabwriter"
@@ -42,7 +41,7 @@ var serviceCmd = &cobra.Command{
 }
 
 func serviceList(c *client.Client) {
-	reply, err := c.ServiceList(context.TODO())
+	reply, err := c.ServiceList(globalCtx)
 	if err != nil {
 		helpers.Fatalf("error listing services: %s", err)
 	}
@@ -59,7 +58,7 @@ func serviceList(c *client.Client) {
 }
 
 func serviceInfo(c *client.Client, id string) {
-	s, err := c.ServiceInfo(context.TODO(), id)
+	s, err := c.ServiceInfo(globalCtx, id)
 	if err != nil {
 		helpers.Fatalf("error listing services: %s", err)
 	}

cmd/osctl/cmd/shutdown.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -26,7 +25,7 @@ var shutdownCmd = &cobra.Command{
 		}
 
 		setupClient(func(c *client.Client) {
-			if err := c.Shutdown(context.TODO()); err != nil {
+			if err := c.Shutdown(globalCtx); err != nil {
 				helpers.Fatalf("error executing shutdown: %s", err)
 			}
 		})

cmd/osctl/cmd/stats.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"text/tabwriter"
@@ -37,7 +36,7 @@ var statsCmd = &cobra.Command{
 			} else {
 				namespace = constants.SystemContainerdNamespace
 			}
-			reply, err := c.Stats(context.TODO(), namespace)
+			reply, err := c.Stats(globalCtx, namespace)
 			if err != nil {
 				helpers.Fatalf("error getting stats: %s", err)
 			}

cmd/osctl/cmd/top.go

@@ -41,7 +41,7 @@ var topCmd = &cobra.Command{
 
 			if oneTime {
 				var output string
-				output, err = topOutput(context.TODO(), c)
+				output, err = topOutput(globalCtx, c)
 				if err != nil {
 					log.Fatal(err)
 				}
@@ -56,7 +56,7 @@ var topCmd = &cobra.Command{
 			}
 			defer ui.Close()
 
-			topUI(context.TODO(), c)
+			topUI(globalCtx, c)
 		})
 	},
 }

cmd/osctl/cmd/upgrade.go

@@ -6,7 +6,6 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 
 	"github.com/spf13/cobra"
@@ -55,7 +54,7 @@ func remoteUpgrade() error {
 	setupClient(func(c *client.Client) {
 		// TODO: See if we can validate version and prevent
 		// starting upgrades to an unknown version
-		ack, err = c.Upgrade(context.TODO(), assetURL)
+		ack, err = c.Upgrade(globalCtx, assetURL)
 	})
 
 	if err == nil {

cmd/osctl/cmd/version.go

@@ -5,7 +5,6 @@
 package cmd
 
 import (
-	"context"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -37,7 +36,7 @@ var versionCmd = &cobra.Command{
 			}
 		}
 		setupClient(func(c *client.Client) {
-			version, err := c.Version(context.TODO())
+			version, err := c.Version(globalCtx)
 			if err != nil {
 				helpers.Fatalf("error getting version: %s", err)
 			}