content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  -- |
  -- Module      : Control.Scheduler
  -- Copyright   : (c) Alexey Kuleshevich 2018-2021
  -- License     : BSD3
  -- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
  -- Stability   : experimental
  -- Portability : non-portable
  --
  module Control.Scheduler
    ( -- * Scheduler
      Scheduler
    , SchedulerWS
    , Results(..)
      -- ** Regular
    , withScheduler
    , withScheduler_
    , withSchedulerR
      -- ** Stateful workers
    , withSchedulerWS
    , withSchedulerWS_
    , withSchedulerWSR
    , unwrapSchedulerWS
      -- ** Trivial (no parallelism)
    , trivialScheduler_
    , withTrivialScheduler
    , withTrivialSchedulerR
    -- * Scheduling computation
    , scheduleWork
    , scheduleWork_
    , scheduleWorkId
    , scheduleWorkId_
    , scheduleWorkState
    , scheduleWorkState_
    , replicateWork
    , replicateWork_
    -- * Batches
    , Batch
    , runBatch
    , runBatch_
    , runBatchR
    , cancelBatch
    , cancelBatch_
    , cancelBatchWith
    , hasBatchFinished
    , getCurrentBatch
    -- * Early termination
    , terminate
    , terminate_
    , terminateWith
    -- * Workers
    , WorkerId(..)
    , WorkerStates
    , numWorkers
    , workerStatesComp
    , initWorkerStates
    -- * Computation strategies
    , Comp(..)
    , getCompWorkers
    -- * Useful functions
    , replicateConcurrently
    , replicateConcurrently_
    , traverseConcurrently
    , traverseConcurrently_
    , traverse_
    -- * Exceptions
    -- $exceptions
    , MutexException(..)
    ) where
  
  import Control.Monad
  import Control.Monad.ST
  import Control.Monad.IO.Unlift
  import Control.Monad.Primitive
  import Control.Scheduler.Computation
  import Control.Scheduler.Internal
  import Control.Scheduler.Types
  import Control.Scheduler.Queue
  import qualified Data.Foldable as F (traverse_, toList)
  import Data.Primitive.SmallArray
  import Data.Traversable
  
  
  -- | Get the underlying `Scheduler`, which cannot access `WorkerStates`.
  --
  -- @since 1.4.0
  unwrapSchedulerWS :: SchedulerWS ws a -> Scheduler RealWorld a
  unwrapSchedulerWS = _getScheduler
  
  
  -- | Get the computation strategy the states where initialized with.
  --
  -- @since 1.4.0
  workerStatesComp :: WorkerStates ws -> Comp
  workerStatesComp = _workerStatesComp
  
  
  -- | Run a scheduler with stateful workers. Throws `MutexException` if an attempt is made
  -- to concurrently use the same `WorkerStates` with another `SchedulerWS`.
  --
  -- ==== __Examples__
  --
  -- A good example of using stateful workers would be generation of random number
  -- in parallel. A lof of times random number generators are not thread safe, so
  -- we can work around this problem with a separate stateful generator for
  -- each of the workers.
  --
  -- >>> import Control.Monad as M ((>=>), replicateM)
  -- >>> import Control.Concurrent (yield, threadDelay)
  -- >>> import Data.List (sort)
  -- >>> -- ^ Above imports are used to make sure output is deterministic, which is needed for doctest
  -- >>> import System.Random.MWC as MWC
  -- >>> import Data.Vector.Unboxed as V (singleton)
  -- >>> states <- initWorkerStates (ParN 4) (MWC.initialize . V.singleton . fromIntegral . getWorkerId)
  -- >>> let scheduleGen scheduler = scheduleWorkState scheduler (MWC.uniform >=> \r -> yield >> threadDelay 200000 >> pure r)
  -- >>> sort <$> withSchedulerWS states (M.replicateM 4 . scheduleGen) :: IO [Double]
  -- [0.21734983682025255,0.5000843862105709,0.5759825622603018,0.8587171114177893]
  -- >>> sort <$> withSchedulerWS states (M.replicateM 4 . scheduleGen) :: IO [Double]
  -- [2.3598617298033475e-2,9.949679290089553e-2,0.38223134248645885,0.7408640677124702]
  --
  -- In the above example we use four different random number generators from
  -- [`mwc-random`](https://www.stackage.org/package/mwc-random) in order to generate 4
  -- numbers, all in separate threads. The subsequent call to the `withSchedulerWS` function
  -- with the same @states@ is allowed to reuse the same generators, thus avoiding expensive
  -- initialization.
  --
  -- /Side note/ - The example presented was crafted with slight trickery in order to
  -- guarantee that the output is deterministic, so if you run instructions exactly the same
  -- way in GHCI you will get the exact same output. Non-determinism comes from thread
  -- scheduling, rather than from random number generator, because we use exactly the same
  -- seed for each worker, but workers run concurrently. Exact output is not really needed,
  -- except for the doctests to pass.
  --
  -- @since 1.4.0
  withSchedulerWS ::
       MonadUnliftIO m => WorkerStates ws -> (SchedulerWS ws a -> m b) -> m [a]
  withSchedulerWS = withSchedulerWSInternal withScheduler
  
  -- | Run a scheduler with stateful workers, while discarding computation results.
  --
  -- @since 1.4.0
  withSchedulerWS_ ::
       MonadUnliftIO m => WorkerStates ws -> (SchedulerWS ws () -> m b) -> m ()
  withSchedulerWS_ = withSchedulerWSInternal withScheduler_
  
  -- | Same as `withSchedulerWS`, except instead of a list it produces `Results`, which
  -- allows for distinguishing between the ways computation was terminated.
  --
  -- @since 1.4.2
  withSchedulerWSR ::
       MonadUnliftIO m
    => WorkerStates ws
    -> (SchedulerWS ws a -> m b)
    -> m (Results a)
  withSchedulerWSR = withSchedulerWSInternal withSchedulerR
  
  
  -- | Schedule a job that will get a worker state passed as an argument
  --
  -- @since 1.4.0
  scheduleWorkState :: (PrimBase m, RealWorld ~ PrimState m) => SchedulerWS ws a -> (ws -> m a) -> m ()
  scheduleWorkState schedulerS withState =
    scheduleWorkId (_getScheduler schedulerS) $ \(WorkerId i) ->
      withState (indexSmallArray (_workerStatesArray (_workerStates schedulerS)) i)
  
  -- | Same as `scheduleWorkState`, but dont' keep the result of computation.
  --
  -- @since 1.4.0
  scheduleWorkState_ :: (PrimBase m, RealWorld ~ PrimState m) => SchedulerWS ws () -> (ws -> m ()) -> m ()
  scheduleWorkState_ schedulerS withState =
    scheduleWorkId_ (_getScheduler schedulerS) $ \(WorkerId i) ->
      withState (indexSmallArray (_workerStatesArray (_workerStates schedulerS)) i)
  
  
  -- | Get the number of workers. Will mainly depend on the computation strategy and/or number of
  -- capabilities you have. Related function is `getCompWorkers`.
  --
  -- @since 1.0.0
  numWorkers :: Scheduler s a -> Int
  numWorkers = _numWorkers
  
  
  -- | Schedule an action to be picked up and computed by a worker from a pool of
  -- jobs. Argument supplied to the job will be the id of the worker doing the job. This is
  -- useful for identification of a thread that will be doing the work, since there is
  -- one-to-one mapping from `Control.Concurrent.ThreadId` to `WorkerId` for a particular
  -- scheduler.
  --
  -- @since 1.2.0
  scheduleWorkId :: (PrimBase m, s ~ PrimState m) => Scheduler s a -> (WorkerId -> m a) -> m ()
  scheduleWorkId s f = stToPrim (_scheduleWorkId s (primToPrim . f))
  
  -- | As soon as possible try to terminate any computation that is being performed by all
  -- workers managed by this scheduler and collect whatever results have been computed, with
  -- supplied element guaranteed to being the last one. In case when `Results` type is
  -- returned this function will cause the scheduler to produce `FinishedEarly`
  --
  -- /Important/ - With `Seq` strategy this will not stop other scheduled tasks from being computed,
  -- although it will make sure their results are discarded.
  --
  -- @since 1.1.0
  terminate :: (PrimMonad m, s ~ PrimState m) => Scheduler s a -> a -> m a
  terminate scheduler a = stToPrim (_terminate scheduler (Early a))
  
  -- | Same as `terminate`, but returning a single element list containing the supplied
  -- argument. This can be very useful for parallel search algorithms. In case when
  -- `Results` is the return type this function will cause the scheduler to produce
  -- `FinishedEarlyWith`
  --
  -- /Important/ - Same as with `terminate`, when `Seq` strategy is used, this will not prevent
  -- computation from continuing, but the scheduler will return only the result supplied to this
  -- function.
  --
  -- @since 1.1.0
  terminateWith :: (PrimMonad m, s ~ PrimState m) => Scheduler s a -> a -> m a
  terminateWith scheduler a = stToPrim $ _terminate scheduler (EarlyWith a)
  
  -- | Schedule an action to be picked up and computed by a worker from a pool of
  -- jobs. Similar to `scheduleWorkId`, except the job doesn't get the worker id.
  --
  -- @since 1.0.0
  scheduleWork :: (PrimBase m, s ~ PrimState m) => Scheduler s a -> m a -> m ()
  scheduleWork scheduler f = stToPrim $ _scheduleWorkId scheduler (const (primToPrim f))
  
  
  -- FIXME: get rid of scheduleJob and decide at `scheduleWork` level if we should use Job or Job_
  -- Type here should be `scheduleWork_ :: Scheduler s a -> m () -> m ()
  -- | Same as `scheduleWork`, but only for a `Scheduler` that doesn't keep the results.
  --
  -- @since 1.1.0
  scheduleWork_ :: (PrimBase m, s ~ PrimState m) => Scheduler s () -> m () -> m ()
  scheduleWork_ s = stToPrim . scheduleWork s . primToPrim
  
  -- | Same as `scheduleWorkId`, but only for a `Scheduler` that doesn't keep the results.
  --
  -- @since 1.2.0
  scheduleWorkId_ :: (PrimBase m, s ~ PrimState m) => Scheduler s () -> (WorkerId -> m ()) -> m ()
  scheduleWorkId_ scheduler f = stToPrim $ _scheduleWorkId scheduler (primToPrim . f)
  
  -- | Schedule the same action to run @n@ times concurrently. This differs from
  -- `replicateConcurrently` by allowing the caller to use the `Scheduler` freely,
  -- or to allow early termination via `terminate` across all (identical) threads.
  -- To be called within a `withScheduler` block.
  --
  -- @since 2.0.0
  replicateWork :: (PrimBase m, s ~ PrimState m) => Scheduler s a -> Int -> m a -> m ()
  replicateWork scheduler n f = go n
    where
      go !k
        | k <= 0 = pure ()
        | otherwise = stToPrim (scheduleWork scheduler (primToPrim f)) *> go (k - 1)
  
  
  -- | Same as `replicateWork`, but it does not retain the results of scheduled jobs
  --
  -- @since 2.0.0
  replicateWork_ :: (PrimBase m, s ~ PrimState m) => Scheduler s () -> Int -> m a -> m ()
  replicateWork_ scheduler n f = go n
    where
      go !k
        | k <= 0 = pure ()
        | otherwise = stToPrim (scheduleWork_ scheduler (primToPrim (void f))) *> go (k - 1)
  
  -- | Similar to `terminate`, but for a `Scheduler` that does not keep any results of computation.
  --
  -- /Important/ - In case of `Seq` computation strategy this function has no affect.
  --
  -- @since 1.1.0
  terminate_ :: (PrimMonad m, s ~ PrimState m) => Scheduler s () -> m ()
  terminate_ s = stToPrim $ _terminate s (Early ())
  
  
  -- | This trivial scheduler will behave in the same way as `withScheduler` with `Seq`
  -- computation strategy, except it is restricted to `PrimMonad`, instead of `MonadUnliftIO`.
  --
  -- @since 1.4.2
  withTrivialScheduler :: (PrimMonad m, s ~ PrimState m) => (Scheduler s a -> m b) -> m [a]
  withTrivialScheduler action = F.toList <$> withTrivialSchedulerR action
  
  
  
  -- | Map an action over each element of the `Traversable` @t@ acccording to the supplied computation
  -- strategy.
  --
  -- @since 1.0.0
  traverseConcurrently :: (MonadUnliftIO m, Traversable t) => Comp -> (a -> m b) -> t a -> m (t b)
  traverseConcurrently comp f xs = withRunInIO $ \run -> do
    ys <- withScheduler comp $ \s -> traverse_ (scheduleWork s . run . f) xs
    pure $ transList ys xs
  
  transList :: Traversable t => [a] -> t b -> t a
  transList xs' = snd . mapAccumL withR xs'
    where
      withR (x:xs) _ = (xs, x)
      withR _      _ = errorWithoutStackTrace "Impossible<traverseConcurrently> - Mismatched sizes"
  
  -- | Just like `traverseConcurrently`, but restricted to `Foldable` and discards the results of
  -- computation.
  --
  -- @since 1.0.0
  traverseConcurrently_ :: (MonadUnliftIO m, Foldable t) => Comp -> (a -> m b) -> t a -> m ()
  traverseConcurrently_ comp f xs =
    withRunInIO $ \run ->
      withScheduler_ comp $ \s -> F.traverse_ (scheduleWork_ s . void . run . f) xs
  {-# INLINE traverseConcurrently_ #-}
  
  -- | Replicate an action @n@ times and schedule them acccording to the supplied computation
  -- strategy.
  --
  -- @since 1.1.0
  replicateConcurrently :: MonadUnliftIO m => Comp -> Int -> m a -> m [a]
  replicateConcurrently comp n f =
    withRunInIO $ \run ->
      withScheduler comp $ \s -> replicateM_ n $ scheduleWork s (run f)
  {-# INLINE replicateConcurrently #-}
  
  -- | Just like `replicateConcurrently`, but discards the results of computation.
  --
  -- @since 1.1.0
  replicateConcurrently_ :: MonadUnliftIO m => Comp -> Int -> m a -> m ()
  replicateConcurrently_ comp n f =
    withRunInIO $ \run -> do
      withScheduler_ comp $ \s -> replicateM_ n (scheduleWork_ s $ void $ run f)
  {-# INLINE replicateConcurrently_ #-}
  
  
  
  -- | Initialize a scheduler and submit jobs that will be computed sequentially or in parallelel,
  -- which is determined by the `Comp`utation strategy.
  --
  -- Here are some cool properties about the `withScheduler`:
  --
  -- * This function will block until all of the submitted jobs have finished or at least one of them
  --   resulted in an exception, which will be re-thrown at the callsite.
  --
  -- * It is totally fine for nested jobs to submit more jobs for the same or other scheduler
  --
  -- * It is ok to initialize multiple schedulers at the same time, although that will likely result
  --   in suboptimal performance, unless workers are pinned to different capabilities.
  --
  -- * __Warning__ It is pretty dangerous to schedule jobs that can block, because it might
  --   lead to a potential deadlock, if you are not careful. Consider this example. First
  --   execution works fine, since there are two scheduled workers, and one can unblock the
  --   other, but the second scenario immediately results in a deadlock.
  --
  -- >>> withScheduler (ParOn [1,2]) $ \s -> newEmptyMVar >>= (\ mv -> scheduleWork s (readMVar mv) >> scheduleWork s (putMVar mv ()))
  -- [(),()]
  -- >>> import System.Timeout
  -- >>> timeout 1000000 $ withScheduler (ParOn [1]) $ \s -> newEmptyMVar >>= (\ mv -> scheduleWork s (readMVar mv) >> scheduleWork s (putMVar mv ()))
  -- Nothing
  --
  -- __Important__: In order to get work done truly in parallel, program needs to be compiled with
  -- @-threaded@ GHC flag and executed with @+RTS -N -RTS@ to use all available cores.
  --
  -- @since 1.0.0
  withScheduler ::
       MonadUnliftIO m
    => Comp -- ^ Computation strategy
    -> (Scheduler RealWorld a -> m b)
       -- ^ Action that will be scheduling all the work.
    -> m [a]
  withScheduler Seq f =
    withRunInIO $ \run -> do
      reverse . resultsToList <$> withTrivialSchedulerRIO (run . f)
  withScheduler comp f =
    withRunInIO $ \run -> do
      reverse . resultsToList <$> withSchedulerInternal comp scheduleJobs readResults (run . f)
  {-# INLINE withScheduler #-}
  
  -- | Same as `withScheduler`, except instead of a list it produces `Results`, which allows
  -- for distinguishing between the ways computation was terminated.
  --
  -- @since 1.4.2
  withSchedulerR ::
       MonadUnliftIO m
    => Comp -- ^ Computation strategy
    -> (Scheduler RealWorld a -> m b)
       -- ^ Action that will be scheduling all the work.
    -> m (Results a)
  withSchedulerR Seq f =
    withRunInIO $ \run -> do
      reverseResults <$> withTrivialSchedulerRIO (run . f)
  withSchedulerR comp f =
    withRunInIO $ \run -> do
      reverseResults <$> withSchedulerInternal comp scheduleJobs readResults (run .f)
  {-# INLINE withSchedulerR #-}
  
  
  -- | Same as `withScheduler`, but discards results of submitted jobs.
  --
  -- @since 1.0.0
  withScheduler_ ::
       MonadUnliftIO m
    => Comp -- ^ Computation strategy
    -> (Scheduler RealWorld a -> m b)
       -- ^ Action that will be scheduling all the work.
    -> m ()
  withScheduler_ Seq f =
    withRunInIO $ \run -> do
      void $ withTrivialSchedulerRIO (run . f)
  withScheduler_ comp f =
    withRunInIO $ \run -> do
      void $ withSchedulerInternal comp scheduleJobs_ (const (pure [])) (run . f)
  {-# INLINE withScheduler_ #-}
  
  
  
  -- | Check if the supplied batch has already finished.
  --
  -- @since 1.5.0
  hasBatchFinished :: (PrimMonad m, s ~ PrimState m) => Batch s a -> m Bool
  hasBatchFinished = stToPrim . batchHasFinished
  {-# INLINE hasBatchFinished #-}
  
  
  -- | Cancel batch with supplied identifier, which will lead to scheduler to return
  -- `FinishedEarly` result. This is an idempotent operation and has no affect if currently
  -- running batch does not match supplied identifier. Returns `False` when cancelling did
  -- not succeed due to mismatched identifier or does not return at all since all jobs get
  -- cancelled immediately. For trivial schedulers however there is no way to perform
  -- concurrent cancelation and it will return `True`.
  --
  -- @since 1.5.0
  cancelBatch :: (PrimMonad m, s ~ PrimState m) => Batch s a -> a -> m Bool
  cancelBatch b = stToPrim . batchCancel b
  {-# INLINE cancelBatch #-}
  
  -- | Same as `cancelBatch`, but only works with schedulers that don't care about results
  --
  -- @since 1.5.0
  cancelBatch_ :: (PrimMonad m, s ~ PrimState m) => Batch s () -> m Bool
  cancelBatch_ b = stToPrim $ batchCancel b ()
  {-# INLINE cancelBatch_ #-}
  
  -- | Same as `cancelBatch_`, but the result of computation will be set to `FinishedEarlyWith`
  --
  -- @since 1.5.0
  cancelBatchWith :: (PrimMonad m, s ~ PrimState m) => Batch s a -> a -> m Bool
  cancelBatchWith b = stToPrim . batchCancelWith b
  {-# INLINE cancelBatchWith #-}
  
  
  -- | This function gives a way to get access to the main batch that started implicitely.
  --
  -- @since 1.5.0
  getCurrentBatch ::
       (PrimMonad m, s ~ PrimState m) => Scheduler s a -> m (Batch s a)
  getCurrentBatch scheduler = stToPrim $ do
    batchId <- _currentBatchId scheduler
    pure $ Batch
      { batchCancel = _cancelBatch scheduler batchId . Early
      , batchCancelWith = _cancelBatch scheduler batchId . EarlyWith
      , batchHasFinished = (batchId /=) <$> _currentBatchId scheduler
      }
  {-# INLINE getCurrentBatch #-}
  
  
  -- | Run a single batch of jobs. Supplied action will not return until all jobs placed on
  -- the queue are done or the whole batch is cancelled with one of these `cancelBatch`,
  -- `cancelBatch_` or `cancelBatchWith`.
  --
  -- It waits for all scheduled jobs to finish and collects the computed results into a
  -- list. It is a blocking operation, but if there are no jobs in progress it will return
  -- immediately. It is safe to continue using the supplied scheduler after this function
  -- returns. However, if any of the jobs resulted in an exception it will be rethrown by this
  -- function, which, unless caught, will further put the scheduler in a terminated state.
  --
  -- It is important to note that any job that hasn't had its results collected from the
  -- scheduler prior to starting the batch it will end up on the batch result list.
  --
  -- @since 1.5.0
  runBatch :: (PrimBase m, s ~ PrimState m) => Scheduler s a -> (Batch s a -> m c) -> m [a]
  runBatch scheduler f = stToPrim $ do
    _ <- primToPrim . f =<< getCurrentBatch scheduler
    reverse . resultsToList <$> _waitForCurrentBatch scheduler
  {-# INLINE runBatch #-}
  
  -- | Same as `runBatch`, except it ignores results of computation
  --
  -- @since 1.5.0
  runBatch_ ::
       (PrimBase m, s ~ PrimState m) => Scheduler s () -> (Batch s () -> m c) -> m ()
  runBatch_ scheduler f = stToPrim $ do
    _ <- primToPrim . f =<< getCurrentBatch scheduler
    void (_waitForCurrentBatch scheduler)
  {-# INLINE runBatch_ #-}
  
  
  -- | Same as `runBatch`, except it produces `Results` instead of a list.
  --
  -- @since 1.5.0
  runBatchR ::
       (PrimBase m, s ~ PrimState m) => Scheduler s a -> (Batch s a -> m c) -> m (Results a)
  runBatchR scheduler f = stToPrim $ do
    _ <- primToPrim . f =<< getCurrentBatch scheduler
    reverseResults <$> _waitForCurrentBatch scheduler
  {-# INLINE runBatchR #-}
  
  {- $setup
  
  >>> import Control.Exception
  >>> import Control.Concurrent
  >>> import Control.Concurrent.MVar
  
  -}
  
  
  {- $exceptions
  
  If any one of the workers dies with an exception, even if that exceptions is asynchronous, it will be
  re-thrown in the scheduling thread.
  
  
  >>> let didAWorkerDie = handleJust asyncExceptionFromException (return . (== ThreadKilled)) . fmap or
  >>> :t didAWorkerDie
  didAWorkerDie :: Foldable t => IO (t Bool) -> IO Bool
  >>> didAWorkerDie $ withScheduler Par $ \ s -> scheduleWork s $ pure False
  False
  >>> didAWorkerDie $ withScheduler Par $ \ s -> scheduleWork s $ myThreadId >>= killThread >> pure False
  True
  
  -}
  