JobControl.hs

module UnitTests.Distribution.Client.JobControl (tests) where

import Distribution.Client.JobControl

import Distribution.Client.Compat.Prelude
import Prelude ()

import Control.Concurrent (threadDelay)
import Control.Exception (try)
import Control.Monad (replicateM, replicateM_)
import Data.IORef (atomicModifyIORef, newIORef)
import qualified Data.Set as Set

import Test.Tasty
import Test.Tasty.QuickCheck hiding (collect)

tests :: [TestTree]
tests =
  [ testGroup
      "serial"
      [ testProperty "submit batch" prop_submit_serial
      , testProperty "submit batch" prop_remaining_serial
      , testProperty "submit interleaved" prop_interleaved_serial
      , testProperty "concurrent jobs" prop_concurrent_serial
      , testProperty "cancel" prop_cancel_serial
      , testProperty "exceptions" prop_exception_serial
      ]
  , testGroup
      "parallel"
      [ testProperty "submit batch" prop_submit_parallel
      , testProperty "submit batch" prop_remaining_parallel
      , testProperty "submit interleaved" prop_interleaved_parallel
      , testProperty "concurrent jobs" prop_concurrent_parallel
      , testProperty "cancel" prop_cancel_parallel
      , testProperty "exceptions" prop_exception_parallel
      ]
  ]

prop_submit_serial :: [Int] -> Property
prop_submit_serial xs =
  ioProperty $ do
    jobCtl <- newSerialJobControl
    prop_submit jobCtl xs

prop_submit_parallel :: Positive (Small Int) -> [Int] -> Property
prop_submit_parallel (Positive (Small maxJobLimit)) xs =
  ioProperty $ do
    jobCtl <- newParallelJobControl maxJobLimit
    prop_submit jobCtl xs

prop_remaining_serial :: [Int] -> Property
prop_remaining_serial xs =
  ioProperty $ do
    jobCtl <- newSerialJobControl
    prop_remaining jobCtl xs

prop_remaining_parallel :: Positive (Small Int) -> [Int] -> Property
prop_remaining_parallel (Positive (Small maxJobLimit)) xs =
  ioProperty $ do
    jobCtl <- newParallelJobControl maxJobLimit
    prop_remaining jobCtl xs

prop_interleaved_serial :: [Int] -> Property
prop_interleaved_serial xs =
  ioProperty $ do
    jobCtl <- newSerialJobControl
    prop_submit_interleaved jobCtl xs

prop_interleaved_parallel :: Positive (Small Int) -> [Int] -> Property
prop_interleaved_parallel (Positive (Small maxJobLimit)) xs =
  ioProperty $ do
    jobCtl <- newParallelJobControl maxJobLimit
    prop_submit_interleaved jobCtl xs

