From f5fce902e8c42b80b8e35e49d136f46134f93663 Mon Sep 17 00:00:00 2001
From: Brian Picciano <me@mediocregopher.com>
Date: Fri, 19 Jul 2024 15:49:13 +0200
Subject: [PATCH] Introduce Group config option

---
 pmux-example.yml   |  8 ++++++++
 pmuxlib/pmuxlib.go | 35 ++++++++++++++++++++++++++---------
 pmuxlib/process.go |  8 ++++++++
 3 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/pmux-example.yml b/pmux-example.yml
index 72fbe73..ec827d2 100644
--- a/pmux-example.yml
+++ b/pmux-example.yml
@@ -35,6 +35,14 @@ processes:
     # process to exit before sending it a SIGKILL (aka a kill -9).
     sigKillWait: 10s
 
+    # group can be used to control the order and grouping of processes as they
+    # shut down.
+    #
+    # Processes will not be shut down until all processes with a higher group
+    # number are already shut down. Processes with the same group number will be
+    # shut down simultaneously.
+    group: 1
+
   # This process will not immediately exit when pmux tells it to do so, but pmux
   # will SIGKILL it after sigKillWait has elapsed.
   stubborn-pinger:
diff --git a/pmuxlib/pmuxlib.go b/pmuxlib/pmuxlib.go
index 9081277..6d77998 100644
--- a/pmuxlib/pmuxlib.go
+++ b/pmuxlib/pmuxlib.go
@@ -68,17 +68,34 @@ func (p *Pmux) Restart(name string) {
 func (p *Pmux) Stop() {
 	var wg sync.WaitGroup
 
-	p.sysLogger.Println("killing child processes")
-	for _, proc := range p.processes {
-		proc := proc
+	var minGroup, maxGroup int
 
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			proc.Stop()
-		}()
+	for _, proc := range p.processes {
+		if maxGroup < proc.cfg.Group {
+			maxGroup = proc.cfg.Group
+		}
+
+		if minGroup > proc.cfg.Group {
+			minGroup = proc.cfg.Group
+		}
+	}
+
+	for group := maxGroup; group >= minGroup; group-- {
+		p.sysLogger.Printf("killing child processes (group %d)", group)
+		for _, proc := range p.processes {
+			proc := proc
+			if proc.cfg.Group != group {
+				continue
+			}
+
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				proc.Stop()
+			}()
+		}
+		wg.Wait()
 	}
 
-	wg.Wait()
 	p.sysLogger.Println("exited gracefully, ciao!")
 }
diff --git a/pmuxlib/process.go b/pmuxlib/process.go
index 760ba07..3d216b7 100644
--- a/pmuxlib/process.go
+++ b/pmuxlib/process.go
@@ -49,6 +49,14 @@ type ProcessConfig struct {
 	// this function returns. If used with RunProcessOnce and an error is
 	// returned, RunProcessOnce will return that error.
 	StartAfterFunc func(context.Context) error
+
+	// Group can be used to control the order and grouping of processes as they
+	// shut down.
+	//
+	// Processes will not be shut down until all processes with a higher group
+	// number are already shut down. Processes with the same group number will
+	// be shut down simultaneously.
+	Group int
 }
 
 func (cfg ProcessConfig) withDefaults() ProcessConfig {