From 46c5db92da8ff971db26a3f11282fba5594a7f28 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 | 15 +++++++++++++--
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/caddy.go b/caddy.go
index 92a1fdf2..1e0eacd2 100644
--- a/caddy.go
+++ b/caddy.go
@@ -715,6 +715,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 {
@@ -736,6 +737,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
@@ -774,6 +785,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 e90475b1..d73af770 100644
--- a/context.go
+++ b/context.go
@@ -44,8 +44,9 @@ type Context struct {
 
 	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
@@ -86,7 +87,8 @@ func (ctx *Context) OnCancel(f func()) {
 	ctx.cleanupFuncs = append(ctx.cleanupFuncs, f)
 }
 
-// Filesystems returns a ref to the FilesystemMap
+// Filesystems returns a ref to the FilesystemMap.
+// EXPERIMENTAL: This API is subject to change.
 func (ctx *Context) Filesystems() FileSystems {
 	// if no config is loaded, we use a default filesystemmap, which includes the osfs
 	if ctx.cfg == nil {
@@ -95,6 +97,15 @@ func (ctx *Context) Filesystems() FileSystems {
 	return ctx.cfg.filesystems
 }
 
+// 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