avoid spawning threads for simple lookups

Author: pepeiborra

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

@@ -1,20 +1,24 @@
 -- We deliberately want to ensure the function we add to the rule database
 -- has the constraints we need on it when we get it out.
 {-# OPTIONS_GHC -Wno-redundant-constraints #-}
-{-# LANGUAGE LambdaCase          #-}
-{-# LANGUAGE NamedFieldPuns      #-}
-{-# LANGUAGE RecordWildCards     #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TupleSections       #-}
-{-# LANGUAGE TypeFamilies        #-}
-{-# LANGUAGE ViewPatterns        #-}
+{-# LANGUAGE DeriveFunctor              #-}
+{-# LANGUAGE DerivingStrategies         #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase                 #-}
+{-# LANGUAGE NamedFieldPuns             #-}
+{-# LANGUAGE RecordWildCards            #-}
+{-# LANGUAGE ScopedTypeVariables        #-}
+{-# LANGUAGE TupleSections              #-}
+{-# LANGUAGE TypeFamilies               #-}
+{-# LANGUAGE ViewPatterns               #-}
 
-module Development.IDE.Graph.Internal.Database where
+module Development.IDE.Graph.Internal.Database (newDatabase, incDatabase, build) where
 
 import           Control.Concurrent.Async
 import           Control.Concurrent.Extra
 import           Control.Exception
 import           Control.Monad
+import           Control.Monad.IO.Class                (MonadIO (liftIO))
 import           Control.Monad.Trans.Class             (lift)
 import           Control.Monad.Trans.Reader
 import qualified Control.Monad.Trans.State.Strict      as State
@@ -76,81 +80,84 @@ build
     :: forall key value . (Shake.RuleResult key ~ value, Typeable key, Show key, Hashable key, Eq key, Typeable value)
     => Database -> [key] -> IO ([Id], [value])
 build db keys = do
-    (ids, vs) <- fmap unzip $ builder db $ map (Right . Key) keys
+    (ids, vs) <- runAIO $ fmap unzip $ either return liftIO =<< builder db (map (Right . Key) keys)
     pure (ids, map (asV . resultValue) vs)
     where
         asV :: Value -> value
         asV (Value x) = unwrapDynamic x
 
--- | Build a list of keys in parallel
+-- | Build a list of keys and return their results.
+--  If none of the keys are dirty, we can return the results immediately.
+--  Otherwise, a blocking computation is returned *which must be evaluated asynchronously* to avoid deadlock.
 builder
-    :: Database -> [Either Id Key] -> IO [(Id, Result)]
+    :: Database -> [Either Id Key] -> AIO (Either [(Id, Result)] (IO [(Id, Result)]))
 builder db@Database{..} keys = do
-    -- Async things that I own and am responsible for killing
-    ownedAsync <- newIORef []
-    flip onException (cleanupAsync ownedAsync) $ do
-
         -- Things that I need to force before my results are ready
-        toForce <- newIORef []
+        toForce <- liftIO $ newIORef []
 
-        results <- withLock databaseLock $ do
-            forM keys $ \idKey -> do
+        results <- withLockAIO databaseLock $ do
+            flip traverse keys $ \idKey -> do
                 -- Resolve the id
                 id <- case idKey of
                     Left id -> pure id
                     Right key -> do
-                        ids <- readIORef databaseIds
+                        ids <- liftIO $ readIORef databaseIds
                         case Intern.lookup key ids of
                             Just v -> pure v
                             Nothing -> do
                                 (ids, id) <- pure $ Intern.add key ids
-                                writeIORef' databaseIds ids
+                                liftIO $ writeIORef' databaseIds ids
                                 return id
 
                 -- Spawn the id if needed
-                status <- Ids.lookup databaseValues id
+                status <- liftIO $ Ids.lookup databaseValues id
                 val <- case fromMaybe (fromRight undefined idKey, Dirty Nothing) status of
                     (_, Clean r) -> pure r
                     (_, Running act _) -> do
-                        -- we promise to force everything in todo before reading the results
-                        -- so the following unsafePerformIO isn't actually unsafe
                         let (force, val) = splitIO act
-                        modifyIORef toForce (force:)
+                        liftIO $ modifyIORef toForce (force:)
                         pure val
                     (key, Dirty s) -> do
-                        -- Important we don't lose any Async things we create
-                        act <- uninterruptibleMask $ \restore -> do
-                            -- the child actions should always be spawned unmasked
-                            -- or they can't be killed
-                            async <- async $ restore $ check db key id s
-                            modifyIORef ownedAsync (async:)
-                            pure $ wait async
-                        Ids.insert databaseValues id (key, Running act s)
+                        act <- join <$> unliftAIO (refresh db key id s)
+                        liftIO $ Ids.insert databaseValues id (key, Running act s)
                         let (force, val) = splitIO act
-                        modifyIORef toForce (force:)
+                        liftIO $ modifyIORef toForce (force:)
                         pure val
-
                 pure (id, val)
 
-        sequence_ =<< readIORef toForce
-        pure results
+        toForceList <- liftIO $ readIORef toForce
+        case toForceList of
+            [] -> return $ Left results
+            _ -> return $ Right $ do
+                    sequence_ toForceList
+                    pure results
 
-cleanupAsync :: IORef [Async a] -> IO ()
-cleanupAsync ref = mapConcurrently_ uninterruptibleCancel =<< readIORef ref
-
--- | Check if we need to run the database.
-check :: Database -> Key -> Id -> Maybe Result -> IO Result
-check db key id result@(Just me@Result{resultDeps=Just deps}) = do
+-- | Refresh a key:
+--     * If no dirty dependencies and we have evaluated the key previously, then we refresh it in the current thread.
+--       This assumes that the implementation will be a lookup
+--     * Otherwise, we spawn a new thread to refresh the dirty deps (if any) and the key itself
+refresh :: Database -> Key -> Id -> Maybe Result -> AIO (IO Result)
+refresh db key id result@(Just me@Result{resultDeps=Just deps}) = do
     res <- builder db $ map Left deps
-    let dirty = any (\(_,dep) -> resultBuilt me < resultChanged dep) res
-    let mode = if dirty then Shake.RunDependenciesChanged else Shake.RunDependenciesSame
-    spawn db key id mode result
-check db key id result = spawn db key id Shake.RunDependenciesChanged result
+    case res of
+      Left res ->
+        if isDirty res
+            then asyncWithCleanUp $ liftIO $ compute db key id Shake.RunDependenciesChanged result
+            else pure $ compute db key id Shake.RunDependenciesSame result
+      Right iores -> asyncWithCleanUp $ liftIO $ do
+        res <- iores
+        let mode = if isDirty res then Shake.RunDependenciesChanged else Shake.RunDependenciesSame
+        compute db key id mode result
+    where
+        isDirty = any (\(_,dep) -> resultBuilt me < resultChanged dep)
+
+refresh db key id result =
+    asyncWithCleanUp $ liftIO $ compute db key id Shake.RunDependenciesChanged result
 
 
--- | Spawn a new computation to run the action.
-spawn :: Database -> Key -> Id -> Shake.RunMode -> Maybe Result -> IO Result
-spawn db@Database{..} key id mode result = do
+-- | Compute a key.
+compute :: Database -> Key -> Id -> Shake.RunMode -> Maybe Result -> IO Result
+compute db@Database{..} key id mode result = do
     let act = runRule databaseRules key (fmap resultData result) mode
     deps <- newIORef $ Just []
     (execution, Shake.RunResult{..}) <-
@@ -218,6 +225,34 @@ transitiveDirtySet database = flip State.execStateT Set.empty . traverse_ loop
             next <- lift $ getReverseDependencies database x
             traverse_ loop (maybe mempty Set.toList next)
 
+-- | IO extended to track created asyncs to clean them up when the thread is killed,
+--   generalizing 'withAsync'
+newtype AIO a = AIO { unAIO :: ReaderT (IORef [Async ()]) IO a }
+  deriving newtype (Applicative, Functor, Monad, MonadIO)
 
-idFromInt :: Set.Key -> Id
-idFromInt = id
+runAIO :: AIO a -> IO a
+runAIO (AIO act) = do
+    asyncs <- newIORef []
+    runReaderT act asyncs `onException` cleanupAsync asyncs
+
+asyncWithCleanUp :: AIO a -> AIO (IO a)
+asyncWithCleanUp act = do
+    st <- AIO ask
+    io <- unliftAIO act
+    liftIO $ uninterruptibleMask $ \restore -> do
+        a <- async $ restore io
+        modifyIORef st (void a :)
+        return $ wait a
+
+withLockAIO :: Lock -> AIO a -> AIO a
+withLockAIO lock act = do
+    io <- unliftAIO act
+    liftIO $ withLock lock io
+
+unliftAIO :: AIO a -> AIO (IO a)
+unliftAIO act = do
+    st <- AIO ask
+    return $ runReaderT (unAIO act) st
+
+cleanupAsync :: IORef [Async a] -> IO ()
+cleanupAsync ref = mapConcurrently_ uninterruptibleCancel =<< readIORef ref