Enable profiling options when 'pprof' tag is set (#2813)

This commit enables cpu, mem, block, and mutex profling of the FZF
executable. To support flushing the profiles at program exit it adds
util.AtExit to register "at exit" functions and mandates that util.Exit
is used instead of os.Exit to stop the program.

Co-authored-by: Junegunn Choi <junegunn.c@gmail.com>

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