Reject messages that come in after 'shutdown' (#564)

* Reject messages that come in after 'shutdown'

This is mandated by the spec so we can do it. We also expose the
shutdown barrier, which I think can be convenient. For more
sophisticated usecases people should just install a proper shutdown
handler.

I tried to write some tests for this, but I gave up because it requires
some singificant surgery on `lsp-test`, which stops recording messages
when it receives `shutdown` :(

* Fix formatting

Author: michaelpj

lsp/ChangeLog.md

@@ -1,5 +1,11 @@
 # Revision history for lsp
 
+## Unreleased
+
+- The server will now reject messages sent after `shutdown` has been received.
+- There is a `shutdownBarrier` member in the server state which can be used to
+  conveniently run actions when shutdown is triggered.
+
 ## 2.4.0.0
 
 - Server-created progress now will not send reports until and unless the client 

lsp/lsp.cabal

@@ -59,6 +59,7 @@ library
     , data-default          ^>=0.7
     , directory             ^>=1.3
     , exceptions            ^>=0.10
+    , extra                 ^>=1.7
     , filepath              >=1.4 && < 1.6
     , hashable              ^>=1.4
     , lens                  >=5.1    && <5.3

lsp/src/Language/LSP/Server.hs

@@ -35,6 +35,10 @@ module Language.LSP.Server (
   requestConfigUpdate,
   tryChangeConfig,
 
+  -- * Shutdown
+  isShuttingDown,
+  waitShuttingDown,
+
   -- * VFS
   getVirtualFile,
   getVirtualFiles,

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

@@ -21,6 +21,7 @@ import Colog.Core (
  )
 import Control.Applicative
 import Control.Concurrent.Async
+import Control.Concurrent.Extra as C
 import Control.Concurrent.MVar
 import Control.Concurrent.STM
 import Control.Exception qualified as E
@@ -211,6 +212,10 @@ data LanguageContextState config = LanguageContextState
   , resRegistrationsNot :: !(TVar (RegistrationMap Notification))
   , 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
   }
 
 type ResponseMap = IxMap LspId (Product SMethod ServerResponseCallback)
@@ -903,6 +908,25 @@ requestConfigUpdate logger = do
         Left err -> logger <& BadConfigurationResponse err `WithSeverity` Error
     else logger <& ConfigurationNotSupported `WithSeverity` Debug
 
+--------------------------------------------------------------------------------
+-- CONFIG
+--------------------------------------------------------------------------------
+
+-- | Checks if the server has received a 'shutdown' request.
+isShuttingDown :: (m ~ LspM config) => m Bool
+isShuttingDown = do
+  b <- resShutdown . 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
+  b <- resShutdown . resState <$> getLspEnv
+  liftIO $ C.waitBarrier b
+
 {- Note [LSP configuration]
 LSP configuration is a huge mess.
 - The configuration model of the client is not specified
@@ -988,3 +1012,16 @@ The 'cancellable' property that we can set when making progress reports just
 affects whether the client should show a 'Cancel' button to the user in the UI.
 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'.
+
+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`.
+
+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.
+-}

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

@@ -21,6 +21,7 @@ import Colog.Core (
   (<&),
  )
 
+import Control.Concurrent.Extra as C
 import Control.Concurrent.STM
 import Control.Exception qualified as E
 import Control.Lens hiding (Empty)
@@ -69,6 +70,8 @@ data LspProcessingLog
   | MessageProcessingError BSL.ByteString String
   | forall m. MissingHandler Bool (SClientMethod m)
   | ProgressCancel ProgressToken
+  | forall m. MessageDuringShutdown (SClientMethod m)
+  | ShuttingDown
   | Exiting
 
 deriving instance Show LspProcessingLog
@@ -85,7 +88,9 @@ instance Pretty LspProcessingLog where
       ]
   pretty (MissingHandler _ m) = "LSP: no handler for:" <+> pretty m
   pretty (ProgressCancel tid) = "LSP: cancelling action for token:" <+> pretty tid
-  pretty Exiting = "LSP: Got exit, exiting"
+  pretty (MessageDuringShutdown m) = "LSP: received message during shutdown:" <+> pretty m
+  pretty ShuttingDown = "LSP: received shutdown"
+  pretty Exiting = "LSP: received exit"
 
 processMessage :: (m ~ LspM config) => LogAction m (WithSeverity LspProcessingLog) -> BSL.ByteString -> m ()
 processMessage logger jsonStr = do
@@ -164,6 +169,7 @@ initializeRequestHandler logger ServerDefinition{..} vfs sendFunc req = do
       resRegistrationsNot <- newTVarIO mempty
       resRegistrationsReq <- newTVarIO mempty
       resLspId <- newTVarIO 0
+      resShutdown <- C.newBarrier
       pure LanguageContextState{..}
 
     -- Call the 'duringInitialization' callback to let the server kick stuff up
@@ -414,13 +420,21 @@ inferServerCapabilities _clientCaps o h =
 {- | Invokes the registered dynamic or static handlers for the given message and
  method, as well as doing some bookkeeping.
 -}
-handle :: (m ~ LspM config) => LogAction m (WithSeverity LspProcessingLog) -> SClientMethod meth -> TClientMessage meth -> m ()
+handle :: forall m config meth. (m ~ LspM config) => LogAction m (WithSeverity LspProcessingLog) -> SClientMethod meth -> TClientMessage meth -> m ()
 handle logger m msg =
   case m of
     SMethod_WorkspaceDidChangeWorkspaceFolders -> handle' logger (Just updateWorkspaceFolders) 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_Shutdown -> handle' logger (Just $ \_ -> signalShutdown) m msg
+     where
+      -- See Note [Shutdown]
+      signalShutdown :: LspM config ()
+      signalShutdown = do
+        logger <& ShuttingDown `WithSeverity` Info
+        b <- resShutdown . resState <$> getLspEnv
+        liftIO $ signalBarrier b ()
     SMethod_TextDocumentDidOpen -> handle' logger (Just $ vfsFunc logger openVFS) m msg
     SMethod_TextDocumentDidChange -> handle' logger (Just $ vfsFunc logger changeFromClientVFS) m msg
     SMethod_TextDocumentDidClose -> handle' logger (Just $ vfsFunc logger closeVFS) m msg
@@ -445,48 +459,40 @@ handle' logger mAction m msg = do
 
   env <- getLspEnv
   let Handlers{reqHandlers, notHandlers} = resHandlers env
-
-  let mkRspCb :: TRequestMessage (m1 :: Method ClientToServer Request) -> Either ResponseError (MessageResult m1) -> IO ()
-      mkRspCb req (Left err) =
-        runLspT env $
-          sendToClient $
-            FromServerRsp (req ^. L.method) $
-              TResponseMessage "2.0" (Just (req ^. L.id)) (Left err)
-      mkRspCb req (Right rsp) =
-        runLspT env $
-          sendToClient $
-            FromServerRsp (req ^. L.method) $
-              TResponseMessage "2.0" (Just (req ^. L.id)) (Right rsp)
+  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 -> do
-            reportMissingHandler
+        | 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 (mkRspCb msg)
+      Just h -> liftIO $ h msg (runLspT env . sendResponse msg)
       Nothing
-        | SMethod_Shutdown <- m -> liftIO $ shutdownRequestHandler msg (mkRspCb msg)
-        | otherwise -> do
-            let errorMsg = T.pack $ unwords ["lsp:no handler for: ", show m]
-                err = ResponseError (InR ErrorCodes_MethodNotFound) errorMsg Nothing
-            sendToClient $
-              FromServerRsp (msg ^. L.method) $
-                TResponseMessage "2.0" (Just (msg ^. L.id)) (Left err)
+        | SMethod_Shutdown <- m -> liftIO $ shutdownRequestHandler msg (runLspT env . sendResponse msg)
+        | otherwise -> missingRequestHandler msg
     IsClientEither -> case msg of
+      -- See Note [Shutdown]
+      NotMess _ | shutdown -> notificationDuringShutdown
       NotMess noti -> case pickHandler dynNotHandlers notHandlers of
         Just h -> liftIO $ h noti
-        Nothing -> reportMissingHandler
+        Nothing -> missingNotificationHandler
+      -- See Note [Shutdown]
+      ReqMess req | shutdown -> requestDuringShutdown req
       ReqMess req -> case pickHandler dynReqHandlers reqHandlers of
-        Just h -> liftIO $ h req (mkRspCb req)
-        Nothing -> do
-          let errorMsg = T.pack $ unwords ["lsp:no handler for: ", show m]
-              err = ResponseError (InR ErrorCodes_MethodNotFound) errorMsg Nothing
-          sendToClient $
-            FromServerRsp (req ^. L.method) $
-              TResponseMessage "2.0" (Just (req ^. L.id)) (Left err)
+        Just h -> liftIO $ h req (runLspT env . sendResponse req)
+        Nothing -> missingRequestHandler req
  where
   -- \| Checks to see if there's a dynamic handler, and uses it in favour of the
   -- static handler, if it exists.
@@ -496,14 +502,32 @@ handle' logger mAction m msg = do
     (Nothing, Just (ClientMessageHandler h)) -> Just h
     (Nothing, Nothing) -> Nothing
 
+  sendResponse :: forall m1. TRequestMessage (m1 :: Method ClientToServer Request) -> Either ResponseError (MessageResult m1) -> m ()
+  sendResponse req res = sendToClient $ FromServerRsp (req ^. L.method) $ TResponseMessage "2.0" (Just (req ^. L.id)) res
+
+  requestDuringShutdown :: forall m1. TRequestMessage (m1 :: Method ClientToServer Request) -> m ()
+  requestDuringShutdown req = do
+    logger <& MessageDuringShutdown m `WithSeverity` Warning
+    sendResponse req (Left (ResponseError (InR ErrorCodes_InvalidRequest) "Server is shutdown" Nothing))
+
+  notificationDuringShutdown :: m ()
+  notificationDuringShutdown = logger <& MessageDuringShutdown m `WithSeverity` Warning
+
   -- '$/' notifications should/could be ignored by server.
   -- Don't log errors in that case.
   -- See https://microsoft.github.io/language-server-protocol/specifications/specification-current/#-notifications-and-requests.
-  reportMissingHandler :: m ()
-  reportMissingHandler =
+  missingNotificationHandler :: m ()
+  missingNotificationHandler =
     let optional = isOptionalMethod (SomeMethod m)
      in logger <& MissingHandler optional m `WithSeverity` if optional then Warning else Error
 
+  missingRequestHandler :: TRequestMessage (m1 :: Method ClientToServer Request) -> m ()
+  missingRequestHandler req = do
+    logger <& MissingHandler False m `WithSeverity` Error
+    let errorMsg = T.pack $ unwords ["No handler for: ", show m]
+        err = ResponseError (InR ErrorCodes_MethodNotFound) errorMsg Nothing
+    sendResponse req (Left err)
+
 progressCancelHandler :: (m ~ LspM config) => LogAction m (WithSeverity LspProcessingLog) -> TMessage Method_WindowWorkDoneProgressCancel -> m ()
 progressCancelHandler logger (TNotificationMessage _ _ (WorkDoneProgressCancelParams tid)) = do
   pdata <- getsState (progressCancel . resProgressData)