Make the cache thread-safe

Author: michaelpj

src/Ormolu/Utils/Cabal.hs

@@ -14,7 +14,6 @@ where
 import Control.Exception
 import Control.Monad.IO.Class
 import Data.ByteString qualified as B
-import Data.IORef
 import Data.Map.Lazy (Map)
 import Data.Map.Lazy qualified as M
 import Data.Maybe (maybeToList)
@@ -29,7 +28,7 @@ import Language.Haskell.Extension
 import Ormolu.Config
 import Ormolu.Exception
 import Ormolu.Fixity
-import Ormolu.Utils.IO (findClosestFileSatisfying, withIORefCache)
+import Ormolu.Utils.IO (Cache, findClosestFileSatisfying, newCache, withCache)
 import System.Directory
 import System.FilePath
 import System.IO.Unsafe (unsafePerformIO)
@@ -101,8 +100,8 @@ data CachedCabalFile = CachedCabalFile
   deriving (Show)
 
 -- | Cache ref that stores 'CachedCabalFile' per Cabal file.
-cacheRef :: IORef (Map FilePath CachedCabalFile)
-cacheRef = unsafePerformIO $ newIORef M.empty
+cacheRef :: Cache FilePath CachedCabalFile
+cacheRef = unsafePerformIO newCache
 {-# NOINLINE cacheRef #-}
 
 -- | Parse 'CabalInfo' from a @.cabal@ file at the given 'FilePath'.
@@ -118,7 +117,7 @@ parseCabalInfo ::
 parseCabalInfo cabalFileAsGiven sourceFileAsGiven = liftIO $ do
   cabalFile <- makeAbsolute cabalFileAsGiven
   sourceFileAbs <- makeAbsolute sourceFileAsGiven
-  CachedCabalFile {..} <- withIORefCache cacheRef cabalFile $ do
+  CachedCabalFile {..} <- withCache cacheRef cabalFile $ do
     cabalFileBs <- B.readFile cabalFile
     genericPackageDescription <-
       whenLeft (snd . runParseResult $ parseGenericPackageDescription cabalFileBs) $

src/Ormolu/Utils/Fixity.hs

@@ -10,18 +10,15 @@ where
 import Control.Exception (throwIO)
 import Control.Monad.IO.Class
 import Data.Bifunctor (first)
-import Data.IORef
 import Data.List.NonEmpty (NonEmpty)
-import Data.Map.Strict (Map)
-import Data.Map.Strict qualified as Map
 import Data.Text qualified as T
 import Data.Text.IO.Utf8 qualified as T.Utf8
 import Distribution.ModuleName (ModuleName)
 import Distribution.Types.PackageName (PackageName)
 import Ormolu.Exception
 import Ormolu.Fixity
 import Ormolu.Fixity.Parser
-import Ormolu.Utils.IO (findClosestFileSatisfying, withIORefCache)
+import Ormolu.Utils.IO (Cache, findClosestFileSatisfying, newCache, withCache)
 import System.Directory
 import System.IO.Unsafe (unsafePerformIO)
 import Text.Megaparsec (errorBundlePretty)
@@ -37,7 +34,7 @@ getDotOrmoluForSourceFile ::
   m (FixityOverrides, ModuleReexports)
 getDotOrmoluForSourceFile sourceFile =
   liftIO (findDotOrmoluFile sourceFile) >>= \case
-    Just dotOrmoluFile -> liftIO $ withIORefCache cacheRef dotOrmoluFile $ do
+    Just dotOrmoluFile -> liftIO $ withCache cacheRef dotOrmoluFile $ do
       dotOrmoluRelative <- makeRelativeToCurrentDirectory dotOrmoluFile
       contents <- T.Utf8.readFile dotOrmoluFile
       case parseDotOrmolu dotOrmoluRelative contents of
@@ -58,8 +55,8 @@ findDotOrmoluFile = findClosestFileSatisfying $ \x ->
   x == ".ormolu"
 
 -- | Cache ref that maps names of @.ormolu@ files to their contents.
-cacheRef :: IORef (Map FilePath (FixityOverrides, ModuleReexports))
-cacheRef = unsafePerformIO (newIORef Map.empty)
+cacheRef :: Cache FilePath (FixityOverrides, ModuleReexports)
+cacheRef = unsafePerformIO newCache
 {-# NOINLINE cacheRef #-}
 
 -- | A wrapper around 'parseFixityDeclaration' for parsing individual fixity

src/Ormolu/Utils/IO.hs

@@ -3,13 +3,15 @@
 
 module Ormolu.Utils.IO
   ( findClosestFileSatisfying,
-    withIORefCache,
+    Cache,
+    newCache,
+    withCache,
   )
 where
 
+import Control.Concurrent (MVar, modifyMVar_, newMVar, readMVar)
 import Control.Exception (catch, throwIO)
 import Control.Monad.IO.Class
-import Data.IORef
 import Data.Map.Lazy (Map)
 import Data.Map.Lazy qualified as M
 import System.Directory
@@ -48,14 +50,26 @@ findClosestFileSatisfying isRightFile rootOfSearch = liftIO $ do
         then pure Nothing
         else findClosestFileSatisfying isRightFile parentDir
 
+newtype Cache k v = Cache (MVar (Map k v))
+
+newCache :: (Ord k) => IO (Cache k v)
+newCache = do
+  var <- newMVar mempty
+  pure (Cache var)
+
 -- | Execute an 'IO' action but only if the given key is not found in the
--- 'IORef' cache.
-withIORefCache :: (Ord k) => IORef (Map k v) -> k -> IO v -> IO v
-withIORefCache cacheRef k action = do
-  cache <- readIORef cacheRef
+-- cache.
+withCache :: (Ord k) => Cache k v -> k -> IO v -> IO v
+withCache (Cache cacheVar) k action = do
+  -- Note we do not use modifyMVar here. The reason is that
+  -- 'action' may be long-running and we don't want to lock
+  -- the cache for the whole time. This means that potentially
+  -- we might end up running the same action twice if they
+  -- both start before one has populated the cache.
+  cache <- readMVar cacheVar
   case M.lookup k cache of
     Just v -> pure v
     Nothing -> do
       v <- action
-      modifyIORef' cacheRef (M.insert k v)
+      modifyMVar_ cacheVar (pure . M.insert k v)
       pure v