Gracefully exit the server (#622)

* 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

---------

Co-authored-by: fendor <fendor@users.noreply.github.com>

Author: soulomoon

.github/workflows/format.yaml

@@ -1,6 +1,6 @@
 name: Format
 
-on: 
+on:
   push:
     branches:
       - master

.github/workflows/nix.yaml

@@ -1,6 +1,6 @@
 name: Nix
 
-on: 
+on:
   push:
     branches:
       - master

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
 

lsp/src/Language/LSP/Server/Control.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module Language.LSP.Server.Control (
   -- * Running
@@ -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, finally, throwIO)
 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

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,17 @@ 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 +857,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`.
 -}

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,11 @@ 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 () ->
   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 +173,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 +189,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 +443,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 +507,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 +571,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