diff --git a/alt_exit.go b/alt_exit.go index 8af90637a..8fd189e1c 100644 --- a/alt_exit.go +++ b/alt_exit.go @@ -51,9 +51,9 @@ func Exit(code int) { os.Exit(code) } -// RegisterExitHandler adds a Logrus Exit handler, call logrus.Exit to invoke -// all handlers. The handlers will also be invoked when any Fatal log entry is -// made. +// RegisterExitHandler appends a Logrus Exit handler to the list of handlers, +// call logrus.Exit to invoke all handlers. The handlers will also be invoked when +// any Fatal log entry is made. // // This method is useful when a caller wishes to use logrus to log a fatal // message but also needs to gracefully shutdown. An example usecase could be @@ -62,3 +62,15 @@ func Exit(code int) { func RegisterExitHandler(handler func()) { handlers = append(handlers, handler) } + +// DeferExitHandler prepends a Logrus Exit handler to the list of handlers, +// call logrus.Exit to invoke all handlers. The handlers will also be invoked when +// any Fatal log entry is made. +// +// This method is useful when a caller wishes to use logrus to log a fatal +// message but also needs to gracefully shutdown. An example usecase could be +// closing database connections, or sending a alert that the application is +// closing. +func DeferExitHandler(handler func()) { + handlers = append([]func(){handler}, handlers...) +} diff --git a/alt_exit_test.go b/alt_exit_test.go index 0a2ff5650..54d503cb4 100644 --- a/alt_exit_test.go +++ b/alt_exit_test.go @@ -14,9 +14,61 @@ import ( func TestRegister(t *testing.T) { current := len(handlers) - RegisterExitHandler(func() {}) - if len(handlers) != current+1 { - t.Fatalf("expected %d handlers, got %d", current+1, len(handlers)) + + var results []string + + h1 := func() { results = append(results, "first") } + h2 := func() { results = append(results, "second") } + + RegisterExitHandler(h1) + RegisterExitHandler(h2) + + if len(handlers) != current+2 { + t.Fatalf("expected %d handlers, got %d", current+2, len(handlers)) + } + + runHandlers() + + if len(results) != 2 { + t.Fatalf("expected 2 handlers to be run, ran %d", len(results)) + } + + if results[0] != "first" { + t.Fatal("expected handler h1 to be run first, but it wasn't") + } + + if results[1] != "second" { + t.Fatal("expected handler h2 to be run second, but it wasn't") + } +} + +func TestDefer(t *testing.T) { + current := len(handlers) + + var results []string + + h1 := func() { results = append(results, "first") } + h2 := func() { results = append(results, "second") } + + DeferExitHandler(h1) + DeferExitHandler(h2) + + if len(handlers) != current+2 { + t.Fatalf("expected %d handlers, got %d", current+2, len(handlers)) + } + + runHandlers() + + if len(results) != 2 { + t.Fatalf("expected 2 handlers to be run, ran %d", len(results)) + } + + if results[0] != "second" { + t.Fatal("expected handler h2 to be run first, but it wasn't") + } + + if results[1] != "first" { + t.Fatal("expected handler h1 to be run second, but it wasn't") } }