{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Test.Hspec.Core.Runner.JobQueue (
  Concurrency(..)
, JobQueue
, Open
, Final
, new
, finalize
, run
, enqueue
, Report(..)
, AbortEarly(..)
, Result
) where

import           Prelude ()
import           Test.Hspec.Core.Compat hiding (all)

import           Foreign.Marshal.Utils (withMany)
import           Control.Concurrent
import           Control.Concurrent.Async (Async, withAsync, link, wait, cancelMany)

import           Test.Hspec.Core.Util (safeTry)
import           Test.Hspec.Core.Clock (Seconds, measure)

data Concurrency = Sequential | Concurrent

data Open
data Final

type QueuedJob progress a = IO (Continuation progress a)

newtype JobQueue final progress a = JobQueue (IORef [QueuedJob progress a])

new :: IO (JobQueue Open progress a)
new :: forall progress a. IO (JobQueue Open progress a)
new = IORef [QueuedJob progress a] -> JobQueue Open progress a
forall final progress a.
IORef [QueuedJob progress a] -> JobQueue final progress a
JobQueue (IORef [QueuedJob progress a] -> JobQueue Open progress a)
-> IO (IORef [QueuedJob progress a])
-> IO (JobQueue Open progress a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [QueuedJob progress a] -> IO (IORef [QueuedJob progress a])
forall a. a -> IO (IORef a)
newIORef []

enqueueJob :: JobQueue Open progress a -> QueuedJob progress a -> IO ()
enqueueJob :: forall progress a.
JobQueue Open progress a -> QueuedJob progress a -> IO ()
enqueueJob (JobQueue IORef [QueuedJob progress a]
ref) QueuedJob progress a
job = IORef [QueuedJob progress a]
-> ([QueuedJob progress a] -> [QueuedJob progress a]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef [QueuedJob progress a]
ref (QueuedJob progress a
job QueuedJob progress a
-> [QueuedJob progress a] -> [QueuedJob progress a]
forall a. a -> [a] -> [a]
:)

finalize :: JobQueue Open progress a -> IO (JobQueue Final progress a)
finalize :: forall progress a.
JobQueue Open progress a -> IO (JobQueue Final progress a)
finalize (JobQueue IORef [QueuedJob progress a]
ref) = do
  [QueuedJob progress a]
jobs <- [QueuedJob progress a] -> [QueuedJob progress a]
forall a. [a] -> [a]
reverse ([QueuedJob progress a] -> [QueuedJob progress a])
-> IO [QueuedJob progress a] -> IO [QueuedJob progress a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [QueuedJob progress a]
-> [QueuedJob progress a] -> IO [QueuedJob progress a]
forall a. IORef a -> a -> IO a
atomicSwapIORef IORef [QueuedJob progress a]
ref []
  IORef [QueuedJob progress a] -> JobQueue Final progress a
forall final progress a.
IORef [QueuedJob progress a] -> JobQueue final progress a
JobQueue (IORef [QueuedJob progress a] -> JobQueue Final progress a)
-> IO (IORef [QueuedJob progress a])
-> IO (JobQueue Final progress a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [QueuedJob progress a] -> IO (IORef [QueuedJob progress a])
forall a. a -> IO (IORef a)
newIORef [QueuedJob progress a]
jobs

dequeue :: JobQueue Final progress a -> IO (Maybe (QueuedJob progress a))
dequeue :: forall progress a.
JobQueue Final progress a -> IO (Maybe (QueuedJob progress a))
dequeue (JobQueue IORef [QueuedJob progress a]
ref) = IORef [QueuedJob progress a]
-> ([QueuedJob progress a]
    -> ([QueuedJob progress a], Maybe (QueuedJob progress a)))
-> IO (Maybe (QueuedJob progress a))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef [QueuedJob progress a]
ref (([QueuedJob progress a]
  -> ([QueuedJob progress a], Maybe (QueuedJob progress a)))
 -> IO (Maybe (QueuedJob progress a)))
-> ([QueuedJob progress a]
    -> ([QueuedJob progress a], Maybe (QueuedJob progress a)))
-> IO (Maybe (QueuedJob progress a))
forall a b. (a -> b) -> a -> b
$ \ case
  QueuedJob progress a
job : [QueuedJob progress a]
jobs -> ([QueuedJob progress a]
jobs, QueuedJob progress a -> Maybe (QueuedJob progress a)
forall a. a -> Maybe a
Just QueuedJob progress a
job)
  [] -> ([], Maybe (QueuedJob progress a)
forall a. Maybe a
Nothing)

data Done = Abort | Done

run :: forall progress a. Int -> JobQueue Final progress a -> [Report progress a] -> IO ()
run :: forall progress a.
Int -> JobQueue Final progress a -> [Report progress a] -> IO ()
run Int
concurrency JobQueue Final progress a
jobs [Report progress a]
items = do
  MVar Done
parent <- IO (MVar Done)
forall a. IO (MVar a)
newEmptyMVar
  let
    worker :: Continuation progress a -> IO ()
    worker :: Continuation progress a -> IO ()
worker = MVar Done
-> JobQueue Final progress a -> Continuation progress a -> IO ()
forall progress a.
MVar Done
-> JobQueue Final progress a -> Continuation progress a -> IO ()
workerThread MVar Done
parent JobQueue Final progress a
jobs

    workers :: [IO ()]
    workers :: [IO ()]
workers = Continuation progress a -> IO ()
worker ([Report progress a] -> Continuation progress a
forall progress a. [Report progress a] -> Continuation progress a
ReportingThread [Report progress a]
items) IO () -> [IO ()] -> [IO ()]
forall a. a -> [a] -> [a]
: Int -> IO () -> [IO ()]
forall a. Int -> a -> [a]
replicate (Int
concurrency Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Continuation progress a -> IO ()
worker Continuation progress a
forall progress a. Continuation progress a
WorkerThread)

  [IO ()] -> ([Async ()] -> IO ()) -> IO ()
forall a r. [IO a] -> ([Async a] -> IO r) -> IO r
withAsyncs [IO ()]
workers (([Async ()] -> IO ()) -> IO ()) -> ([Async ()] -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ [Async ()]
threads -> do
    (Async () -> IO ()) -> [Async ()] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Async () -> IO ()
forall a. Async a -> IO ()
link [Async ()]
threads
    MVar Done -> IO Done
forall a. MVar a -> IO a
takeMVar MVar Done
parent IO Done -> (Done -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
      Done
Abort -> [Async ()] -> IO ()
forall a. [Async a] -> IO ()
cancelMany [Async ()]
threads
      Done
Done -> (Async () -> IO ()) -> [Async ()] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Async () -> IO ()
forall a. Async a -> IO a
wait [Async ()]
threads

type Job progress a = (progress -> IO ()) -> IO a

data Result progress a =
    SequentialResult (Job progress a)
  | ConcurrentResult (IORef (ConcurrentJobState progress a))

enqueue :: JobQueue Open progress a -> Concurrency -> Job progress a -> IO (Result progress a)
enqueue :: forall progress a.
JobQueue Open progress a
-> Concurrency -> Job progress a -> IO (Result progress a)
enqueue JobQueue Open progress a
queue Concurrency
concurrency Job progress a
job = case Concurrency
concurrency of
  Concurrency
Sequential -> Result progress a -> IO (Result progress a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Result progress a -> IO (Result progress a))
-> Result progress a -> IO (Result progress a)
forall a b. (a -> b) -> a -> b
$ Job progress a -> Result progress a
forall progress a. Job progress a -> Result progress a
SequentialResult Job progress a
job
  Concurrency
Concurrent -> JobQueue Open progress a
-> Job progress a -> IO (Result progress a)
forall progress a.
JobQueue Open progress a
-> Job progress a -> IO (Result progress a)
enqueueConcurrent JobQueue Open progress a
queue Job progress a
job

data ConcurrentJobState progress a =
    Running
  | Reporting [Report progress a] (progress -> IO ())
  | Completed (Seconds, Either SomeException a)

enqueueConcurrent :: forall progress a. JobQueue Open progress a -> Job progress a -> IO (Result progress a)
enqueueConcurrent :: forall progress a.
JobQueue Open progress a
-> Job progress a -> IO (Result progress a)
enqueueConcurrent JobQueue Open progress a
queue Job progress a
job = do
  IORef (ConcurrentJobState progress a)
ref <- ConcurrentJobState progress a
-> IO (IORef (ConcurrentJobState progress a))
forall a. a -> IO (IORef a)
newIORef ConcurrentJobState progress a
forall progress a. ConcurrentJobState progress a
Running

  let
    reportProgress :: progress -> IO ()
    reportProgress :: progress -> IO ()
reportProgress progress
progress = IORef (ConcurrentJobState progress a)
-> IO (ConcurrentJobState progress a)
forall b. IORef b -> IO b
atomicReadIORef IORef (ConcurrentJobState progress a)
ref IO (ConcurrentJobState progress a)
-> (ConcurrentJobState progress a -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
      Reporting [Report progress a]
_ progress -> IO ()
report -> progress -> IO ()
report progress
progress
      ConcurrentJobState progress a
_ -> IO ()
forall (m :: * -> *). Applicative m => m ()
pass

    reportResult :: (Seconds, Either SomeException a) -> IO (Continuation progress a)
    reportResult :: (Seconds, Either SomeException a) -> IO (Continuation progress a)
reportResult (Seconds, Either SomeException a)
result = IORef (ConcurrentJobState progress a)
-> ConcurrentJobState progress a
-> IO (ConcurrentJobState progress a)
forall a. IORef a -> a -> IO a
atomicSwapIORef IORef (ConcurrentJobState progress a)
ref ((Seconds, Either SomeException a) -> ConcurrentJobState progress a
forall progress a.
(Seconds, Either SomeException a) -> ConcurrentJobState progress a
Completed (Seconds, Either SomeException a)
result) IO (ConcurrentJobState progress a)
-> (ConcurrentJobState progress a -> IO (Continuation progress a))
-> IO (Continuation progress a)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
      ConcurrentJobState progress a
Running -> Continuation progress a -> IO (Continuation progress a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Continuation progress a
forall progress a. Continuation progress a
WorkerThread
      Reporting [Report progress a]
items progress -> IO ()
_ -> Continuation progress a -> IO (Continuation progress a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Continuation progress a -> IO (Continuation progress a))
-> Continuation progress a -> IO (Continuation progress a)
forall a b. (a -> b) -> a -> b
$ [Report progress a] -> Continuation progress a
forall progress a. [Report progress a] -> Continuation progress a
ReportingThread [Report progress a]
items
      Completed (Seconds, Either SomeException a)
_ -> String -> IO (Continuation progress a)
forall a. HasCallStack => String -> a
thisCanNeverHappen String
"job completed twice"

  JobQueue Open progress a -> IO (Continuation progress a) -> IO ()
forall progress a.
JobQueue Open progress a -> QueuedJob progress a -> IO ()
enqueueJob JobQueue Open progress a
queue (IO (Continuation progress a) -> IO ())
-> IO (Continuation progress a) -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    IO a -> IO (Seconds, Either SomeException a)
forall a. IO a -> IO (Seconds, Either SomeException a)
eval (Job progress a
job progress -> IO ()
reportProgress) IO (Seconds, Either SomeException a)
-> ((Seconds, Either SomeException a)
    -> IO (Continuation progress a))
-> IO (Continuation progress a)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Seconds, Either SomeException a) -> IO (Continuation progress a)
reportResult

  Result progress a -> IO (Result progress a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Result progress a -> IO (Result progress a))
-> Result progress a -> IO (Result progress a)
forall a b. (a -> b) -> a -> b
$ IORef (ConcurrentJobState progress a) -> Result progress a
forall progress a.
IORef (ConcurrentJobState progress a) -> Result progress a
ConcurrentResult IORef (ConcurrentJobState progress a)
ref

data Report progress a =
    Report (IO ())
  | ReportResult {
      forall progress a. Report progress a -> Result progress a
_resultAction :: Result progress a
    , forall progress a. Report progress a -> progress -> IO ()
_reportProgress :: progress -> IO ()
    , forall progress a.
Report progress a
-> (Seconds, Either SomeException a) -> IO AbortEarly
_reportResult :: (Seconds, Either SomeException a) -> IO AbortEarly
    }

data AbortEarly = NoAbortEarly | AbortEarly

data Continuation progress a = WorkerThread | ReportingThread [Report progress a]

workerThread :: forall progress a. MVar Done -> JobQueue Final progress a -> Continuation progress a -> IO ()
workerThread :: forall progress a.
MVar Done
-> JobQueue Final progress a -> Continuation progress a -> IO ()
workerThread MVar Done
parent JobQueue Final progress a
jobs = Continuation progress a -> IO ()
loop
  where
    loop :: Continuation progress a -> IO ()
    loop :: Continuation progress a -> IO ()
loop = \ case
      Continuation progress a
WorkerThread -> JobQueue Final progress a -> IO (Maybe (QueuedJob progress a))
forall progress a.
JobQueue Final progress a -> IO (Maybe (QueuedJob progress a))
dequeue JobQueue Final progress a
jobs IO (Maybe (QueuedJob progress a))
-> (Maybe (QueuedJob progress a) -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
        Maybe (QueuedJob progress a)
Nothing -> IO ()
forall (m :: * -> *). Applicative m => m ()
pass
        Just QueuedJob progress a
job -> QueuedJob progress a
job QueuedJob progress a -> (Continuation progress a -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Continuation progress a -> IO ()
loop
      ReportingThread [] -> IO ()
done
      ReportingThread all :: [Report progress a]
all@(Report progress a
item : [Report progress a]
items) -> case Report progress a
item of
        Report IO ()
action -> do
          IO ()
action IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
keepReporting
        ReportResult Result progress a
result progress -> IO ()
reportProgress (((Seconds, Either SomeException a) -> IO AbortEarly)
-> (Seconds, Either SomeException a) -> IO () -> IO ()
forall r. (r -> IO AbortEarly) -> r -> IO () -> IO ()
checkAbort -> (Seconds, Either SomeException a) -> IO () -> IO ()
reportResult) -> case Result progress a
result of
          SequentialResult Job progress a
action -> do
            (Seconds, Either SomeException a)
r <- IO a -> IO (Seconds, Either SomeException a)
forall a. IO a -> IO (Seconds, Either SomeException a)
eval (Job progress a
action progress -> IO ()
reportProgress)
            (Seconds, Either SomeException a) -> IO () -> IO ()
reportResult (Seconds, Either SomeException a)
r IO ()
keepReporting
          ConcurrentResult IORef (ConcurrentJobState progress a)
ref -> IORef (ConcurrentJobState progress a)
-> ConcurrentJobState progress a
-> IO (ConcurrentJobState progress a)
forall a. IORef a -> a -> IO a
passReportingResponsibility IORef (ConcurrentJobState progress a)
ref ([Report progress a]
-> (progress -> IO ()) -> ConcurrentJobState progress a
forall progress a.
[Report progress a]
-> (progress -> IO ()) -> ConcurrentJobState progress a
Reporting [Report progress a]
all progress -> IO ()
reportProgress) IO (ConcurrentJobState progress a)
-> (ConcurrentJobState progress a -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
            ConcurrentJobState progress a
Running -> IO ()
continueAsWorkerThread
            Completed (Seconds, Either SomeException a)
r -> (Seconds, Either SomeException a) -> IO () -> IO ()
reportResult (Seconds, Either SomeException a)
r IO ()
keepReporting
            Reporting [Report progress a]
_ progress -> IO ()
_ -> String -> IO ()
forall a. HasCallStack => String -> a
thisCanNeverHappen String
"more than one reporting thread"
        where
          passReportingResponsibility :: IORef s -> s -> IO s
          passReportingResponsibility :: forall a. IORef a -> a -> IO a
passReportingResponsibility = IORef s -> s -> IO s
forall a. IORef a -> a -> IO a
atomicSwapIORef

          continueAsWorkerThread :: IO ()
          continueAsWorkerThread :: IO ()
continueAsWorkerThread = Continuation progress a -> IO ()
loop Continuation progress a
forall progress a. Continuation progress a
WorkerThread

          keepReporting :: IO ()
          keepReporting :: IO ()
keepReporting = Continuation progress a -> IO ()
loop ([Report progress a] -> Continuation progress a
forall progress a. [Report progress a] -> Continuation progress a
ReportingThread [Report progress a]
items)

    checkAbort :: (r -> IO AbortEarly) -> r -> IO () -> IO ()
    checkAbort :: forall r. (r -> IO AbortEarly) -> r -> IO () -> IO ()
checkAbort r -> IO AbortEarly
report r
r IO ()
continue = r -> IO AbortEarly
report r
r IO AbortEarly -> (AbortEarly -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
      AbortEarly
NoAbortEarly -> IO ()
continue
      AbortEarly
AbortEarly -> IO ()
abort

    abort :: IO ()
    abort :: IO ()
abort = MVar Done -> Done -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Done
parent Done
Abort

    done :: IO ()
    done :: IO ()
done = MVar Done -> Done -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Done
parent Done
Done

eval :: IO a -> IO (Seconds, Either SomeException a)
eval :: forall a. IO a -> IO (Seconds, Either SomeException a)
eval = IO (Either SomeException a) -> IO (Seconds, Either SomeException a)
forall a. IO a -> IO (Seconds, a)
measure (IO (Either SomeException a)
 -> IO (Seconds, Either SomeException a))
-> (IO a -> IO (Either SomeException a))
-> IO a
-> IO (Seconds, Either SomeException a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> IO (Either SomeException a)
forall a. IO a -> IO (Either SomeException a)
safeTry

withAsyncs :: forall a r. [IO a] -> ([Async a] -> IO r) -> IO r
withAsyncs :: forall a r. [IO a] -> ([Async a] -> IO r) -> IO r
withAsyncs = (IO a -> (Async a -> IO r) -> IO r)
-> [IO a] -> ([Async a] -> IO r) -> IO r
forall a b res.
(a -> (b -> res) -> res) -> [a] -> ([b] -> res) -> res
withMany IO a -> (Async a -> IO r) -> IO r
forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync

thisCanNeverHappen :: HasCallStack => String -> a
thisCanNeverHappen :: forall a. HasCallStack => String -> a
thisCanNeverHappen = String -> a
forall a. HasCallStack => String -> a
error