zeromq3-haskell-0.2: examples/threads.hs

{-# LANGUAGE DeriveDataTypeable #-}
module Main where

import Control.Concurrent (threadDelay)
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.CatchIO
import Control.Scope
import Data.ByteString.Char8 (pack, unpack)
import System.ZMQ3.Safe
import System.Random
import System.Console.CmdArgs.Implicit
import Control.Concurrent.Thread.Group (ThreadGroup)
import qualified Control.Concurrent.Thread.Group as TG

runner :: (InScope m1 m2, MonadIO m2) => StdGen -> Int -> Socket m1 Req -> m2 ()
runner gen runs s = go 0
  where
    go ctr = unless (ctr == runs) $ do
        let (r, g) = randomR (100, 100000) gen
        liftIO $ setStdGen g
        send s [] (pack . show $ ctr)
        ret <- receive s
        liftIO $ print ctr
        if ((read . unpack $ ret) /= ctr)
            then error "unexpected return value"
            else (liftIO $ threadDelay r) >> go (ctr + 1)

data Threads = Threads
  { address :: String
  , threads :: Int
  , runs :: Int
  } deriving (Show, Data, Typeable)

opts :: Threads
opts = Threads { address = def &= help "address to connect"
               , threads = def &= help "number of client threads"
               , runs    = def &= help "number of runs" &= opt (1000 :: Int)
               }
               &= summary "Threads Test"

setup :: (InScope m1 m2, MonadCatchIO m2) => Threads -> Context m1 -> m2 ()
setup options ctx = do
    rnd <- liftIO $ getStdGen
    scope $ socket ctx Req >>= \s -> do
        connect s (address options)
        runner rnd (runs options) s
    return ()

main :: IO ()
main = do
    o <- cmdArgs opts
    g <- TG.new
    scope $ context 2 >>= \c -> do
        setup o c