test.hs

{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
import           Data.Pool

import           Control.Concurrent
import           Control.Concurrent.Async
import           Data.Monoid
import           Data.Time.Clock

import           Control.Monad

import           Data.List (delete, splitAt)

main = do
  (spacing, count, resources) <- return (10000, 10000, 50)

  pool <- createPool (pure ()) (const $ pure ()) 8 5 resources

  let list = replicate count spacing

  time <- elapsed $
    mapConcurrentlyN_ resources (\i -> withResource pool $ \_ -> threadDelay i) list
--    mapConcurrently_ (\i -> withResource pool $ \_ -> threadDelay i) list

  capabilities <- getNumCapabilities
  print $ "Took " <> show time <> " using " <> show capabilities <> " capabilities."

elapsed :: IO a -> IO NominalDiffTime
elapsed action = do
  current <- getCurrentTime
  action
  updated <- getCurrentTime

  return $ diffUTCTime updated current

mapConcurrentlyN_ :: Int -> (a -> IO b) -> [a] -> IO ()
mapConcurrentlyN_ n act xs = do
    let (initial, rest) = splitAt n (map (async . act) xs)
    as <- sequence initial
    waitAndContinue as rest
  where
    waitAndContinue as [] = mapM_ wait as
    waitAndContinue as (next:rest) = do
        (a, _) <- waitAny as
        n <- next
        waitAndContinue (n : delete a as) rest

-- $ ghc -O2 -threaded -rtsopts test.hs -o test -debug
-- $ ./test +RTS -N1
-- "Took 2.169734537s using 1 capabilities."
-- $ ./test +RTS -N
-- "Took 2.168515792s using 8 capabilities."