Lock.hs

-- | Serial Haskell module compilation
module Language.Haskell.TH.Lock
  ( ensureSerialCompilation
  , ensureSerialCompilationQuietly
  , ensureSerialCompilationVerbose
  ) where

import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar, tryPutMVar, tryReadMVar, MVar)
import Control.Monad (void)
import Language.Haskell.TH.Syntax
import Prelude hiding (log)
import System.IO.Unsafe (unsafePerformIO)

-- | 'MVar' holds a Haskell module name expecting sequentian
-- compilation and which is currently goes through the compilation
-- process.
moduleCompilationLock :: MVar String
moduleCompilationLock = unsafePerformIO (newEmptyMVar)

-- | Call this function right after import section to prevent
-- concurrent TH code exection for modules with mutable compile time
-- shared state.  The function acquires a lock in the scope of GHC
-- process and the lock is released once type checker completes module
-- verification.
ensureSerialCompilation :: (String -> IO ()) -> Q [Dec]
ensureSerialCompilation log = do
  m <- loc_module <$> location
  go moduleCompilationLock m
  addModFinalizer (goAway moduleCompilationLock m)
  pure []
  where
    goAway l m = do
      runIO (tryReadMVar l) >>= \case
        Just holdingModule ->
          if holdingModule == m
          then runIO $ do
            log $ "Module [" <> m <> "] released lock"
            void $ takeMVar l
          else
            reportError $
              "Module [" <> m <> "] attempted to release lock of TH MVar holding by [" <>
                holdingModule <> "]"
        Nothing ->
          reportError $ "Module [" <> m <> "] attempted to release unlocked TH MVar"

    go l m = runIO (goIo l m)
    goIo l m =
      tryPutMVar l m >>= \case
        True -> do
          log $ "Module [" <> m <> "] acquired TH lock"
        False -> tryReadMVar l >>= \case
          Just holdingModule -> do
            log $ "Module [" <> m <> "] is waiting for [" <> holdingModule <> "]"
            putMVar l m
            log $ "Module [" <> m <> "] acquired TH lock"
          Nothing -> do
            log $ "Module [" <> m <> "] retries to acquire TH lock"
            goIo l m

-- | 'ensureSerialCompilation' without lock logging
ensureSerialCompilationQuietly :: Q [Dec]
ensureSerialCompilationQuietly = ensureSerialCompilation (void . pure)

-- | 'ensureSerialCompilation' with lock logging
ensureSerialCompilationVerbose :: Q [Dec]
ensureSerialCompilationVerbose = ensureSerialCompilation putStrLn