Stream.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

-- | Processing streams with a fixed number of worker threads
module Control.Concurrent.Stream
  ( stream
  , streamBound
  , streamWithInput
  , streamWithOutput
  , streamWithInputOutput
  , mapConcurrentlyBounded
  , forConcurrentlyBounded
  ) where

import Control.Concurrent.Async
import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Data.Maybe
import Data.IORef

data ShouldBindThreads = BoundThreads | UnboundThreads

-- | Maps a fixed number of workers concurrently over a stream of values
-- produced by a producer function. The producer is passed a function to
-- call for each work item. If a worker throws a synchronous exception, it
-- will be propagated to the caller.
stream
  :: Int -- ^ Maximum Concurrency
  -> ((a -> IO ()) -> IO ()) -- ^ Producer
  -> (a -> IO ()) -- ^ Worker
  -> IO ()
stream maxConcurrency producer worker =
  streamWithInput producer (replicate maxConcurrency ()) $ const worker

-- | Like stream, but uses bound threads for the workers.  See
-- 'Control.Concurrent.forkOS' for details on bound threads.
streamBound
  :: Int -- ^ Maximum Concurrency
  -> ((a -> IO ()) -> IO ()) -- ^ Producer
  -> (a -> IO ()) -- ^ Worker
  -> IO ()
streamBound maxConcurrency producer worker =
  stream_ BoundThreads producer (replicate maxConcurrency ()) $ const worker

-- | Like stream, but each worker is passed an element of an input list.
streamWithInput
  :: ((a -> IO ()) -> IO ()) -- ^ Producer
  -> [b] -- ^ Worker state
  -> (b -> a -> IO ()) -- ^ Worker
  -> IO ()
streamWithInput = stream_ UnboundThreads

-- | Like 'stream', but collects the results of each worker
streamWithOutput
  :: Int 
  -> ((a -> IO ()) -> IO ()) -- ^ Producer
  -> (a -> IO c) -- ^ Worker
  -> IO [c]
streamWithOutput maxConcurrency producer worker =
  streamWithInputOutput producer (replicate maxConcurrency ()) $ 
    const worker

-- | Like 'streamWithInput', but collects the results of each worker
streamWithInputOutput
  :: ((a -> IO ()) -> IO ()) -- ^ Producer
  -> [b] -- ^ Worker input
  -> (b -> a -> IO c) -- ^ Worker
  -> IO [c]
streamWithInputOutput producer workerInput worker = do
  results <- newIORef []
  let prod write = producer $ \a -> do
        res <- newIORef Nothing
        modifyIORef results (res :)
        write (a, res)
  stream_ UnboundThreads prod workerInput $ \s (a,ref) -> do
    worker s a >>= writeIORef ref . Just
  readIORef results >>= mapM readIORef >>= return . catMaybes . reverse
    
stream_
  :: ShouldBindThreads -- use bound threads?
  -> ((a -> IO ()) -> IO ()) -- ^ Producer
  -> [b] -- Worker input
  -> (b -> a -> IO ()) -- ^ Worker
  -> IO ()
stream_ useBoundThreads producer workerInput worker = do
  let maxConcurrency = length workerInput
  q <- atomically $ newTBQueue (fromIntegral maxConcurrency)
  let write x = atomically $ writeTBQueue q (Just x)
  mask $ \unmask ->
    concurrently_ (runWorkers unmask q) $ unmask $ do
      -- run the producer
      producer write
      -- write end-markers for all workers
      replicateM_ maxConcurrency $
        atomically $ writeTBQueue q Nothing
  where
    runWorkers unmask q = case useBoundThreads of
      BoundThreads ->
        foldr1 concurrentlyBound $
          map (runWorker unmask q) workerInput
      UnboundThreads ->
        mapConcurrently_ (runWorker unmask q) workerInput

    concurrentlyBound l r =
      withAsyncBound l $ \a ->
      withAsyncBound r $ \b ->
      void $ waitBoth a b

    runWorker unmask q s = do
      v <- atomically $ readTBQueue q
      case v of
        Nothing -> return ()
        Just t -> do
          unmask (worker s t)
          runWorker unmask q s

-- | Concurrent map over a list of values, using a bounded number of threads.
mapConcurrentlyBounded
  :: Int -- ^ Maximum concurrency
  -> (a -> IO b) -- ^ Function to map over the input values
  -> [a] -- ^ List of input values
  -> IO [b] -- ^ List of output values
mapConcurrentlyBounded maxConcurrency f input =
  streamWithOutput maxConcurrency (forM_ input) f
  
-- | 'mapConcurrentlyBounded' but with its arguments reversed
forConcurrentlyBounded
  :: Int -- ^ Maximum concurrency
  -> [a] -- ^ List of input values
  -> (a -> IO b) -- ^ Function to map over the input values
  -> IO [b] -- ^ List of output values
forConcurrentlyBounded = flip . mapConcurrentlyBounded