Review masking and add traces when things don't cancel timely (#2768)

* Review masking and add traces when things don't cancel timely

* fixup

* use sleep consistently

* redundant imports

* hlints

* fix 9.2 build

* Fix 9.2 build for real

* remove unnecessary polymorphism

* Avoid spawning loop async unnecessrily

* flush asyncs ref

* Add comments and apply @michaelpj suggestions

Author: pepeiborra

ghcide/src/Development/IDE/Core/FileExists.hs

@@ -109,8 +109,7 @@ modifyFileExists :: IdeState -> [(NormalizedFilePath, FileChangeType)] -> IO ()
 modifyFileExists state changes = do
   FileExistsMapVar var <- getIdeGlobalState state
   -- Masked to ensure that the previous values are flushed together with the map update
-    -- update the map
-  mask_ $ join $ atomicallyNamed "modifyFileExists" $ do
+  join $ mask_ $ atomicallyNamed "modifyFileExists" $ do
     forM_ changes $ \(f,c) ->
         case fromChange c of
             Just c' -> STM.focus (Focus.insert c') f var

ghcide/src/Development/IDE/Core/Shake.hs

@@ -151,7 +151,6 @@ import           Development.IDE.Types.Options
 import           Development.IDE.Types.Shake
 import qualified Focus
 import           GHC.Fingerprint
-import           GHC.Generics
 import           HieDb.Types
 import           Ide.Plugin.Config
 import qualified Ide.PluginUtils                        as HLS
@@ -173,6 +172,7 @@ data Log
   = LogCreateHieDbExportsMapStart
   | LogCreateHieDbExportsMapFinish !Int
   | LogBuildSessionRestart !String ![DelayedActionInternal] !(HashSet Key) !Seconds !(Maybe FilePath)
+  | LogBuildSessionRestartTakingTooLong !Seconds
   | LogDelayedAction !(DelayedAction ()) !Seconds
   | LogBuildSessionFinish !(Maybe SomeException)
   | LogDiagsDiffButNoLspEnv ![FileDiagnostic]
@@ -192,6 +192,8 @@ instance Pretty Log where
         , "Action Queue:" <+> pretty (map actionName actionQueue)
         , "Keys:" <+> pretty (map show $ HSet.toList keyBackLog)
         , "Aborting previous build session took" <+> pretty (showDuration abortDuration) <+> pretty shakeProfilePath ]
