exiter/exiter.go

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
}