{-# LANGUAGE LambdaCase #-}
module MockQueue where
import Control.Concurrent ( threadDelay )
import Control.Concurrent.STM
import Control.Immortal.Queue
data Task
= Log Integer Int
-- ^ Always succeeds with Integer after Int milliseconds
| Fail String
-- ^ Always fails
queueConfig :: TVar ([Integer], [String]) -> TQueue Task -> ImmortalQueue Task
queueConfig output q = ImmortalQueue { qThreadCount = 2
, qPollWorkerTime = 200
, qPop = atomically $ readTQueue q
, qPush = atomically . writeTQueue q
, qHandler = performTask
, qFailure = handleError
}
where
performTask = \case
Log i t -> threadDelay (1000 * t) >> atomically (addSuccess i)
Fail _ -> error "failed"
handleError t _ = atomically $ case t of
Log _ _ -> return ()
Fail s -> addFailure s
addSuccess i = modifyTVar output $ \(s, f) -> (s ++ [i], f)
addFailure m = modifyTVar output $ \(s, f) -> (s, f ++ [m])
-- | Run pool that processes all the given tasks, splitting successes and
-- failures.
runPool :: [Task] -> IO ([Integer], [String])
runPool = runPool_ True Nothing
-- | Run a pool that processes the given tasks. `cleanClose` indicates if the
-- pool should be closed cleanly and `waitTime` will wait the specified
-- time before closing/killing or wait until the queue is empty if Nothing.
runPool_ :: Bool -> Maybe Int -> [Task] -> IO ([Integer], [String])
runPool_ cleanClose waitTime tasks = do
output <- newTVarIO ([], [])
tqueue <- newTQueueIO
workers <- processImmortalQueue $ queueConfig output tqueue
atomically $ mapM_ (writeTQueue tqueue) tasks
maybe (waitEmpty tqueue) (threadDelay . (* 1000)) waitTime
if cleanClose then closeImmortalQueue workers else killImmortalQueue workers
readTVarIO output
-- | Wait until the given queue is empty
waitEmpty :: TQueue a -> IO ()
waitEmpty q = do
isEmpty <- atomically $ isEmptyTQueue q
if isEmpty then return () else threadDelay 1000000 >> waitEmpty q