module ZooKeeper.Recipe.Lock
  ( withLock

    -- * Internal
    -- Use them seperately is not safe!
  , lock
  , unlock
  ) where

import           Control.Exception      (bracket, catch, try)
import           Control.Monad
import           Z.Data.CBytes          (CBytes)

import           ZooKeeper
import           ZooKeeper.Exception    (ZNODEEXISTS, ZooException)
import           ZooKeeper.Recipe.Utils (SequenceNumWithGUID (..),
                                         createSeqEphemeralZNode,
                                         mkSequenceNumWithGUID)
import           ZooKeeper.Types

-- | To acquire a distributed lock.
-- Warning: do not forget to unlock it! Use 'withLock' instead if possible.
lock :: ZHandle
     -- ^ The zookeeper handle obtained by a call to 'zookeeperResInit'
     -> CBytes
     -- ^ The path to get the lock. Ephemeral znodes will be put on it
     -> CBytes
     -- ^ The GUID for this zookeeper session. To handle recoverable execptions
     -- correctly, it should be distinct from different sessions
     -> IO CBytes
     -- ^ The real path of the lock that acquired. It will be used when unlocking
     -- the same lock
lock :: ZHandle -> CBytes -> CBytes -> IO CBytes
lock ZHandle
zk CBytes
lockPath CBytes
guid = do
  -- Check persistent paths
  do Maybe StatCompletion
electionExists <- HasCallStack => ZHandle -> CBytes -> IO (Maybe StatCompletion)
ZHandle -> CBytes -> IO (Maybe StatCompletion)
zooExists ZHandle
zk CBytes
lockPath
     case Maybe StatCompletion
electionExists of
       Just StatCompletion
_  -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
       Maybe StatCompletion
Nothing -> IO (Either ZooException StringCompletion) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO StringCompletion -> IO (Either ZooException StringCompletion)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO StringCompletion -> IO (Either ZooException StringCompletion))
-> IO StringCompletion -> IO (Either ZooException StringCompletion)
forall a b. (a -> b) -> a -> b
$ HasCallStack =>
ZHandle
-> CBytes
-> Maybe Bytes
-> AclVector
-> CreateMode
-> IO StringCompletion
ZHandle
-> CBytes
-> Maybe Bytes
-> AclVector
-> CreateMode
-> IO StringCompletion
zooCreate ZHandle
zk CBytes
lockPath Maybe Bytes
forall a. Maybe a
Nothing AclVector
zooOpenAclUnsafe CreateMode
ZooPersistent :: IO (Either ZooException StringCompletion))
     IO () -> (ZNODEEXISTS -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(ZNODEEXISTS
_ :: ZNODEEXISTS) -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())

  -- Create Ephemeral and Sequece znode, and get the seq number i
  (StringCompletion CBytes
this) <- ZHandle -> CBytes -> CBytes -> IO StringCompletion
createSeqEphemeralZNode ZHandle
zk CBytes
lockPath CBytes
guid
  let thisSeqNumWithGUID :: SequenceNumWithGUID
thisSeqNumWithGUID = CBytes -> SequenceNumWithGUID
mkSequenceNumWithGUID CBytes
this

  ZHandle -> SequenceNumWithGUID -> IO CBytes
callback ZHandle
zk SequenceNumWithGUID
thisSeqNumWithGUID
  where
    callback :: ZHandle -> SequenceNumWithGUID -> IO CBytes
callback ZHandle
zk_ SequenceNumWithGUID
self = do
      (StringsCompletion (StringVector [CBytes]
children)) <- HasCallStack => ZHandle -> CBytes -> IO StringsCompletion
ZHandle -> CBytes -> IO StringsCompletion
zooGetChildren ZHandle
zk_ CBytes
lockPath
      let childrenSeqNumWithGUID :: [SequenceNumWithGUID]
childrenSeqNumWithGUID = CBytes -> SequenceNumWithGUID
mkSequenceNumWithGUID (CBytes -> SequenceNumWithGUID)
-> [CBytes] -> [SequenceNumWithGUID]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CBytes]
children
      case (SequenceNumWithGUID -> Bool)
-> [SequenceNumWithGUID] -> [SequenceNumWithGUID]
forall a. (a -> Bool) -> [a] -> [a]
filter (SequenceNumWithGUID -> SequenceNumWithGUID -> Bool
forall a. Ord a => a -> a -> Bool
< SequenceNumWithGUID
self) [SequenceNumWithGUID]
childrenSeqNumWithGUID of
        [] -> do
          CBytes -> IO CBytes
