{-# LANGUAGE NamedFieldPuns #-}

module Test.Util.FileLock (mockFileLock) where

import Control.Monad (join, void)
import Control.Monad.IOSim
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import GHC.Stack (HasCallStack)
import Ouroboros.Consensus.Util.FileLock
import Ouroboros.Consensus.Util.IOLike
import Test.Util.Orphans.IOLike ()

{-------------------------------------------------------------------------------
  mockFileLock
-------------------------------------------------------------------------------}

-- | Mock in-memory implementation of 'FileLock'.
--
-- Supports an optional delay in release to simulate lazy, non-synchronous
-- unlocking as done by Linux (near instantaneous but not instant) and
-- Windows.
mockFileLock ::
  -- | Optional release delay
  Maybe DiffTime ->
  IOSim s (FileLock (IOSim s))
mockFileLock :: forall s. Maybe DiffTime -> IOSim s (FileLock (IOSim s))
mockFileLock Maybe DiffTime
releaseDelay = do
  locks <- Maybe DiffTime -> IOSim s (MockFileLocks (IOSim s))
forall (m :: * -> *).
IOLike m =>
Maybe DiffTime -> m (MockFileLocks m)
newMockFileLocks Maybe DiffTime
releaseDelay
  return
    FileLock
      { lockFile = \FilePath
fp -> MockFileLocks (IOSim s) -> FilePath -> IOSim s ()
forall (m :: * -> *).
(IOLike m, HasCallStack) =>
MockFileLocks m -> FilePath -> m ()
mockUnlockFile MockFileLocks (IOSim s)
locks (FilePath -> IOSim s ())
-> IOSim s FilePath -> IOSim s (IOSim s ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MockFileLocks (IOSim s) -> FilePath -> IOSim s FilePath
forall (m :: * -> *).
IOLike m =>
MockFileLocks m -> FilePath -> m FilePath
mockLockFile MockFileLocks (IOSim s)
locks FilePath
fp
      }

{-------------------------------------------------------------------------------
  MockFileLocks
-------------------------------------------------------------------------------}

data MockFileLocks m = MockFileLocks
  { forall (m :: * -> *).
MockFileLocks m -> StrictTVar m (Map FilePath LockStatus)
varLocks :: StrictTVar m (Map FilePath LockStatus)
  , forall (m :: * -> *). MockFileLocks m -> Maybe DiffTime
releaseDelay :: Maybe DiffTime
  }

-- | The status of a file lock, required to account for lazy releases. Note
-- that we don't have to model \"unlocked\", as the absence in 'varLocks'
-- already means that the lock is not held.
data LockStatus = Held | LazyRelease

newMockFileLocks :: IOLike m => Maybe DiffTime -> m (MockFileLocks m)
newMockFileLocks :: forall (m :: * -> *).
IOLike m =>
Maybe DiffTime -> m (MockFileLocks m)
newMockFileLocks Maybe DiffTime
releaseDelay = do
  varLocks <- Map FilePath LockStatus
-> m (StrictTVar m (Map FilePath LockStatus))
forall (m :: * -> *) a. MonadSTM m => a -> m (StrictTVar m a)
uncheckedNewTVarM Map FilePath LockStatus
forall k a. Map k a
Map.empty
  return MockFileLocks{varLocks, releaseDelay}

mockLockFile :: IOLike m => MockFileLocks m -> FilePath -> m FilePath
mockLockFile :: forall (m :: * -> *).
IOLike m =>
MockFileLocks m -> FilePath -> m FilePath
mockLockFile MockFileLocks{StrictTVar m (Map FilePath LockStatus)
varLocks :: forall (m :: * -> *).
MockFileLocks m -> StrictTVar m (Map FilePath LockStatus)
varLocks :: StrictTVar m (Map FilePath LockStatus)
varLocks} FilePath
path = STM m FilePath -> m FilePath
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m FilePath -> m FilePath) -> STM m FilePath -> m FilePath
forall a b. (a -> b) -> a -> b
$ do
  locks <- StrictTVar m (Map FilePath LockStatus)
-> STM m (Map FilePath LockStatus)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar StrictTVar m (Map FilePath LockStatus)
varLocks
  if Map.member path locks
    then retry
    else writeTVar varLocks $ Map.insert path Held locks
  return path

-- | We simulate lazy lock releases by changing the status of the lock to
-- 'LazyRelease' and spawning a thread that waits for 'releaseDelay' before
-- removing the lock from 'varLocks'.
mockUnlockFile ::
  (IOLike m, HasCallStack) =>
  MockFileLocks m -> FilePath -> m ()
mockUnlockFile :: forall (m :: * -> *).
(IOLike m, HasCallStack) =>
MockFileLocks m -> FilePath -> m ()
mockUnlockFile MockFileLocks{StrictTVar m (Map FilePath LockStatus)
varLocks :: forall (m :: * -> *).
MockFileLocks m -> StrictTVar m (Map FilePath LockStatus)
varLocks :: StrictTVar m (Map FilePath LockStatus)
varLocks, Maybe DiffTime
releaseDelay :: forall (m :: * -> *). MockFileLocks m -> Maybe DiffTime
releaseDelay :: Maybe DiffTime
releaseDelay} FilePath
path =
  m (m ()) -> m ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (m (m ()) -> m ()) -> m (m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ STM m (m ()) -> m (m ())
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (m ()) -> m (m ())) -> STM m (m ()) -> m (m ())
forall a b. (a -> b) -> a -> b
$ do
    locks <- StrictTVar m (Map FilePath LockStatus)
-> STM m (Map FilePath LockStatus)
forall (m :: * -> *) a. MonadSTM m => StrictTVar m a -> STM m a
readTVar StrictTVar m (Map FilePath LockStatus)
varLocks
    case Map.lookup path locks of
      Maybe LockStatus
Nothing ->
        FilePath -> STM m (m ())
forall a. HasCallStack => FilePath -> a
error (FilePath -> STM m (m ())) -> FilePath -> STM m (m ())
forall a b. (a -> b) -> a -> b
$ FilePath
"unlocking an unlocked file: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath -> FilePath
forall a. Show a => a -> FilePath
show FilePath
path
      Just LockStatus
LazyRelease ->
        FilePath -> STM m (m ())
forall a. HasCallStack => FilePath -> a
error (FilePath -> STM m (m ())) -> FilePath -> STM m (m ())
forall a b. (a -> b) -> a -> b
$ FilePath
"unlocking a file that is still being unlocked: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath -> FilePath
forall a. Show a => a -> FilePath
show FilePath
path
      Just LockStatus
Held -> case Maybe DiffTime
releaseDelay of
        Maybe DiffTime
Nothing -> do
          StrictTVar m (Map FilePath LockStatus)
-> Map FilePath LockStatus -> STM m ()
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
StrictTVar m a -> a -> STM m ()
writeTVar StrictTVar m (Map FilePath LockStatus)
varLocks (Map FilePath LockStatus -> STM m ())
-> Map FilePath LockStatus -> STM m ()
forall a b. (a -> b) -> a -> b
$ FilePath -> Map FilePath LockStatus -> Map FilePath LockStatus
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete FilePath
path Map FilePath LockStatus
locks
          m () -> STM m (m ())
forall a. a -> STM m a
forall (m :: * -> *) a. Monad m => a -> m a
return (m () -> STM m (m ())) -> m () -> STM m (m ())
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Just DiffTime
delay -> do
          StrictTVar m (Map FilePath LockStatus)
-> Map FilePath LockStatus -> STM m ()
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
StrictTVar m a -> a -> STM m ()
writeTVar StrictTVar m (Map FilePath LockStatus)
varLocks (Map FilePath LockStatus -> STM m ())
-> Map FilePath LockStatus -> STM m ()
forall a b. (a -> b) -> a -> b
$ FilePath
-> LockStatus -> Map FilePath LockStatus -> Map FilePath LockStatus
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert FilePath
path LockStatus
LazyRelease Map FilePath LockStatus
locks
          m () -> STM m (m ())
forall a. a -> STM m a
forall (m :: * -> *) a. Monad m => a -> m a
return (m () -> STM m (m ())) -> m () -> STM m (m ())
forall a b. (a -> b) -> a -> b
$ m (ThreadId m) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (ThreadId m) -> m ()) -> m (ThreadId m) -> m ()
forall a b. (a -> b) -> a -> b
$ m () -> m (ThreadId m)
forall (m :: * -> *). MonadFork m => m () -> m (ThreadId m)
forkIO (m () -> m (ThreadId m)) -> m () -> m (ThreadId m)
forall a b. (a -> b) -> a -> b
$ do
            DiffTime -> m ()
forall (m :: * -> *). MonadDelay m => DiffTime -> m ()
threadDelay DiffTime
delay
            STM m () -> m ()
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m () -> m ()) -> STM m () -> m ()
forall a b. (a -> b) -> a -> b
$ StrictTVar m (Map FilePath LockStatus)
-> Map FilePath LockStatus -> STM m ()
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
StrictTVar m a -> a -> STM m ()
writeTVar StrictTVar m (Map FilePath LockStatus)
varLocks (Map FilePath LockStatus -> STM m ())
-> Map FilePath LockStatus -> STM m ()
forall a b. (a -> b) -> a -> b
$ FilePath -> Map FilePath LockStatus -> Map FilePath LockStatus
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete FilePath
path Map FilePath LockStatus
locks