v0.1.4 fixed blocking shutdown, better isolation of the streams

src/cmd/logwisp/main.go

@@ -20,10 +20,7 @@ func main() {
 	// Parse CLI flags
 	var (
 		configFile = flag.String("config", "", "Config file path")
-		// Legacy compatibility flags
-		port       = flag.Int("port", 0, "HTTP port (legacy, maps to --http-port)")
-		bufferSize = flag.Int("buffer-size", 0, "Buffer size (legacy, maps to --http-buffer-size)")
-		// New explicit flags
+		// Flags
 		httpPort      = flag.Int("http-port", 0, "HTTP server port")
 		httpBuffer    = flag.Int("http-buffer-size", 0, "HTTP server buffer size")
 		tcpPort       = flag.Int("tcp-port", 0, "TCP server port")
@@ -41,15 +38,7 @@ func main() {
 	// Build CLI args for config
 	var cliArgs []string
 
-	// Legacy mapping
-	if *port > 0 {
-		cliArgs = append(cliArgs, fmt.Sprintf("--httpserver.port=%d", *port))
-	}
-	if *bufferSize > 0 {
-		cliArgs = append(cliArgs, fmt.Sprintf("--httpserver.buffer_size=%d", *bufferSize))
-	}
-
-	// New flags
+	// Flags
 	if *httpPort > 0 {
 		cliArgs = append(cliArgs, fmt.Sprintf("--httpserver.port=%d", *httpPort))
 	}
@@ -92,8 +81,6 @@ func main() {
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
 
-	var wg sync.WaitGroup
-
 	// Create monitor
 	mon := monitor.New()
 	mon.SetCheckInterval(time.Duration(cfg.Monitor.CheckIntervalMs) * time.Millisecond)
@@ -119,14 +106,23 @@ func main() {
 		tcpChan := mon.Subscribe()
 		tcpServer = stream.NewTCPStreamer(tcpChan, cfg.TCPServer)
 
-		wg.Add(1)
+		// Start TCP server in separate goroutine without blocking wg.Wait()
+		tcpStarted := make(chan error, 1)
 		go func() {
-			defer wg.Done()
-			if err := tcpServer.Start(); err != nil {
-				fmt.Fprintf(os.Stderr, "TCP server error: %v\n", err)
-			}
+			tcpStarted <- tcpServer.Start()
 		}()
 
+		// Check if TCP server started successfully
+		select {
+		case err := <-tcpStarted:
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "TCP server failed to start: %v\n", err)
+				os.Exit(1)
+			}
+		case <-time.After(1 * time.Second):
+			// Server is running
+		}
+
 		fmt.Printf("TCP streaming on port %d\n", cfg.TCPServer.Port)
 	}
 
@@ -135,14 +131,23 @@ func main() {
 		httpChan := mon.Subscribe()
 		httpServer = stream.NewHTTPStreamer(httpChan, cfg.HTTPServer)
 
-		wg.Add(1)
+		// Start HTTP server in separate goroutine without blocking wg.Wait()
+		httpStarted := make(chan error, 1)
 		go func() {
-			defer wg.Done()
-			if err := httpServer.Start(); err != nil {
-				fmt.Fprintf(os.Stderr, "HTTP server error: %v\n", err)
-			}
+			httpStarted <- httpServer.Start()
 		}()
 
+		// Check if HTTP server started successfully
+		select {
+		case err := <-httpStarted:
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "HTTP server failed to start: %v\n", err)
+				os.Exit(1)
+			}
+		case <-time.After(1 * time.Second):
+			// Server is running
+		}
+
 		fmt.Printf("HTTP/SSE streaming on http://localhost:%d/stream\n", cfg.HTTPServer.Port)
 		fmt.Printf("Status available at http://localhost:%d/status\n", cfg.HTTPServer.Port)
 	}