forall (m :: * -> *) a. Monad m => a -> m a
return (CBytes -> IO CBytes) -> CBytes -> IO CBytes
forall a b. (a -> b) -> a -> b
$ SequenceNumWithGUID -> CBytes
unSequenceNumWithGUID SequenceNumWithGUID
self
        [SequenceNumWithGUID]
xs -> do
          let toWatch :: CBytes
toWatch = CBytes
lockPath CBytes -> CBytes -> CBytes
forall a. Semigroup a => a -> a -> a
<> CBytes
"/" CBytes -> CBytes -> CBytes
forall a. Semigroup a => a -> a -> a
<> SequenceNumWithGUID -> CBytes
unSequenceNumWithGUID ([SequenceNumWithGUID] -> SequenceNumWithGUID
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [SequenceNumWithGUID]
xs)
          HasCallStack =>
ZHandle
-> CBytes
-> (HsWatcherCtx -> IO ())
-> (Maybe StatCompletion -> IO ())
-> IO ()
ZHandle
-> CBytes
-> (HsWatcherCtx -> IO ())
-> (Maybe StatCompletion -> IO ())
-> IO ()
zooWatchExists ZHandle
zk_ CBytes
toWatch (\HsWatcherCtx{Maybe CBytes
ZooEvent
ZooState
ZHandle
watcherCtxPath :: HsWatcherCtx -> Maybe CBytes
watcherCtxState :: HsWatcherCtx -> ZooState
watcherCtxType :: HsWatcherCtx -> ZooEvent
watcherCtxZHandle :: HsWatcherCtx -> ZHandle
watcherCtxPath :: Maybe CBytes
watcherCtxState :: ZooState
watcherCtxType :: ZooEvent
watcherCtxZHandle :: ZHandle
..} -> IO CBytes -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO CBytes -> IO ()) -> IO CBytes -> IO ()
forall a b. (a -> b) -> a -> b
$ ZHandle -> SequenceNumWithGUID -> IO CBytes
callback ZHandle
watcherCtxZHandle SequenceNumWithGUID
self)
            (\Maybe StatCompletion
_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
          CBytes -> IO CBytes
forall (m :: * -> *) a. Monad m => a -> m a
return (CBytes -> IO CBytes) -> CBytes -> IO CBytes
forall a b. (a -> b) -> a -> b
$ SequenceNumWithGUID -> CBytes
unSequenceNumWithGUID SequenceNumWithGUID
self

-- | To release a distributed lock. Note that the real lock path
-- should be the one acquired by 'lock', otherwise, a 'ZooException'
-- will be thrown.
unlock :: ZHandle
       -- ^ The zookeeper handle obtained by a call to 'zookeeperResInit'
       -> CBytes
       -- ^ The real lock path acquired by 'lock'. An exception will be
       -- thrown if it is bad (for example, does not exist)
       -> IO ()
unlock :: ZHandle -> CBytes -> IO ()
unlock ZHandle
zk CBytes
thisLock = HasCallStack => ZHandle -> CBytes -> Maybe CInt -> IO ()
ZHandle -> CBytes -> Maybe CInt -> IO ()
zooDelete ZHandle
zk CBytes
thisLock Maybe CInt
forall a. Maybe a
Nothing

-- | To do an action with a distributed lock. Only one caller with the same
-- 'lockPath' can execute the action at the same time. If the action throws
-- any exception during the locking period, the lock will be released and the
-- exception will be thrown again.
withLock :: ZHandle
         -- ^ The zookeeper handle obtained by a call to 'zookeeperResInit'
         -> CBytes
         -- ^ The path to get the lock. Ephemeral znodes will be put on it.
         -> CBytes
         -- ^ The GUID for this zookeeper session. To handle recoverable execptions
         -- correctly, it should be distinct from different sessions.
         -> IO a
         -- ^ The action to be executed within the lock.
         -> IO a
withLock :: ZHandle -> CBytes -> CBytes -> IO a -> IO a
withLock ZHandle
zk CBytes
lockPath CBytes
guid IO a
action =
  IO CBytes -> (CBytes -> IO ()) -> (CBytes -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (ZHandle -> CBytes -> CBytes -> IO CBytes
lock ZHandle
zk CBytes
lockPath CBytes
guid) (ZHandle -> CBytes -> IO ()
unlock ZHandle
zk) (IO a -> CBytes -> IO a
forall a b. a -> b -> a
const IO a
action)