Add a DeferExitHandler function

Useful for running exit handlers in the same order as defer statements

Author: Georgi Dimitrov

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...)
+}

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")
 	}
 }