diff --git a/pkg/util/atexit/atexit.go b/pkg/util/atexit/atexit.go new file mode 100644 index 0000000000..38efa8c65d --- /dev/null +++ b/pkg/util/atexit/atexit.go @@ -0,0 +1,96 @@ +package atexit + +import ( + "os" + "runtime/debug" + "sort" + "sync" +) + +// ExitHandlerFunc is the type of handler func +type ExitHandlerFunc func(ExitHandler) + +// ExitHandler defines the spec of handler +// +// Reason and Func are mandatory and must not be empty or nil +// +// Handlers with smaller Prio will be executed earlier than those with bigger +// Prio at exit time. Handler func will receive a copy of the ExitHandler +// struct previously registered +type ExitHandler struct { + Prio int + Reason string + Func ExitHandlerFunc + Value interface{} +} + +var ( + handlers = map[int][]ExitHandler{} + handlersLock = &sync.Mutex{} + once = &sync.Once{} +) + +// Register registers ExitHandler +// +// Smaller prio number mean higher priority and exit handlers with higher +// priority will be executed first. For handlers with equal priorities, those +// registered first will be executed earlier at exit time +func Register(eh ExitHandler) { + if eh.Reason == "" { + panic("handler reason must not be empty") + } + if eh.Func == nil { + panic("handler func must not be nil") + } + + handlersLock.Lock() + defer handlersLock.Unlock() + + ehs, ok := handlers[eh.Prio] + if ok { + ehs = append(ehs, eh) + } else { + ehs = []ExitHandler{eh} + } + handlers[eh.Prio] = ehs +} + +// Handle calls registered handlers sequentially according to priority and +// registration order +// +// Panics caused by handler func will be caught, recorded, then next func will +// be run +func Handle() { + once.Do(func() { + handlersLock.Lock() + defer handlersLock.Unlock() + + prios := make([]int, 0, len(handlers)) + for prio := range handlers { + prios = append(prios, prio) + } + sort.Ints(prios) + for _, prio := range prios { + ehs := handlers[prio] + for _, eh := range ehs { + print("atexit: prio=", prio, ", reason=", eh.Reason, "\n") + func() { + defer func() { + val := recover() + if val != nil { + print("panic ", val, "\n") + debug.PrintStack() + } + }() + eh.Func(eh) + }() + } + } + }) +} + +// Exit calls handlers then does os.Exit(code) +func Exit(code int) { + defer os.Exit(code) + Handle() +} diff --git a/pkg/util/atexit/atexit_test.go b/pkg/util/atexit/atexit_test.go new file mode 100644 index 0000000000..b2a69c8948 --- /dev/null +++ b/pkg/util/atexit/atexit_test.go @@ -0,0 +1,56 @@ +package atexit + +import ( + "testing" +) + +func TestAtExit(t *testing.T) { + t.Run("empty reason", func(t *testing.T) { + func() { + defer func() { + val := recover() + if val == nil { + t.Errorf("should panic") + } + }() + Register(ExitHandler{}) + }() + }) + t.Run("empty func", func(t *testing.T) { + func() { + defer func() { + val := recover() + if val == nil { + t.Errorf("should panic") + } + }() + Register(ExitHandler{ + Reason: "have reason", + }) + }() + }) + t.Run("order & prio", func(t *testing.T) { + verdict := "" + handler := func(eh ExitHandler) { verdict += eh.Reason } + Register(ExitHandler{ + Prio: 2, + Reason: "2", + Func: handler, + }) + Register(ExitHandler{ + Prio: 0, + Reason: "0", + Func: handler, + }) + Register(ExitHandler{ + Prio: 0, + Reason: "1", + Func: handler, + }) + Handle() + Handle() + if verdict != "012" { + t.Errorf("expecting %q, got %q", "012", verdict) + } + }) +} diff --git a/pkg/util/atexit/const.go b/pkg/util/atexit/const.go new file mode 100644 index 0000000000..8bed426fc3 --- /dev/null +++ b/pkg/util/atexit/const.go @@ -0,0 +1,5 @@ +package atexit + +const ( + PRIO_LOG_CLOSE = 40000 +)