{-# 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 ::
     Maybe DiffTime  -- ^ Optional release delay
  -> IOSim s (FileLock (IOSim s))
mockFileLock :: forall s. Maybe DiffTime -> IOSim s (FileLock (IOSim s))
mockFileLock Maybe DiffTime
releaseDelay = do
    MockFileLocks (IOSim s)
locks <- Maybe DiffTime -> IOSim s (MockFileLocks (IOSim s))
forall (m :: * -> *).
IOLike m =>
Maybe DiffTime -> m (MockFileLocks m)
newMockFileLocks Maybe DiffTime
releaseDelay
    FileLock (IOSim s) -> IOSim s (FileLock (IOSim s))
forall a. a -> IOSim s a
forall (m :: * -> *) a. Monad m => a -> m a
return FileLock {
        lockFile :: FilePath -> IOSim s (IOSim s ())
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
    StrictTVar m (Map FilePath LockStatus)
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
    MockFileLocks m -> m (MockFileLocks m)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return MockFileLocks { StrictTVar m (Map FilePath LockStatus)
varLocks :: StrictTVar m (Map FilePath LockStatus)
varLocks :: StrictTVar m (Map FilePath LockStatus)
varLocks, Maybe DiffTime
releaseDelay :: Maybe DiffTime
releaseDelay :: Maybe DiffTime
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
    Map FilePath LockStatus
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 FilePath -> Map FilePath LockStatus -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member FilePath
path Map FilePath LockStatus
locks
      then STM m ()
forall a. STM m a
forall (m :: * -> *) a. MonadSTM m => STM m a
retry
      else 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
Held Map FilePath LockStatus
locks
    FilePath -> STM m FilePath
forall a. a -> STM m a
forall (m :: * -> *) a. Monad m => a -> m a
return FilePath
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
      Map FilePath LockStatus
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 FilePath -> Map FilePath LockStatus -> Maybe LockStatus
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup FilePath
path Map FilePath LockStatus
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