+    LogBuildSessionRestartTakingTooLong seconds ->
+        "Build restart is taking too long (" <> pretty seconds <> " seconds)"
     LogDelayedAction delayedAction duration ->
       hsep
         [ "Finished:" <+> pretty (actionName delayedAction)
@@ -322,7 +324,7 @@ getVirtualFile nf = do
 -- Take a snapshot of the current LSP VFS
 vfsSnapshot :: Maybe (LSP.LanguageContextEnv a) -> IO VFS
 vfsSnapshot Nothing       = pure $ VFS mempty ""
-vfsSnapshot (Just lspEnv) = LSP.runLspT lspEnv $ LSP.getVirtualFiles
+vfsSnapshot (Just lspEnv) = LSP.runLspT lspEnv LSP.getVirtualFiles
 
 
 addIdeGlobal :: IsIdeGlobal a => a -> Rules ()
@@ -596,7 +598,7 @@ shakeOpen recorder lspEnv defaultConfig logger debouncer
 
         dirtyKeys <- newTVarIO mempty
         -- Take one VFS snapshot at the start
-        vfs <- atomically . newTVar =<< vfsSnapshot lspEnv
+        vfs <- newTVarIO =<< vfsSnapshot lspEnv
         pure ShakeExtras{..}
     shakeDb  <-
         shakeNewDatabase
@@ -683,7 +685,7 @@ shakeRestart recorder IdeState{..} reason acts =
         shakeSession
         (\runner -> do
               let log = logWith recorder
-              (stopTime,()) <- duration (cancelShakeSession runner)
+              (stopTime,()) <- duration $ logErrorAfter 10 recorder $ cancelShakeSession runner
               res <- shakeDatabaseProfile shakeDb
               backlog <- readTVarIO $ dirtyKeys shakeExtras
               queue <- atomicallyNamed "actionQueue - peek" $ peekInProgress $ actionQueue shakeExtras
@@ -706,6 +708,11 @@ shakeRestart recorder IdeState{..} reason acts =
         -- See https://github.com/haskell/ghcide/issues/79
         (\() -> do
           (,()) <$> newSession recorder shakeExtras shakeDb acts reason)
+    where
+        logErrorAfter :: Seconds -> Recorder (WithPriority Log) -> IO () -> IO ()
+        logErrorAfter seconds recorder action = flip withAsync (const action) $ do
+            sleep seconds
+            logWith recorder Error (LogBuildSessionRestartTakingTooLong seconds)
 
 notifyTestingLogMessage :: ShakeExtras -> T.Text -> IO ()
 notifyTestingLogMessage extras msg = do
@@ -1100,7 +1107,7 @@ defineEarlyCutoff' doDiagnostics cmp key file old mode action = do
                     -- No changes in the dependencies and we have
                     -- an existing successful result.
                     Just (v@(Succeeded _ x), diags) -> do
-                        ver <- estimateFileVersionUnsafely state key (Just x) file
+                        ver <- estimateFileVersionUnsafely key (Just x) file
                         doDiagnostics (vfsVersion =<< ver) $ Vector.toList diags
                         return $ Just $ RunResult ChangedNothing old $ A v
                     _ -> return Nothing
@@ -1121,7 +1128,7 @@ defineEarlyCutoff' doDiagnostics cmp key file old mode action = do
                     \(e :: SomeException) -> do
                         pure (Nothing, ([ideErrorText file $ T.pack $ show e | not $ isBadDependency e],Nothing))
 
-                ver <- estimateFileVersionUnsafely state key res file
+                ver <- estimateFileVersionUnsafely key res file
                 (bs, res) <- case res of
                     Nothing -> do
                         pure (toShakeValue ShakeStale bs, staleV)
@@ -1147,12 +1154,11 @@ defineEarlyCutoff' doDiagnostics cmp key file old mode action = do
     estimateFileVersionUnsafely
         :: forall k v
          . IdeRule k v
-        => Values
-        -> k
+        => k
         -> Maybe v
         -> NormalizedFilePath
         -> Action (Maybe FileVersion)
-    estimateFileVersionUnsafely state _k v fp
+    estimateFileVersionUnsafely _k v fp
         | fp == emptyFilePath = pure Nothing
         | Just Refl <- eqT @k @GetModificationTime = pure v
         -- GetModificationTime depends on these rules, so avoid creating a cycle

hls-graph/src/Development/IDE/Graph/Internal/Database.hs

@@ -8,6 +8,7 @@
 {-# LANGUAGE RecordWildCards            #-}
 {-# LANGUAGE ScopedTypeVariables        #-}
 {-# LANGUAGE TypeFamilies               #-}
+{-# LANGUAGE TupleSections #-}
 
 module Development.IDE.Graph.Internal.Database (newDatabase, incDatabase, build, getDirtySet, getKeysAndVisitAge) where
 
@@ -32,14 +33,15 @@ import           Data.IORef.Extra
 import           Data.Maybe
 import           Data.Traversable                     (for)
 import           Data.Tuple.Extra
+import           Debug.Trace (traceM)
 import           Development.IDE.Graph.Classes
 import           Development.IDE.Graph.Internal.Rules
 import           Development.IDE.Graph.Internal.Types
 import qualified Focus
 import qualified ListT
 import qualified StmContainers.Map                    as SMap
+import           System.Time.Extra                    (duration, sleep)
 import           System.IO.Unsafe
-import           System.Time.Extra                    (duration)
 
 newDatabase :: Dynamic -> TheRules -> IO Database
 newDatabase databaseExtra databaseRules = do
@@ -120,7 +122,7 @@ builder db@Database{..} stack keys = withRunInIO $ \(RunInIO run) -> do
             pure (id, val)
 
     toForceList <- liftIO $ readTVarIO toForce
-    let waitAll = run $ mapConcurrentlyAIO_ id toForceList
+    let waitAll = run $ waitConcurrently_ toForceList
     case toForceList of
         [] -> return $ Left results
         _ -> return $ Right $ do
@@ -170,6 +172,10 @@ compute db@Database{..} stack key mode result = do
         deps | not(null deps)
             && runChanged /= ChangedNothing
                     -> do
+            -- IMPORTANT: record the reverse deps **before** marking the key Clean.
+            -- If an async exception strikes before the deps have been recorded,
+            -- we won't be able to accurately propagate dirtiness for this key
+            -- on the next build.
             void $
                 updateReverseDeps key db
                     (getResultDepsDefault [] previousDeps)
@@ -224,7 +230,8 @@ updateReverseDeps
     -> [Key] -- ^ Previous direct dependencies of Id
     -> HashSet Key -- ^ Current direct dependencies of Id
     -> IO ()
-updateReverseDeps myId db prev new = uninterruptibleMask_ $ do
+-- mask to ensure that all the reverse dependencies are updated
+updateReverseDeps myId db prev new = do
     forM_ prev $ \d ->
         unless (d `HSet.member` new) $
             doOne (HSet.delete myId) d
@@ -252,20 +259,27 @@ transitiveDirtySet database = flip State.execStateT HSet.empty . traverse_ loop
             next <- lift $ atomically $ getReverseDependencies database x
             traverse_ loop (maybe mempty HSet.toList next)
 
--- | IO extended to track created asyncs to clean them up when the thread is killed,
---   generalizing 'withAsync'
+--------------------------------------------------------------------------------
+-- Asynchronous computations with cancellation
+
+-- | A simple monad to implement cancellation on top of 'Async',
+--   generalizing 'withAsync' to monadic scopes.
 newtype AIO a = AIO { unAIO :: ReaderT (IORef [Async ()]) IO a }
   deriving newtype (Applicative, Functor, Monad, MonadIO)
 
+-- | Run the monadic computation, cancelling all the spawned asyncs if an exception arises
 runAIO :: AIO a -> IO a
 runAIO (AIO act) = do
     asyncs <- newIORef []
     runReaderT act asyncs `onException` cleanupAsync asyncs
 
+-- | Like 'async' but with built-in cancellation.
+--   Returns an IO action to wait on the result.
 asyncWithCleanUp :: AIO a -> AIO (IO a)
 asyncWithCleanUp act = do
     st <- AIO ask
     io <- unliftAIO act
+    -- mask to make sure we keep track of the spawned async
     liftIO $ uninterruptibleMask $ \restore -> do
         a <- async $ restore io
         atomicModifyIORef'_ st (void a :)
@@ -284,27 +298,40 @@ withRunInIO k = do
     k $ RunInIO (\aio -> runReaderT (unAIO aio) st)
 
 cleanupAsync :: IORef [Async a] -> IO ()
-cleanupAsync ref = uninterruptibleMask_ $ do
-    asyncs <- readIORef ref
+-- mask to make sure we interrupt all the asyncs
+cleanupAsync ref = uninterruptibleMask $ \unmask -> do
+    asyncs <- atomicModifyIORef' ref ([],)
+    -- interrupt all the asyncs without waiting
     mapM_ (\a -> throwTo (asyncThreadId a) AsyncCancelled) asyncs
-    mapM_ waitCatch asyncs
+    -- Wait until all the asyncs are done
+    -- But if it takes more than 10 seconds, log to stderr
+    unless (null asyncs) $ do
+        let warnIfTakingTooLong = unmask $ forever $ do
+                sleep 10
+                traceM "cleanupAsync: waiting for asyncs to finish"
+        withAsync warnIfTakingTooLong $ \_ ->
+            mapM_ waitCatch asyncs
+
+data Wait
+    = Wait {justWait :: !(IO ())}
+    | Spawn {justWait :: !(IO ())}
 
-data Wait a
-    = Wait {justWait :: !a}
-    | Spawn {justWait :: !a}
-    deriving Functor
+fmapWait :: (IO () -> IO ()) -> Wait -> Wait
+fmapWait f (Wait io) = Wait (f io)
+fmapWait f (Spawn io) = Spawn (f io)
 
-waitOrSpawn :: Wait (IO a) -> IO (Either (IO a) (Async a))
+waitOrSpawn :: Wait -> IO (Either (IO ()) (Async ()))
 waitOrSpawn (Wait io)  = pure $ Left io
 waitOrSpawn (Spawn io) = Right <$> async io
 
-mapConcurrentlyAIO_ :: (a -> IO ()) -> [Wait a] -> AIO ()
-mapConcurrentlyAIO_ _ [] = pure ()
-mapConcurrentlyAIO_ f [one] = liftIO $ justWait $ fmap f one
-mapConcurrentlyAIO_ f many = do
+waitConcurrently_ :: [Wait] -> AIO ()
+waitConcurrently_ [] = pure ()
+waitConcurrently_ [one] = liftIO $ justWait one
+waitConcurrently_ many = do
     ref <- AIO ask
-    waits <- liftIO $ uninterruptibleMask $ \restore -> do
-        waits <- liftIO $ traverse (waitOrSpawn . fmap (restore . f)) many
+    -- mask to make sure we keep track of all the asyncs
+    waits <- liftIO $ uninterruptibleMask $ \unmask -> do
+        waits <- liftIO $ traverse (waitOrSpawn . fmapWait unmask) many
         let asyncs = rights waits
         liftIO $ atomicModifyIORef'_ ref (asyncs ++)
         return waits