diff --git a/tsc/internal/lsp/server.go b/tsc/internal/lsp/server.go index dcb020b0ca1b7..92ce9e24cf51e 100644 --- a/tsc/internal/lsp/server.go +++ b/tsc/internal/lsp/server.go @@ -113,6 +113,8 @@ type Writer interface { Write(msg *lsproto.Message) error } +var errExitBeforeShutdown = errors.New("exit notification received before shutdown") + type lspReader struct { r *lsproto.BaseReader } @@ -178,6 +180,7 @@ type Server struct { logger *logger initStarted atomic.Bool + shutdownRequested atomic.Bool clientSeq atomic.Int32 requestQueue *dynamicQueue[*lsproto.RequestMessage] outgoingQueue *dynamicQueue[*lsproto.Message] @@ -936,6 +939,13 @@ func (s *Server) readLoop(ctx context.Context) error { } continue } + if errors.Is(err, io.EOF) { + // Queue EOF after pending requests so transport closure cannot preempt dispatch. + if err := s.requestQueue.Put(ctx, nil); err != nil { + return err + } + return nil + } return err } @@ -1008,8 +1018,14 @@ func (s *Server) dispatchLoop(ctx context.Context) error { for { req, err := s.requestQueue.Get(ctx) if err != nil { + if cause := context.Cause(ctx); errors.Is(cause, io.EOF) || errors.Is(cause, errExitBeforeShutdown) { + return cause + } return err } + if req == nil { + return io.EOF + } s.lastRequestTimeMs.Store(time.Now().UnixMilli()) requestCtx := locale.WithLocale(ctx, s.GetLocale()) @@ -1029,8 +1045,8 @@ func (s *Server) dispatchLoop(ctx context.Context) error { if err := s.sendError(req.ID, lsproto.ErrorCodeRequestCancelled); err != nil { lspExit(err) } - } else if errors.Is(err, io.EOF) { - lspExit(nil) + } else if errors.Is(err, io.EOF) || errors.Is(err, errExitBeforeShutdown) { + lspExit(err) } else { if err := s.sendError(req.ID, err); err != nil { lspExit(err) @@ -1336,7 +1352,7 @@ var handlers = sync.OnceValue(func() handlerMap { func (handlers handlerMap) registerNotificationHandler[Req any](info lsproto.NotificationInfo[Req], fn func(*Server, context.Context, Req) error) { handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) (func() error, error) { - if s.session == nil && req.Method != lsproto.MethodInitialized { + if s.session == nil && req.Method != lsproto.MethodInitialized && req.Method != lsproto.MethodExit { return nil, lsproto.ErrorCodeServerNotInitialized } @@ -1834,10 +1850,14 @@ func (s *Server) handleShutdown(ctx context.Context, _ lsproto.NoParams, _ *lspr } s.closeAPISessions() s.session.Close() + s.shutdownRequested.Store(true) return lsproto.ShutdownResponse{}, nil } func (s *Server) handleExit(ctx context.Context, _ lsproto.NoParams) error { + if !s.shutdownRequested.Load() { + return errExitBeforeShutdown + } return io.EOF } diff --git a/tsc/internal/lsp/server_test.go b/tsc/internal/lsp/server_test.go index 826bb41985d6b..514938602e478 100644 --- a/tsc/internal/lsp/server_test.go +++ b/tsc/internal/lsp/server_test.go @@ -23,6 +23,16 @@ type shutdownTestWriter struct{} func (shutdownTestWriter) Write(*lsproto.Message) error { return nil } +type exitThenEOFTestReader struct{ read bool } + +func (r *exitThenEOFTestReader) Read() (*lsproto.Message, error) { + if !r.read { + r.read = true + return (&lsproto.RequestMessage{Method: lsproto.MethodExit}).Message(), nil + } + return nil, io.EOF +} + type cancelingTestWriter struct { writer Writer cancel context.CancelFunc @@ -72,6 +82,78 @@ func TestGetContentMapperWorkspaceRoots(t *testing.T) { }) } +func TestDispatchLoopReturnsEOFOnExit(t *testing.T) { + t.Parallel() + + server := NewServer(&ServerOptions{ + In: shutdownTestReader{}, + Out: shutdownTestWriter{}, + Err: io.Discard, + Cwd: "/test", + }) + server.session = &project.Session{} + server.shutdownRequested.Store(true) + if err := server.requestQueue.Put(t.Context(), &lsproto.RequestMessage{Method: lsproto.MethodExit}); err != nil { + t.Fatal(err) + } + + if err := server.dispatchLoop(t.Context()); !errors.Is(err, io.EOF) { + t.Fatalf("dispatchLoop() error = %v, want %v", err, io.EOF) + } +} + +func TestDispatchLoopReturnsErrorWhenExitBeforeShutdown(t *testing.T) { + t.Parallel() + + server := NewServer(&ServerOptions{ + In: shutdownTestReader{}, + Out: shutdownTestWriter{}, + Err: io.Discard, + Cwd: "/test", + }) + if err := server.requestQueue.Put(t.Context(), &lsproto.RequestMessage{Method: lsproto.MethodExit}); err != nil { + t.Fatal(err) + } + + if err := server.dispatchLoop(t.Context()); !errors.Is(err, errExitBeforeShutdown) { + t.Fatalf("dispatchLoop() error = %v, want %v", err, errExitBeforeShutdown) + } +} + +func TestRunExitsSuccessfullyAfterShutdown(t *testing.T) { + t.Parallel() + + reader := &exitThenEOFTestReader{} + server := NewServer(&ServerOptions{ + In: reader, + Out: shutdownTestWriter{}, + Err: io.Discard, + Cwd: "/test", + }) + server.session = &project.Session{} + server.shutdownRequested.Store(true) + + err := server.Run(t.Context()) + if err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } +} + +func TestRunReturnsErrorWhenExitPrecedesShutdown(t *testing.T) { + t.Parallel() + + server := NewServer(&ServerOptions{ + In: &exitThenEOFTestReader{}, + Out: shutdownTestWriter{}, + Err: io.Discard, + Cwd: "/test", + }) + + if err := server.Run(t.Context()); !errors.Is(err, errExitBeforeShutdown) { + t.Fatalf("Run() error = %v, want %v", err, errExitBeforeShutdown) + } +} + // TestServerShutdownNoDeadlock verifies that operations after shutdown // don't block. func TestServerShutdownNoDeadlock(t *testing.T) {