From 5165da389159bfaab9ad635a69671e113c0fdd25 Mon Sep 17 00:00:00 2001
From: Aditya Thebe <contact@adityathebe.com>
Date: Thu, 24 Oct 2024 20:53:52 +0545
Subject: [PATCH] test: shutdown priority

---
 shutdown/shutdown.go      | 12 +++++++++---
 shutdown/shutdown_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)
 create mode 100644 shutdown/shutdown_test.go

diff --git a/shutdown/shutdown.go b/shutdown/shutdown.go
index 4f470ad2..72f292ad 100644
--- a/shutdown/shutdown.go
+++ b/shutdown/shutdown.go
@@ -27,8 +27,14 @@ func init() {
 }
 
 var Shutdown = sync.OnceFunc(func() {
-	logger.Infof("shutting down")
-	for _, task := range shutdownTaskRegistry {
+	logger.Infof("begin shutdown")
+	for len(shutdownTaskRegistry) > 0 {
+		_task := heap.Pop(&shutdownTaskRegistry)
+		if _task == nil {
+			break
+		}
+
+		task := _task.(ShutdownTask)
 		logger.Infof("shutting down: %s", task.Label)
 		task.Hook()
 	}
@@ -79,7 +85,7 @@ func (st ShutdownTasks) Len() int { return len(st) }
 
 // Less defines higher priority numbers will be processed first
 func (st ShutdownTasks) Less(i, j int) bool {
-	return st[i].Priority > st[j].Priority // Higher priority numbers come first
+	return st[i].Priority < st[j].Priority // Higher priority numbers come first
 }
 
 func (st ShutdownTasks) Swap(i, j int) {
diff --git a/shutdown/shutdown_test.go b/shutdown/shutdown_test.go
new file mode 100644
index 00000000..ee44a781
--- /dev/null
+++ b/shutdown/shutdown_test.go
@@ -0,0 +1,25 @@
+package shutdown
+
+import "testing"
+
+func TestShutdownPriority(t *testing.T) {
+	var closer int
+
+	add := func(label string, priority int) {
+		AddHookWithPriority(label, priority, func() {
+			if closer > priority {
+				t.Fatalf("something higher priority (%d) was closed earlier than (%d)", closer, priority)
+			} else {
+				closer = priority
+			}
+		})
+	}
+
+	add("database", PriorityCritical)
+	add("gRPC", PriorityIngress)
+	add("checkJOB", PriorityJobs)
+	add("echo", PriorityIngress)
+	add("topologyJob", PriorityJobs)
+
+	Shutdown()
+}