@@ -156,29 +161,67 @@ func main() {
 	<-sigChan
 	fmt.Println("\nShutting down...")
 
-	// Stop servers first
+	// Create shutdown group for concurrent server stops
+	var shutdownWg sync.WaitGroup
+
+	// Stop servers first (concurrently)
 	if tcpServer != nil {
-		tcpServer.Stop()
+		shutdownWg.Add(1)
+		go func() {
+			defer shutdownWg.Done()
+			tcpServer.Stop()
+		}()
 	}
 	if httpServer != nil {
-		httpServer.Stop()
+		shutdownWg.Add(1)
+		go func() {
+			defer shutdownWg.Done()
+			httpServer.Stop()
+		}()
 	}
 
-	// Cancel context and stop monitor
+	// Cancel context to stop monitor
 	cancel()
-	mon.Stop()
 
-	// Wait for completion
-	done := make(chan struct{})
+	// Wait for servers to stop with timeout
+	serversDone := make(chan struct{})
 	go func() {
-		wg.Wait()
-		close(done)
+		shutdownWg.Wait()
+		close(serversDone)
 	}()
 
-	select {
-	case <-done:
-		fmt.Println("Shutdown complete")
-	case <-time.After(2 * time.Second):
-		fmt.Println("Shutdown timeout")
+	// Stop monitor after context cancellation
+	monitorDone := make(chan struct{})
+	go func() {
+		mon.Stop()
+		close(monitorDone)
+	}()
+
+	// Wait for all components with proper timeout
+	shutdownTimeout := 5 * time.Second
+	shutdownTimer := time.NewTimer(shutdownTimeout)
+	defer shutdownTimer.Stop()
+
+	serversShutdown := false
+	monitorShutdown := false
+
+	for !serversShutdown || !monitorShutdown {
+		select {
+		case <-serversDone:
+			serversShutdown = true
+		case <-monitorDone:
+			monitorShutdown = true
+		case <-shutdownTimer.C:
+			if !serversShutdown {
+				fmt.Println("Warning: Server shutdown timeout")
+			}
+			if !monitorShutdown {
+				fmt.Println("Warning: Monitor shutdown timeout")
+			}
+			fmt.Println("Forcing exit")
+			os.Exit(1)
+		}
 	}
+
+	fmt.Println("Shutdown complete")
 }

src/internal/stream/http.go

@@ -3,6 +3,7 @@ package stream
 
 import (
 	"bufio"
+	"context"
 	"encoding/json"
 	"fmt"
 	"strings"
@@ -22,6 +23,8 @@ type HTTPStreamer struct {
 	activeClients atomic.Int32
 	mu            sync.RWMutex
 	startTime     time.Time
+	done          chan struct{}
+	wg            sync.WaitGroup
 }
 
 func NewHTTPStreamer(logChan chan monitor.LogEntry, cfg config.HTTPConfig) *HTTPStreamer {
@@ -29,6 +32,7 @@ func NewHTTPStreamer(logChan chan monitor.LogEntry, cfg config.HTTPConfig) *HTTP
 		logChan:   logChan,
 		config:    cfg,
 		startTime: time.Now(),
+		done:      make(chan struct{}),
 	}
 }
 
@@ -41,13 +45,42 @@ func (h *HTTPStreamer) Start() error {
 	}
 
 	addr := fmt.Sprintf(":%d", h.config.Port)
-	return h.server.ListenAndServe(addr)
+
+	// Run server in separate goroutine to avoid blocking
+	errChan := make(chan error, 1)
+	go func() {
+		err := h.server.ListenAndServe(addr)
+		if err != nil {
+			errChan <- err
+		}
+	}()
+
+	// Check if server started successfully
+	select {
+	case err := <-errChan:
+		return err
+	case <-time.After(100 * time.Millisecond):
+		// Server started successfully
+		return nil
+	}
 }
 
 func (h *HTTPStreamer) Stop() {
+	// Signal all client handlers to stop
+	close(h.done)
+
+	// Shutdown HTTP server
 	if h.server != nil {
-		h.server.Shutdown()
+		// Create context with timeout for server shutdown
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		// Use ShutdownWithContext for graceful shutdown
+		h.server.ShutdownWithContext(ctx)
 	}
+
+	// Wait for all active client handlers to finish
+	h.wg.Wait()
 }
 
 func (h *HTTPStreamer) requestHandler(ctx *fasthttp.RequestCtx) {
@@ -72,25 +105,46 @@ func (h *HTTPStreamer) handleStream(ctx *fasthttp.RequestCtx) {
 	ctx.Response.Header.Set("X-Accel-Buffering", "no")
 
 	h.activeClients.Add(1)
-	defer h.activeClients.Add(-1)
+	h.wg.Add(1) // Track this client handler
+	defer func() {
+		h.activeClients.Add(-1)
+		h.wg.Done() // Mark handler as done
+	}()
 
 	// Create subscription for this client
 	clientChan := make(chan monitor.LogEntry, h.config.BufferSize)
+	clientDone := make(chan struct{})
 
 	// Subscribe to monitor's broadcast
 	go func() {
-		for entry := range h.logChan {
+		defer close(clientChan)
+		for {
 			select {
-			case clientChan <- entry:
-			default:
-				// Drop if client buffer full
+			case entry, ok := <-h.logChan:
+				if !ok {
+					return
+				}
+				select {
+				case clientChan <- entry:
+				case <-clientDone:
+					return
+				case <-h.done: // Check for server shutdown
+					return
+				default:
+					// Drop if client buffer full
+				}
+			case <-clientDone:
+				return
+			case <-h.done: // Check for server shutdown
+				return
 			}
 		}
-		close(clientChan)
 	}()
 
 	// Define the stream writer function
 	streamFunc := func(w *bufio.Writer) {
+		defer close(clientDone)
+
 		// Send initial connected event
 		clientID := fmt.Sprintf("%d", time.Now().UnixNano())
 		fmt.Fprintf(w, "event: connected\ndata: {\"client_id\":\"%s\"}\n\n", clientID)
@@ -129,6 +183,12 @@ func (h *HTTPStreamer) handleStream(ctx *fasthttp.RequestCtx) {
 						return
 					}
 				}
+
+			case <-h.done: // ADDED: Check for server shutdown
+				// Send final disconnect event
+				fmt.Fprintf(w, "event: disconnect\ndata: {\"reason\":\"server_shutdown\"}\n\n")
+				w.Flush()
+				return
 			}
 		}
 	}

