Lockless FileExistsMap and position mapping (#2442)

pepeiborra · jneira · mergify[bot] · web-flow · commit a83c850b68a8 · 2021-12-07T18:51:38.000Z
* lock-less position mapping

* lock-less FileExistsMap

* Move to STM

Co-authored-by: Javier Neira &lt;atreyu.bbb@gmail.com&gt;
Co-authored-by: mergify[bot] &lt;37929162+mergify[bot]@users.noreply.github.com&gt;
diff --git a/ghcide/src/Development/IDE/Core/FileExists.hs b/ghcide/src/Development/IDE/Core/FileExists.hs
@@ -10,14 +10,12 @@ module Development.IDE.Core.FileExists
   )
 where
 
-import           Control.Concurrent.STM.Stats
-import           Control.Concurrent.Strict
+import           Control.Concurrent.STM.Stats          (atomically,
+                                                        atomicallyNamed)
 import           Control.Exception
 import           Control.Monad.Extra
 import           Control.Monad.IO.Class
 import qualified Data.ByteString                       as BS
-import           Data.HashMap.Strict                   (HashMap)
-import qualified Data.HashMap.Strict                   as HashMap
 import           Data.List                             (partition)
 import           Data.Maybe
 import           Development.IDE.Core.FileStore
@@ -27,9 +25,11 @@ import           Development.IDE.Core.Shake
 import           Development.IDE.Graph
 import           Development.IDE.Types.Location
 import           Development.IDE.Types.Options
+import qualified Focus
 import           Ide.Plugin.Config                     (Config)
 import           Language.LSP.Server                   hiding (getVirtualFile)
 import           Language.LSP.Types
+import qualified StmContainers.Map                     as STM
 import qualified System.Directory                      as Dir
 import qualified System.FilePath.Glob                  as Glob
 
@@ -75,37 +75,38 @@ fast path by a check that the path also matches our watching patterns.
 -- | A map for tracking the file existence.
 -- If a path maps to 'True' then it exists; if it maps to 'False' then it doesn't exist'; and
 -- if it's not in the map then we don't know.
-type FileExistsMap = (HashMap NormalizedFilePath Bool)
+type FileExistsMap = STM.Map NormalizedFilePath Bool
 
 -- | A wrapper around a mutable 'FileExistsState'
-newtype FileExistsMapVar = FileExistsMapVar (Var FileExistsMap)
+newtype FileExistsMapVar = FileExistsMapVar FileExistsMap
 
 instance IsIdeGlobal FileExistsMapVar
 
 -- | Grab the current global value of 'FileExistsMap' without acquiring a dependency
 getFileExistsMapUntracked :: Action FileExistsMap
 getFileExistsMapUntracked = do
   FileExistsMapVar v <- getIdeGlobalAction
-  liftIO $ readVar v
+  return v
 
 -- | Modify the global store of file exists.
 modifyFileExists :: IdeState -> [(NormalizedFilePath, FileChangeType)] -> IO ()
 modifyFileExists state changes = do
   FileExistsMapVar var <- getIdeGlobalState state
-  changesMap           <- evaluate $ HashMap.fromList changes
   -- Masked to ensure that the previous values are flushed together with the map update
-  mask $ \_ -> do
     -- update the map
-    void $ modifyVar' var $ HashMap.union (HashMap.mapMaybe fromChange changesMap)
+  mask_ $ join $ atomicallyNamed "modifyFileExists" $ do
+    forM_ changes $ \(f,c) ->
+        case fromChange c of
+            Just c' -> STM.focus (Focus.insert c') f var
+            Nothing -> pure ()
     -- See Note [Invalidating file existence results]
     -- flush previous values
     let (fileModifChanges, fileExistChanges) =
-            partition ((== FcChanged) . snd) (HashMap.toList changesMap)
-    join $ atomically $ do
-        mapM_ (deleteValue (shakeExtras state) GetFileExists . fst) fileExistChanges
-        io1 <- recordDirtyKeys (shakeExtras state) GetFileExists $ map fst fileExistChanges
-        io2 <- recordDirtyKeys (shakeExtras state) GetModificationTime $ map fst fileModifChanges
-        return (io1 <> io2)
+            partition ((== FcChanged) . snd) changes
+    mapM_ (deleteValue (shakeExtras state) GetFileExists . fst) fileExistChanges
+    io1 <- recordDirtyKeys (shakeExtras state) GetFileExists $ map fst fileExistChanges
+    io2 <- recordDirtyKeys (shakeExtras state) GetModificationTime $ map fst fileModifChanges
+    return (io1 <> io2)
 
 fromChange :: FileChangeType -> Maybe Bool
 fromChange FcCreated = Just True
@@ -164,7 +165,7 @@ fileExistsRules lspEnv vfs = do
   -- Create the global always, although it should only be used if we have fast rules.
   -- But there's a chance someone will send unexpected notifications anyway,
   -- e.g. https://github.com/haskell/ghcide/issues/599
