Merge pull request #146 from lpabon/ctrl-c-signal

Catch CTRL-C and cleanup

Author: lpabon

cmd/root.go

@@ -16,6 +16,8 @@ limitations under the License.
 package cmd
 
 import (
+	"os"
+
 	"github.com/portworx/pxc/pkg/commander"
 	"github.com/portworx/pxc/pkg/config"
 	"github.com/portworx/pxc/pkg/kubernetes"
@@ -31,8 +33,9 @@ type rootFlags struct {
 
 // rootCmd represents the base command when called without any subcommands
 var (
-	rootCmd     *cobra.Command
-	rootOptions *rootFlags
+	rootCmd           *cobra.Command
+	rootOptions       *rootFlags
+	rootSignalHandler *util.SigIntManager
 
 	// This template allow pxc to override the way it prints out the help. This template
 	// allows pxc to not print out global options unless requested.
@@ -83,6 +86,7 @@ Please see https://docs.portworx.com/reference/ for more information.`,
 		PersistentPostRunE: rootPersistentPostRunE,
 		RunE:               rootCmdExec,
 	}
+
 })
 
 var _ = commander.RegisterCommandInit(func() {
@@ -125,6 +129,14 @@ func rootPersistentPreRunE(cmd *cobra.Command, args []string) error {
 	// Set version
 	logrus.Infof("pxc version: %s", PxVersion)
 
+	// Capture any CTRL-C
+	rootSignalHandler = util.NewSigIntManager(func() {
+		cleanup()
+		util.Printf("\n*** Interrupt\n")
+		os.Exit(1)
+	})
+	rootSignalHandler.Start()
+
 	// The following commands do not need to load configuration
 	switch cmd.Name() {
 	case "version":
@@ -141,7 +153,12 @@ func rootPersistentPreRunE(cmd *cobra.Command, args []string) error {
 
 func rootPersistentPostRunE(cmd *cobra.Command, args []string) error {
 	// Close the global tunnel if any
-	kubernetes.StopTunnel()
+	cleanup()
+	rootSignalHandler.Stop()
 
 	return nil
 }
+
+func cleanup() {
+	kubernetes.StopTunnel()
+}

pkg/kubernetes/portforward.go

@@ -19,8 +19,10 @@ import (
 	"fmt"
 	"os/exec"
 	"strings"
+	"sync"
 
 	"github.com/portworx/pxc/pkg/config"
+	"github.com/portworx/pxc/pkg/util"
 
 	"github.com/sirupsen/logrus"
 )
@@ -34,9 +36,12 @@ type PortForwarder interface {
 
 // KubectlPortForwarder object
 type KubectlPortForwarder struct {
-	kubeconfig string
-	endpoint   string
-	cmd        *exec.Cmd
+	kubeconfig  string
+	endpoint    string
+	cmd         *exec.Cmd
+	signhandler *util.SigIntManager
+	lock        sync.Mutex
+	running     bool
 }
 
 var (
@@ -80,6 +85,13 @@ func newKubectlPortForwarder(kubeconfig string) *KubectlPortForwarder {
 
 // Start creates the portforward using kubectl
 func (p *KubectlPortForwarder) Start() error {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	if p.running {
+		return fmt.Errorf("Tunnel already running")
+	}
+
 	args := config.KM().KubectlFlagsToCliArgs()
 	currentCluster := config.CM().GetCurrentCluster()
 	logrus.Debugf("port-forward: CurrentCluster: %v", *currentCluster)
@@ -98,6 +110,11 @@ func (p *KubectlPortForwarder) Start() error {
 		return fmt.Errorf("Unable to setup kubectl: %v", err)
 	}
 
+	p.signhandler = util.NewSigIntManager(func() {
+		p.Stop()
+	})
+	p.signhandler.Start()
+
 	// Start the port forward process
 	err = cmd.Start()
 	if err != nil {
@@ -129,15 +146,31 @@ func (p *KubectlPortForwarder) Start() error {
 	logrus.Debugf("Read %d bytes", n)
 	logrus.Debugf("Output: %s", sbuf)
 
+	p.running = true
 	return nil
 }
 
 // Stop ends the session
 func (p *KubectlPortForwarder) Stop() error {
-	logrus.Debug("Port forwarding stopped")
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	if !p.running {
+		return nil
+	}
+
 	if p.cmd != nil {
-		return p.cmd.Process.Kill()
+		logrus.Debug("Port forwarding stopped")
+		err := p.cmd.Process.Kill()
+		p.cmd = nil
+		return err
+	}
+
+	if p.signhandler != nil {
+		p.signhandler.Stop()
+		p.signhandler = nil
 	}
+	p.running = false
 	return nil
 }
 
@@ -149,7 +182,6 @@ func (p *KubectlPortForwarder) Endpoint() string {
 func (p *KubectlPortForwarder) getEndpointFromKubectlOutput(sbuf string) (string, error) {
 	index := strings.Index(sbuf, "127.0.0.1:")
 	if index >= 0 {
-		//return strings.Split(sbuf[index:], " ")[0], nil
 		e := strings.Split(sbuf[index:], " ")[0]
 		e = "localhost:" + strings.Split(e, ":")[1]
 		return e, nil

pkg/util/sigint.go

@@ -0,0 +1,87 @@
+/*
+Copyright © 2020 Portworx
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package util
+
+import (
+	"fmt"
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+
+	"github.com/sirupsen/logrus"
+)
+
+type SigIntManager struct {
+	handler func()
+	done    chan bool
+	wg      sync.WaitGroup
+	lock    sync.Mutex
+	running bool
+}
+
+func NewSigIntManager(handler func()) *SigIntManager {
+	return &SigIntManager{
+		handler: handler,
+		done:    make(chan bool, 1),
+	}
+}
+
+func (s *SigIntManager) Start() error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	if s.running {
+		return fmt.Errorf("already running a signal handler")
+	}
+
+	var startwg sync.WaitGroup
+	signalch := make(chan os.Signal, 1)
+	signal.Notify(signalch, os.Interrupt, os.Kill, syscall.SIGINT, syscall.SIGTERM)
+	startwg.Add(1)
+	s.wg.Add(1)
+	go func() {
+		startwg.Done()
+		select {
+		case <-signalch:
+			logrus.Info("Ctrl-C captured...")
+			s.handler()
+		case <-s.done:
+			logrus.Debug("Closing Ctrl-C capturing function")
+		}
+		s.wg.Done()
+	}()
+	startwg.Wait()
+
+	s.running = true
+	return nil
+}
+
+func (s *SigIntManager) Stop() error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	if !s.running {
+		return nil
+	}
+
+	s.done <- true
+	s.wg.Wait()
+	s.running = false
+
+	return nil
+}