ensure we always send begin and end for progress

if startDelay and updateDelay are 0. It ensures
that test cases relying on progress notifications
to be sent are not flaky.

Author: soulomoon

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

@@ -17,16 +17,15 @@ import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.IO.Unlift
 import Data.Aeson qualified as J
-import Data.Foldable
 import Data.Map.Strict qualified as Map
 import Data.Maybe
 import Data.Text (Text)
+import Data.Text qualified as T
 import Language.LSP.Protocol.Lens qualified as L
 import Language.LSP.Protocol.Message
 import Language.LSP.Protocol.Types
 import Language.LSP.Protocol.Types qualified as L
 import Language.LSP.Server.Core
-import UnliftIO qualified as U
 import UnliftIO.Exception qualified as UE
 
 {- | A package indicating the percentage of progress complete and a
@@ -53,16 +52,16 @@ instance E.Exception ProgressCancelledException
 data ProgressCancellable = Cancellable | NotCancellable
 
 -- Get a new id for the progress session and make a new one
-getNewProgressId :: MonadLsp config m => m ProgressToken
-getNewProgressId = do
+getNewProgressId :: (MonadLsp config m) => Text -> m ProgressToken
+getNewProgressId title = do
   stateState (progressNextId . resProgressData) $ \cur ->
     let !next = cur + 1
-     in (L.ProgressToken $ L.InL cur, next)
+     in (L.ProgressToken $ L.InR (title <> T.pack (show cur)), next)
 {-# INLINE getNewProgressId #-}
 
 withProgressBase ::
   forall c m a.
-  MonadLsp c m =>
+  (MonadLsp c m) =>
   Bool ->
   Text ->
   Maybe ProgressToken ->
@@ -102,12 +101,10 @@ withProgressBase indefinite title clientToken cancellable f = do
 
     -- Deregister our 'ProgressToken', specifically its cancellation handler. It is important
     -- to do this reliably or else we will leak handlers.
-    unregisterToken :: m ()
-    unregisterToken = do
+    unregisterToken :: ProgressToken -> m ()
+    unregisterToken token = do
       handlers <- getProgressCancellationHandlers
-      liftIO $ atomically $ do
-        mt <- tryReadTMVar tokenVar
-        for_ mt $ \t -> modifyTVar handlers (Map.delete t)
+      liftIO $ atomically $ modifyTVar handlers (Map.delete token)
 
     -- Find and register our 'ProgressToken', asking the client for it if necessary.
     -- Note that this computation may terminate before we get the token, we need to wait
@@ -120,14 +117,14 @@ withProgressBase indefinite title clientToken cancellable f = do
       -- the title/initial percentage aren't given until the 'begin' mesage. However,
       -- it's neater not to create tokens that we won't use, and clients may find it
       -- easier to clean them up if they receive begin/end reports for them.
-      liftIO $ threadDelay startDelay
+      when (startDelay > 0) $ liftIO $ threadDelay startDelay
       case clientToken of
         -- See Note [Client- versus server-initiated progress]
         -- Client-initiated progress
         Just t -> registerToken t
         -- Try server-initiated progress
         Nothing -> do
-          t <- getNewProgressId
+          t <- getNewProgressId title
           clientCaps <- getClientCapabilities
 
           -- If we don't have a progress token from the client and
@@ -145,43 +142,54 @@ withProgressBase indefinite title clientToken cancellable f = do
               -- Successfully registered the token, we can now use it.
               -- So we go ahead and start. We do this as soon as we get the
               -- token back so the client gets feedback ASAP
-              Right _ -> registerToken t
+              Right _ -> do
+                registerToken t
               -- The client sent us an error, we can't use the token.
-              Left _err -> pure ()
-
-    -- Actually send the progress reports.
-    sendReports :: m ()
-    sendReports = do
-      t <- liftIO $ atomically $ readTMVar tokenVar
-      begin t
-      -- Once we are sending updates, if we get interrupted we should send
-      -- the end notification
-      update t `UE.finally` end t
-     where
-      cancellable' = case cancellable of
-        Cancellable -> Just True
-        NotCancellable -> Just False
-      begin t = do
+              Left _err -> do
+                pure ()
+
+    update t = do
+      forever $ do
+        -- See Note [Delayed progress reporting]
+        when (updateDelay > 0) $ liftIO $ threadDelay updateDelay
         (ProgressAmount pct msg) <- liftIO $ atomically $ takeTMVar reportVar
-        sendProgressReport t $ WorkDoneProgressBegin L.AString title cancellable' msg pct
-      update t =
-        forever $ do
-          -- See Note [Delayed progress reporting]
-          liftIO $ threadDelay updateDelay
-          (ProgressAmount pct msg) <- liftIO $ atomically $ takeTMVar reportVar
-          sendProgressReport t $ WorkDoneProgressReport L.AString Nothing msg pct
-      end t = sendProgressReport t (WorkDoneProgressEnd L.AString Nothing)
+        sendProgressReport t $ WorkDoneProgressReport L.AString Nothing msg pct
+    end t = sendProgressReport t (WorkDoneProgressEnd L.AString Nothing)
+
+    begin t = do
+      sendProgressReport t $ WorkDoneProgressBegin L.AString title cancellable' Nothing Nothing
+      return t
+
+    cancellable' = case cancellable of
+      Cancellable -> Just True
+      NotCancellable -> Just False
+
+    -- if we have no delays then we can use uninterruptibleMask_ to create the token
+    -- to ensure we always get begin and end messages
+    maskTokenCreation =
+      if startDelay == 0 && updateDelay == 0
+        then UE.uninterruptibleMask_
+        else id
 
     -- Create the token and then start sending reports; all of which races with the check for the
     -- progress having ended. In all cases, make sure to unregister the token at the end.
-    progressThreads :: m ()
-    progressThreads =
-      ((createToken >> sendReports) `UE.finally` unregisterToken) `U.race_` liftIO progressEnded
-
+    progressThreads runInBase =
+      runInBase
+        ( UE.bracket
+            ( maskTokenCreation $
+                createToken
+                  >> liftIO (atomically $ readTMVar tokenVar)
+                  >>= begin
+            )
+            ( \t -> end t >> unregisterToken t
+            )
+            update
+        )
+        `race_` progressEnded
   withRunInIO $ \runInBase -> do
     withAsync (runInBase $ f updater) $ \mainAct ->
       -- If the progress gets cancelled then we need to get cancelled too
-      withAsync (runInBase progressThreads) $ \pthreads -> do
+      withAsync (progressThreads runInBase) $ \pthreads -> do
         r <- waitEither mainAct pthreads
         -- TODO: is this weird? I can't see how else to gracefully use the ending barrier
         -- as a guard to cancel the other async