prop_submit :: JobControl IO Int -> [Int] -> IO Bool
prop_submit jobCtl xs = do
  traverse_ (\x -> spawnJob jobCtl (return x)) xs
  xs' <- traverse (\_ -> collectJob jobCtl) xs
  return (sort xs == sort xs')

prop_remaining :: JobControl IO Int -> [Int] -> IO Bool
prop_remaining jobCtl xs = do
  traverse_ (\x -> spawnJob jobCtl (return x)) xs
  xs' <- collectRemainingJobs jobCtl
  return (sort xs == sort xs')

collectRemainingJobs :: Monad m => JobControl m a -> m [a]
collectRemainingJobs jobCtl = go []
  where
    go xs = do
      remaining <- remainingJobs jobCtl
      if remaining
        then do
          x <- collectJob jobCtl
          go (x : xs)
        else return xs

prop_submit_interleaved :: JobControl IO (Maybe Int) -> [Int] -> IO Bool
prop_submit_interleaved jobCtl xs = do
  xs' <-
    sequenceA
      [ spawn >> collect
      | let spawns =
              map (\x -> spawnJob jobCtl (return (Just x))) xs
                ++ repeat (return ())
            collects =
              replicate 5 (return Nothing)
                ++ map (\_ -> collectJob jobCtl) xs
      , (spawn, collect) <- zip spawns collects
      ]
  return (sort xs == sort (catMaybes xs'))

prop_concurrent_serial :: NonNegative (Small Int) -> Property
prop_concurrent_serial (NonNegative (Small ntasks)) =
  ioProperty $ do
    jobCtl <- newSerialJobControl
    countRef <- newIORef (0 :: Int)
    replicateM_ ntasks (spawnJob jobCtl (task countRef))
    counts <- replicateM ntasks (collectJob jobCtl)
    return $
      length counts == ntasks
        && all (\(n0, n1) -> n0 == 0 && n1 == 1) counts
  where
    task countRef = do
      n0 <- atomicModifyIORef countRef (\n -> (n + 1, n))
      threadDelay 100
      n1 <- atomicModifyIORef countRef (\n -> (n - 1, n))
      return (n0, n1)

prop_concurrent_parallel :: Positive (Small Int) -> NonNegative Int -> Property
prop_concurrent_parallel (Positive (Small maxJobLimit)) (NonNegative ntasks) =
  ioProperty $ do
    jobCtl <- newParallelJobControl maxJobLimit
    countRef <- newIORef (0 :: Int)
    replicateM_ ntasks (spawnJob jobCtl (task countRef))
    counts <- replicateM ntasks (collectJob jobCtl)
    return $
      length counts == ntasks
        && all
          ( \(n0, n1) ->
              n0 >= 0
                && n0 < maxJobLimit
                && n1 > 0
                && n1 <= maxJobLimit
          )
          counts
        -- we do hit the concurrency limit (in the right circumstances)
        && if ntasks >= maxJobLimit * 2 -- give us enough of a margin
          then any (\(_, n1) -> n1 == maxJobLimit) counts
          else True
  where
    task countRef = do
      n0 <- atomicModifyIORef countRef (\n -> (n + 1, n))
      threadDelay 100
      n1 <- atomicModifyIORef countRef (\n -> (n - 1, n))
      return (n0, n1)

prop_cancel_serial :: [Int] -> [Int] -> Property
prop_cancel_serial xs ys =
  ioProperty $ do
    jobCtl <- newSerialJobControl
    traverse_ (\x -> spawnJob jobCtl (return x)) (xs ++ ys)
    xs' <- traverse (\_ -> collectJob jobCtl) xs
    cancelJobs jobCtl
    ys' <- collectRemainingJobs jobCtl
    return (sort xs == sort xs' && null ys')

prop_cancel_parallel :: Positive (Small Int) -> [Int] -> [Int] -> Property
prop_cancel_parallel (Positive (Small maxJobLimit)) xs ys = do
  ioProperty $ do
    jobCtl <- newParallelJobControl maxJobLimit
    traverse_ (\x -> spawnJob jobCtl (threadDelay 100 >> return x)) (xs ++ ys)
    xs' <- traverse (\_ -> collectJob jobCtl) xs
    cancelJobs jobCtl
    ys' <- collectRemainingJobs jobCtl
    return $ Set.fromList (xs' ++ ys') `Set.isSubsetOf` Set.fromList (xs ++ ys)

data TestException = TestException Int
  deriving (Show)

instance Exception TestException

prop_exception_serial :: [Either Int Int] -> Property
prop_exception_serial xs =
  ioProperty $ do
    jobCtl <- newSerialJobControl
    prop_exception jobCtl xs

prop_exception_parallel :: Positive (Small Int) -> [Either Int Int] -> Property
prop_exception_parallel (Positive (Small maxJobLimit)) xs =
  ioProperty $ do
    jobCtl <- newParallelJobControl maxJobLimit
    prop_exception jobCtl xs

prop_exception :: JobControl IO Int -> [Either Int Int] -> IO Bool
prop_exception jobCtl xs = do
  traverse_ (\x -> spawnJob jobCtl (either (throwIO . TestException) return x)) xs
  xs' <- replicateM (length xs) $ do
    mx <- try (collectJob jobCtl)
    return $ case mx of
      Left (TestException n) -> Left n
      Right n -> Right n
  return (sort xs == sort xs')