Merge pull request #8 from rebuy-de/context-support

add context support

cmdutil/app.go

@@ -1,10 +1,14 @@
 package cmdutil
 
 import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	graylog "gopkg.in/gemnasium/logrus-graylog-hook.v2"
-
-	log "github.com/sirupsen/logrus"
 )
 
 // ApplicationRunner is an optional interface for NewRootCommand.
@@ -14,6 +18,14 @@ type ApplicationRunner interface {
 	Run(cmd *cobra.Command, args []string)
 }
 
+// ApplicationRunnerWithContext is an optional interface for NewRootCommand.
+type ApplicationRunnerWithContext interface {
+	// Run contains the actual application code. It is equivalent to
+	// the Run command of Cobra plus adding a context. The context gets
+	// cancelled, if the application receives a SIGINT or SIGTERM.
+	Run(ctx context.Context, cmd *cobra.Command, args []string)
+}
+
 // ApplicationBinder is an optional interface for NewRootCommand.
 type ApplicationBinder interface {
 
@@ -22,6 +34,28 @@ type ApplicationBinder interface {
 	Bind(cmd *cobra.Command)
 }
 
+// wrapContext uses a ApplicationRunnerWithContext and implements a
+// ApplicationRunner. It passes a context, that gets cancels on SIGINT or
+// SIGTERM, to the ApplicationRunnerWithContext.Run function.
+type wrapContext struct {
+	runner ApplicationRunnerWithContext
+}
+
+func (w wrapContext) Run(cmd *cobra.Command, args []string) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		sig := <-signals
+		logrus.Warnf("received signal '%v'; cleaning up", sig)
+		cancel()
+	}()
+
+	w.runner.Run(ctx, cmd, args)
+}
+
 // NewRootCommand creates a Cobra command, which reflects our current best
 // practices. It adds a verbose flag, sets up logrus and registers a Graylog
 // hook. Also it registers the NewVersionCommand and prints the version on
@@ -35,20 +69,28 @@ func NewRootCommand(app interface{}) *cobra.Command {
 
 	var run func(cmd *cobra.Command, args []string)
 
+	// Note: since ApplicationRunnerWithContext and ApplicationRunner require
+	// the same function Run with different parameters, they are mutually
+	// exclusive.
 	runner, ok := app.(ApplicationRunner)
 	if ok {
 		run = runner.Run
 	}
 
+	runnerWithContext, ok := app.(ApplicationRunnerWithContext)
+	if ok {
+		run = wrapContext{runner: runnerWithContext}.Run
+	}
+
 	cmd := &cobra.Command{
 		Use: BuildName,
 		Run: run,
 
 		PersistentPreRun: func(cmd *cobra.Command, args []string) {
-			log.SetLevel(log.InfoLevel)
+			logrus.SetLevel(logrus.InfoLevel)
 
 			if verbose {
-				log.SetLevel(log.DebugLevel)
+				logrus.SetLevel(logrus.DebugLevel)
 			}
 
 			if gelfAddress != "" {
@@ -58,19 +100,19 @@ func NewRootCommand(app interface{}) *cobra.Command {
 					"commit-sha": BuildHash,
 				}
 				hook := graylog.NewGraylogHook(gelfAddress, labels)
-				hook.Level = log.DebugLevel
-				log.AddHook(hook)
+				hook.Level = logrus.DebugLevel
+				logrus.AddHook(hook)
 			}
 
-			log.WithFields(log.Fields{
+			logrus.WithFields(logrus.Fields{
 				"Version": BuildVersion,
 				"Date":    BuildDate,
 				"Commit":  BuildHash,
 			}).Infof("%s started", BuildName)
 		},
 
 		PersistentPostRun: func(cmd *cobra.Command, args []string) {
-			log.Infof("%s stopped", BuildName)
+			logrus.Infof("%s stopped", BuildName)
 		},
 	}
 

executil/run.go

@@ -0,0 +1,44 @@
+package executil
+
+import (
+	"context"
+	"os/exec"
+	"strings"
+	"syscall"
+
+	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+)
+
+// Run starts the specified command and waits for it to complete.
+//
+// The difference to Run from exec.CommandContext is that it sends an interrupt
+// instead of a kill, which gives the process time for a graceful shutdown.
+func Run(ctx context.Context, cmd *exec.Cmd) error {
+	commandline := strings.Join(cmd.Args, " ")
+	logrus.WithFields(logrus.Fields{
+		"Args": cmd.Args,
+		"Dir":  cmd.Dir,
+	}).Debugf("running command `%s`", commandline)
+
+	err := cmd.Start()
+	if err != nil {
+		return errors.WithStack(err)
+	}
+
+	done := make(chan struct{}, 1)
+	defer close(done)
+
+	go func() {
+		select {
+		case <-ctx.Done():
+			logrus.Debugf("sending interrupt signal to `%s`", commandline)
+			cmd.Process.Signal(syscall.SIGINT)
+		case <-done:
+			// This mean wait() already exited and we can stop to wait for the
+			// cancelation.
+		}
+	}()
+
+	return errors.Wrapf(cmd.Wait(), "failed to run `%s`", commandline)
+}