88 lines
1.8 KiB
Go
88 lines
1.8 KiB
Go
|
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
|
||
|
|
}
|