{-| An implementation of the mutex from
  <http://zookeeper.apache.org/doc/trunk/recipes.html#sc_recipes_Locks>.

  The mutex node will be automatically created and removed as necessary, and
  must not be used for any other purposes.  The parent node of the mutex node
  must already exist.
-}

module Zookeeper.Mutex ( withMutex ) where

import Control.Concurrent.MVar   ( newEmptyMVar, putMVar, takeMVar )
import Control.Exception as E    ( bracket, catch, catches, Handler(..) )
import Control.Monad             ( liftM )
import Control.Monad.Fix         ( fix )
import Data.List                 ( sort )
import System.FilePath.Posix     ( (</>), takeDirectory, takeFileName )
import qualified Zookeeper.Core as C

acquireMutex :: C.Handle -> String -> IO String
acquireMutex handle path = do
  -- Create the parent node to store the child nodes in.  If the parent node already exists then
  -- this will cause a NodeExists exception;  we catch (and ignore) this exception.
  (C.create handle path "" [] C.openAclUnsafe >> return ()) `E.catch` (\C.NodeExists -> return ())
  -- XXX There is a race here where a release call removes this directory.  Should we instead just
  -- not even attempt to do the automatic removal?
  -- Now create our child node in the directory.
  let child = path </> "lock-"
  c <- takeFileName `liftM` C.create handle child "" [C.Ephemeral, C.Sequence] C.openAclUnsafe
  fix $ \f -> do
    cs <- sort `liftM` C.getChildren handle path Nothing
    if c == head cs  -- Our child node is at the head; we're free to go.
      then return (path </> c)
      else do        -- Else wait on the child node that is lexicographically before us.
        let c' = last $ takeWhile (< c) cs
        m  <- newEmptyMVar
        e' <- C.exists handle (path </> c') $ Just $ \_ _ _ _ -> putMVar m ()
        if e' == Nothing then f else takeMVar m >> return (path </> c)

releaseMutex :: C.Handle -> String -> IO ()
releaseMutex handle child = do
  -- Delete the child node.
  C.delete handle child Nothing
  -- Try deleting the parent node.  If there are still extant child nodes then this will cause
  -- a NotEmpty exception; we catch (and ignore) this exception.  We also catch the case where
  -- another there is another aquire/release that removes the parent node ahead of us.
  C.delete handle (takeDirectory child) Nothing `E.catches`
    [E.Handler (\C.NotEmpty -> return ()), E.Handler (\C.NoNode -> return ())]

-- | Perform an action under mutual exclusion.
withMutex :: C.Handle  -- ^ The ZooKeeper handle.
          -> String    -- ^ The mutex node.
          -> IO ()     -- ^ The action to perform under mutual exclusion.
          -> IO ()
withMutex handle path act = do
  E.bracket (acquireMutex handle path) (releaseMutex handle) $ \_ -> act
