package exiter import ( "container/list" "context" "os" "syscall" "gitlab.com/CRThaze/sigar" ) const ( exitSignalGroupName = "exitSignalGroup" abortSignalGroupName = "abortSignalGroup" ) // ExitTaskFn (`func()`) is a closure represening something that should be // executed right before exiting. type ExitTaskFn sigar.TaskFn // WriteFn is a format function that will write output somewhere. // This could be a custom function or simply a logger function pointer. type WriteFn func(format string, a ...interface{}) // ExiterConfig contains options for the creation of an Exiter. type ExiterConfig struct { ParentCtx context.Context ExitSignals []os.Signal SuppressAllExitMsg bool ShowOKExitMsg bool } // Exiter is a tool for cleanly and consistently exiting the program. // Instead of creating the struct directly, use the NewExiter() or // NewExiterWithConfig() factories. type Exiter struct { ctx context.Context cancel context.CancelFunc sigar *sigar.SignalCatcher preCancelTasks *list.List postCancelTasks *list.List usagePrintFunc ExitTaskFn exitInfoMsgWriteFunc WriteFn exitErrMsgWriteFunc WriteFn suppressAllExitMsg bool showOKExitMsg bool } // DefaultExiterConfig is the ExiterConfig used when calling the // NewExiter() factory. var DefaultExiterConfig = ExiterConfig{ ParentCtx: context.Background(), ExitSignals: []os.Signal{ syscall.SIGINT, syscall.SIGTERM, }, } // NewExiterWithConfig is an Exiter factory that allows you to // customize the configuration. // Be sure to call (Exiter).ListenForSignals() once created to handle // all exits properly. func NewExiterWithConfig(config ExiterConfig) (e *Exiter) { e = &Exiter{ suppressAllExitMsg: config.SuppressAllExitMsg, showOKExitMsg: config.ShowOKExitMsg, } e.ctx, e.cancel = context.WithCancel(config.ParentCtx) e.sigar = sigar.NewSignalCatcher( sigar.SignalGroups{ exitSignalGroupName: config.ExitSignals, abortSignalGroupName: []os.Signal{syscall.SIGABRT}, }, e.ctx, ) e.sigar.RegisterTask(exitSignalGroupName, func() { e.ExitWithOK() }) e.sigar.RegisterTask(abortSignalGroupName, func() { e.Exit(ExitAbort) }) return } // NewExiter is an Exiter factory that uses the DefaultExiterConfig. // Be sure to call (Exiter).ListenForSignals() once created to handle // all exits properly. func NewExiter() *Exiter { return NewExiterWithConfig(DefaultExiterConfig) } // SpawnExiter is a convenience function that calls both NewExiter() // and (Exiter).ListenForSignals(). func SpawnExiter() (e *Exiter) { e = NewExiter() e.ListenForSignals() return } // SpawnExiterWithConfig is a convenience function that calls both // NewExiterWithConfig() and (Exiter).ListenForSignals(). func SpawnExiterWithConfig(config ExiterConfig) (e *Exiter) { e = NewExiterWithConfig(config) e.ListenForSignals() return } // ListenForSignals spawns a goroutine to catch signals so that the // Exiter can gracefully shutdown the program. func (e *Exiter) ListenForSignals() { e.sigar.Listen() } // RegisterPreTask adds a task to the queue that will be executed // prior to unblocking goroutines blocked at (Exiter).Wait(). func (e *Exiter) RegisterPreTask(fn ExitTaskFn) { e.preCancelTasks.PushBack(fn) } // RegisterPostTask adds a task to the queue that will be executed // after unblocking goroutines blocked at (Exiter).Wait(). func (e *Exiter) RegisterPostTask(fn ExitTaskFn) { e.postCancelTasks.PushBack(fn) } // SetUsageFunc tells the Exiter what function to call to print the // program's help/usage information. func (e *Exiter) SetUsageFunc(usage ExitTaskFn) { e.usagePrintFunc = usage } // ExitOnPanic is a convenience function that is meant to be defered // as a means to gracefully exit in the event of a panic. func (e *Exiter) ExitOnPanic() { if r := recover(); r != nil { e.exitErrMsgWriteFunc("Paniced: %v", r) e.ExitWithSoftwareErr() } } func (e *Exiter) exit(code ExitCode) { defer func() { if r := recover(); r != nil { e.exitErrMsgWriteFunc( "Paniced attempting to cleanup tasks (with exit code: %d): %v", code, r, ) os.Exit(int(ExitSoftware)) } }() for x := e.preCancelTasks.Front(); x != nil; x = x.Next() { x.Value.(ExitTaskFn)() } e.cancel() for x := e.postCancelTasks.Front(); x != nil; x = x.Next() { x.Value.(ExitTaskFn)() } os.Exit(int(code)) } func (e *Exiter) writeExitMsg(code ExitCode, fn WriteFn, err ...error) { if !e.suppressAllExitMsg && fn != nil { if code == ExitOK && !e.showOKExitMsg { return } if len(err) != 0 && err[0] != nil { fn("%s: %v", ExitSoftware.Desc(), err[0]) } else { fn("%s", ExitSoftware.Desc()) } } } // ExitWithOK calls (Exiter).Exit() with ExitOK (0). // // It is recommended to use this for all normal exit conditions. func (e *Exiter) ExitWithOK() { e.Exit(ExitOK) } // ExitWithUsageOK prints the help/usage and calls (Exiter).Exit() with // ExitOK (0). If no UsageFunc was set, it will not print anything. // // It is recommended to use this when responding to help arguments. func (e *Exiter) ExitWithUsageOK() { e.RegisterPostTask(e.usagePrintFunc) e.Exit(ExitOK) } // ExitWithError calls (Exiter).Exit() with ExitError (1). // // It is recommended to use this when no other exit code makes sense. func (e *Exiter) ExitWithError(err ...error) { e.Exit(ExitError, err...) } // ExitWithUsageErr prints the help/usage and calls (Exiter).Exit() with // ExitUsage (64). If no UsageFunc was set, it will not print anything. // // It is recommended to use this when bad arguments are passed. func (e *Exiter) ExitWithUsageErr(err ...error) { e.Exit(ExitUsage, err...) } // ExitWithDataErr calls (Exiter).Exit() with ExitDataErr (65). // // It is recommended to use this when data provided through any means // is invalid but NOT when an argument fails basic validation // (though schema errors in complex arguments could be appropriate UNLESS // they are configuartion settings, in which case you should use // (Exiter).ExitWithConfigErr()). func (e *Exiter) ExitWithDataErr(err ...error) { e.Exit(ExitDataErr, err...) } // ExitWithNoInputErr calls (Exiter).Exit() with ExitNoInput (66). // // It is recommended to use this only when an input channel (such as // a pipe, os.STDIN, or virtual device) cannot be opened or read when // needed. // Permissions errors should use (Exiter).ExitWithNoPermErr() instead. func (e *Exiter) ExitWithNoInputErr(err ...error) { e.Exit(ExitNoInput, err...) } // ExitWithNoUserErr calls (Exiter).Exit() with ExitNoUser (67). // // It is recommended to use this only when attempting to address data to // a local or network user, and no suitable match can be found. Mostly // useful for multi-user machines or intranet communication not using // a protocol that has its own suitable statuses for this. func (e *Exiter) ExitWithNoUserErr(err ...error) { e.Exit(ExitNoUser, err...) } // ExitWithNoHostErr calls (Exiter).Exit() with ExitNoHost (68). // // It is recommended to use this when a provided hostname cannot be resolved. func (e *Exiter) ExitWithNoHostErr(err ...error) { e.Exit(ExitNoHost, err...) } // ExitWithUnavailableErr calls (Exiter).Exit() with ExitUnavailable (69) (nice). // // It is recommended to use this for client programs which cannot communicate // with another process or remote service. func (e *Exiter) ExitWithUnavailableErr(err ...error) { e.Exit(ExitUnavailable, err...) } // ExitWithSoftwareErr calls (Exiter).Exit() with ExitSoftware (70). // // It is recommended to use this as a more meaningful alternative to // ExitWithError() or for more graceful exits in the event of a Panic. // For the later case consider deferring (Exiter).ExitOnPanic() which // will call this function in that event. func (e *Exiter) ExitWithSoftwareErr(err ...error) { e.Exit(ExitSoftware, err...) } // ExitWithOSErr calls (Exiter).Exit() with ExitOSErr (71). // // It is recommended to use this only when unable to continue due to // encountering a generalized system error. func (e *Exiter) ExitWithOSErr(err ...error) { e.Exit(ExitOSErr, err...) } // ExitWithOSFile calls (Exiter).Exit() with ExitOSFile (72). // // It is recommended to use this only when unable to find a file or directory // not owned by this or any othe program but is part of the general // system configuration. Such as a standard FHS directory, or a device file // in /dev. func (e *Exiter) ExitWithOSFile(err ...error) { e.Exit(ExitOSFile, err...) } // ExitWithCantCreateErr calls (Exiter).Exit() with ExitCantCreate (73). // // It is recommended to use this only when the creation of a user output // file is impossible for any reason and the program is unable to continue. func (e *Exiter) ExitWithCantCreateErr(err ...error) { e.Exit(ExitCantCreate, err...) } // ExitWithIOErr calls (Exiter).Exit() with ExitIOErr (74). // // It is recommended to use this only when unable to interact with I/O // hardware, failure to open or write to an output stream, or when // encountering an interuption from an already open input stream // (otherwise use (Exiter).ExitWithNoInputErr()). // Permissions errors should use (Exiter).ExitWithNoPermErr() instead. func (e *Exiter) ExitWithIOErr(err ...error) { e.Exit(ExitIOErr, err...) } // ExitWithTempFailErr calls (Exiter).Exit() with ExitTempFail (75). // // It is recommended to use this when the program chooses (by design) to // terminate instead of retrying some operation that is expected to have // failed merely intermittently. func (e *Exiter) ExitWithTempFailErr(err ...error) { e.Exit(ExitTempFail, err...) } // ExitWithProtocolErr calls (Exiter).Exit() with ExitProtocol (76). // // It is recommended to use this only when some communication with a // process or network location fails due to a protocol violation from // the other communication participant. And example would be receiving // a response that is not inline with the protocol (like a SYN or ACK // in response to the same instead of a SYN/ACK when performing a TCP // handshake). func (e *Exiter) ExitWithProtocolErr(err ...error) { e.Exit(ExitTempFail, err...) } // ExitWithNoPermErr calls (Exiter).Exit() with ExitNoPerm (77). // // It is recommended to use this whenever user or credential Permissions // do not allow you to continue. func (e *Exiter) ExitWithNoPermErr(err ...error) { e.Exit(ExitNoPerm, err...) } // ExitWithConfigErr calls (Exiter).Exit() with ExitConfig (78). // // It is recommended to use this whenever a schema validation is encountered // in a config file or other serialiezed configuarion source. func (e *Exiter) ExitWithConfigErr(err ...error) { e.Exit(ExitConfig, err...) } // ExitWithConfigErr exits with the given ExitCode, and prints the // first error provided (if any). func (e *Exiter) Exit(code ExitCode, err ...error) { switch code { case ExitOK: e.RegisterPostTask(func() { e.writeExitMsg(code, e.exitInfoMsgWriteFunc) }) e.exit(code) case ExitUsage: e.RegisterPostTask(func() { e.writeExitMsg(code, e.exitErrMsgWriteFunc, err...) }) e.RegisterPostTask(e.usagePrintFunc) e.exit(code) default: e.RegisterPostTask(func() { e.writeExitMsg(code, e.exitErrMsgWriteFunc, err...) }) e.exit(code) } } // Wait blocks till (Exiter).Exit() is called. // It is recommended to call this function in a select block case in // all goroutines to ensure they are properly terminated. func (e *Exiter) Wait() { <-e.ctx.Done() } // GetCancelableContext returns the underlying cancelable context // used by the Exiter. func (e *Exiter) GetCancelableContext() context.Context { return e.ctx } // GetSignalCatcher returns the underlying *sigar.SignalCatcher // used by the Exiter. func (e *Exiter) GetSignalCatcher() *sigar.SignalCatcher { return e.sigar }