fix: synchronize CLI watchers and server shutdown

Author: piotrlaczkowski

cmd/golanggraph/auto_serve_command.go

@@ -137,6 +137,7 @@ func runAutoServe(cmd *cobra.Command, args []string) error {
 // regenerating routes is intentional: endpoint generation mutates route and
 // agent maps and is unsafe once requests may be in flight.
 func runAutoServeWithContext(ctx context.Context, out io.Writer, config *server.AutoServerConfig, opts autoServeOptions) error {
+	out = synchronizeWriter(out)
 	var changes <-chan struct{}
 	var watchErrors <-chan error
 	var stopWatching func()

cmd/golanggraph/main.go

@@ -19,6 +19,7 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
 
@@ -48,6 +49,28 @@ var (
 	verbose bool
 )
 
+// synchronizedWriter makes command output safe when a command's watcher and
+// serving goroutine both report progress. CLI callers commonly use bytes.Buffer
+// in tests, but the wrapper also protects any non-concurrent io.Writer supplied
+// by an embedding application.
+type synchronizedWriter struct {
+	mu sync.Mutex
+	w  io.Writer
+}
+
+func (w *synchronizedWriter) Write(p []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.w.Write(p)
+}
+
+func synchronizeWriter(w io.Writer) io.Writer {
+	if _, ok := w.(*synchronizedWriter); ok {
+		return w
+	}
+	return &synchronizedWriter{w: w}
+}
+
 // rootCmd represents the base command when called without any subcommands
 var rootCmd = &cobra.Command{
 	Use:   "golanggraph",
@@ -545,6 +568,7 @@ type serverOptions struct {
 // printed "Server started on host:port" -- so a failure to bind was announced
 // as a success. And the dev command's own flags were ignored (see devCmd).
 func runServer(ctx context.Context, out io.Writer, opts serverOptions) error {
+	out = synchronizeWriter(out)
 	if opts.Port <= 0 || opts.Port > 65535 {
 		return fmt.Errorf("invalid port %d", opts.Port)
 	}

pkg/server/server.go

@@ -70,6 +70,7 @@ type Server struct {
 	config   *ServerConfig
 	router   *mux.Router
 	server   *http.Server
+	serverMu sync.RWMutex
 	logger   *logrus.Logger
 	upgrader websocket.Upgrader
 
@@ -91,6 +92,12 @@ type Server struct {
 	// several clients can observe the same agent or graph at once.
 	wsConnections   map[string]map[*websocket.Conn]struct{}
 	wsConnectionsMu sync.RWMutex
+
+	// httpConnections records connection state so Stop can close sockets that
+	// have been accepted but have not yet reached a request. net/http otherwise
+	// gives StateNew connections a five-second grace period during Shutdown.
+	httpConnections   map[net.Conn]http.ConnState
+	httpConnectionsMu sync.Mutex
 }
 
 // NewServer creates a new server
@@ -104,12 +111,13 @@ func NewServer(config *ServerConfig) *Server {
 	}
 
 	server := &Server{
-		config:        config,
-		router:        mux.NewRouter(),
-		logger:        logrus.New(),
-		graphManager:  NewGraphManager(),
-		wsConnections: make(map[string]map[*websocket.Conn]struct{}),
-		startedAt:     time.Now(),
+		config:          config,
+		router:          mux.NewRouter(),
+		logger:          logrus.New(),
+		graphManager:    NewGraphManager(),
+		wsConnections:   make(map[string]map[*websocket.Conn]struct{}),
+		httpConnections: make(map[net.Conn]http.ConnState),
+		startedAt:       time.Now(),
 	}
 
 	// Reject WebSocket upgrades from origins the API does not allow. Accepting
@@ -136,6 +144,16 @@ func NewServer(config *ServerConfig) *Server {
 	return server
 }
 
+func (s *Server) trackHTTPConnection(conn net.Conn, state http.ConnState) {
+	s.httpConnectionsMu.Lock()
+	defer s.httpConnectionsMu.Unlock()
+	if state == http.StateClosed || state == http.StateHijacked {
+		delete(s.httpConnections, conn)
+		return
+	}
+	s.httpConnections[conn] = state
+}
+
 // SetCheckpointer attaches the checkpointer used to serve thread history.
 func (s *Server) SetCheckpointer(cp persistence.Checkpointer) {
 	s.checkpointer = cp
@@ -297,20 +315,24 @@ func (s *Server) setupRoutes() {
 
 // Start starts the server
 func (s *Server) Start() error {
-	s.server = &http.Server{
+	httpServer := &http.Server{
 		Addr:           fmt.Sprintf("%s:%d", s.config.Host, s.config.Port),
 		Handler:        s.router,
 		ReadTimeout:    s.config.ReadTimeout,
 		WriteTimeout:   s.config.WriteTimeout,
 		MaxHeaderBytes: s.config.MaxHeaderBytes,
+		ConnState:      s.trackHTTPConnection,
 	}
+	s.serverMu.Lock()
+	s.server = httpServer
+	s.serverMu.Unlock()
 
 	s.logger.WithFields(logrus.Fields{
 		"host": s.config.Host,
 		"port": s.config.Port,
 	}).Info("Starting GoLangGraph server")
 
-	if err := s.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+	if err := httpServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
 		return err
 	}
 	return nil
@@ -331,10 +353,28 @@ func (s *Server) Stop(ctx context.Context) error {
 	}
 	s.wsConnectionsMu.Unlock()
 
-	if s.server == nil {
+	// Shutdown intentionally waits for active requests. A connection that has
+	// not reached a request is not active work, but net/http treats StateNew as
+	// active for five seconds; close those sockets so a stop/restart is prompt.
+	s.httpConnectionsMu.Lock()
+	newConnections := make([]net.Conn, 0)
+	for conn, state := range s.httpConnections {
+		if state == http.StateNew {
+			newConnections = append(newConnections, conn)
+		}
+	}
+	s.httpConnectionsMu.Unlock()
+	for _, conn := range newConnections {
+		_ = conn.Close()
+	}
+
+	s.serverMu.RLock()
+	httpServer := s.server
+	s.serverMu.RUnlock()
+	if httpServer == nil {
 		return nil
 	}
-	return s.server.Shutdown(ctx)
+	return httpServer.Shutdown(ctx)
 }
 
 // Middleware