poolboy 0.1.0.1 → 0.2.0.0
raw patch · 3 files changed
+45/−39 lines, 3 filesdep +async
Dependencies added: async
Files
- poolboy.cabal +2/−1
- src/Data/Poolboy.hs +41/−32
- test/Spec.hs +2/−6
poolboy.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: poolboy-version: 0.1.0.1+version: 0.2.0.0 author: Gautier DI FOLCO maintainer: gautier.difolco@gmail.com category: Data@@ -18,6 +18,7 @@ base == 4.* , safe-exceptions == 0.1.* , stm == 2.*+ , async == 2.* hs-source-dirs: src exposed-modules: Data.Poolboy
src/Data/Poolboy.hs view
@@ -10,13 +10,18 @@ WorkQueue, withPoolboy, newPoolboy,- stopWorkQueue,- isStopWorkQueue, -- * Driving changeDesiredWorkersCount, waitReadyQueue, + -- * Stopping+ stopWorkQueue,+ isStopedWorkQueue,+ WaitingStopStrategy,+ waitingStopTimeout,+ waitingStopFinishWorkers,+ -- * Enqueueing enqueue, enqueueAfter,@@ -24,11 +29,13 @@ where import Control.Concurrent+import Control.Concurrent.Async import Control.Concurrent.STM.TQueue import Control.Exception.Safe (bracket, tryAny) import Control.Monad import Control.Monad.STM-import Data.IORef+import Data.Maybe (isNothing)+import System.Timeout (timeout) -- | Initial settings data PoolboySettings = PoolboySettings@@ -65,8 +72,8 @@ return () -- | 'backet'-based usage (recommended)-withPoolboy :: PoolboySettings -> (WorkQueue -> IO a) -> IO a-withPoolboy settings = bracket (newPoolboy settings) (\wq -> stopWorkQueue wq >> waitStopWorkQueue wq)+withPoolboy :: PoolboySettings -> WaitingStopStrategy -> (WorkQueue -> IO a) -> IO a+withPoolboy settings waitStopWorkQueue = bracket (newPoolboy settings) (\wq -> stopWorkQueue wq >> waitStopWorkQueue wq) -- | Standalone/manual usage newPoolboy :: PoolboySettings -> IO WorkQueue@@ -75,7 +82,6 @@ WorkQueue <$> newTQueueIO <*> newTQueueIO- <*> newIORef 0 <*> newEmptyMVar <*> return settings.log @@ -85,7 +91,7 @@ FixedWCS x -> return x changeDesiredWorkersCount wq count- void $ forkIO $ controller wq+ void $ forkIO $ controller wq [] return wq @@ -100,15 +106,21 @@ atomically $ writeTQueue wq.commands Stop -- | Non-blocking check of the work queue's running status-isStopWorkQueue :: WorkQueue -> IO Bool-isStopWorkQueue wq =+isStopedWorkQueue :: WorkQueue -> IO Bool+isStopedWorkQueue wq = not <$> isEmptyMVar wq.stopped +type WaitingStopStrategy = WorkQueue -> IO ()+ -- | Block until the queue is totally stopped (no more running worker)-waitStopWorkQueue :: WorkQueue -> IO ()-waitStopWorkQueue wq =+waitingStopFinishWorkers :: WaitingStopStrategy+waitingStopFinishWorkers wq = void $ tryAny $ readMVar wq.stopped +-- | Block until the queue is totally stopped or deadline (in micro seconds) is reached+waitingStopTimeout :: Int -> WaitingStopStrategy+waitingStopTimeout delay wq = void $ timeout delay $ waitingStopFinishWorkers wq+ -- | Enqueue one action in the work queue (non-blocking) enqueue :: WorkQueue -> IO () -> IO () enqueue wq =@@ -132,7 +144,6 @@ data WorkQueue = WorkQueue { commands :: TQueue Commands, queue :: TQueue (Either () (IO ())),- currentWorkersCount :: IORef Int, stopped :: MVar (), log :: String -> IO () }@@ -142,42 +153,40 @@ | Stop deriving stock (Show) -controller :: WorkQueue -> IO ()-controller wq = do+controller :: WorkQueue -> [Async ()] -> IO ()+controller wq workers = do command <- atomically $ readTQueue wq.commands let stopOneWorker = atomically $ writeTQueue wq.queue $ Left ()+ getLiveWorkers = filterM (fmap isNothing . poll) workers wq.log $ "Command: " <> show command case command of ChangeDesiredWorkersCount n -> do- currentCount <- readIORef wq.currentWorkersCount- let diff = currentCount - n- if diff > 0- then replicateM_ diff stopOneWorker- else replicateM_ (abs diff) $ do- wq.log "Pre-fork"- forkIO $ worker wq- controller wq+ liveWorkers <- getLiveWorkers+ let diff = length liveWorkers - n+ newWorkers <-+ if diff > 0+ then do+ replicateM_ diff stopOneWorker+ return []+ else replicateM (abs diff) $ do+ wq.log "Pre-fork"+ async $ worker wq+ controller wq $ newWorkers <> liveWorkers Stop -> do- currentCount <- readIORef wq.currentWorkersCount+ liveWorkers <- getLiveWorkers+ let currentCount = length liveWorkers wq.log $ "Stopping " <> show currentCount <> " workers" replicateM_ currentCount stopOneWorker+ forM_ liveWorkers waitCatch+ void $ tryPutMVar wq.stopped () worker :: WorkQueue -> IO () worker wq = do wq.log "New worker"- newCount <- atomicModifyIORef' wq.currentWorkersCount $ \n -> (n + 1, n + 1)- wq.log $ "New worker count " <> show newCount let loop = do command <- atomically $ readTQueue wq.queue case command of Left () -> do wq.log "Stopping"- remaining <-- atomicModifyIORef' wq.currentWorkersCount $ \n ->- let count = max 0 (n - 1) in (count, count)- wq.log $ "Remaining: " <> show remaining- when (remaining == 0) $- void $- tryPutMVar wq.stopped () Right act -> wq.log "pop" >> void (tryAny act) >> wq.log "poped" >> loop loop
test/Spec.hs view
@@ -18,20 +18,16 @@ it "threadDelay should be absorbed in mulitple threads" $ do computations <- timeout 2500 $- withPoolboy (poolboySettingsWith 100) $ \wq -> do+ withPoolboy (poolboySettingsWith 100) waitingStopFinishWorkers $ \wq -> do replicateM_ 100 $ enqueue wq $ threadDelay 1000- waitReadyQueue wq- threadDelay 1000 computations `shouldSatisfy` isJust replicateM_ 1 $ it "should be resilient to errors and Exceptions" $ do witness <- newIORef False computations <- timeout 10000 $- withPoolboy (poolboySettingsWith 5) $ \wq -> do+ withPoolboy (poolboySettingsWith 5) waitingStopFinishWorkers $ \wq -> do mapM_ (enqueue wq) [error "an error", throw RandomException, writeIORef witness True]- waitReadyQueue wq- threadDelay 100 computations `shouldSatisfy` isJust readIORef witness `shouldReturn` True