Merge pull request #11 from NETWAYS/feature/timeout

cmd/check_example/main.go

@@ -3,7 +3,7 @@ package main
 import (
 	"github.com/NETWAYS/go-check"
 	log "github.com/sirupsen/logrus"
-	"os"
+	"time"
 )
 
 func main() {
@@ -12,18 +12,18 @@ func main() {
 	flags.Name = "check_test"
 	flags.Readme = `Test Plugin`
 	flags.Version = "1.0.0"
+	flags.Timeout = 10
 
-	value := flags.Set.IntP("value", "t", 10, "test value")
+	value := flags.Set.IntP("value", "", 10, "test value")
 	warning := flags.Set.IntP("warning", "w", 20, "warning threshold")
 	critical := flags.Set.IntP("critical", "c", 50, "critical threshold")
 
-	// value should be calculated
-
-	flags.Parse(os.Args[1:])
-	flags.SetupLogging()
+	flags.ParseArguments()
 
 	log.Info("Start logging")
 
+	time.Sleep(20 * time.Second)
+
 	if *value > *critical {
 		check.Exit(check.Critical, "value is %d", *value)
 	} else if *value > *warning {

flags.go

@@ -11,13 +11,16 @@ import (
 )
 
 type Flags struct {
-	Name         string
-	Readme       string
-	Version      string
-	Verbose      bool
-	Debug        bool
-	PrintVersion bool
-	Set          *flag.FlagSet
+	Name          string
+	Readme        string
+	Version       string
+	Timeout       int
+	Verbose       bool
+	Debug         bool
+	PrintVersion  bool
+	DefaultFlags  bool
+	DefaultHelper bool
+	Set           *flag.FlagSet
 }
 
 func NewFlags() *Flags {
@@ -38,16 +41,25 @@ func NewFlags() *Flags {
 		flagSet.PrintDefaults()
 	}
 
-	flagSet.BoolVarP(&flags.Debug, "debug", "d", false, "Enable debug mode")
-	flagSet.BoolVarP(&flags.Verbose, "verbose", "v", false, "Enable verbose mode")
-	flagSet.BoolVarP(&flags.PrintVersion, "version", "V", false, "Print version and exit")
-
 	flags.Set = flagSet
 
+	// set some defaults
+	flags.DefaultFlags = true
+	flags.Timeout = 30
+	flags.DefaultHelper = true
+
 	return flags
 }
 
-func (f *Flags) Parse(arguments []string) {
+func (f *Flags) ParseArguments() {
+	f.ParseArray(os.Args[1:])
+}
+
+func (f *Flags) ParseArray(arguments []string) {
+	if f.DefaultFlags {
+		f.addDefaultFlags()
+	}
+
 	err := f.Set.Parse(arguments)
 	if err != nil {
 		if err != flag.ErrHelp {
@@ -60,6 +72,20 @@ func (f *Flags) Parse(arguments []string) {
 		fmt.Println(f.Name, "version", f.Version)
 		BaseExit(3)
 	}
+
+	if f.DefaultHelper {
+		f.SetupLogging()
+		f.EnableTimeoutHandler()
+	}
+}
+
+func (f *Flags) addDefaultFlags() {
+	f.Set.IntVarP(&f.Timeout, "timeout", "t", f.Timeout, "Abort the check after n seconds")
+	f.Set.BoolVarP(&f.Debug, "debug", "d", false, "Enable debug mode")
+	f.Set.BoolVarP(&f.Verbose, "verbose", "v", false, "Enable verbose mode")
+	f.Set.BoolVarP(&f.PrintVersion, "version", "V", false, "Print version and exit")
+
+	f.DefaultFlags = false
 }
 
 func (f *Flags) SetupLogging() {

flags_test.go

@@ -12,8 +12,8 @@ func ExampleFlags() {
 
 	_ = flags.Set.StringP("hostname", "H", "localhost", "Hostname to check")
 
-	// flags.Parse(os.Args[1:])
-	flags.Parse([]string{"--help"})
+	// flags.ParseArray(os.Args[1:])
+	flags.ParseArray([]string{"--help"})
 	flags.SetupLogging()
 
 	log.Info("test")

timeout.go

@@ -0,0 +1,39 @@
+package check
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+)
+
+var timeoutEnabled bool
+
+// Start the timeout and signal handler in a goroutine
+func (f *Flags) EnableTimeoutHandler() {
+	go HandleTimeout(f.Timeout)
+}
+
+// Helper for a goroutine, to wait for signals and timeout, and exit with a proper code
+func HandleTimeout(timeout int) {
+	if timeoutEnabled {
+		// signal handling has already been set up
+		return
+	}
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
+
+	if timeout < 1 {
+		Exit(Unknown, "Invalid timeout: %d", timeout)
+	}
+
+	timedOut := time.After(time.Duration(timeout) * time.Second)
+	timeoutEnabled = true
+
+	select {
+	case s := <-signals:
+		Exit(Unknown, "Received signal: %s", s)
+	case <-timedOut:
+		Exit(Unknown, "Timeout reached")
+	}
+}