From dc63bcdc49d6b345b19fb5dbb4d0ef8f99e8c72c Mon Sep 17 00:00:00 2001
From: Dustin Pianalto <dustin@djpianalto.com>
Date: Fri, 7 Mar 2025 14:13:30 -0900
Subject: [PATCH] fix shutdown maybe

---
 examples/main.go | 20 ++++++++--
 web/server.go    | 99 ++++++++++++++++++++++++++++++------------------
 2 files changed, 79 insertions(+), 40 deletions(-)

diff --git a/examples/main.go b/examples/main.go
index f9d009f..b04800e 100644
--- a/examples/main.go
+++ b/examples/main.go
@@ -94,11 +94,23 @@ func main() {
 	// Cancel context to signal all components to shut down
 	cancel()
 
-	// Give a moment for any pending operations to complete
-	time.Sleep(500 * time.Millisecond)
+	// Create a timeout context for shutdown
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer shutdownCancel()
 
-	// Wait for all interfaces to shut down
-	wg.Wait()
+	// Wait for all interfaces to shut down with a timeout
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		fmt.Println("All components shut down successfully")
+	case <-shutdownCtx.Done():
+		fmt.Println("Shutdown timed out, forcing exit")
+	}
 }
 
 // startWebInterface initializes and runs the web interface
diff --git a/web/server.go b/web/server.go
index 30a659f..ae26043 100644
--- a/web/server.go
+++ b/web/server.go
@@ -256,46 +256,61 @@ func (s *Server) Start() error {
 		IdleTimeout:  120 * time.Second,
 	}
 
-	// Check if certificate files exist
-	certFile := "server.crt"
-	keyFile := "server.key"
-
-	certExists := fileExists(certFile)
-	keyExists := fileExists(keyFile)
-
-	s.logger.Info("Starting web server", "port", s.port, "http2", true)
-
-	if certExists && keyExists {
-		// Start HTTPS server with HTTP/2 support
-		s.logger.Info("Using TLS with HTTP/2", "certFile", certFile, "keyFile", keyFile)
-		return s.server.ListenAndServeTLS(certFile, keyFile)
-	} else {
-		// Generate self-signed certificate for development
-		s.logger.Info("Generating self-signed certificate for HTTP/2")
-		cert, key, err := generateSelfSignedCert()
-		if err != nil {
-			s.logger.Error("Failed to generate self-signed certificate", "error", err)
-
-			// Fall back to HTTP/1.1
-			s.logger.Info("Falling back to HTTP/1.1 (no TLS)")
+	// Start server in a goroutine so we can return from this function
+	go func() {
+		s.logger.Info("Starting web server", "port", s.port, "http2", s.useHTTP2)
+
+		var err error
+		if s.useHTTP2 {
+			// Check if certificate files exist
+			certFile := "server.crt"
+			keyFile := "server.key"
+
+			certExists := fileExists(certFile)
+			keyExists := fileExists(keyFile)
+
+			if certExists && keyExists {
+				// Start HTTPS server with HTTP/2 support
+				s.logger.Info("Using TLS with HTTP/2", "certFile", certFile, "keyFile", keyFile)
+				err = s.server.ListenAndServeTLS(certFile, keyFile)
+			} else {
+				// Generate self-signed certificate for development
+				s.logger.Info("Generating self-signed certificate for HTTP/2")
+				cert, key, genErr := generateSelfSignedCert()
+				if genErr != nil {
+					s.logger.Error("Failed to generate self-signed certificate", "error", genErr)
+
+					// Fall back to HTTP/1.1
+					s.logger.Info("Falling back to HTTP/1.1 (no TLS)")
+					s.server.TLSConfig = nil
+					err = s.server.ListenAndServe()
+				} else {
+					// Write certificate files
+					if writeErr := os.WriteFile(certFile, cert, 0600); writeErr != nil {
+						s.logger.Error("Failed to write certificate file", "error", writeErr)
+						err = writeErr
+					} else if writeErr := os.WriteFile(keyFile, key, 0600); writeErr != nil {
+						s.logger.Error("Failed to write key file", "error", writeErr)
+						err = writeErr
+					} else {
+						s.logger.Info("Self-signed certificate generated", "certFile", certFile, "keyFile", keyFile)
+						err = s.server.ListenAndServeTLS(certFile, keyFile)
+					}
+				}
+			}
+		} else {
+			// Start HTTP/1.1 server
+			s.logger.Info("Using HTTP/1.1 (no TLS)")
 			s.server.TLSConfig = nil
-			return s.server.ListenAndServe()
+			err = s.server.ListenAndServe()
 		}
 
-		// Write certificate files
-		if err := os.WriteFile(certFile, cert, 0600); err != nil {
-			s.logger.Error("Failed to write certificate file", "error", err)
-			return err
+		if err != nil && err != http.ErrServerClosed {
+			s.logger.Error("HTTP server error", "error", err)
 		}
+	}()
 
-		if err := os.WriteFile(keyFile, key, 0600); err != nil {
-			s.logger.Error("Failed to write key file", "error", err)
-			return err
-		}
-
-		s.logger.Info("Self-signed certificate generated", "certFile", certFile, "keyFile", keyFile)
-		return s.server.ListenAndServeTLS(certFile, keyFile)
-	}
+	return nil
 }
 
 // Helper function to check if a file exists
@@ -347,6 +362,10 @@ func (s *Server) Shutdown(ctx context.Context) error {
 	// Signal event forwarder to stop
 	close(s.shutdown)
 
+	// Close all client connections with a timeout
+	shutdownCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
 	// Close all client connections
 	s.clientsMux.Lock()
 	for clientChan := range s.clients {
@@ -363,7 +382,15 @@ func (s *Server) Shutdown(ctx context.Context) error {
 	s.adminclientsMux.Unlock()
 
 	// Shutdown the HTTP server
-	return s.server.Shutdown(ctx)
+	err := s.server.Shutdown(shutdownCtx)
+	if err != nil {
+		// If shutdown times out, force close
+		s.logger.Warn("Server shutdown timed out, forcing close", "error", err)
+		return s.server.Close()
+	}
+
+	s.logger.Info("Web server shutdown complete")
+	return nil
 }
 
 // forwardEvents forwards derby events to SSE clients