network-transport-zeromq-0.2: benchmarks/Channels.hs

{-# LANGUAGE LambdaCase, OverloadedStrings #-}
module Channels where

-- | Like Latency, but creating lots of channels

import System.Environment
import Control.Monad (void, forM_, forever, replicateM_)
import Control.Concurrent.MVar
import Control.Concurrent (forkOS, threadDelay)
import Control.Applicative
import Control.Distributed.Process
import Control.Distributed.Process.Node
import Criterion.Measurement
import Data.Binary (encode, decode)
import Data.ByteString.Char8 (pack)
import Network.Transport.ZMQ (createTransport, defaultZMQParameters)
import qualified Data.ByteString.Lazy as BSL
import Text.Printf

pingServer :: Process ()
pingServer = forever $ do
  them <- expect
  sendChan them ()
  -- TODO: should this be automatic?
  reconnectPort them

pingClient :: Int -> ProcessId -> Process ()
pingClient n them = do
  replicateM_ n $ do
    (sc, rc) <- newChan :: Process (SendPort (), ReceivePort ())
    send them sc
    receiveChan rc

initialServer :: Process ()
initialServer = do
  us <- getSelfPid
  liftIO $ BSL.writeFile "pingServer.pid" (encode us)
  pingServer

initialClient :: Int -> Process ()
initialClient n = do
  them <- liftIO $ decode <$> BSL.readFile "pingServer.pid"
  pingClient n them

main :: IO ()
main = getArgs >>= \case
    [] -> defaultBench 
    [role, host] -> do
       Right transport <- createTransport defaultZMQParameters (pack host)
       node <- newLocalNode transport initRemoteTable
       case role of
         "SERVER" -> runProcess node initialServer
         "CLIENT" -> fmap read getLine >>= runProcess node .  initialClient
         _       -> error "Role should be either SERVER or CLIENT"
    _ -> error "either call benchmark with [SERVER|CLIENT] host or without arguments"
  where
    defaultBench = do
      void . forkOS $ do
        Right transport <- createTransport defaultZMQParameters "127.0.0.1"
        node <- newLocalNode transport initRemoteTable
        runProcess node $ initialServer
      threadDelay 1000000
      e <- newEmptyMVar
      void . forkOS $ do
        putStrLn "pings        time\n---          ---\n"
        forM_ [100,200,600,800,1000,2000,5000,8000,10000] $ \i -> do
            Right transport <- createTransport defaultZMQParameters "127.0.0.1"
            node <- newLocalNode transport initRemoteTable
            d <- time_ (runProcess node $ initialClient i)
            printf "%-8i %10.4f\n" i d
        putMVar e ()
      takeMVar e