From 3c877c504b6102daf5dcc1083b1f1a7db88d304c Mon Sep 17 00:00:00 2001
From: Charlie Vieth <charlie.vieth@gmail.com>
Date: Sat, 13 Apr 2024 01:58:11 -0400
Subject: 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>
---
 src/util/atexit.go      | 38 ++++++++++++++++++++++++++++++++++++++
 src/util/atexit_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 src/util/atexit.go
 create mode 100644 src/util/atexit_test.go

(limited to 'src/util')

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")
+	}
+}
-- 
cgit v1.2.3