From df52cc025d15b2a9f68cc55f3d3eba822e541910 Mon Sep 17 00:00:00 2001 From: soulomoon Date: Tue, 25 Nov 2025 17:20:06 +0800 Subject: [PATCH 1/4] Ensure graceful LSP server shutdown and robust IO handling - Track shutdown/exit separately and wait for sender to flush before exit - Make runServerWithHandles resilient to broken pipes and log truncated messages - Have runServerWith/ioLoop return proper exit codes based on shutdown state - Adjust processing to use sender wait action during shutdown and log exit - Fix functional test harness to assert zero exit code instead of ExitSuccess - Tidy minor CI workflow whitespace in format.yaml and nix.yaml --- .github/workflows/format.yaml | 2 +- .github/workflows/nix.yaml | 2 +- lsp-test/func-test/FuncTest.hs | 8 ++- lsp/src/Language/LSP/Server/Control.hs | 84 +++++++++++++++-------- lsp/src/Language/LSP/Server/Core.hs | 57 +++++++++++---- lsp/src/Language/LSP/Server/Processing.hs | 23 +++++-- 6 files changed, 122 insertions(+), 54 deletions(-) diff --git a/.github/workflows/format.yaml b/.github/workflows/format.yaml index bf7370fc..85300e05 100644 --- a/.github/workflows/format.yaml +++ b/.github/workflows/format.yaml @@ -1,6 +1,6 @@ name: Format -on: +on: push: branches: - master diff --git a/.github/workflows/nix.yaml b/.github/workflows/nix.yaml index 02f6f6f5..798614c8 100644 --- a/.github/workflows/nix.yaml +++ b/.github/workflows/nix.yaml @@ -1,6 +1,6 @@ name: Nix -on: +on: push: branches: - master diff --git a/lsp-test/func-test/FuncTest.hs b/lsp-test/func-test/FuncTest.hs index 818b1696..ff805de6 100644 --- a/lsp-test/func-test/FuncTest.hs +++ b/lsp-test/func-test/FuncTest.hs @@ -39,13 +39,15 @@ runSessionWithServer logger defn testConfig caps root session = do (hinRead, hinWrite) <- createPipe (houtRead, houtWrite) <- createPipe - server <- async $ void $ runServerWithHandles logger (L.hoistLogAction liftIO logger) hinRead houtWrite defn + server <- async $ runServerWithHandles logger (L.hoistLogAction liftIO logger) hinRead houtWrite defn res <- Test.runSessionWithHandles hinWrite houtRead testConfig caps root session timeout 3000000 $ do - Left (fromException -> Just ExitSuccess) <- waitCatch server - pure () + return_code <- wait server + case return_code of + 0 -> pure () + _ -> error $ "Server exited with non-zero code: " ++ show return_code pure res diff --git a/lsp/src/Language/LSP/Server/Control.hs b/lsp/src/Language/LSP/Server/Control.hs index 3a7041c9..51eed2c3 100644 --- a/lsp/src/Language/LSP/Server/Control.hs +++ b/lsp/src/Language/LSP/Server/Control.hs @@ -1,5 +1,6 @@ {-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE RecordWildCards #-} module Language.LSP.Server.Control ( @@ -29,7 +30,7 @@ import Control.Applicative ((<|>)) import Control.Concurrent import Control.Concurrent.Async import Control.Concurrent.STM.TChan -import Control.Exception (finally) +import Control.Exception (catchJust, throwIO, finally) import Control.Monad import Control.Monad.IO.Class import Control.Monad.STM @@ -51,19 +52,23 @@ import Language.LSP.VFS import Network.WebSockets qualified as WS import Prettyprinter import System.IO +import System.IO.Error (isResourceVanishedError) data LspServerLog = LspProcessingLog Processing.LspProcessingLog | DecodeInitializeError String | HeaderParseFail [String] String | EOF + | BrokenPipeWhileSending TL.Text -- truncated outgoing message (including header) | Starting + | ServerStopped | ParsedMsg T.Text | SendMsg TL.Text | WebsocketLog WebsocketLog deriving (Show) instance Pretty LspServerLog where + pretty ServerStopped = "Server stopped" pretty (LspProcessingLog l) = pretty l pretty (DecodeInitializeError err) = vsep @@ -76,7 +81,12 @@ instance Pretty LspServerLog where , pretty (intercalate " > " ctxs) <> ": " <+> pretty err ] pretty EOF = "Got EOF" - pretty Starting = "Starting server" + pretty (BrokenPipeWhileSending msg) = + vsep + [ "Broken pipe while sending (client likely closed output handle):" + , indent 2 (pretty msg) + ] + pretty Starting = "Server starting" pretty (ParsedMsg msg) = "---> " <> pretty msg pretty (SendMsg msg) = "<--2-- " <> pretty msg pretty (WebsocketLog msg) = "Websocket:" <+> pretty msg @@ -130,9 +140,15 @@ runServerWithHandles ioLogger logger hin hout serverDefinition = do let clientIn = BS.hGetSome hin defaultChunkSize - clientOut out = do - BSL.hPut hout out - hFlush hout + clientOut out = + catchJust + (\e -> if isResourceVanishedError e then Just e else Nothing) + (BSL.hPut hout out >> hFlush hout) + ( \e -> do + let txt = TL.toStrict $ TL.take 400 $ TL.decodeUtf8 out -- limit size + ioLogger <& BrokenPipeWhileSending (TL.fromStrict txt) `WithSeverity` Error + throwIO e + ) runServerWith ioLogger logger clientIn clientOut serverDefinition @@ -179,15 +195,12 @@ runServerWithConfig :: IO Int runServerWithConfig ServerConfig{..} serverDefinition = do ioLogger <& Starting `WithSeverity` Info - - cout <- atomically newTChan :: IO (TChan J.Value) - _rhpid <- forkIO $ sendServer ioLogger cout outwards prepareOutwards - - let sendMsg msg = atomically $ writeTChan cout $ J.toJSON msg - - ioLoop ioLogger lspLogger inwards parseInwards serverDefinition emptyVFS sendMsg - - return 1 + cout <- atomically newTChan :: IO (TChan FromServerMessage) + withAsync (sendServer ioLogger cout outwards prepareOutwards) $ \sendServerAsync -> do + let sendMsg = atomically . writeTChan cout + res <- ioLoop ioLogger lspLogger inwards parseInwards serverDefinition emptyVFS sendMsg (wait sendServerAsync) + ioLogger <& ServerStopped `WithSeverity` Info + return res -- --------------------------------------------------------------------- @@ -200,33 +213,39 @@ ioLoop :: ServerDefinition config -> VFS -> (FromServerMessage -> IO ()) -> - IO () -ioLoop ioLogger logger clientIn parser serverDefinition vfs sendMsg = do + IO () -> + IO Int +ioLoop ioLogger logger clientIn parser serverDefinition vfs sendMsg waitSenderFinish = do minitialize <- parseOne ioLogger clientIn (parse parser "") case minitialize of - Nothing -> pure () + Nothing -> pure 1 Just (msg, remainder) -> do case J.eitherDecode $ BSL.fromStrict msg of - Left err -> ioLogger <& DecodeInitializeError err `WithSeverity` Error + Left err -> do + ioLogger <& DecodeInitializeError err `WithSeverity` Error + return 1 Right initialize -> do - mInitResp <- Processing.initializeRequestHandler pioLogger serverDefinition vfs sendMsg initialize + mInitResp <- Processing.initializeRequestHandler pioLogger serverDefinition vfs sendMsg waitSenderFinish initialize case mInitResp of - Nothing -> pure () + Nothing -> pure 1 Just env -> runLspT env $ loop (parse parser remainder) where pioLogger = L.cmap (fmap LspProcessingLog) ioLogger pLogger = L.cmap (fmap LspProcessingLog) logger - loop :: Result BS.StrictByteString -> LspM config () loop = go where go r = do - res <- parseOne logger clientIn r - case res of - Nothing -> pure () - Just (msg, remainder) -> do - Processing.processMessage pLogger $ BSL.fromStrict msg - go (parse parser remainder) + b <- isExiting + if b + then pure 0 + else do + res <- parseOne logger clientIn r + case res of + Nothing -> pure 1 + Just (msg, remainder) -> do + Processing.processMessage pLogger $ BSL.fromStrict msg + go (parse parser remainder) parseOne :: MonadIO m => @@ -366,9 +385,10 @@ withWebsocketRunServer wsConf withLspDefinition ioLogger lspLogger = -- --------------------------------------------------------------------- -- | Simple server to make sure all output is serialised -sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.LazyByteString -> IO ()) -> (BSL.LazyByteString -> BSL.LazyByteString) -> IO () -sendServer _logger msgChan clientOut prepareMessage = do - forever $ do +sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan FromServerMessage -> (BSL.LazyByteString -> IO ()) -> (BSL.LazyByteString -> BSL.LazyByteString) -> IO () +sendServer _logger msgChan clientOut prepareMessage = go + where + go = do msg <- atomically $ readTChan msgChan -- We need to make sure we only send over the content of the message, @@ -377,6 +397,10 @@ sendServer _logger msgChan clientOut prepareMessage = do let out = prepareMessage str clientOut out + -- close the client sender when we send out the shutdown request's response + case msg of + FromServerRsp SMethod_Shutdown _ -> pure () + _ -> go -- TODO: figure out how to re-enable -- This can lead to infinite recursion in logging, see https://github.com/haskell/lsp/issues/447 diff --git a/lsp/src/Language/LSP/Server/Core.hs b/lsp/src/Language/LSP/Server/Core.hs index 9c57992c..ccc2476a 100644 --- a/lsp/src/Language/LSP/Server/Core.hs +++ b/lsp/src/Language/LSP/Server/Core.hs @@ -137,6 +137,9 @@ data LanguageContextEnv config = LanguageContextEnv -- ^ The delay before starting a progress reporting session, in microseconds , resProgressUpdateDelay :: Int -- ^ The delay between sending progress updates, in microseconds + , resWaitSender :: !(IO ()) + -- ^ An IO action that waits for the sender thread to finish sending all pending messages. + -- This is used to ensure all responses are sent before the server exits. See Note [Shutdown] } -- --------------------------------------------------------------------- @@ -211,9 +214,9 @@ data LanguageContextState config = LanguageContextState , resRegistrationsReq :: !(TVar (RegistrationMap Request)) , resLspId :: !(TVar Int32) , resShutdown :: !(C.Barrier ()) - -- ^ Has the server received 'shutdown'? Can be used to conveniently trigger e.g. thread termination, - -- but if you need a cleanup action to terminate before exiting, then you should install a full - -- 'shutdown' handler + -- ^ Barrier signaled when the server receives the 'shutdown' request. See Note [Shutdown] + , resExit :: !(C.Barrier ()) + -- ^ Barrier signaled when the server receives the 'exit' notification. See Note [Shutdown] } type ResponseMap = IxMap LspId (Product SMethod ServerResponseCallback) @@ -754,6 +757,16 @@ isShuttingDown = do Just _ -> True Nothing -> False +-- | Check if the server has received the 'exit' notification. +-- See Note [Shutdown] +isExiting :: (m ~ LspM config) => m Bool +isExiting = do + b <- resExit . resState <$> getLspEnv + r <- liftIO $ C.waitBarrierMaybe b + pure $ case r of + Just () -> True + Nothing -> False + -- | Blocks until the server receives a 'shutdown' request. waitShuttingDown :: (m ~ LspM config) => m () waitShuttingDown = do @@ -843,18 +856,34 @@ The client can still always choose to cancel for another reason. -} {- Note [Shutdown] -The 'shutdown' request basically tells the server to clean up and stop doing things. -In particular, it allows us to ignore or reject all further messages apart from 'exit'. +~~~~~~~~~~~~~~~~~~ +The LSP protocol has a two-phase shutdown sequence: + +1. `shutdown` request: ask the server to stop doing work and finish + any in-flight operations. +2. `exit` notification: tell the server to terminate the process. + +We expose two `Barrier`s to track this state: + +- `resShutdown`: signalled when we receive the `shutdown` request. + Use `isShuttingDown` to check this. +- `resExit`: signalled when we receive the `exit` notification. + Use `isExiting` to check this. + +Shutdown is itself a request, and we assume the client will not send +`exit` before `shutdown`. If you want to be sure that some cleanup has +run before the server exits, make that cleanup part of your customize +`shutdown` handler. -We also provide a `Barrier` that indicates whether or not we are shutdown, this can -be convenient, e.g. you can race a thread against `waitBarrier` to have it automatically -be cancelled when we receive `shutdown`. +We use a dedicated sender thread to serialise all messages that go to +the client. That thread is set up to stop sending messages after the +`shutdown` response has been sent. -Shutdown is a request, and the client won't send `exit` until a server responds, so if you -want to be sure that some cleanup happens, you need to ensure we don't respond to `shutdown` -until it's done. The best way to do this is just to install a specific `shutdown` handler. +While handling the `shutdown` request we call `resWaitSender` to wait +for the sender thread to flush and finish. Otherwise, we might get a +"broken pipe" error from trying to send messages after the client has +closed our output handle. -After the `shutdown` request, we don't handle any more requests and notifications other than -`exit`. We also don't handle any more responses to requests we have sent but just throw the -responses away. +After the `shutdown` request has been processed, we do not handle any +more requests or notifications except for `exit`. -} diff --git a/lsp/src/Language/LSP/Server/Processing.hs b/lsp/src/Language/LSP/Server/Processing.hs index 3d744ba5..049eeee8 100644 --- a/lsp/src/Language/LSP/Server/Processing.hs +++ b/lsp/src/Language/LSP/Server/Processing.hs @@ -59,7 +59,6 @@ import Language.LSP.Protocol.Utils.SMethodMap qualified as SMethodMap import Language.LSP.Server.Core import Language.LSP.VFS as VFS import Prettyprinter -import System.Exit data LspProcessingLog = VfsLog VfsLog @@ -119,9 +118,10 @@ initializeRequestHandler :: ServerDefinition config -> VFS -> (FromServerMessage -> IO ()) -> + IO () -> -- ^ Action that waits for the sender thread to finish. We use it to set LanguageContextEnv.resWaitSender, See Note [Shutdown]. TMessage Method_Initialize -> IO (Maybe (LanguageContextEnv config)) -initializeRequestHandler logger ServerDefinition{..} vfs sendFunc req = do +initializeRequestHandler logger ServerDefinition{..} vfs sendFunc waitSender req = do let sendResp = sendFunc . FromServerRsp SMethod_Initialize handleErr (Left err) = do sendResp $ makeResponseError (req ^. L.id) err @@ -172,6 +172,7 @@ initializeRequestHandler logger ServerDefinition{..} vfs sendFunc req = do resRegistrationsReq <- newTVarIO mempty resLspId <- newTVarIO 0 resShutdown <- C.newBarrier + resExit <- C.newBarrier pure LanguageContextState{..} -- Call the 'duringInitialization' callback to let the server kick stuff up @@ -187,6 +188,7 @@ initializeRequestHandler logger ServerDefinition{..} vfs sendFunc req = do rootDir (optProgressStartDelay options) (optProgressUpdateDelay options) + waitSender configChanger config = forward interpreter (onConfigChange config) handlers = transmuteHandlers interpreter (staticHandlers clientCaps) interpreter = interpretHandler initializationResult @@ -440,6 +442,13 @@ handle logger m msg = SMethod_WorkspaceDidChangeConfiguration -> handle' logger (Just $ handleDidChangeConfiguration logger) m msg -- See Note [LSP configuration] SMethod_Initialized -> handle' logger (Just $ \_ -> initialDynamicRegistrations logger >> requestConfigUpdate (cmap (fmap LspCore) logger)) m msg + SMethod_Exit -> handle' logger (Just $ \_ -> signalExit) m msg + where + signalExit :: LspM config () + signalExit = do + logger <& Exiting `WithSeverity` Info + b <- resExit . resState <$> getLspEnv + liftIO $ signalBarrier b () SMethod_Shutdown -> handle' logger (Just $ \_ -> signalShutdown) m msg where -- See Note [Shutdown] @@ -497,6 +506,10 @@ handle' logger mAction m msg = do -- See Note [Shutdown] IsClientReq | shutdown, not (allowedMethod m) -> requestDuringShutdown msg IsClientReq -> case pickHandler dynReqHandlers reqHandlers of + Just h | SMethod_Shutdown <- m -> do + waitSender <- resWaitSender <$> getLspEnv + liftIO $ h msg (runLspT env . sendResponse msg) + liftIO waitSender Just h -> liftIO $ h msg (runLspT env . sendResponse msg) Nothing | SMethod_Shutdown <- m -> liftIO $ shutdownRequestHandler msg (runLspT env . sendResponse msg) @@ -557,9 +570,9 @@ progressCancelHandler logger (TNotificationMessage _ _ (WorkDoneProgressCancelPa liftIO cancelAction exitNotificationHandler :: (MonadIO m) => LogAction m (WithSeverity LspProcessingLog) -> Handler m Method_Exit -exitNotificationHandler logger _ = do - logger <& Exiting `WithSeverity` Info - liftIO exitSuccess +exitNotificationHandler _logger _ = do + -- default exit handler do nothing + return () -- | Default Shutdown handler shutdownRequestHandler :: Handler IO Method_Shutdown From 6d3bacaa89e126a815c9ccc762a63bc93701fcb2 Mon Sep 17 00:00:00 2001 From: soulomoon Date: Tue, 25 Nov 2025 17:32:49 +0800 Subject: [PATCH 2/4] format --- lsp/src/Language/LSP/Server/Control.hs | 4 ++-- lsp/src/Language/LSP/Server/Core.hs | 5 +++-- lsp/src/Language/LSP/Server/Processing.hs | 3 ++- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/lsp/src/Language/LSP/Server/Control.hs b/lsp/src/Language/LSP/Server/Control.hs index 51eed2c3..792d52f1 100644 --- a/lsp/src/Language/LSP/Server/Control.hs +++ b/lsp/src/Language/LSP/Server/Control.hs @@ -1,7 +1,7 @@ {-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeFamilies #-} module Language.LSP.Server.Control ( -- * Running @@ -30,7 +30,7 @@ import Control.Applicative ((<|>)) import Control.Concurrent import Control.Concurrent.Async import Control.Concurrent.STM.TChan -import Control.Exception (catchJust, throwIO, finally) +import Control.Exception (catchJust, finally, throwIO) import Control.Monad import Control.Monad.IO.Class import Control.Monad.STM diff --git a/lsp/src/Language/LSP/Server/Core.hs b/lsp/src/Language/LSP/Server/Core.hs index ccc2476a..c705a26a 100644 --- a/lsp/src/Language/LSP/Server/Core.hs +++ b/lsp/src/Language/LSP/Server/Core.hs @@ -757,8 +757,9 @@ isShuttingDown = do Just _ -> True Nothing -> False --- | Check if the server has received the 'exit' notification. --- See Note [Shutdown] +{- | Check if the server has received the 'exit' notification. +See Note [Shutdown] +-} isExiting :: (m ~ LspM config) => m Bool isExiting = do b <- resExit . resState <$> getLspEnv diff --git a/lsp/src/Language/LSP/Server/Processing.hs b/lsp/src/Language/LSP/Server/Processing.hs index 049eeee8..2005d91b 100644 --- a/lsp/src/Language/LSP/Server/Processing.hs +++ b/lsp/src/Language/LSP/Server/Processing.hs @@ -118,7 +118,8 @@ initializeRequestHandler :: ServerDefinition config -> VFS -> (FromServerMessage -> IO ()) -> - IO () -> -- ^ Action that waits for the sender thread to finish. We use it to set LanguageContextEnv.resWaitSender, See Note [Shutdown]. + IO () -> + -- ^ Action that waits for the sender thread to finish. We use it to set LanguageContextEnv.resWaitSender, See Note [Shutdown]. TMessage Method_Initialize -> IO (Maybe (LanguageContextEnv config)) initializeRequestHandler logger ServerDefinition{..} vfs sendFunc waitSender req = do From ca6d2e4ae122288286c8bcc539f8d8b4e9736c94 Mon Sep 17 00:00:00 2001 From: soulomoon Date: Tue, 25 Nov 2025 18:06:54 +0800 Subject: [PATCH 3/4] format --- lsp/src/Language/LSP/Server/Processing.hs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lsp/src/Language/LSP/Server/Processing.hs b/lsp/src/Language/LSP/Server/Processing.hs index 2005d91b..963997cc 100644 --- a/lsp/src/Language/LSP/Server/Processing.hs +++ b/lsp/src/Language/LSP/Server/Processing.hs @@ -118,8 +118,8 @@ initializeRequestHandler :: ServerDefinition config -> VFS -> (FromServerMessage -> IO ()) -> + -- | Action that waits for the sender thread to finish. We use it to set LanguageContextEnv.resWaitSender, See Note [Shutdown]. IO () -> - -- ^ Action that waits for the sender thread to finish. We use it to set LanguageContextEnv.resWaitSender, See Note [Shutdown]. TMessage Method_Initialize -> IO (Maybe (LanguageContextEnv config)) initializeRequestHandler logger ServerDefinition{..} vfs sendFunc waitSender req = do From 3883db0db9da457d4be3ebf6a6e734dabb387e27 Mon Sep 17 00:00:00 2001 From: patrick Date: Tue, 25 Nov 2025 18:56:38 +0800 Subject: [PATCH 4/4] Apply suggestion from @fendor Co-authored-by: fendor --- lsp/src/Language/LSP/Server/Processing.hs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lsp/src/Language/LSP/Server/Processing.hs b/lsp/src/Language/LSP/Server/Processing.hs index 963997cc..b78cc921 100644 --- a/lsp/src/Language/LSP/Server/Processing.hs +++ b/lsp/src/Language/LSP/Server/Processing.hs @@ -118,7 +118,7 @@ initializeRequestHandler :: ServerDefinition config -> VFS -> (FromServerMessage -> IO ()) -> - -- | Action that waits for the sender thread to finish. We use it to set LanguageContextEnv.resWaitSender, See Note [Shutdown]. + -- | Action that waits for the sender thread to finish. We use it to set 'LanguageContextEnv.resWaitSender', See Note [Shutdown]. IO () -> TMessage Method_Initialize -> IO (Maybe (LanguageContextEnv config))