-- This Source Code Form is subject to the terms of the Mozilla Public
-- License, v. 2.0. If a copy of the MPL was not distributed with this
-- file, You can obtain one at http://mozilla.org/MPL/2.0/.

module Database.Redis.IO.Timeouts
    ( TimeoutManager
    , create
    , destroy
    , Action
    , Milliseconds (..)
    , add
    , cancel
    , withTimeout
    ) where

import Control.Applicative
import Control.Concurrent.STM
import Control.Exception (mask_, bracket)
import Control.Reaper
import Control.Monad
import Database.Redis.IO.Types (Milliseconds (..), ignore)
import Prelude

data TimeoutManager = TimeoutManager
    { TimeoutManager -> Int
roundtrip :: !Int
    , TimeoutManager -> Reaper [Action] Action
reaper    :: !(Reaper [Action] Action)
    }

data Action = Action
    { Action -> IO ()
action :: !(IO ())
    , Action -> TVar State
state  :: !(TVar State)
    }

data State = Running !Int | Canceled

create :: Milliseconds -> IO TimeoutManager
create :: Milliseconds -> IO TimeoutManager
create (Ms Int
n) = Int -> Reaper [Action] Action -> TimeoutManager
TimeoutManager Int
n (Reaper [Action] Action -> TimeoutManager)
-> IO (Reaper [Action] Action) -> IO TimeoutManager
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReaperSettings [Action] Action -> IO (Reaper [Action] Action)
forall workload item.
ReaperSettings workload item -> IO (Reaper workload item)
mkReaper ReaperSettings [Action] Action
forall item. ReaperSettings [item] item
defaultReaperSettings
    { reaperAction :: [Action] -> IO ([Action] -> [Action])
reaperAction = (Action -> IO (Maybe Action))
-> [Action] -> IO ([Action] -> [Action])
forall item item'.
(item -> IO (Maybe item')) -> [item] -> IO ([item'] -> [item'])
mkListAction Action -> IO (Maybe Action)
prune
    , reaperDelay :: Int
reaperDelay  = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000
    }
  where
    prune :: Action -> IO (Maybe Action)
prune Action
a = do
        State
s <- STM State -> IO State
forall a. STM a -> IO a
atomically (STM State -> IO State) -> STM State -> IO State
forall a b. (a -> b) -> a -> b
$ do
            State
x <- TVar State -> STM State
forall a. TVar a -> STM a
readTVar (Action -> TVar State
state Action
a)
            TVar State -> State -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (Action -> TVar State
state Action
a) (State -> State
newState State
x)
            State -> STM State
forall (m :: * -> *) a. Monad m => a -> m a
return State
x
        case State
s of
            Running Int
0 -> do
                IO () -> IO ()
ignore (Action -> IO ()
action Action
a)
                Maybe Action -> IO (Maybe Action)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Action
forall a. Maybe a
Nothing
            State
Canceled -> Maybe Action -> IO (Maybe Action)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Action
forall a. Maybe a
Nothing
            State
_        -> Maybe Action -> IO (Maybe Action)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Action -> IO (Maybe Action))
-> Maybe Action -> IO (Maybe Action)
forall a b. (a -> b) -> a -> b
$ Action -> Maybe Action
forall a. a -> Maybe a
Just Action
a

    newState :: State -> State
newState (Running Int
k) = Int -> State
Running (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    newState State
s           = State
s

destroy :: TimeoutManager -> Bool -> IO ()
destroy :: TimeoutManager -> Bool -> IO ()
destroy TimeoutManager
tm Bool
exec = IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    [Action]
a <- Reaper [Action] Action -> IO [Action]
forall workload item. Reaper workload item -> IO workload
reaperStop (TimeoutManager -> Reaper [Action] Action
reaper TimeoutManager
tm)
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
exec (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ (Action -> IO ()) -> [Action] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Action -> IO ()
f [Action]
a
  where
    f :: Action -> IO ()
f Action
e = TVar State -> IO State
forall a. TVar a -> IO a
readTVarIO (Action -> TVar State
state Action
e) IO State -> (State -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \State
s -> case State
s of
        Running Int
_ -> IO () -> IO ()
ignore (Action -> IO ()
action Action
e)
        State
Canceled  -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

add :: TimeoutManager -> Milliseconds -> IO () -> IO Action
add :: TimeoutManager -> Milliseconds -> IO () -> IO Action
add TimeoutManager
tm (Ms Int
n) IO ()
a = do
    Action
r <- IO () -> TVar State -> Action
Action IO ()
a (TVar State -> Action) -> IO (TVar State) -> IO Action
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State -> IO (TVar State)
forall a. a -> IO (TVar a)
newTVarIO (Int -> State
Running (Int -> State) -> Int -> State
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` TimeoutManager -> Int
roundtrip TimeoutManager
tm)
    Reaper [Action] Action -> Action -> IO ()
forall workload item. Reaper workload item -> item -> IO ()
reaperAdd (TimeoutManager -> Reaper [Action] Action
reaper TimeoutManager
tm) Action
r
    Action -> IO Action
forall (m :: * -> *) a. Monad m => a -> m a
return Action
r

cancel :: Action -> IO ()
cancel :: Action -> IO ()
cancel Action
a = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar State -> State -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (Action -> TVar State
state Action
a) State
Canceled

withTimeout :: TimeoutManager -> Milliseconds -> IO () -> IO a -> IO a
withTimeout :: TimeoutManager -> Milliseconds -> IO () -> IO a -> IO a
withTimeout TimeoutManager
tm Milliseconds
m IO ()
x IO a
a = IO Action -> (Action -> IO ()) -> (Action -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (TimeoutManager -> Milliseconds -> IO () -> IO Action
add TimeoutManager
tm Milliseconds
m IO ()
x) Action -> IO ()
cancel ((Action -> IO a) -> IO a) -> (Action -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ IO a -> Action -> IO a
forall a b. a -> b -> a
const IO a
a