-  addIdeGlobal . FileExistsMapVar =<< liftIO (newVar [])
+  addIdeGlobal . FileExistsMapVar =<< liftIO STM.newIO
 
   extras <- getShakeExtrasRules
   opts <- liftIO $ getIdeOptionsIO extras
@@ -213,7 +214,7 @@ fileExistsFast vfs file = do
     -- Could in principle use 'alwaysRerun' here, but it's too slwo, See Note [Invalidating file existence results]
     mp <- getFileExistsMapUntracked
 
-    let mbFilesWatched = HashMap.lookup file mp
+    mbFilesWatched <- liftIO $ atomically $ STM.lookup file mp
     exist <- case mbFilesWatched of
       Just exist -> pure exist
       -- We don't know about it: use the slow route.
diff --git a/ghcide/src/Development/IDE/Core/Shake.hs b/ghcide/src/Development/IDE/Core/Shake.hs
@@ -196,7 +196,7 @@ data ShakeExtras = ShakeExtras
     ,publishedDiagnostics :: STM.Map NormalizedUri [Diagnostic]
     -- ^ This represents the set of diagnostics that we have published.
     -- Due to debouncing not every change might get published.
-    ,positionMapping :: Var (HMap.HashMap NormalizedUri (Map TextDocumentVersion (PositionDelta, PositionMapping)))
+    ,positionMapping :: STM.Map NormalizedUri (Map TextDocumentVersion (PositionDelta, PositionMapping))
     -- ^ Map from a text document version to a PositionMapping that describes how to map
     -- positions in a version of that document to positions in the latest version
     -- First mapping is delta from previous version and second one is an
@@ -328,7 +328,6 @@ getIdeOptionsIO ide = do
 -- for the version of that value.
 lastValueIO :: IdeRule k v => ShakeExtras -> k -> NormalizedFilePath -> IO (Maybe (v, PositionMapping))
 lastValueIO s@ShakeExtras{positionMapping,persistentKeys,state} k file = do
-    allMappings <- readVar positionMapping
 
     let readPersistent
           | IdeTesting testing <- ideTesting s -- Don't read stale persistent values in tests
@@ -340,13 +339,13 @@ lastValueIO s@ShakeExtras{positionMapping,persistentKeys,state} k file = do
             f <- MaybeT $ pure $ HMap.lookup (Key k) pmap
             (dv,del,ver) <- MaybeT $ runIdeAction "lastValueIO" s $ f file
             MaybeT $ pure $ (,del,ver) <$> fromDynamic dv
-          case mv of
+          atomically $ case mv of
             Nothing -> do
-                void $ atomically $ STM.focus (Focus.alter (alterValue $ Failed True)) (toKey k file) state
+                STM.focus (Focus.alter (alterValue $ Failed True)) (toKey k file) state
                 return Nothing
             Just (v,del,ver) -> do
-                void $ atomically $ STM.focus (Focus.alter (alterValue $ Stale (Just del) ver (toDyn v))) (toKey k file) state
-                return $ Just (v,addDelta del $ mappingForVersion allMappings file ver)
+                STM.focus (Focus.alter (alterValue $ Stale (Just del) ver (toDyn v))) (toKey k file) state
+                Just . (v,) . addDelta del <$> mappingForVersion positionMapping file ver
 
         -- We got a new stale value from the persistent rule, insert it in the map without affecting diagnostics
         alterValue new Nothing = Just (ValueWithDiagnostics new mempty) -- If it wasn't in the map, give it empty diagnostics
@@ -359,8 +358,10 @@ lastValueIO s@ShakeExtras{positionMapping,persistentKeys,state} k file = do
     atomically (STM.lookup (toKey k file) state) >>= \case
       Nothing -> readPersistent
       Just (ValueWithDiagnostics v _) -> case v of
-        Succeeded ver (fromDynamic -> Just v) -> pure (Just (v, mappingForVersion allMappings file ver))
-        Stale del ver (fromDynamic -> Just v) -> pure (Just (v, maybe id addDelta del $ mappingForVersion allMappings file ver))
+        Succeeded ver (fromDynamic -> Just v) ->
+            atomically $ Just . (v,) <$> mappingForVersion positionMapping file ver
+        Stale del ver (fromDynamic -> Just v) ->
+            atomically $ Just . (v,) . maybe id addDelta del <$> mappingForVersion positionMapping file ver
         Failed p | not p -> readPersistent
         _ -> pure Nothing
 
@@ -372,14 +373,13 @@ lastValue key file = do
     liftIO $ lastValueIO s key file
 
 mappingForVersion
-    :: HMap.HashMap NormalizedUri (Map TextDocumentVersion (a, PositionMapping))
+    :: STM.Map NormalizedUri (Map TextDocumentVersion (a, PositionMapping))
     -> NormalizedFilePath
     -> TextDocumentVersion