src/internal/stream/tcp.go

@@ -2,6 +2,7 @@
 package stream
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"sync"
@@ -17,9 +18,11 @@ type TCPStreamer struct {
 	logChan     chan monitor.LogEntry
 	config      config.TCPConfig
 	server      *tcpServer
-	done        chan struct{}
 	activeConns atomic.Int32
 	startTime   time.Time
+	done        chan struct{}
+	engine      *gnet.Engine
+	wg          sync.WaitGroup
 }
 
 type tcpServer struct {
@@ -32,8 +35,8 @@ func NewTCPStreamer(logChan chan monitor.LogEntry, cfg config.TCPConfig) *TCPStr
 	return &TCPStreamer{
 		logChan:   logChan,
 		config:    cfg,
-		done:      make(chan struct{}),
 		startTime: time.Now(),
+		done:      make(chan struct{}),
 	}
 }
 
@@ -41,23 +44,53 @@ func (t *TCPStreamer) Start() error {
 	t.server = &tcpServer{streamer: t}
 
 	// Start log broadcast loop
-	go t.broadcastLoop()
+	t.wg.Add(1)
+	go func() {
+		defer t.wg.Done()
+		t.broadcastLoop()
+	}()
 
 	// Configure gnet with no-op logger
 	addr := fmt.Sprintf("tcp://:%d", t.config.Port)
 
-	err := gnet.Run(t.server, addr,
-		gnet.WithLogger(noopLogger{}), // No-op logger: discard everything
-		gnet.WithMulticore(true),
-		gnet.WithReusePort(true),
-	)
+	// Run gnet in separate goroutine to avoid blocking
+	errChan := make(chan error, 1)
+	go func() {
+		err := gnet.Run(t.server, addr,
+			gnet.WithLogger(noopLogger{}), // No-op logger: discard everything
+			gnet.WithMulticore(true),
+			gnet.WithReusePort(true),
+		)
+		errChan <- err
+	}()
 
-	return err
+	// Wait briefly for server to start or fail
+	select {
+	case err := <-errChan:
+		// Server failed immediately
+		close(t.done)
+		t.wg.Wait()
+		return err
+	case <-time.After(100 * time.Millisecond):
+		// Server started successfully
+		return nil
+	}
 }
 
 func (t *TCPStreamer) Stop() {
+	// Signal broadcast loop to stop
 	close(t.done)
-	// No engine to stop with gnet v2
+
+	// Stop gnet engine if running
+	if t.engine != nil {
+		// Use Stop() method to gracefully shutdown gnet
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+		t.engine.Stop(ctx)
+	}
+
+	// Wait for broadcast loop to finish
+	t.wg.Wait()
 }
 
 func (t *TCPStreamer) broadcastLoop() {
@@ -72,7 +105,10 @@ func (t *TCPStreamer) broadcastLoop() {
 
 	for {
 		select {
-		case entry := <-t.logChan:
+		case entry, ok := <-t.logChan:
+			if !ok {
+				return // Channel closed
+			}
 			data, err := json.Marshal(entry)
 			if err != nil {
 				continue
@@ -122,6 +158,7 @@ func (t *TCPStreamer) formatHeartbeat() []byte {
 }
 
 func (s *tcpServer) OnBoot(eng gnet.Engine) gnet.Action {
+	s.streamer.engine = &eng
 	return gnet.None
 }