diff --git a/conduit-concurrent-map.cabal b/conduit-concurrent-map.cabal --- a/conduit-concurrent-map.cabal +++ b/conduit-concurrent-map.cabal @@ -1,13 +1,13 @@ name: conduit-concurrent-map -version: 0.1.1 +version: 0.1.2 license: MIT copyright: 2017 Niklas Hambüchen author: Niklas Hambüchen maintainer: Niklas Hambüchen category: Data, Conduit build-type: Simple -tested-with: GHC==8.2.2, GHC==8.4.3 -cabal-version: >= 1.8 +tested-with: GHC==8.10.7 +cabal-version: >= 1.10 homepage: https://github.com/nh2/conduit-concurrent-map bug-Reports: https://github.com/nh2/conduit-concurrent-map/issues synopsis: Concurrent, order-preserving mapping Conduit @@ -25,7 +25,7 @@ src build-depends: base >= 4 && < 5 - , conduit + , conduit >= 1.3.0 , containers , mtl , resourcet @@ -33,6 +33,7 @@ , unliftio-core , vector + default-language: Haskell2010 ghc-options: -Wall @@ -49,6 +50,8 @@ , conduit , hspec >= 1.3.0.1 , HUnit >= 1.2 + , QuickCheck , say + default-language: Haskell2010 ghc-options: -Wall -threaded -with-rtsopts=-N4 diff --git a/src/Data/Conduit/ConcurrentMap.hs b/src/Data/Conduit/ConcurrentMap.hs --- a/src/Data/Conduit/ConcurrentMap.hs +++ b/src/Data/Conduit/ConcurrentMap.hs @@ -12,7 +12,7 @@ import Control.Monad (when) import Control.Monad.IO.Class (liftIO) -import Control.Monad.IO.Unlift (MonadUnliftIO, askRunInIO) +import Control.Monad.IO.Unlift (MonadUnliftIO, UnliftIO, unliftIO, askUnliftIO) import Control.Monad.Trans (lift) import Control.Monad.Trans.Resource (MonadResource) import Data.Conduit (ConduitT, await, bracketP) @@ -45,8 +45,10 @@ a :< _ -> Just a --- | Concurrent, order-preserving conduit mapping function. +-- | @concurrentMapM_ numThreads workerOutputBufferSize f@ -- +-- Concurrent, order-preserving conduit mapping function. +-- -- Like `Data.Conduit.mapM`, but runs in parallel with the given number of threads, -- returns outputs in the order of inputs (like @mapM@, no reordering), -- and allows defining a bounded size output buffer for elements of type @b@ to @@ -59,15 +61,15 @@ -- Unless we buffer the queued result somewhere, the thread that finished the -- short-running computation is now blocked and sits idle (low utilisation). -- --- To cope with this, @concurrentMapM_ numThreads workerOutputBufferSize f@ gives each +-- To cope with this, this function gives each -- thread @workerOutputBufferSize@ output slots to store @b@s while they are blocked. -- --- Use the convenience `concurrentMapM_` when @f@ is CPU-bound. +-- Use the convenience `concurrentMapM_numCaps` when @f@ is CPU-bound. -- -- @workerOutputBufferSize@ must be given >= 1. -- -- The @workerOutputBufferSize@ keeps the memory usage of the conduit bounded, --- namely to @getNumCapabilities * (workerOutputBufferSize + 1)@ many @b@s at any +-- namely to @numThreads * (workerOutputBufferSize + 1)@ many @b@s at any -- given time (the @+ 1@ is for the currently processing ones). -- -- To achieve maximum parallelism/utilisation, you should choose @@ -93,7 +95,7 @@ -- the same order as their corresponding `a`s came in (the parallelism -- doesn't change the order). -- * Bounded memory: The conduit will only hold to --- @getNumCapabilities * (workerOutputBufferSize + 1)@ as many @b@s. +-- @numThreads * (workerOutputBufferSize + 1)@ as many @b@s. -- * Full utilisation: The conduit will try to keep all cores busy as much as -- it can. This means that it will always try to `await` if there's a free -- core, and will only `yield` once it has to to make a core free. @@ -125,7 +127,7 @@ -- ------------------------- [ workerOutVar(N-1)a workerOutVar(N-1)b ... ] <- f / -- [ workerOutVar(N )a workerOutVar(N )b ... ] <- f / -- o <- button to signal - -- inVarInqueued + -- inVarEnqueued -- -- Any worker that's not busy is hanging onto `inVar`, grabbing -- its contents as soon as `inVar` is filled. @@ -147,7 +149,68 @@ -- block on putting their `b`s in, so there are maximally -- `N * (workerOutputBufferSize + 1)` many `b`s held alive -- by this function. - + -- + -- Note that as per this "+ 1" logic, for each worker there may up to 1 + -- `workerOutVar` that is in in the `outQueue` twice. + -- For example, for `numThreads = 2` and `workerOutputBufferSize = 2`, + -- we may have: + -- + -- ------------------------- [ worker1OutVarSlotA worker1OutVarSlotB ] <- f \ + -- outQueue of workerOutVars - inVar + -- ------------------------- [ worker2OutVarSlotA worker2OutVarSlotB ] <- f / + -- + -- with an input conduit streaming elements + -- [A, B, C, D] + -- with processing times + -- [9, 0, 0, 0] + -- this may lead to an `outQueue` as follows: + -- + -- +-----------------------------------+ + -- | | + -- V | + -- ------------------------- [ worker1OutVarSlot_a worker1OutVarSlot_a ] <- f \ + -- A B C - inVar (containing element D) + -- ------------------------- [ worker2OutVarSlot_b worker2OutVarSlot_b ] <- f / + -- ^ ^ | | + -- +--|-----------------------------+ | + -- | | + -- +--------------------------------------------------+ + -- + -- where worker 1 is still processing work item A, and worker 2 has just finished + -- processing work items B and C. + -- Now worker 2 is idle, pops element D as the next work item from the `inVar`, + -- and enqueues enqueues MVar `worker2OutVarSlot_b` into `outQueue`, + -- processes element D, and runs `putMVar worker2OutVarSlot_b (f D)`; + -- it is at this time that worker 2 blocks until `worker2OutVarSlot_b` + -- is emptied when the conduit `yield`s the result. + -- Thus we have this situation: + -- + -- +-----------------------------------+ + -- | | + -- V | + -- ------------------------- [ worker1OutVarSlot_a worker1OutVarSlot_a ] <- f \ + -- A B C D - inVar + -- ------------------------- [ worker2OutVarSlot_b worker2OutVarSlot_b ] <- f / + -- ^ ^ ^ | | | + -- +--|--|--------------------------+ | | + -- | | | | + -- +--|-----------------------------|-----------------+ + -- | | + -- +-----------------------------+ + -- + -- It is thus NOT an invariant that every `outVar` is in the `outQueue` only once. + -- + -- TODO: This whole design has producing the "+ 1" logic has a bit of an ugliness + -- in that it's not possible to make each worker use at max 1 `b`; only + -- 2 or more `b`s are possible. + -- The whole design might be simplified by changing it so that instead + -- of each worker having a fixed number of `workerOutVar`s, + -- workers make up new `workerOutVar`s on demand (enqueuing them + -- into `outQueue` as before), and the conduit keeping track of + -- how many work items are between `inVar` and being yielded + -- (this is currently `numInQueue`), and ensuring that this number + -- is < than some maximum number M (blocking on `takeMVar` of the + -- front MVar in `outQueue` when the M limit is reached). inVar :: MVar (Maybe a) <- newEmptyMVar inVarEnqueued :: MVar () <- newEmptyMVar outQueueRef :: IORef (Seq (MVar b)) <- newIORef Seq.empty @@ -161,37 +224,41 @@ -- we can use it in conduit `bracketP`'s IO-based resource acquisition -- function (where we have to spawn our workers to guarantee they shut down -- when somebody async-kills the conduit). - runInIO :: (m b -> IO b) <- lift askRunInIO -- lift brings us into `m` + u :: UnliftIO m <- lift askUnliftIO -- `lift` here brings us into `m` -- `spawnWorkers` uses `async` and thus MUST be run with interrupts disabled -- (e.g. as initialisation function of `bracket`) to be async exception safe. + -- + -- Note `async` does not unmask, but `unliftIO u` will restore the original + -- masking state (thus typically unmask). let spawnWorkers :: IO (Async ()) spawnWorkers = do workersAsync <- async $ do -- see comment above for exception safety - forConcurrently_ [1..numThreads] $ \_ -> do - -- Each worker has `workerOutputBufferSize` many `workerOutVar`s - -- in a ring buffer; until the shutdown signal is received, a worker - -- loops to: grab an `a` from the `inVar`, pick its next `workerOutVar, - -- put it into the `outQueue`, signal that it has atomically done these - -- 2 actions, process `b <- f x`, and write the `b` to the `workerOutVar`. - workerOutVars <- V.replicateM workerOutputBufferSize newEmptyMVar - let loop :: Int -> IO () - loop !i = do + unliftIO u $ liftIO $ do -- use `runInIO` to restore masking state + forConcurrently_ [1..numThreads] $ \_i_worker -> do + -- Each worker has `workerOutputBufferSize` many `workerOutVar`s + -- in a ring buffer; until the shutdown signal is received, a worker + -- loops to: grab an `a` from the `inVar`, pick its next `workerOutVar, + -- put it into the `outQueue`, signal that it has atomically done these + -- 2 actions, process `b <- f x`, and write the `b` to the `workerOutVar`. + workerOutVars <- V.replicateM workerOutputBufferSize newEmptyMVar + let loop :: Int -> IO () + loop !i_outVarSlot = do - m'a <- takeMVar inVar - case m'a of - Nothing -> return () -- shutdown signal, worker quits - Just a -> do - let workerOutVar = workerOutVars ! i - atomicModifyIORef_' outQueueRef (|> workerOutVar) - signal inVarEnqueued - -- Important: Force WHNF here so that f gets evaluated inside the - -- worker; it's `f`'s job to decide whether to deepseq or not. - !b <- runInIO (f a) - putMVar workerOutVar b - loop ((i + 1) `rem` workerOutputBufferSize) + m'a <- takeMVar inVar + case m'a of + Nothing -> return () -- shutdown signal, worker quits + Just a -> do + let workerOutVar = workerOutVars ! i_outVarSlot + atomicModifyIORef_' outQueueRef (|> workerOutVar) + signal inVarEnqueued + -- Important: Force WHNF here so that f gets evaluated inside the + -- worker; it's `f`'s job to decide whether to deepseq or not. + !b <- unliftIO u (f a) + putMVar workerOutVar b + loop ((i_outVarSlot + 1) `rem` workerOutputBufferSize) - loop 0 + loop 0 link workersAsync @@ -234,8 +301,10 @@ -- 2) Cruise phase, -- during which we always have at least `numWorkersRampedUp` many -- `workerOutVar`s in the output queue (this is an invariant). - -- At all times `numInQueue` keeps track of how many `workerOutVar`s - -- are in the output queue. + -- At all times `numInQueue` keeps track of how many work units + -- are under processing (that is, are after being read off the `inVar` + -- and before being read off an `outVar`; + -- so <= `N * (workerOutputBufferSize + 1)` many). -- Cruise phase doesn't happen if the conduit terminates before -- `numThreads` elements are awaited. -- 3) Drain phase, @@ -248,10 +317,10 @@ loop numWorkersRampedUp numInQueue = do await >>= \case - Nothing -> do -- upstream conduit is done, tell all workers to finish + Nothing -> do -- Drain phase: Upstream conduit is done, tell all workers to finish. for_ [1..numWorkersRampedUp] $ \_ -> do - putInVar Nothing yieldQueueHead -- This will succeed due to the "Cruise phase invariant", see above. + putInVar Nothing -- This will not block forever because we just freed an `outVar` in the line above. for_ [1..(numThreads - numWorkersRampedUp)] $ \_ -> do -- need to quit workers that were never ramped up too putInVar Nothing let numInQueueAfterStopping = numInQueue - numWorkersRampedUp diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -1,15 +1,58 @@ +{-# LANGUAGE ScopedTypeVariables #-} + module Main where +import Control.Monad (replicateM) import Control.Concurrent (threadDelay) import Control.Monad.IO.Class (liftIO) import Data.Conduit +import qualified Data.Conduit.Combinators as CC import qualified Data.Conduit.List as CL import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Monadic (monadicIO, run, pick, assert) import Say (sayString) import Data.Conduit.ConcurrentMap +prop_concurrentMapM_is_like_mapM :: Property +prop_concurrentMapM_is_like_mapM = monadicIO $ do + ints :: [Int] <- pick (replicateM 10 (choose (1, 10))) + bufferSize :: Int <- pick (choose (1, 30)) + numThreads :: Int <- pick (choose (1, 20)) + + isEquals <- run $ do + + let serial = runConduit + ( CC.yieldMany ints + .| CC.mapM (liftIO . f) + .| CC.sinkList + ) + let buffered = + runConduitRes + ( CC.yieldMany ints + .| concurrentMapM_ numThreads bufferSize (liftIO . f) + .| CC.sinkList + ) + + outSerial <- serial + outBuffered <- buffered + + return (outSerial == outBuffered) + + -- print ints -- for debugging failures + assert isEquals + + where + f :: Int -> IO Int + f i = do + -- sayString (show i ++ " before") -- for debugging + threadDelay i -- microseconds + -- sayString (show i ++ " after") -- for debugging + return (i*2) + + main :: IO () main = hspec $ do @@ -28,3 +71,6 @@ .| CL.consume l `shouldBe` [2,4,6,8,10,12] + + describe "concurrentMapM_" $ do + it "is like mapM" $ prop_concurrentMapM_is_like_mapM