revert some changes

Author: soulomoon

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

@@ -13,7 +13,7 @@ import Colog.Core (LogAction (..), Severity (..), WithSeverity (..), (<&))
 import Colog.Core qualified as L
 import Control.Applicative ((<|>))
 import Control.Concurrent (threadDelay)
-import Control.Concurrent.Async (cancel, race, wait, withAsync)
+import Control.Concurrent.Async (race, wait, withAsync)
 import Control.Concurrent.STM.TChan
 import Control.Exception (catchJust, throwIO)
 import Control.Monad.IO.Class

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

@@ -441,7 +441,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 -> exitNotificationHandler logger 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]
@@ -468,9 +474,18 @@ 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 -> f msg
+    Just f | not shutdown || allowedMethod m -> f msg
     _ -> pure ()
 
   dynReqHandlers <- getsState resRegistrationsReq
@@ -481,12 +496,14 @@ handle' logger mAction m msg = do
 
   case splitClientMethod m of
     -- See Note [Shutdown]
-    IsClientNot | shutdown -> notificationDuringShutdown
+    IsClientNot | shutdown, not (allowedMethod m) -> notificationDuringShutdown
     IsClientNot -> case pickHandler dynNotHandlers notHandlers of
       Just h -> liftIO $ h msg
-      Nothing | otherwise -> missingNotificationHandler
+      Nothing
+        | SMethod_Exit <- m -> exitNotificationHandler logger msg
+        | otherwise -> missingNotificationHandler
     -- See Note [Shutdown]
-    IsClientReq | shutdown -> requestDuringShutdown msg
+    IsClientReq | shutdown, not (allowedMethod m) -> requestDuringShutdown msg
     IsClientReq -> case pickHandler dynReqHandlers reqHandlers of
       Just h -> liftIO $ h msg (runLspT env . sendResponse msg)
       Nothing
@@ -547,11 +564,10 @@ progressCancelHandler logger (TNotificationMessage _ _ (WorkDoneProgressCancelPa
       logger <& ProgressCancel tid `WithSeverity` Debug
       liftIO cancelAction
 
-exitNotificationHandler :: (MonadIO m, MonadLsp config0 m) => LogAction m (WithSeverity LspProcessingLog) -> Handler m Method_Exit
-exitNotificationHandler logger _ = do
-  logger <& Exiting `WithSeverity` Info
-  b <- resExit . resState <$> getLspEnv
-  liftIO $ signalBarrier b ()
+exitNotificationHandler :: (MonadIO m) => LogAction m (WithSeverity LspProcessingLog) -> Handler m Method_Exit
+exitNotificationHandler _logger _ = do
+  -- default exit handler do nothing
+  return ()
 
 -- | Default Shutdown handler
 shutdownRequestHandler :: Handler IO Method_Shutdown