From e7ae898ddb2752ae5fc30cd811a58e955a4a4042 Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Thu, 7 Oct 2021 14:23:49 +0200
Subject: [PATCH] Add global shutdown hook function to modules

---
 modules/start.go |  3 ++-
 modules/stop.go  | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/modules/start.go b/modules/start.go
index 2245d6f..8ee35a0 100644
--- a/modules/start.go
+++ b/modules/start.go
@@ -16,8 +16,9 @@ var (
 	globalPrepFn          func() error
 )
 
-// SetGlobalPrepFn sets a global prep function that is run before all modules. This can be used to pre-initialize modules, such as setting the data root or database path.
 // SetGlobalPrepFn sets a global prep function that is run before all modules.
+// This can be used to pre-initialize modules, such as setting the data root
+// or database path.
 func SetGlobalPrepFn(fn func() error) {
 	if globalPrepFn == nil {
 		globalPrepFn = fn
diff --git a/modules/stop.go b/modules/stop.go
index cbfefad..0ad30cd 100644
--- a/modules/stop.go
+++ b/modules/stop.go
@@ -14,8 +14,17 @@ var (
 	shutdownFlag   = abool.NewBool(false)
 
 	shutdownCompleteSignal = make(chan struct{})
+
+	globalShutdownFn func()
 )
 
+// SetGlobalShutdownFn sets a global shutdown function that is called first when shutting down.
+func SetGlobalShutdownFn(fn func()) {
+	if globalShutdownFn == nil {
+		globalShutdownFn = fn
+	}
+}
+
 // IsShuttingDown returns whether the global shutdown is in progress.
 func IsShuttingDown() bool {
 	return shutdownFlag.IsSet()
@@ -39,6 +48,11 @@ func Shutdown() error {
 		return errors.New("shutdown already initiated")
 	}
 
+	// Execute global shutdown function.
+	if globalShutdownFn != nil {
+		globalShutdownFn()
+	}
+
 	if initialStartCompleted.IsSet() {
 		log.Warning("modules: starting shutdown...")
 	} else {