-    -> PositionMapping
-mappingForVersion allMappings file ver =
-    maybe zeroMapping snd $
-    Map.lookup ver =<<
-    HMap.lookup (filePathToUri' file) allMappings
+    -> STM PositionMapping
+mappingForVersion allMappings file ver = do
+    mapping <- STM.lookup (filePathToUri' file) allMappings
+    return $ maybe zeroMapping snd $ Map.lookup ver =<< mapping
 
 type IdeRule k v =
   ( Shake.RuleResult k ~ v
@@ -513,7 +513,7 @@ shakeOpen lspEnv defaultConfig logger debouncer
         diagnostics <- STM.newIO
         hiddenDiagnostics <- STM.newIO
         publishedDiagnostics <- STM.newIO
-        positionMapping <- newVar HMap.empty
+        positionMapping <- STM.newIO
         knownTargetsVar <- newVar $ hashed HMap.empty
         let restartShakeSession = shakeRestart ideState
         persistentKeys <- newVar HMap.empty
@@ -1222,18 +1222,17 @@ getAllDiagnostics ::
 getAllDiagnostics =
     fmap (concatMap (\(k,v) -> map (fromUri k,ShowDiag,) $ getDiagnosticsFromStore v)) . ListT.toList . STM.listT
 
-updatePositionMapping :: IdeState -> VersionedTextDocumentIdentifier -> List TextDocumentContentChangeEvent -> IO ()
-updatePositionMapping IdeState{shakeExtras = ShakeExtras{positionMapping}} VersionedTextDocumentIdentifier{..} (List changes) = do
-    modifyVar_ positionMapping $ \allMappings -> do
-        let uri = toNormalizedUri _uri
-        let mappingForUri = HMap.lookupDefault Map.empty uri allMappings
-        let (_, updatedMapping) =
+updatePositionMapping :: IdeState -> VersionedTextDocumentIdentifier -> List TextDocumentContentChangeEvent -> STM ()
+updatePositionMapping IdeState{shakeExtras = ShakeExtras{positionMapping}} VersionedTextDocumentIdentifier{..} (List changes) =
+    STM.focus (Focus.alter f) uri positionMapping
+      where
+        uri = toNormalizedUri _uri
+        f = Just . f' . fromMaybe mempty
+        f' mappingForUri = snd $
                 -- Very important to use mapAccum here so that the tails of
                 -- each mapping can be shared, otherwise quadratic space is
                 -- used which is evident in long running sessions.
                 Map.mapAccumRWithKey (\acc _k (delta, _) -> let new = addDelta delta acc in (new, (delta, acc)))
                   zeroMapping
                   (Map.insert _version (shared_change, zeroMapping) mappingForUri)
-        pure $ HMap.insert uri updatedMapping allMappings
-  where
-    shared_change = mkDelta changes
+        shared_change = mkDelta changes
diff --git a/ghcide/src/Development/IDE/LSP/Notifications.hs b/ghcide/src/Development/IDE/LSP/Notifications.hs
@@ -14,6 +14,7 @@ module Development.IDE.LSP.Notifications
 import           Language.LSP.Types
 import qualified Language.LSP.Types                    as LSP
 
+import           Control.Concurrent.STM.Stats          (atomically)
 import           Control.Monad.Extra
 import           Control.Monad.IO.Class
 import qualified Data.HashMap.Strict                   as HM
@@ -42,7 +43,7 @@ descriptor :: PluginId -> PluginDescriptor IdeState
 descriptor plId = (defaultPluginDescriptor plId) { pluginNotificationHandlers = mconcat
   [ mkPluginNotificationHandler LSP.STextDocumentDidOpen $
       \ide _ (DidOpenTextDocumentParams TextDocumentItem{_uri,_version}) -> liftIO $ do
-      updatePositionMapping ide (VersionedTextDocumentIdentifier _uri (Just _version)) (List [])
+      atomically $ updatePositionMapping ide (VersionedTextDocumentIdentifier _uri (Just _version)) (List [])
       whenUriFile _uri $ \file -> do
           -- We don't know if the file actually exists, or if the contents match those on disk
           -- For example, vscode restores previously unsaved contents on open
@@ -52,7 +53,7 @@ descriptor plId = (defaultPluginDescriptor plId) { pluginNotificationHandlers =
 
   , mkPluginNotificationHandler LSP.STextDocumentDidChange $
       \ide _ (DidChangeTextDocumentParams identifier@VersionedTextDocumentIdentifier{_uri} changes) -> liftIO $ do
-        updatePositionMapping ide identifier changes
+        atomically $ updatePositionMapping ide identifier changes
         whenUriFile _uri $ \file -> do
           addFileOfInterest ide file Modified{firstOpen=False}
           setFileModified ide False file
