diff --git a/cmdutil/app.go b/cmdutil/app.go index 1086517..c7b70a4 100644 --- a/cmdutil/app.go +++ b/cmdutil/app.go @@ -2,9 +2,6 @@ package cmdutil import ( "context" - "os" - "os/signal" - "syscall" "github.com/sirupsen/logrus" "github.com/spf13/cobra" @@ -34,28 +31,6 @@ 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 @@ -79,7 +54,9 @@ func NewRootCommand(app interface{}) *cobra.Command { runnerWithContext, ok := app.(ApplicationRunnerWithContext) if ok { - run = wrapContext{runner: runnerWithContext}.Run + run = func(cmd *cobra.Command, args []string) { + runnerWithContext.Run(SignalRootContext(), cmd, args) + } } cmd := &cobra.Command{ diff --git a/cmdutil/context.go b/cmdutil/context.go new file mode 100644 index 0000000..3381408 --- /dev/null +++ b/cmdutil/context.go @@ -0,0 +1,33 @@ +package cmdutil + +import ( + "context" + "os" + "os/signal" + "syscall" + + "github.com/sirupsen/logrus" +) + +// SignalRootContext returns a new empty context, that gets canneld on SIGINT +// or SIGTEM. +func SignalRootContext() context.Context { + return SignalContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) +} + +// SignalContext returns a copy of the parent context that gets cancelled if +// the application gets any of the given signals. +func SignalContext(ctx context.Context, signals ...os.Signal) context.Context { + ctx, cancel := context.WithCancel(ctx) + + c := make(chan os.Signal, 1) + signal.Notify(c, signals...) + + go func() { + sig := <-c + logrus.Warnf("received signal '%v'; cleaning up", sig) + cancel() + }() + + return ctx +}