2 files changed, 62 insertions, 0 deletions
diff --git a/src/util/atexit.go b/src/util/atexit.go
new file mode 100644
index 00000000..a22a3a96
--- /dev/null
+++ b/src/util/atexit.go
@@ -0,0 +1,38 @@
+package util
+
+import (
+	"os"
+	"sync"
+)
+
+var atExitFuncs []func()
+
+// AtExit registers the function fn to be called on program termination.
+// The functions will be called in reverse order they were registered.
+func AtExit(fn func()) {
+	if fn == nil {
+		panic("AtExit called with nil func")
+	}
+	once := &sync.Once{}
+	atExitFuncs = append(atExitFuncs, func() {
+		once.Do(fn)
+	})
+}
+
+// RunAtExitFuncs runs any functions registered with AtExit().
+func RunAtExitFuncs() {
+	fns := atExitFuncs
+	for i := len(fns) - 1; i >= 0; i-- {
+		fns[i]()
+	}
+}
+
+// Exit executes any functions registered with AtExit() then exits the program
+// with os.Exit(code).
+//
+// NOTE: It must be used instead of os.Exit() since calling os.Exit() terminates
+// the program before any of the AtExit functions can run.
+func Exit(code int) {
+	defer os.Exit(code)
+	RunAtExitFuncs()
+}
diff --git a/src/util/atexit_test.go b/src/util/atexit_test.go
new file mode 100644
index 00000000..1ff85be8
--- /dev/null
+++ b/src/util/atexit_test.go
@@ -0,0 +1,24 @@
+package util
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestAtExit(t *testing.T) {
+	want := []int{3, 2, 1, 0}
+	var called []int
+	for i := 0; i < 4; i++ {
+		n := i
+		AtExit(func() { called = append(called, n) })
+	}
+	RunAtExitFuncs()
+	if !reflect.DeepEqual(called, want) {
+		t.Errorf("AtExit: want call order: %v got: %v", want, called)
+	}
+
+	RunAtExitFuncs()
+	if !reflect.DeepEqual(called, want) {
+		t.Error("AtExit: should only call exit funcs once")
+	}
+}