ConcurrentReadOrderedWrite.hs

{-# LANGUAGE BangPatterns #-}
import Control.Monad
import Control.Concurrent
import Control.Concurrent.STM
import Numeric.Natural

type QJob = TBQueue (A, MVar B)

worker :: (A -> B) -> QJob -> IO ()
worker f qjob = forever $ do
  -- Read the new element from the job queue
  (a,res) <- atomically $ readTBQueue qjob

  -- Compute the result
  let !b = f a

  -- And write it to the output
  putMVar res b

makeJobQ :: Natural -> TBQueue A -> TBQueue B -> IO QJob
makeJobQ size qa qb = do
  -- Queue that workers take their job from
  qjob <- newTBQueueIO size

  -- Queue that is used to order the results
  qresult <- newTBQueueIO size

  -- Worker that takes elements from the input and
  -- adds them to the two auxillary queues
  void $ forkIO $ forever $ do
    a <- atomically $ readTBQueue qa
    res <- newEmptyMVar
    atomically $ writeTBQueue qjob (a, res)
    atomically $ writeTBQueue qresult res

  -- Worker that takes elements from the result
  -- queue and puts then in the output queue
  void $ forkIO $ forever $ do
    res <- atomically $ readTBQueue qresult
    b <- readMVar res
    atomically $ writeTBQueue qb b

  return qjob


-- testing
type A = Int
type B = Int

main :: IO ()
main = do
  -- Make some input buffers
  qa <- newTBQueueIO 500
  qb <- newTBQueueIO 500

  -- Spawn the job queue with its special workers
  qjob <- makeJobQ 500 qa qb

  -- Spawn some workers
  let workerf a = head [b | b <- [0..], b * b * b == a * a * a] -- identity function, but then a bit slower
  replicateM_ 20 $ forkIO $ worker workerf qjob

  -- Fill the input queue
  let inp = [0..100]
  forM_ inp $ atomically . writeTBQueue qa

  -- Read the output queue
  res <- replicateM (length inp) $ atomically $ readTBQueue qb

  print $ res == inp