runServerWith gracefully exit by returning exit code 1

Author: soulomoon

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

@@ -149,9 +149,9 @@ runServerWith ioLogger logger clientIn clientOut serverDefinition = do
   cout <- atomically newTChan :: IO (TChan FromServerMessage)
   withAsync (sendServer ioLogger cout clientOut) $ \_sendAsync -> do
     let sendMsg = atomically . writeTChan cout
-    ioLoop ioLogger logger clientIn serverDefinition emptyVFS sendMsg
+    res <- ioLoop ioLogger logger clientIn serverDefinition emptyVFS sendMsg
     ioLogger <& ServerStopped `WithSeverity` Info
-    return 0
+    return res
 
 -- ---------------------------------------------------------------------
 
@@ -163,33 +163,37 @@ ioLoop ::
   ServerDefinition config ->
   VFS ->
   (FromServerMessage -> IO ()) ->
-  IO ()
+  IO Int
 ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = 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
           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 ()

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

@@ -211,6 +211,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 +755,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

@@ -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
@@ -440,6 +441,7 @@ 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 -> exitNotificationHandler logger msg
     SMethod_Shutdown -> handle' logger (Just $ \_ -> signalShutdown) m msg
      where
       -- See Note [Shutdown]
@@ -466,18 +468,9 @@ handle' ::
   m ()
 handle' logger mAction m msg = do
   shutdown <- isShuttingDown
-  -- These are the methods that we are allowed to process during shutdown.
-  -- The reason that we do not include 'shutdown' itself here is because
-  -- by the time we get the first 'shutdown' message, isShuttingDown will
-  -- still be false, so we would still be able to process it.
-  -- This ensures we won't process the second 'shutdown' message and only
-  -- process 'exit' during shutdown.
-  let allowedMethod m = case (splitClientMethod m, m) of
-        (IsClientNot, SMethod_Exit) -> True
-        _ -> False
 
   case mAction of
-    Just f | not shutdown || allowedMethod m -> f msg
+    Just f | not shutdown -> f msg
     _ -> pure ()
 
   dynReqHandlers <- getsState resRegistrationsReq
@@ -488,14 +481,12 @@ handle' logger mAction m msg = do
 
   case splitClientMethod m of
     -- See Note [Shutdown]
-    IsClientNot | shutdown, not (allowedMethod m) -> notificationDuringShutdown
+    IsClientNot | shutdown -> notificationDuringShutdown
     IsClientNot -> case pickHandler dynNotHandlers notHandlers of
       Just h -> liftIO $ h msg
-      Nothing
-        | SMethod_Exit <- m -> exitNotificationHandler logger msg
-        | otherwise -> missingNotificationHandler
+      Nothing | otherwise -> missingNotificationHandler
     -- See Note [Shutdown]
-    IsClientReq | shutdown, not (allowedMethod m) -> requestDuringShutdown msg
+    IsClientReq | shutdown -> requestDuringShutdown msg
     IsClientReq -> case pickHandler dynReqHandlers reqHandlers of
       Just h -> liftIO $ h msg (runLspT env . sendResponse msg)
       Nothing
@@ -556,10 +547,11 @@ progressCancelHandler logger (TNotificationMessage _ _ (WorkDoneProgressCancelPa
       logger <& ProgressCancel tid `WithSeverity` Debug
       liftIO cancelAction
 
-exitNotificationHandler :: (MonadIO m) => LogAction m (WithSeverity LspProcessingLog) -> Handler m Method_Exit
+exitNotificationHandler :: (MonadIO m, MonadLsp config0 m) => LogAction m (WithSeverity LspProcessingLog) -> Handler m Method_Exit
 exitNotificationHandler logger _ = do
-  logger <& Exiting `WithSeverity` Info
-  liftIO exitSuccess
+  logger <& ShuttingDown `WithSeverity` Info
+  b <- resExit . resState <$> getLspEnv
+  liftIO $ signalBarrier b ()
 
 -- | Default Shutdown handler
 shutdownRequestHandler :: Handler IO Method_Shutdown