Gracefully exit the server

.github/workflows/format.yaml

@@ -1,6 +1,6 @@
 name: Format
 
-on: 
+on:
   push:
     branches:
       - master
@@ -12,7 +12,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v3
-      - uses: cachix/install-nix-action@v22
+      - uses: cachix/install-nix-action@v31
         with:
           nix_path: nixpkgs=channel:nixos-unstable
       - run: nix-shell --run "fourmolu -m check ."

.github/workflows/nix.yaml

@@ -1,6 +1,6 @@
 name: Nix
 
-on: 
+on:
   push:
     branches:
       - master
@@ -17,7 +17,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v3
-      - uses: cachix/install-nix-action@v22
+      - uses: cachix/install-nix-action@v31
         with:
           nix_path: nixpkgs=channel:nixos-unstable
       - run: nix-shell --run "cabal update && cabal build all"

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,4 +1,5 @@
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module Language.LSP.Server.Control (
   -- * Running
@@ -11,9 +12,9 @@ module Language.LSP.Server.Control (
 import Colog.Core (LogAction (..), Severity (..), WithSeverity (..), (<&))
 import Colog.Core qualified as L
 import Control.Applicative ((<|>))
-import Control.Concurrent
+import Control.Concurrent.Async (wait, withAsync)
 import Control.Concurrent.STM.TChan
-import Control.Monad
+import Control.Exception (catchJust, throwIO)
 import Control.Monad.IO.Class
 import Control.Monad.STM
 import Data.Aeson qualified as J
@@ -33,18 +34,22 @@ import Language.LSP.Server.Processing qualified as Processing
 import Language.LSP.VFS
 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
   deriving (Show)
 
 instance Pretty LspServerLog where
+  pretty ServerStopped = "Server stopped"
   pretty (LspProcessingLog l) = pretty l
   pretty (DecodeInitializeError err) =
     vsep
@@ -57,7 +62,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
 
@@ -108,9 +118,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
 
@@ -130,15 +146,12 @@ runServerWith ::
   IO Int -- exit code
 runServerWith ioLogger logger clientIn clientOut serverDefinition = do
   ioLogger <& Starting `WithSeverity` Info
-
-  cout <- atomically newTChan :: IO (TChan J.Value)
-  _rhpid <- forkIO $ sendServer ioLogger cout clientOut
-
-  let sendMsg msg = atomically $ writeTChan cout $ J.toJSON msg
-
-  ioLoop ioLogger logger clientIn serverDefinition emptyVFS sendMsg
-
-  return 1
+  cout <- atomically newTChan :: IO (TChan FromServerMessage)
+  withAsync (sendServer ioLogger cout clientOut) $ \_sendAsync -> do
+    let sendMsg = atomically . writeTChan cout
+    res <- ioLoop ioLogger logger clientIn serverDefinition emptyVFS sendMsg (wait _sendAsync)
+    ioLogger <& ServerStopped `WithSeverity` Info
+    return res
 
 -- ---------------------------------------------------------------------
 
@@ -150,33 +163,40 @@ ioLoop ::
   ServerDefinition config ->
   VFS ->
   (FromServerMessage -> IO ()) ->
-  IO ()
-ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = do
+  IO () ->
+  IO Int
+ioLoop ioLogger logger clientIn 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.ByteString -> LspM config ()
+  loop :: Result BS.ByteString -> LspM config Int
   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)
 
   parser = do
     try contentType <|> return ()
@@ -224,9 +244,10 @@ parseOne logger clientIn = go
 -- ---------------------------------------------------------------------
 
 -- | Simple server to make sure all output is serialised
-sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.ByteString -> IO ()) -> IO ()
-sendServer _logger msgChan clientOut = do
-  forever $ do
+sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan FromServerMessage -> (BSL.ByteString -> IO ()) -> IO ()
+sendServer _logger msgChan clientOut = go
+ where
+  go = do
     msg <- atomically $ readTChan msgChan
 
     -- We need to make sure we only send over the content of the message,
@@ -241,6 +262,10 @@ sendServer _logger msgChan clientOut = do
             ]
 
     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,7 @@ 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 ())
   }
 
 -- ---------------------------------------------------------------------
@@ -211,6 +212,7 @@ data LanguageContextState config = LanguageContextState
   , resRegistrationsReq :: !(TVar (RegistrationMap Request))
   , resLspId :: !(TVar Int32)
   , resShutdown :: !(C.Barrier ())
+  , resExit :: !(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
@@ -754,6 +756,14 @@ isShuttingDown = do
     Just _ -> True
     Nothing -> False
 
+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

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

@@ -59,7 +59,7 @@ 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
+import UnliftIO (race)
 
 data LspProcessingLog
   = VfsLog VfsLog
@@ -69,13 +69,19 @@ data LspProcessingLog
   | ProgressCancel ProgressToken
   | forall m. MessageDuringShutdown (SClientMethod m)
   | ShuttingDown
+  | ShuttingDownSenderTimeout !Int
   | Exiting
 
 deriving instance Show LspProcessingLog
 
 instance Pretty LspProcessingLog where
   pretty (VfsLog l) = pretty l
   pretty (LspCore l) = pretty l
+  pretty (ShuttingDownSenderTimeout sec) =
+    vsep
+      [ "LSP: sender did not shut down within" <+> pretty sec <+> "seconds"
+      , "This likely indicates a broken pipe, and the server should exit now"
+      ]
   pretty (MessageProcessingError bs err) =
     vsep
       [ "LSP: incoming message parse error:"
@@ -119,9 +125,10 @@ initializeRequestHandler ::
   ServerDefinition config ->
   VFS ->
   (FromServerMessage -> IO ()) ->
+  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 +179,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 +195,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 +449,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 +513,17 @@ 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)
+        -- wait for the sender to finish before we return
+        -- todo magic number
+        let waitTime = 3 -- seconds
+        r <- liftIO $ race waitSender (threadDelay $ waitTime * 1_000_000)
+        case r of
+          Left _ -> pure ()
+          Right _ -> logger <& ShuttingDownSenderTimeout waitTime `WithSeverity` Warning
+        liftIO waitSender
       Just h -> liftIO $ h msg (runLspT env . sendResponse msg)
       Nothing
         | SMethod_Shutdown <- m -> liftIO $ shutdownRequestHandler msg (runLspT env . sendResponse msg)
@@ -557,9 +584,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