-- | Polling helpers for tests.
module Arbiter.Test.Poll
  ( waitUntil
  , withLinkedAsync
  , bumpRef
  ) where

import Control.Concurrent (threadDelay)
import Control.Monad (unless)
import Data.IORef (IORef, atomicModifyIORef')
import GHC.Stack (HasCallStack)
import Test.Hspec (expectationFailure)
import UnliftIO (MonadUnliftIO)
import UnliftIO.Async (Async, link, withAsync)

-- | Milliseconds between polls of 'waitUntil'.
pollIntervalMs :: Int
pollIntervalMs :: Int
pollIntervalMs = Int
100

-- | Poll every 'pollIntervalMs' until the predicate returns 'True'.
-- Fails with 'expectationFailure' after @timeoutMs@ milliseconds.
waitUntil :: (HasCallStack) => Int -> IO Bool -> IO ()
waitUntil :: HasCallStack => Int -> IO Bool -> IO ()
waitUntil Int
timeoutMs IO Bool
check = HasCallStack => Int -> IO ()
Int -> IO ()
go (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int
timeoutMs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
pollIntervalMs))
  where
    go :: (HasCallStack) => Int -> IO ()
    go :: HasCallStack => Int -> IO ()
go Int
0 = HasCallStack => String -> IO ()
String -> IO ()
expectationFailure String
"waitUntil: timed out waiting for condition"
    go Int
remaining = do
      satisfied <- IO Bool
check
      unless satisfied $ do
        threadDelay (pollIntervalMs * microsPerMilli)
        go (remaining - 1)

microsPerMilli :: Int
microsPerMilli :: Int
microsPerMilli = Int
1_000

-- | Run 'withAsync' with a linked handle. Raise a child-thread failure in the
-- caller and stop the associated 'waitUntil'.
withLinkedAsync :: (MonadUnliftIO m) => m a -> (Async a -> m b) -> m b
withLinkedAsync :: forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withLinkedAsync m a
run Async a -> m b
body = m a -> (Async a -> m b) -> m b
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync m a
run ((Async a -> m b) -> m b) -> (Async a -> m b) -> m b
forall a b. (a -> b) -> a -> b
$ \Async a
handle -> Async a -> m ()
forall (m :: * -> *) a. MonadIO m => Async a -> m ()
link Async a
handle m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Async a -> m b
body Async a
handle

-- | Add one to a counter.
bumpRef :: (Num a) => IORef a -> IO ()
bumpRef :: forall a. Num a => IORef a -> IO ()
bumpRef IORef a
ref = IORef a -> (a -> (a, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef a
ref (\a
count -> (a
count a -> a -> a
forall a. Num a => a -> a -> a
+ a
1, ()))