poolboy 0.3.0.0 → 0.4.0.0
raw patch · 4 files changed
+153/−42 lines, 4 files
Files
- CHANGELOG.md +8/−0
- poolboy.cabal +1/−1
- src/Data/Poolboy.hs +98/−41
- test/Spec.hs +46/−0
CHANGELOG.md view
@@ -1,3 +1,11 @@+## 0.4.0.0++* Fix race conditions+* Add command log for debugging+ * Add `PoolboyCommand`+ * Add `poolboySettingsLog`+ * Add `Monad` type parameters to `PoolboySettings` and `WorkQueue`+ ## 0.3.0.0 * Use `MonadUnliftIO` in functions
poolboy.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: poolboy-version: 0.3.0.0+version: 0.4.0.0 author: Gautier DI FOLCO maintainer: foss@difolco.dev category: Data
src/Data/Poolboy.hs view
@@ -4,9 +4,11 @@ ( -- * Configuration PoolboySettings (..), WorkersCountSettings (..),+ PoolboyCommand (..), defaultPoolboySettings, poolboySettingsWith, poolboySettingsName,+ poolboySettingsLog, -- * Running WorkQueue,@@ -37,7 +39,7 @@ import Control.Monad import Data.Functor (($>)) import qualified Data.HashMap.Strict as HM-import Data.Maybe (listToMaybe)+import Data.Maybe (isJust, listToMaybe) import GHC.Conc (labelThread) import GHC.Stack (HasCallStack, callStack, getCallStack, prettySrcLoc, withFrozenCallStack) import UnliftIO (MonadIO (liftIO), MonadUnliftIO)@@ -49,11 +51,11 @@ import UnliftIO.Timeout (timeout) -- | Initial settings-data PoolboySettings = PoolboySettings+data PoolboySettings m = PoolboySettings { workersCount :: WorkersCountSettings,- workQueueName :: String+ workQueueName :: String,+ logger :: PoolboyCommand -> m () }- deriving stock (Eq, Show) -- | Initial number of threads data WorkersCountSettings@@ -62,8 +64,31 @@ | FixedWCS Int -- arbitrary number deriving stock (Eq, Show) +-- | Commands performed on the queue+data PoolboyCommand+ = CreatePool ThreadId+ | ChangeDesiredWorkersCount ThreadId Int+ | SetMaxWorkersCount ThreadId Int+ | AddAvailableWorkers ThreadId Int+ | WaitAvailableWorkers ThreadId Int+ | SetStoppedWorkQueue ThreadId+ | WaitingStopFinishWorkers ThreadId+ | WaitingWorkersCompletion ThreadId [ThreadId]+ | ResetMaxWorkers ThreadId+ | EmptyAvailableWorkers ThreadId Int+ | WaitReady ThreadId+ | EnqueueAfter ThreadId+ | Enqueue ThreadId+ | EnsureRunning ThreadId Bool+ | SpawnTask ThreadId+ | WaitAvailableWorker ThreadId+ | StartTask ThreadId+ | CompleteTask ThreadId+ | EnqueueRegisterTask ThreadId ThreadId+ deriving stock (Eq, Show)+ -- | Usual configuration 'CapabilitiesWCS' and no name-defaultPoolboySettings :: (HasCallStack) => PoolboySettings+defaultPoolboySettings :: (HasCallStack, Monad m) => PoolboySettings m defaultPoolboySettings = withFrozenCallStack $ PoolboySettings@@ -71,33 +96,40 @@ workQueueName = case listToMaybe $ getCallStack callStack of Nothing -> "<no name>"- Just (f, loc) -> "<no name> created in " <> f <> ", called at " <> prettySrcLoc loc+ Just (f, loc) -> "<no name> created in " <> f <> ", called at " <> prettySrcLoc loc,+ logger = const $ return () } -- | Arbitrary-numbered of workers settings-poolboySettingsWith :: Int -> PoolboySettings+poolboySettingsWith :: (HasCallStack, Monad m) => Int -> PoolboySettings m poolboySettingsWith c = defaultPoolboySettings {workersCount = FixedWCS c} -- | Name of the work queue settings (used in error and debugging)-poolboySettingsName :: String -> PoolboySettings -> PoolboySettings+poolboySettingsName :: String -> PoolboySettings m -> PoolboySettings m poolboySettingsName n s = s {workQueueName = n} +-- | Name of the work queue settings (used in error and debugging)+poolboySettingsLog :: (Functor m) => (PoolboyCommand -> m a) -> PoolboySettings m -> PoolboySettings m+poolboySettingsLog f s = s {logger = void . f}+ -- | 'backet'-based usage (recommended)-withPoolboy :: (MonadUnliftIO m) => PoolboySettings -> WaitingStopStrategy m -> (WorkQueue -> m a) -> m a+withPoolboy :: (MonadUnliftIO m) => PoolboySettings m -> WaitingStopStrategy m -> (WorkQueue m -> m a) -> m a withPoolboy settings waitStopWorkQueue = bracket (newPoolboy settings) (\wq -> stopWorkQueue wq >> waitStopWorkQueue wq) -- | Standalone/manual usage-newPoolboy :: (MonadUnliftIO m) => PoolboySettings -> m WorkQueue+newPoolboy :: (MonadIO m) => PoolboySettings m -> m (WorkQueue m) newPoolboy settings = do+ settings.logger . CreatePool =<< myThreadId+ count <- case settings.workersCount of CapabilitiesWCS -> getNumCapabilities FixedWCS x -> return x - WorkQueue settings.workQueueName+ WorkQueue settings.workQueueName (\c -> settings.logger . c =<< myThreadId) <$> newQSemN count <*> newIORef count <*> newEmptyMVar@@ -106,50 +138,59 @@ -- | Request a worker number adjustment -- -- Warning: non-concurrent operation-changeDesiredWorkersCount :: (MonadUnliftIO m) => WorkQueue -> Int -> m ()+changeDesiredWorkersCount :: (MonadUnliftIO m) => WorkQueue m -> Int -> m () changeDesiredWorkersCount wq n = do+ wq.onCommand $ flip ChangeDesiredWorkersCount n ensureRunning wq refWorkers <- readIORef wq.maxWorkers when (n > 0) $ do+ wq.onCommand $ flip SetMaxWorkersCount n writeIORef wq.maxWorkers n if refWorkers < n- then signalQSemN wq.availableWorkers $ n - refWorkers- else- void $- async' wq $- waitQSemN wq.availableWorkers $- refWorkers - n+ then do+ wq.onCommand $ flip AddAvailableWorkers $ n - refWorkers+ signalQSemN wq.availableWorkers $ n - refWorkers+ else void $+ async' wq $ do+ wq.onCommand $ flip WaitAvailableWorkers $ refWorkers - n+ waitQSemN wq.availableWorkers $+ refWorkers - n -- | Request stopping wokers-stopWorkQueue :: (MonadUnliftIO m) => WorkQueue -> m ()+stopWorkQueue :: (MonadUnliftIO m) => WorkQueue m -> m () stopWorkQueue wq = do stopped <- isStopedWorkQueue wq unless stopped $ do+ wq.onCommand SetStoppedWorkQueue void $ tryPutMVar wq.stopped () -- | Non-blocking check of the work queue's running status-isStopedWorkQueue :: (MonadUnliftIO m) => WorkQueue -> m Bool+isStopedWorkQueue :: (MonadUnliftIO m) => WorkQueue m -> m Bool isStopedWorkQueue wq = not <$> isEmptyMVar wq.stopped -type WaitingStopStrategy m = WorkQueue -> m ()+type WaitingStopStrategy m = WorkQueue m -> m () -- | Block until the queue is totally stopped (no more running worker) waitingStopFinishWorkers :: (MonadUnliftIO m) => WaitingStopStrategy m waitingStopFinishWorkers wq =- void $+ void $ do+ wq.onCommand WaitingStopFinishWorkers retryBlocked 10 $ do readMVar wq.stopped let waitWorkerCompletion = do workers <- readIORef wq.inflightWorkers unless (HM.null workers) $ do- forM_ (HM.elems workers) waitCatch+ wq.onCommand $ flip WaitingWorkersCompletion $ HM.keys workers+ forM_ (HM.elems workers) $ mapM_ waitCatch atomicModifyIORef wq.inflightWorkers $ \ws ->- (foldr HM.delete ws $ HM.keys workers, ())+ (foldr HM.delete ws $ HM.keys $ HM.filter isJust workers, ()) threadDelay 10000 waitWorkerCompletion waitWorkerCompletion+ wq.onCommand ResetMaxWorkers workersCount <- atomicModifyIORef' wq.maxWorkers (0,)+ wq.onCommand $ flip EmptyAvailableWorkers workersCount waitQSemN wq.availableWorkers workersCount -- | Block until the queue is totally stopped or deadline (in micro seconds) is reached@@ -159,20 +200,21 @@ -- | Enqueue one action in the work queue (non-blocking) -- -- Throws 'WorkQueueStoppedException' if the work queue is stopped-enqueue :: (MonadUnliftIO m) => WorkQueue -> m () -> m ()+enqueue :: (MonadUnliftIO m) => WorkQueue m -> m () -> m () enqueue wq = void . enqueueTracking wq -- | Enqueue one action in the work queue (non-blocking) -- -- Throws 'WorkQueueStoppedException' if the work queue is stopped-enqueueTracking :: (MonadUnliftIO m) => WorkQueue -> m a -> m (Async a)+enqueueTracking :: (MonadUnliftIO m) => WorkQueue m -> m a -> m (Async a) enqueueTracking wq f = do ensureRunning wq enqueueTrackingAfterUnsafe wq (return ()) f -- | Block until one worker is available-waitReadyQueue :: (MonadUnliftIO m) => WorkQueue -> m ()+waitReadyQueue :: (MonadUnliftIO m) => WorkQueue m -> m () waitReadyQueue wq = do+ wq.onCommand WaitReady ensureRunning wq waitQSemN wq.availableWorkers 1 signalQSemN wq.availableWorkers 1@@ -180,34 +222,42 @@ -- | Enqueue action and some actions to be run after it -- -- Throws 'WorkQueueStoppedException' if the work queue is stopped-enqueueAfter :: (Traversable f, MonadUnliftIO m) => WorkQueue -> m () -> f (m ()) -> m ()+enqueueAfter :: (Traversable f, MonadUnliftIO m) => WorkQueue m -> m () -> f (m ()) -> m () enqueueAfter wq x xs = void $ enqueueAfterTracking wq x xs -- | Enqueue action and some actions to be run after it -- -- Throws 'WorkQueueStoppedException' if the work queue is stopped-enqueueAfterTracking :: (Traversable f, MonadUnliftIO m) => WorkQueue -> m a -> f (m b) -> m (Async a, f (Async b))+enqueueAfterTracking :: (Traversable f, MonadUnliftIO m) => WorkQueue m -> m a -> f (m b) -> m (Async a, f (Async b)) enqueueAfterTracking wq x xs = do+ wq.onCommand EnqueueAfter+ mainRef <- newEmptyMVar+ st <- forM xs $ enqueueTrackingAfterUnsafe wq (wait =<< readMVar mainRef) rm <- enqueueTracking wq x- st <- forM xs $ enqueueTrackingAfterUnsafe wq (wait rm)+ putMVar mainRef rm return (rm, st) -- Support (internal)-data WorkQueue = WorkQueue+data WorkQueue m = WorkQueue { name :: String,+ onCommand :: (ThreadId -> PoolboyCommand) -> m (), availableWorkers :: QSemN, maxWorkers :: IORef Int, stopped :: MVar (),- inflightWorkers :: IORef (HM.HashMap ThreadId (Async ()))+ inflightWorkers :: IORef (HM.HashMap ThreadId (Maybe (Async ()))) } -- | Ensure the queue is stopped -- -- Throws 'WorkQueueStoppedException' if not-ensureRunning :: (MonadUnliftIO m) => WorkQueue -> m ()+ensureRunning :: (MonadUnliftIO m) => WorkQueue m -> m () ensureRunning wq = do stopped <- isStopedWorkQueue wq- when stopped $+ threadId <- myThreadId+ workers <- readIORef wq.inflightWorkers+ let isRunning = not stopped || HM.member threadId workers+ wq.onCommand $ flip EnsureRunning isRunning+ unless isRunning $ throwIO $ WorkQueueStoppedException wq.name @@ -219,25 +269,32 @@ -- | Enqueue one action in the work queue (non-blocking) -- -- Does not check if the queue is stopping-enqueueTrackingAfterUnsafe :: forall a m i. (MonadUnliftIO m) => WorkQueue -> m i -> m a -> m (Async a)+enqueueTrackingAfterUnsafe :: forall a m i. (MonadUnliftIO m) => WorkQueue m -> m i -> m a -> m (Async a) enqueueTrackingAfterUnsafe wq prerequisite f = do+ wq.onCommand Enqueue+ let register threadId taskM = atomicModifyIORef wq.inflightWorkers (\ws -> (HM.insert threadId taskM ws, ())) task <- async' wq $ do+ wq.onCommand SpawnTask threadId <- myThreadId+ register threadId Nothing void prerequisite bracket_- (waitQSemN wq.availableWorkers 1)- ( signalQSemN wq.availableWorkers 1+ ( wq.onCommand WaitAvailableWorker+ >> waitQSemN wq.availableWorkers 1+ >> wq.onCommand StartTask+ )+ ( wq.onCommand CompleteTask >> atomicModifyIORef wq.inflightWorkers (\ws -> (HM.delete threadId ws, ()))+ >> signalQSemN wq.availableWorkers 1 ) f- atomicModifyIORef wq.inflightWorkers (\ws -> (HM.insert (asyncThreadId task) (task $> ()) ws, ()))+ wq.onCommand $ flip EnqueueRegisterTask (asyncThreadId task)+ register (asyncThreadId task) (Just $ task $> ()) return task --- f- -- | Start and label a thread-async' :: (MonadUnliftIO m) => WorkQueue -> m a -> m (Async a)+async' :: (MonadUnliftIO m) => WorkQueue m -> m a -> m (Async a) async' wq f = do t <- async f liftIO $ labelThread (asyncThreadId t) wq.name
test/Spec.hs view
@@ -52,6 +52,52 @@ computations `shouldSatisfy` isJust duration `shouldSatisfy` (< 2000) readIORef counter `shouldReturn` 110+ replicateM_ 1 $+ it "nested enqueueing should work on all jobs" $ do+ counter <- newIORef @Int 0+ let inc = atomicModifyIORef' counter $ \n -> (n + 1, ())+ (duration, computations) <-+ timeItT $+ timeout 100000000 $+ withPoolboy (poolboySettingsWith 30) waitingStopFinishWorkers $ \wq -> do+ replicateM_ 10 $+ enqueue wq $ do+ inc+ enqueue wq $ do+ inc+ enqueue wq $ do+ inc+ enqueue wq $ do+ inc+ computations `shouldSatisfy` isJust+ duration `shouldSatisfy` (< 2000)+ readIORef counter `shouldReturn` 40+ replicateM_ 1 $+ it "nested enqueueing on small pool should work on all jobs" $ do+ counter <- newIORef @Int 0+ tracesRef <- newIORef mempty+ let inc = atomicModifyIORef' counter $ \n -> (n + 1, ())+ addTrace c = atomicModifyIORef' tracesRef $ \cs -> (c : cs, ())+ (duration, computations) <-+ timeItT $+ timeout 100000000 $+ withPoolboy (poolboySettingsLog addTrace $ poolboySettingsWith 5) waitingStopFinishWorkers $ \wq -> do+ replicateM_ 5 $+ enqueue wq $ do+ inc+ enqueue wq $ do+ inc+ enqueue wq $ do+ inc+ enqueue wq $ do+ inc+ computations `shouldSatisfy` isJust+ duration `shouldSatisfy` (< 2000)+ c <- readIORef counter+ when (c /= 20) $ do+ cs <- readIORef tracesRef+ forM_ (reverse cs) print+ readIORef counter `shouldReturn` 20 data RandomException = RandomException deriving (Show)