From 7dd1bfa797b462791dc398690f599a2c5e5d1962 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Mon, 16 Mar 2026 04:21:55 +0200 Subject: Track fire-and-forget goroutines with sync.WaitGroup for clean shutdown Adds inflight WaitGroup to Server and wraps inline-prompt, chat-response, and deferShowDocument goroutines. Run() waits for all in-flight work before returning. Co-Authored-By: Claude Opus 4.6 --- internal/lsp/handlers_document.go | 14 ++++++++++++-- internal/lsp/server.go | 7 ++++++- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/internal/lsp/handlers_document.go b/internal/lsp/handlers_document.go index a411f4e..f69c626 100644 --- a/internal/lsp/handlers_document.go +++ b/internal/lsp/handlers_document.go @@ -124,7 +124,11 @@ func (s *Server) maybeRunInlinePrompt(uri string, lineIdx int, raw string, openS } if s.currentLLMClient() != nil { pos := Position{Line: lineIdx, Character: len(raw)} - go s.runInlinePrompt(uri, pos) + s.inflight.Add(1) + go func() { + defer s.inflight.Done() + s.runInlinePrompt(uri, pos) + }() } return true } @@ -193,7 +197,11 @@ func (s *Server) handleChatPrompt(uri string, lineIdx int, match chatPromptLine) } return } - go s.requestChatResponse(uri, lineIdx, match) + s.inflight.Add(1) + go func() { + defer s.inflight.Done() + s.requestChatResponse(uri, lineIdx, match) + }() } func (s *Server) requestChatResponse(uri string, lineIdx int, match chatPromptLine) { @@ -432,7 +440,9 @@ func (s *Server) clientShowDocument(uri string, sel *Range) { // The goroutine respects s.serverCtx so it won't write after shutdown. func (s *Server) deferShowDocument(uri string, sel Range) { ctx := s.serverCtx + s.inflight.Add(1) go func() { + defer s.inflight.Done() if ctx == nil { time.Sleep(120 * time.Millisecond) s.clientShowDocument(uri, &sel) diff --git a/internal/lsp/server.go b/internal/lsp/server.go index 6442342..c039255 100644 --- a/internal/lsp/server.go +++ b/internal/lsp/server.go @@ -31,6 +31,7 @@ type Server struct { serverCancel context.CancelFunc statusSink StatusSink exited atomic.Bool + inflight sync.WaitGroup // tracks background goroutines (inline prompt, chat, etc.) mu sync.RWMutex docs map[string]*document logContext bool @@ -332,8 +333,12 @@ func (s *Server) emitGlobalStatus(reqs int64, rpm float64, sent int64, recv int6 } // Run starts the server's main loop, reading and dispatching LSP messages until EOF or exit. +// On shutdown it cancels the server context and waits for in-flight goroutines. func (s *Server) Run() error { - defer s.cancelRequests() + defer func() { + s.cancelRequests() + s.inflight.Wait() + }() for { body, err := s.readMessage() if errors.Is(err, io.EOF) { -- cgit v1.2.3