Lockless iorefs (#2460)

* lock-less KnownTargets

* lock-less exportsMap

* lock-less globals

* lock-less persistentKeys

* switch to TVar

* fix build in plugins

* add comments

Co-authored-by: Javier Neira <atreyu.bbb@gmail.com>

Author: pepeiborra

ghcide/session-loader/Development/IDE/Session.hs

@@ -73,11 +73,12 @@ import           System.IO
 import           System.Info
 
 import           Control.Applicative                  (Alternative ((<|>)))
-import           Control.Exception                    (evaluate)
 import           Data.Void
 
-import           Control.Concurrent.STM               (atomically)
+import           Control.Concurrent.STM.Stats         (atomically, modifyTVar',
+                                                       readTVar, writeTVar)
 import           Control.Concurrent.STM.TQueue
+import           Data.Foldable                        (for_)
 import qualified Data.HashSet                         as Set
 import           Database.SQLite.Simple
 import           Development.IDE.Core.Tracing         (withTrace)
@@ -265,13 +266,17 @@ loadSessionWithOptions SessionLoadingOptions{..} dir = do
               TargetModule _ -> do
                 found <- filterM (IO.doesFileExist . fromNormalizedFilePath) targetLocations
                 return (targetTarget, found)
-          join $ atomically $ recordDirtyKeys extras GetKnownTargets  [emptyFilePath]
-          modifyVarIO' knownTargetsVar $ traverseHashed $ \known -> do
-            let known' = HM.unionWith (<>) known $ HM.fromList $ map (second Set.fromList) knownTargets
-            when (known /= known') $
+          hasUpdate <- join $ atomically $ do
+            known <- readTVar knownTargetsVar
+            let known' = flip mapHashed known $ \k ->
+                            HM.unionWith (<>) k $ HM.fromList $ map (second Set.fromList) knownTargets
+                hasUpdate = if known /= known' then Just (unhashed known') else Nothing
+            writeTVar knownTargetsVar known'
+            logDirtyKeys <- recordDirtyKeys extras GetKnownTargets [emptyFilePath]
+            return (logDirtyKeys >> pure hasUpdate)
+          for_ hasUpdate $ \x ->
                 logDebug logger $ "Known files updated: " <>
-                    T.pack(show $ (HM.map . Set.map) fromNormalizedFilePath known')
-            pure known'
+                    T.pack(show $ (HM.map . Set.map) fromNormalizedFilePath x)
 
     -- Create a new HscEnv from a hieYaml root and a set of options
     -- If the hieYaml file already has an HscEnv, the new component is
@@ -405,7 +410,7 @@ loadSessionWithOptions SessionLoadingOptions{..} dir = do
                     -- update exports map
                     extras <- getShakeExtras
                     let !exportsMap' = createExportsMap $ mapMaybe (fmap hirModIface) modIfaces
-                    liftIO $ modifyVar_ (exportsMap extras) $ evaluate . (exportsMap' <>)
+                    liftIO $ atomically $ modifyTVar' (exportsMap extras) (exportsMap' <>)
 
           return (second Map.keys res)
 

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

@@ -25,7 +25,8 @@ import qualified Data.HashMap.Strict                    as HashMap
 import qualified Data.Text                              as T
 import           Development.IDE.Graph
 
-import           Control.Concurrent.STM.Stats           (atomically)
+import           Control.Concurrent.STM.Stats           (atomically,
+                                                         modifyTVar')
 import qualified Data.ByteString                        as BS
 import           Data.Maybe                             (catMaybes)
 import           Development.IDE.Core.ProgressReporting
@@ -114,7 +115,7 @@ kick = do
     -- Update the exports map
     results <- uses GenerateCore files <* uses GetHieAst files
     let mguts = catMaybes results
-    void $ liftIO $ modifyVar' exportsMap (updateExportsMapMg mguts)
+    void $ liftIO $ atomically $ modifyTVar' exportsMap (updateExportsMapMg mguts)
 
     liftIO $ progressUpdate progress KickCompleted
 

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

@@ -189,7 +189,9 @@ data ShakeExtras = ShakeExtras
      lspEnv :: Maybe (LSP.LanguageContextEnv Config)
     ,debouncer :: Debouncer NormalizedUri
     ,logger :: Logger
-    ,globals :: Var (HMap.HashMap TypeRep Dynamic)
+    ,globals :: TVar (HMap.HashMap TypeRep Dynamic)
+      -- ^ Registry of global state used by rules.
+      -- Small and immutable after startup, so not worth using an STM.Map.
     ,state :: Values
     ,diagnostics :: STMDiagnosticStore
     ,hiddenDiagnostics :: STMDiagnosticStore
@@ -210,17 +212,18 @@ data ShakeExtras = ShakeExtras
         -> IO ()
     ,ideNc :: IORef NameCache
     -- | A mapping of module name to known target (or candidate targets, if missing)
-    ,knownTargetsVar :: Var (Hashed KnownTargets)
+    ,knownTargetsVar :: TVar (Hashed KnownTargets)
     -- | A mapping of exported identifiers for local modules. Updated on kick
-    ,exportsMap :: Var ExportsMap
+    ,exportsMap :: TVar ExportsMap
     -- | A work queue for actions added via 'runInShakeSession'
     ,actionQueue :: ActionQueue
     ,clientCapabilities :: ClientCapabilities
     , hiedb :: HieDb -- ^ Use only to read.
     , hiedbWriter :: HieDbWriter -- ^ use to write
-    , persistentKeys :: Var (HMap.HashMap Key GetStalePersistent)
+    , persistentKeys :: TVar (HMap.HashMap Key GetStalePersistent)
       -- ^ Registery for functions that compute/get "stale" results for the rule
       -- (possibly from disk)
+      -- Small and immutable after startup, so not worth using an STM.Map.
     , vfs :: VFSHandle
     , defaultConfig :: Config
       -- ^ Default HLS config, only relevant if the client does not provide any Config
@@ -258,7 +261,7 @@ getPluginConfig plugin = do
 addPersistentRule :: IdeRule k v => k -> (NormalizedFilePath -> IdeAction (Maybe (v,PositionDelta,TextDocumentVersion))) -> Rules ()
 addPersistentRule k getVal = do
   ShakeExtras{persistentKeys} <- getShakeExtrasRules
-  void $ liftIO $ modifyVar' persistentKeys $ HMap.insert (Key k) (fmap (fmap (first3 toDyn)) . getVal)
+  void $ liftIO $ atomically $ modifyTVar' persistentKeys $ HMap.insert (Key k) (fmap (fmap (first3 toDyn)) . getVal)
 
 class Typeable a => IsIdeGlobal a where
 
@@ -282,15 +285,15 @@ addIdeGlobal x = do
 
 addIdeGlobalExtras :: IsIdeGlobal a => ShakeExtras -> a -> IO ()
 addIdeGlobalExtras ShakeExtras{globals} x@(typeOf -> ty) =
-    void $ liftIO $ modifyVarIO' globals $ \mp -> case HMap.lookup ty mp of
-        Just _ -> errorIO $ "Internal error, addIdeGlobalExtras, got the same type twice for " ++ show ty
-        Nothing -> return $! HMap.insert ty (toDyn x) mp
+    void $ liftIO $ atomically $ modifyTVar' globals $ \mp -> case HMap.lookup ty mp of
+        Just _ -> error $ "Internal error, addIdeGlobalExtras, got the same type twice for " ++ show ty
+        Nothing -> HMap.insert ty (toDyn x) mp
 
 
 getIdeGlobalExtras :: forall a . IsIdeGlobal a => ShakeExtras -> IO a
 getIdeGlobalExtras ShakeExtras{globals} = do
     let typ = typeRep (Proxy :: Proxy a)
-    x <- HMap.lookup (typeRep (Proxy :: Proxy a)) <$> readVar globals
+    x <- HMap.lookup (typeRep (Proxy :: Proxy a)) <$> readTVarIO globals
     case x of
         Just x
             | Just x <- fromDynamic x -> pure x
@@ -333,7 +336,7 @@ lastValueIO s@ShakeExtras{positionMapping,persistentKeys,state} k file = do
           | IdeTesting testing <- ideTesting s -- Don't read stale persistent values in tests
           , testing = pure Nothing
           | otherwise = do
-          pmap <- readVar persistentKeys
+          pmap <- readTVarIO persistentKeys
           mv <- runMaybeT $ do
             liftIO $ Logger.logDebug (logger s) $ T.pack $ "LOOKUP UP PERSISTENT FOR: " ++ show k
             f <- MaybeT $ pure $ HMap.lookup (Key k) pmap
@@ -477,7 +480,7 @@ getValues state key file = do
 knownTargets :: Action (Hashed KnownTargets)
 knownTargets = do
   ShakeExtras{knownTargetsVar} <- getShakeExtras
-  liftIO $ readVar knownTargetsVar
+  liftIO $ readTVarIO knownTargetsVar
 
 -- | Seq the result stored in the Shake value. This only
 -- evaluates the value to WHNF not NF. We take care of the latter
@@ -508,25 +511,25 @@ shakeOpen lspEnv defaultConfig logger debouncer
     us <- mkSplitUniqSupply 'r'
     ideNc <- newIORef (initNameCache us knownKeyNames)
     shakeExtras <- do
-        globals <- newVar HMap.empty
+        globals <- newTVarIO HMap.empty
         state <- STM.newIO
         diagnostics <- STM.newIO
         hiddenDiagnostics <- STM.newIO
         publishedDiagnostics <- STM.newIO
         positionMapping <- STM.newIO
-        knownTargetsVar <- newVar $ hashed HMap.empty
+        knownTargetsVar <- newTVarIO $ hashed HMap.empty
         let restartShakeSession = shakeRestart ideState
-        persistentKeys <- newVar HMap.empty
+        persistentKeys <- newTVarIO HMap.empty
         indexPending <- newTVarIO HMap.empty
         indexCompleted <- newTVarIO 0
         indexProgressToken <- newVar Nothing
         let hiedbWriter = HieDbWriter{..}
-        exportsMap <- newVar mempty
+        exportsMap <- newTVarIO mempty
         -- lazily initialize the exports map with the contents of the hiedb
         _ <- async $ do
             logDebug logger "Initializing exports map from hiedb"
             em <- createExportsMapHieDb hiedb
-            _ <- modifyVar' exportsMap (<> em)
+            atomically $ modifyTVar' exportsMap (<> em)
             logDebug logger $ "Done initializing exports map from hiedb (" <> pack(show (ExportsMap.size em)) <> ")"
 
         progress <- do

ghcide/src/Development/IDE/Plugin/CodeAction/Args.hs

@@ -14,7 +14,7 @@ module Development.IDE.Plugin.CodeAction.Args
   )
 where
 
-import           Control.Concurrent.Extra
+import           Control.Concurrent.STM.Stats                 (readTVarIO)
 import           Control.Monad.Reader
 import           Control.Monad.Trans.Maybe
 import           Data.Either                                  (fromRight)
@@ -59,7 +59,7 @@ runGhcideCodeAction state (CodeActionParams _ _ (TextDocumentIdentifier uri) _ra
       runRule GhcSession >>= \case
         Just env -> do
           pkgExports <- envPackageExports env
-          localExports <- readVar (exportsMap $ shakeExtras state)
+          localExports <- readTVarIO (exportsMap $ shakeExtras state)
           pure $ localExports <> pkgExports
         _ -> pure mempty
   caaIdeOptions <- onceIO $ runAction "GhcideCodeActions.getIdeOptions" state getIdeOptions

ghcide/src/Development/IDE/Plugin/Completions.hs

@@ -8,7 +8,7 @@ module Development.IDE.Plugin.Completions
     ) where
 
 import           Control.Concurrent.Async                     (concurrently)
-import           Control.Concurrent.Extra
+import           Control.Concurrent.STM.Stats                 (readTVarIO)
 import           Control.Monad.Extra
 import           Control.Monad.IO.Class
 import           Control.Monad.Trans.Maybe
@@ -138,7 +138,7 @@ getCompletionsLSP ide plId
             -- set up the exports map including both package and project-level identifiers
             packageExportsMapIO <- fmap(envPackageExports . fst) <$> useWithStaleFast GhcSession npath
             packageExportsMap <- mapM liftIO packageExportsMapIO
-            projectExportsMap <- liftIO $ readVar (exportsMap $ shakeExtras ide)
+            projectExportsMap <- liftIO $ readTVarIO (exportsMap $ shakeExtras ide)
             let exportsMap = fromMaybe mempty packageExportsMap <> projectExportsMap
 
             let moduleExports = getModuleExportsMap exportsMap

plugins/hls-retrie-plugin/hls-retrie-plugin.cabal

@@ -30,6 +30,7 @@ library
     , lsp-types
     , retrie                >=0.1.1.0
     , safe-exceptions
+    , stm
     , text
     , transformers
     , unordered-containers

plugins/hls-retrie-plugin/src/Ide/Plugin/Retrie.hs

@@ -17,6 +17,7 @@
 module Ide.Plugin.Retrie (descriptor) where
 
 import           Control.Concurrent.Extra             (readVar)
+import           Control.Concurrent.STM               (readTVarIO)
 import           Control.Exception.Safe               (Exception (..),
                                                        SomeException, catch,
                                                        throwIO, try)
@@ -356,7 +357,7 @@ callRetrie ::
   Bool ->
   IO ([CallRetrieError], WorkspaceEdit)
 callRetrie state session rewrites origin restrictToOriginatingFile = do
-  knownFiles <- toKnownFiles . unhashed <$> readVar (knownTargetsVar $ shakeExtras state)
+  knownFiles <- toKnownFiles . unhashed <$> readTVarIO (knownTargetsVar $ shakeExtras state)
   let reuseParsedModule f = do
         pm <-
           useOrFail "GetParsedModule" NoParse GetParsedModule f