Stop doing action for handles and ignore responses after shutdown (#567)

* stop doing action before handler after shutdown * allowed method can run the before-action * fix * add log * swap log location * block the response call back * stop tryChangeConfig * Update lsp/src/Language/LSP/Server/Processing.hs Co-authored-by: Michael Peyton Jones <[email protected]> * revert barrior in tryChangeConfig * fix typo * revert tryChangeConfig * add documentation * Update lsp/src/Language/LSP/Server/Processing.hs Co-authored-by: Michael Peyton Jones <[email protected]> * refine shutdown blocking * refine shutdown blocking * fix format * fix format * fix doc * fix doc * Update Processing.hs --------- Co-authored-by: Michael Peyton Jones <[email protected]>
haskell · Apr 10, 2024 · eab156c · eab156c
1 parent 1e5940b
commit eab156c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/lsp/src/Language/LSP/Server/Core.hs b/lsp/src/Language/LSP/Server/Core.hs
@@ -1024,4 +1024,8 @@ be cancelled when we receive `shutdown`.
 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.
+
+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.
 -}
diff --git a/lsp/src/Language/LSP/Server/Processing.hs b/lsp/src/Language/LSP/Server/Processing.hs
@@ -92,6 +92,7 @@ instance Pretty LspProcessingLog where
 processMessage :: (m ~ LspM config) => LogAction m (WithSeverity LspProcessingLog) -> BSL.ByteString -> m ()
 processMessage logger jsonStr = do
   pendingResponsesVar <- LspT $ asks $ resPendingResponses . resState
+  shutdown <- isShuttingDown
   join $ liftIO $ atomically $ fmap handleErrors $ runExceptT $ do
     val <- except $ eitherDecode jsonStr
     pending <- lift $ readTVar pendingResponsesVar
@@ -100,8 +101,10 @@ processMessage logger jsonStr = do
       FromClientMess m mess ->
         pure $ handle logger m mess
       FromClientRsp (P.Pair (ServerResponseCallback f) (Const !newMap)) res -> do
+        -- see Note [Shutdown]
         writeTVar pendingResponsesVar newMap
-        pure $ liftIO $ f (res ^. L.result)
+        unless shutdown <$> do
+          pure $ liftIO $ f (res ^. L.result)
  where
   parser :: ResponseMap -> Value -> Parser (FromClientMessage' (P.Product ServerResponseCallback (Const ResponseMap)))
   parser rm = parseClientMessage $ \i ->
@@ -449,31 +452,37 @@ handle' ::
   TClientMessage meth ->
   m ()
 handle' logger mAction m msg = do
-  maybe (return ()) (\f -> f msg) mAction
+  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
+    _ -> pure ()
 
   dynReqHandlers <- getsState resRegistrationsReq
   dynNotHandlers <- getsState resRegistrationsNot
 
   env <- getLspEnv
   let Handlers{reqHandlers, notHandlers} = resHandlers env
-  shutdown <- isShuttingDown
 
   case splitClientMethod m of
     -- See Note [Shutdown]
     IsClientNot | shutdown, not (allowedMethod m) -> notificationDuringShutdown
-     where
-      allowedMethod SMethod_Exit = True
-      allowedMethod _ = False
     IsClientNot -> case pickHandler dynNotHandlers notHandlers of
       Just h -> liftIO $ h msg
       Nothing
         | SMethod_Exit <- m -> exitNotificationHandler logger msg
         | otherwise -> missingNotificationHandler
     -- See Note [Shutdown]
     IsClientReq | shutdown, not (allowedMethod m) -> requestDuringShutdown msg
-     where
-      allowedMethod SMethod_Shutdown = True
-      allowedMethod _ = False
     IsClientReq -> case pickHandler dynReqHandlers reqHandlers of
       Just h -> liftIO $ h msg (runLspT env . sendResponse msg)
       Nothing