{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Thread and process write lock.
--
-- Allows synchronisation between threads and processes.
-- Uses an 'MVar' for synchronisation between threads
-- and fcntl write locks for synchronisation between processes.
--
-- Only ever have one 'Lock' object per lock file per process!
module Data.CAS.Lock
  ( Lock,
    openLock,
    closeLock,
    withLock,
  )
where

import Control.Concurrent (threadDelay)
import Control.Exception.Safe
import Control.Monad (unless)
import Network.HostName (getHostName)
import Path
import Path.IO
import System.Posix.Files
import System.Posix.IO
import System.Posix.Process
import System.Random
import UnliftIO (MonadUnliftIO, withRunInIO)
import UnliftIO.MVar

-- | Thread and process write lock.
--
-- Only ever have one 'Lock' object per lock file per process!
data Lock = Lock
  { Lock -> MVar ()
lockMVar :: MVar (),
    Lock -> Path Abs Dir
lockDir :: Path Abs Dir
  }

-- | Open the lock file and create a lock object.
--
-- This does not acquire the lock.
--
-- Only ever have one 'Lock' object per lock file per process!
openLock :: Path Abs Dir -> IO Lock
openLock :: Path Abs Dir -> IO Lock
openLock Path Abs Dir
dir = do
  MVar ()
mvar <- () -> IO (MVar ())
forall (m :: * -> *) a. MonadIO m => a -> m (MVar a)
newMVar ()
  Bool -> Path Abs Dir -> IO ()
forall (m :: * -> *) b. MonadIO m => Bool -> Path b Dir -> m ()
createDirIfMissing Bool
True Path Abs Dir
dir
  Lock -> IO Lock
forall (m :: * -> *) a. Monad m => a -> m a
return
    (Lock -> IO Lock) -> Lock -> IO Lock
forall a b. (a -> b) -> a -> b
$! Lock :: MVar () -> Path Abs Dir -> Lock
Lock
      { lockMVar :: MVar ()
lockMVar = MVar ()
mvar,
        lockDir :: Path Abs Dir
lockDir = Path Abs Dir
dir
      }

-- | Close the lock file.
--
-- Does not release the lock.
--
-- Blocks if the lock is taken.
closeLock :: Lock -> IO ()
closeLock :: Lock -> IO ()
closeLock Lock
lock = do
  MVar () -> IO ()
forall (m :: * -> *) a. MonadIO m => MVar a -> m a
takeMVar (Lock -> MVar ()
lockMVar Lock
lock)

-- | Acquire the lock for the duration of the given action and release after.
withLock :: MonadUnliftIO m => Lock -> m a -> m a
withLock :: Lock -> m a -> m a
withLock Lock
lock m a
action = ((forall a. m a -> IO a) -> IO a) -> m a
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO a) -> m a)
-> ((forall a. m a -> IO a) -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
unliftIO ->
  MVar () -> (() -> IO a) -> IO a
forall (m :: * -> *) a b.
MonadUnliftIO m =>
MVar a -> (a -> m b) -> m b
withMVar (Lock -> MVar ()
lockMVar Lock
lock) ((() -> IO a) -> IO a) -> (() -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \() ->
    IO () -> IO () -> IO a -> IO a
forall (m :: * -> *) a b c. MonadMask m => m a -> m b -> m c -> m c
bracket_ (Path Abs Dir -> IO ()
acquireDirLock (Path Abs Dir -> IO ()) -> Path Abs Dir -> IO ()
forall a b. (a -> b) -> a -> b
$ Lock -> Path Abs Dir
lockDir Lock
lock) (Path Abs Dir -> IO ()
releaseDirLock (Path Abs Dir -> IO ()) -> Path Abs Dir -> IO ()
forall a b. (a -> b) -> a -> b
$ Lock -> Path Abs Dir
lockDir Lock
lock) (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$
      m a -> IO a
forall a. m a -> IO a
unliftIO m a
action

----------------------------------------------------------------------
-- Internals

-- | Generate unique (per process) filename.
--
-- Combines the host name and process ID.
getUniqueFileName :: IO (Path Rel File)
getUniqueFileName :: IO (Path Rel File)
getUniqueFileName = do
  HostName
hostName <- IO HostName
getHostName
  ProcessID
pid <- IO ProcessID
getProcessID
  HostName -> IO (Path Rel File)
forall (m :: * -> *). MonadThrow m => HostName -> m (Path Rel File)
parseRelFile (HostName -> IO (Path Rel File)) -> HostName -> IO (Path Rel File)
forall a b. (a -> b) -> a -> b
$ HostName
hostName HostName -> HostName -> HostName
forall a. [a] -> [a] -> [a]
++ ProcessID -> HostName
forall a. Show a => a -> HostName
show ProcessID
pid

lockFileName :: Path Rel File
lockFileName :: Path Rel File
lockFileName = [relfile|lock|]

-- | Acquire the lock.
--
-- Uses an algorithm that is described in the man-page of open (2) in the
-- last paragraph to @O_EXCL@ in release 4.14 of the Linux man-pages project.
--
-- Creates a file under a unique (per process) filename.
-- Attempts to hard-link that file to a common lock path.
-- If the operation succeeds, then the lock was acquired.
-- If not, but if the link count of the file under the unique filename
-- increased to two, then the lock was acquired.
-- Otherwise, another process holds the lock and this process waits
-- and retries.
acquireDirLock :: Path Abs Dir -> IO ()
acquireDirLock :: Path Abs Dir -> IO ()
acquireDirLock Path Abs Dir
dir = do
  Path Rel File
file <- IO (Path Rel File)
getUniqueFileName
  let path :: Path Abs File
path = Path Abs Dir
dir Path Abs Dir -> Path Rel File -> Path Abs File
forall b t. Path b Dir -> Path Rel t -> Path b t
</> Path Rel File
file
  Fd
fd <- HostName -> FileMode -> IO Fd
createFile (Path Abs File -> HostName
fromAbsFile Path Abs File
path) FileMode
ownerWriteMode
  Fd -> IO ()
closeFd Fd
fd
  Either IOError ()
r <- IO () -> IO (Either IOError ())
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
try (IO () -> IO (Either IOError ()))
-> IO () -> IO (Either IOError ())
forall a b. (a -> b) -> a -> b
$ HostName -> HostName -> IO ()
createLink (Path Abs File -> HostName
fromAbsFile Path Abs File
path) (Path Abs File -> HostName
fromAbsFile (Path Abs File -> HostName) -> Path Abs File -> HostName
forall a b. (a -> b) -> a -> b
$ Path Abs Dir
dir Path Abs Dir -> Path Rel File -> Path Abs File
forall b t. Path b Dir -> Path Rel t -> Path b t
</> Path Rel File
lockFileName)
  case Either IOError ()
r of
    Right () -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Left (IOError
_ :: IOError) -> do
      LinkCount
count <- FileStatus -> LinkCount
linkCount (FileStatus -> LinkCount) -> IO FileStatus -> IO LinkCount
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HostName -> IO FileStatus
getFileStatus (Path Abs File -> HostName
fromAbsFile Path Abs File
path)
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (LinkCount
count LinkCount -> LinkCount -> Bool
forall a. Eq a => a -> a -> Bool
== LinkCount
2) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        Int
delay <- (Int, Int) -> IO Int
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Int
50000, Int
100000)
        Int -> IO ()
threadDelay Int
delay
        Path Abs Dir -> IO ()
acquireDirLock Path Abs Dir
dir

-- | Release the lock.
--
-- Unlinks the file under the unique file name and the common lock file.
releaseDirLock :: Path Abs Dir -> IO ()
releaseDirLock :: Path Abs Dir -> IO ()
releaseDirLock Path Abs Dir
dir = do
  Path Rel File
file <- IO (Path Rel File)
getUniqueFileName
  let path :: Path Abs File
path = Path Abs Dir
dir Path Abs Dir -> Path Rel File -> Path Abs File
forall b t. Path b Dir -> Path Rel t -> Path b t
</> Path Rel File
file
  HostName -> IO ()
removeLink (Path Abs File -> HostName
fromAbsFile (Path Abs File -> HostName) -> Path Abs File -> HostName
forall a b. (a -> b) -> a -> b
$ Path Abs Dir
dir Path Abs Dir -> Path Rel File -> Path Abs File
forall b t. Path b Dir -> Path Rel t -> Path b t
</> Path Rel File
lockFileName)
  HostName -> IO ()
removeLink (Path Abs File -> HostName
fromAbsFile Path Abs File
path)