From 391c4f20d54da05629e0bcec76ae599c5f26abf2 Mon Sep 17 00:00:00 2001
From: sawka <mike.sawka@gmail.com>
Date: Wed, 12 Jun 2024 18:17:56 -0700
Subject: [PATCH] add a stdin watcher that also leads to a shutdown

---
 cmd/server/main-server.go | 32 +++++++++++++++++++++++++-------
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/cmd/server/main-server.go b/cmd/server/main-server.go
index f81d1d1a..b6e5239c 100644
--- a/cmd/server/main-server.go
+++ b/cmd/server/main-server.go
@@ -11,6 +11,7 @@ import (
 	"os/signal"
 	"runtime"
 	"strconv"
+	"sync"
 	"syscall"
 	"time"
 
@@ -23,14 +24,18 @@ import (
 
 const ReadySignalPidVarName = "WAVETERM_READY_SIGNAL_PID"
 
+var shutdownOnce sync.Once
+
 func doShutdown(reason string) {
-	log.Printf("shutting down: %s\n", reason)
-	ctx, cancelFn := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancelFn()
-	// TODO deal with flush in progress
-	filestore.WFS.FlushCache(ctx)
-	time.Sleep(200 * time.Millisecond)
-	os.Exit(0)
+	shutdownOnce.Do(func() {
+		log.Printf("shutting down: %s\n", reason)
+		ctx, cancelFn := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancelFn()
+		// TODO deal with flush in progress
+		filestore.WFS.FlushCache(ctx)
+		time.Sleep(200 * time.Millisecond)
+		os.Exit(0)
+	})
 }
 
 func installShutdownSignalHandlers() {
@@ -44,6 +49,18 @@ func installShutdownSignalHandlers() {
 	}()
 }
 
+// watch stdin, kill server if stdin is closed
+func stdinReadWatch() {
+	buf := make([]byte, 1024)
+	for {
+		_, err := os.Stdin.Read(buf)
+		if err != nil {
+			doShutdown(fmt.Sprintf("stdin closed/error (%v)", err))
+			break
+		}
+	}
+}
+
 func main() {
 	err := service.ValidateServiceMap()
 	if err != nil {
@@ -78,6 +95,7 @@ func main() {
 		return
 	}
 	installShutdownSignalHandlers()
+	go stdinReadWatch()
 
 	go web.RunWebSocketServer()
 	go func() {