thread pool with thread-local state

ThreadPool.hs

module Control.Concurrent.ThreadPool
  ( Pool
  , create
  , queue
  , state
  , kill
  ) where

import           Control.Concurrent     (ThreadId, forkIO, killThread)
import           Control.Concurrent.STM (TQueue, atomically, newTQueueIO,
                                         readTQueue, writeTQueue)
import           Control.Monad          (replicateM)
import           Data.IORef             (IORef, newIORef, readIORef, writeIORef)
import           Data.Traversable       (for)

data Pool a = Pool (TQueue (a -> IO a)) [(ThreadId, IORef a)]

create :: Word -> IO a -> IO (Pool a)
create n ini = do
  q <- newTQueueIO
  ts <- replicateM (fromIntegral n) $ do
    s <- ini
    ref <- newIORef s
    tid <- forkIO $ go q ref s
    pure (tid, ref)
  pure $ Pool q ts
  where
    go q r s = do
      task <- atomically $ readTQueue q
      s' <- task s
      writeIORef r s'
      go q r s'

queue :: Pool a -> (a -> IO a) -> IO ()
queue (Pool q _) task = atomically $ writeTQueue q task

state :: Pool a -> IO [a]
state (Pool _ ts) = for ts $ \(_, ref) -> readIORef ref

kill :: Pool a -> IO [a]
kill (Pool _ ts) =
  for ts $ \(tid, ref) -> do
    killThread tid
    readIORef ref