From 4071dd4e20ca05587c0b4d349f1f39f2cde7e95b Mon Sep 17 00:00:00 2001
From: Matt Holt <mholt@users.noreply.github.com>
Date: Fri, 1 Mar 2024 09:57:05 -0700
Subject: [PATCH] core: OnExit hooks (#6128)

* core: OnExit callbacks

* core: Process-global OnExit callbacks
---
 caddy.go   | 28 ++++++++++++++++++++++++++++
 context.go | 12 +++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/caddy.go b/caddy.go
index 8e7dce50..4086b4a1 100644
--- a/caddy.go
+++ b/caddy.go
@@ -707,6 +707,7 @@ func exitProcess(ctx context.Context, logger *zap.Logger) {
 	logger.Warn("exiting; byeee!! 👋")
 
 	exitCode := ExitCodeSuccess
+	lastContext := ActiveContext()
 
 	// stop all apps
 	if err := Stop(); err != nil {
@@ -728,6 +729,16 @@ func exitProcess(ctx context.Context, logger *zap.Logger) {
 		}
 	}
 
+	// execute any process-exit callbacks
+	for _, exitFunc := range lastContext.exitFuncs {
+		exitFunc(ctx)
+	}
+	exitFuncsMu.Lock()
+	for _, exitFunc := range exitFuncs {
+		exitFunc(ctx)
+	}
+	exitFuncsMu.Unlock()
+
 	// shut down admin endpoint(s) in goroutines so that
 	// if this function was called from an admin handler,
 	// it has a chance to return gracefully
@@ -766,6 +777,23 @@ var exiting = new(int32) // accessed atomically
 // EXPERIMENTAL API: subject to change or removal.
 func Exiting() bool { return atomic.LoadInt32(exiting) == 1 }
 
+// OnExit registers a callback to invoke during process exit.
+// This registration is PROCESS-GLOBAL, meaning that each
+// function should only be registered once forever, NOT once
+// per config load (etc).
+//
+// EXPERIMENTAL API: subject to change or removal.
+func OnExit(f func(context.Context)) {
+	exitFuncsMu.Lock()
+	exitFuncs = append(exitFuncs, f)
+	exitFuncsMu.Unlock()
+}
+
+var (
+	exitFuncs   []func(context.Context)
+	exitFuncsMu sync.Mutex
+)
+
 // Duration can be an integer or a string. An integer is
 // interpreted as nanoseconds. If a string, it is a Go
 // time.Duration value such as `300ms`, `1.5h`, or `2h45m`;
diff --git a/context.go b/context.go
index 85978d49..efc2c560 100644
--- a/context.go
+++ b/context.go
@@ -39,8 +39,9 @@ type Context struct {
 	context.Context
 	moduleInstances map[string][]Module
 	cfg             *Config
-	cleanupFuncs    []func()
 	ancestry        []Module
+	cleanupFuncs    []func()                // invoked at every config unload
+	exitFuncs       []func(context.Context) // invoked at config unload ONLY IF the process is exiting (EXPERIMENTAL)
 }
 
 // NewContext provides a new context derived from the given
@@ -81,6 +82,15 @@ func (ctx *Context) OnCancel(f func()) {
 	ctx.cleanupFuncs = append(ctx.cleanupFuncs, f)
 }
 
+// OnExit executes f when the process exits gracefully.
+// The function is only executed if the process is gracefully
+// shut down while this context is active.
+//
+// EXPERIMENTAL API: subject to change or removal.
+func (ctx *Context) OnExit(f func(context.Context)) {
+	ctx.exitFuncs = append(ctx.exitFuncs, f)
+}
+
 // LoadModule loads the Caddy module(s) from the specified field of the parent struct
 // pointer and returns the loaded module(s). The struct pointer and its field name as
 // a string are necessary so that reflection can be used to read the struct tag on the