fix: handle graceful shutdown

1 files changed, 45 insertions, 11 deletions

diff --git a/main.go b/main.go
index f9e7057..b8cbdac 100644
--- a/main.go
+++ b/main.go
@@ -4,7 +4,6 @@ import (
 	"flag"
 	"log"
 	"os"
-	"os/exec"
 	"os/signal"
 	"strings"
 	"sync"
@@ -15,7 +14,11 @@ import (
 	"github.com/xlgmokha/minit/pkg/procfile"
 )
 
-var procfilePath *string
+var (
+	pidMutex     sync.Mutex
+	pids         []int
+	procfilePath *string
+)
 
 func init() {
 	procfilePath = flag.String("f", "Procfile", "path to Procfile")
@@ -23,6 +26,34 @@ func init() {
 	log.SetFlags(0)
 }
 
+func addPid(pid int) {
+	pidMutex.Lock()
+	defer pidMutex.Unlock()
+	pids = append(pids, pid)
+}
+
+func removePid(pid int) {
+	pidMutex.Lock()
+	defer pidMutex.Unlock()
+
+	for i, p := range pids {
+		if p == pid {
+			pids = append(pids[:i], pids[i+1:]...)
+			break
+		}
+	}
+}
+
+func forwardSignalToAll(sig os.Signal) {
+	pidMutex.Lock()
+	defer pidMutex.Unlock()
+
+	signal := sig.(syscall.Signal)
+	for _, pid := range pids {
+		syscall.Kill(-pid, signal)
+	}
+}
+
 func main() {
 	var wg sync.WaitGroup
 	var shutdown int32
@@ -37,32 +68,35 @@ func main() {
 			wg.Add(1)
 			go func(proc *procfile.Proc) {
 				defer wg.Done()
-				var cmd *exec.Cmd
 
 				for atomic.LoadInt32(&shutdown) == 0 {
-					cmd = proc.NewCommand()
+					cmd := proc.NewCommand()
 
 					if cmd.Start() != nil {
 						time.Sleep(2 * time.Second)
 						continue
 					}
+
+					addPid(cmd.Process.Pid)
 					cmd.Wait()
+					removePid(cmd.Process.Pid)
 					time.Sleep(time.Second)
 				}
-
-				if cmd != nil && cmd.Process != nil {
-					syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM)
-				}
 			}(proc)
 		}
 	}
 
 	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	signal.Notify(sigChan)
 
 	go func() {
-		<-sigChan
-		atomic.StoreInt32(&shutdown, 1)
+		for sig := range sigChan {
+			if sig == syscall.SIGINT || sig == syscall.SIGTERM {
+				atomic.StoreInt32(&shutdown, 1)
+			}
+
+			forwardSignalToAll(sig)
+		}
 	}()
 
 	wg.Wait()