setup signalhandler and exiter

exiter/exiter.go

@@ -0,0 +1,97 @@
+package exiter
+
+import (
+	"container/list"
+	"context"
+	"os"
+	"syscall"
+
+	"gitlab.com/CRThaze/sugar/sigar"
+)
+
+const exitSignalGroupName = "exitSignalGroup"
+
+type ExitTaskFn sigar.TaskFn
+
+type ExiterConfig struct {
+	ParentCtx   context.Context
+	ExitSignals []os.Signal
+}
+
+type Exiter struct {
+	ctx             context.Context
+	cancel          context.CancelFunc
+	sigar           *sigar.SignalCatcher
+	preCancelTasks  *list.List
+	postCancelTasks *list.List
+}
+
+var DefaultExiterConfig = ExiterConfig{
+	ParentCtx: context.Background(),
+	ExitSignals: []os.Signal{
+		syscall.SIGINT,
+		syscall.SIGTERM,
+		syscall.SIGABRT,
+	},
+}
+
+func NewExiterWithConfig(config ExiterConfig) (e *Exiter) {
+	e = &Exiter{}
+	e.ctx, e.cancel = context.WithCancel(config.ParentCtx)
+	e.sigar = sigar.NewSignalCatcher(
+		sigar.SignalGroups{
+			exitSignalGroupName: config.ExitSignals,
+		},
+		e.ctx,
+	)
+	e.sigar.RegisterTask(exitSignalGroupName, func() {
+		e.Exit(0)
+	})
+	return
+}
+
+func NewExiter() *Exiter { return NewExiterWithConfig(DefaultExiterConfig) }
+
+func SpawnExiter() (e *Exiter) {
+	e = NewExiter()
+	e.ListenForSignals()
+	return
+}
+
+func (e *Exiter) ListenForSignals() {
+	e.sigar.Listen()
+}
+
+func (e *Exiter) RegisterPreTask(fn ExitTaskFn) {
+	e.preCancelTasks.PushBack(fn)
+}
+
+func (e *Exiter) RegisterPostTask(fn ExitTaskFn) {
+	e.postCancelTasks.PushBack(fn)
+}
+
+func (e *Exiter) Exit(code ...int) {
+	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)()
+	}
+	if len(code) == 0 {
+		code = make([]int, 1, 1)
+	}
+	os.Exit(code[0])
+}
+
+func (e *Exiter) Wait() {
+	<-e.ctx.Done()
+}
+
+func (e *Exiter) GetCancelableContext() context.Context {
+	return e.ctx
+}
+
+func (e *Exiter) GetSignalCatcher() *sigar.SignalCatcher {
+	return e.sigar
+}

sigar/signalHandler.go

@@ -0,0 +1,87 @@
+package sigar
+
+import (
+	"container/list"
+	"context"
+	"errors"
+	"os"
+	"os/signal"
+)
+
+type taskMap map[string]*list.List
+
+type SignalGroups map[string][]os.Signal
+
+func (sg SignalGroups) initTaskMap(tm taskMap) {
+	for groupName := range sg {
+		tm[groupName] = list.New()
+	}
+}
+
+type TaskFn func()
+
+type SignalCatcher struct {
+	signals SignalGroups
+	tasks   map[string]*list.List
+	exitCtx context.Context
+}
+
+func NewSignalCatcher(sg SignalGroups, exitCtx context.Context) (sc *SignalCatcher) {
+	if sg == nil {
+		sg = SignalGroups{}
+	}
+	sc = &SignalCatcher{
+		signals: sg,
+		exitCtx: exitCtx,
+		tasks:   taskMap{},
+	}
+	sg.initTaskMap(sc.tasks)
+	if exitCtx == nil {
+		// Create a cancelable context and discard the cancel function, as we will never use it.
+		sc.exitCtx, _ = context.WithCancel(context.Background())
+	}
+	return
+}
+
+func SpawnSignalCatcher(sg SignalGroups, exitCtx context.Context) (sc *SignalCatcher) {
+	sc = NewSignalCatcher(sg, exitCtx)
+	sc.Listen()
+	return
+}
+
+func (sc *SignalCatcher) runTasks(group string) {
+	taskList := sc.tasks[group]
+	for e := taskList.Front(); e != nil; e = e.Next() {
+		e.Value.(TaskFn)()
+	}
+}
+
+func (sc *SignalCatcher) Listen() {
+	for groupName, signalSet := range sc.signals {
+		// Prevent loop closure by copying variables.
+		group := groupName
+		signals := signalSet
+		go func() {
+			sigChan := make(chan os.Signal, 1)
+			signal.Notify(sigChan, signals...)
+			select {
+			case <-sigChan:
+				sc.runTasks(group)
+			case <-sc.exitCtx.Done():
+			}
+		}()
+	}
+}
+
+func (sc *SignalCatcher) RegisterTaskGroup(group string, signals []os.Signal) {
+	sc.signals[group] = signals
+	sc.tasks[group] = list.New()
+}
+
+func (sc *SignalCatcher) RegisterTask(group string, fn TaskFn) error {
+	if _, ok := sc.tasks[group]; !ok {
+		return errors.New("No such task group")
+	}
+	sc.tasks[group].PushBack(fn)
+	return nil
+}