@@ -17,9 +17,11 @@ import qualified Language.Haskell.LSP.Core as LSP
1717import Control.Concurrent.Chan
1818import Control.Concurrent.Extra
1919import Control.Concurrent.Async
20+import Control.Concurrent.STM
2021import Control.Exception.Safe
2122import Data.Default
2223import Data.Maybe
24+import qualified Data.Set as Set
2325import qualified Data.Text as T
2426import GHC.IO.Handle (hDuplicate, hDuplicateTo)
2527import System.IO
@@ -65,13 +67,38 @@ runLanguageServer options userHandlers getIdeState = do
6567 -- dies and can be restarted instead of losing threads silently.
6668 clientMsgBarrier <- newBarrier
6769
68− let withResponse wrap f = Just $ \r -> writeChan clientMsgChan $ Response r wrap f
70+ -- The set of requests ids that we have received but not finished processing
71+ pendingRequests <- newTVarIO Set.empty
72+ -- The set of requests that have been cancelled and are also in pendingRequests
73+ cancelledRequests <- newTVarIO Set.empty
74+
75+ let withResponse wrap f = Just $ \r@RequestMessage{_id} -> do
76+ atomically $ modifyTVar pendingRequests (Set.insert _id)
77+ writeChan clientMsgChan $ Response r wrap f
6978 let withNotification old f = Just $ \r -> writeChan clientMsgChan $ Notification r (\lsp ide x -> f lsp ide x >> whenJust old ($ r))
79+ let cancelRequest reqId = atomically $ do
80+ queued <- readTVar pendingRequests
81+ -- We want to avoid that the list of cancelled requests
82+ -- keeps growing if we receive cancellations for requests
83+ -- that do not exist or have already been processed.
84+ when (reqId `elem` queued) $
85+ modifyTVar cancelledRequests (Set.insert reqId)
86+ let clearReqId reqId = atomically $ do
87+ modifyTVar pendingRequests (Set.delete reqId)
88+ modifyTVar cancelledRequests (Set.delete reqId)
89+ -- We implement request cancellation by racing waitForCancel against
90+ -- the actual request handler.
91+ let waitForCancel reqId = atomically $ do
92+ cancelled <- readTVar cancelledRequests
93+ unless (reqId `Set.member` cancelled) retry
7094 let PartialHandlers parts =
7195 setHandlersIgnore <> -- least important
7296 setHandlersDefinition <> setHandlersHover <> setHandlersCodeAction <> -- useful features someone may override
7397 userHandlers <>
74− setHandlersNotifications -- absolutely critical, join them with user notifications
98+ setHandlersNotifications <> -- absolutely critical, join them with user notifications
99+ cancelHandler cancelRequest
100+ -- Cancel requests are special since they need to be handled
101+ -- out of order to be useful. Existing handlers are run afterwards.
75102 handlers <- parts WithMessage{withResponse, withNotification} def
76103
77104 void $ waitAnyCancel =<< traverse async
@@ -79,7 +106,7 @@ runLanguageServer options userHandlers getIdeState = do
79106 stdin
80107 newStdout
81108 ( const $ Right ()
82− , handleInit (signalBarrier clientMsgBarrier ()) clientMsgChan
109+ , handleInit (signalBarrier clientMsgBarrier ()) clearReqId waitForCancel clientMsgChan
83110 )
84111 handlers
85112 (modifyOptions options)
@@ -87,8 +114,8 @@ runLanguageServer options userHandlers getIdeState = do
87114 , void $ waitBarrier clientMsgBarrier
88115 ]
89116 where
90− handleInit :: IO () -> Chan Message -> LSP.LspFuncs () -> IO (Maybe err)
91− handleInit exitClientMsg clientMsgChan lspFuncs@LSP.LspFuncs{..} = do
117+ handleInit :: IO () -> (LspId -> IO ()) -> (LspId -> IO ()) -> Chan Message -> LSP.LspFuncs () -> IO (Maybe err)
118+ handleInit exitClientMsg clearReqId waitForCancel clientMsgChan lspFuncs@LSP.LspFuncs{..} = do
92119 ide <- getIdeState sendFunc (makeLSPVFSHandle lspFuncs)
93120 _ <- flip forkFinally (const exitClientMsg) $ forever $ do
94121 msg <- readChan clientMsgChan
@@ -100,9 +127,21 @@ runLanguageServer options userHandlers getIdeState = do
100127 "Message: " ++ show x ++ "\n" ++
101128 "Exception: " ++ show e
102129 Response x@RequestMessage{_id, _params} wrap act ->
130+ flip finally (clearReqId _id) $
103131 catch (do
104− res <- act lspFuncs ide _params
105− sendFunc $ wrap $ ResponseMessage "2.0" (responseId _id) (Just res) Nothing
132+ -- We could optimize this by first checking if the id
133+ -- is in the cancelled set. However, this is unlikely to be a
134+ -- bottleneck and the additional check might hide
135+ -- issues with async exceptions that need to be fixed.
136+ cancelOrRes <- race (waitForCancel _id) $ act lspFuncs ide _params
137+ case cancelOrRes of
138+ Left () -> do
139+ logDebug (ideLogger ide) $ T.pack $
140+ "Cancelled request " <> show _id
141+ sendFunc $ wrap $ ResponseMessage "2.0" (responseId _id) Nothing $
142+ Just $ ResponseError RequestCancelled "" Nothing
143+ Right res ->
144+ sendFunc $ wrap $ ResponseMessage "2.0" (responseId _id) (Just res) Nothing
106145 ) $ \(e :: SomeException) -> do
107146 logError (ideLogger ide) $ T.pack $
108147 "Unexpected exception on request, please report!\n" ++
@@ -117,11 +156,17 @@ runLanguageServer options userHandlers getIdeState = do
117156-- Set them to avoid a warning in VS Code output.
118157setHandlersIgnore :: PartialHandlers
119158setHandlersIgnore = PartialHandlers $ \_ x -> return x
120− {LSP.cancelNotificationHandler = none
121− ,LSP.initializedHandler = none
159+ {LSP.initializedHandler = none
122160 }
123161 where none = Just $ const $ return ()
124162
163+cancelHandler :: (LspId -> IO ()) -> PartialHandlers
164+cancelHandler cancelRequest = PartialHandlers $ \_ x -> return x
165+ {LSP.cancelNotificationHandler = Just $ \msg@NotificationMessage {_params = CancelParams {_id}} -> do
166+ cancelRequest _id
167+ whenJust (LSP.cancelNotificationHandler x) ($ msg)
168+ }
169+
125170
126171-- | A message that we need to deal with - the pieces are split up with existentials to gain additional type safety
127172-- and defer precise processing until later (allows us to keep at a higher level of abstraction slightly longer)