test.hs

{-# LANGUAGE TemplateHaskell #-}

module Main
where

import System.Environment (getArgs)
import Data.ByteString.Char8 (pack)
import Control.Distributed.Process (say, Process(..), spawn, RemoteTable(..))
import Control.Distributed.Process.Node (newLocalNode, initRemoteTable, runProcess, localNodeId, LocalNode(..))
import Control.Distributed.Process.Internal.Types (NodeId(..))
import Control.Distributed.Process.Closure (mkClosure, remotable)
import Control.Concurrent (threadDelay)
import Network.Transport (EndPointAddress(..))
import Network.Transport.TCP (defaultTCPParameters, createTransport)
import Control.Monad (replicateM_)
import Control.Monad.Trans  (liftIO)

sleepSeconds :: Int -> IO ()
sleepSeconds n = threadDelay (1000000 * n)

printLoop s n = replicateM_ n $ do
  say s
  liftIO $ sleepSeconds 1;

printProcess :: (String, Int) -> Process ()
printProcess (s, n) = printLoop s n

remotable ['printProcess]

startLocalNode :: String -> String -> IO LocalNode
startLocalNode host port = do
  Right transport <- createTransport host port defaultTCPParameters
  newLocalNode transport rtable

nodeFromEndpoint :: String -> NodeId
nodeFromEndpoint = NodeId . EndPointAddress . pack

leaderProcess :: LocalNode -> NodeId -> String -> Int -> IO ()
leaderProcess localNode remoteNode message numTimes =
  runProcess localNode $ do
    remotePid <- spawn remoteNode ($(mkClosure 'printProcess) (message, numTimes))
    liftIO $ print remotePid

    liftIO $ print "finished spawning all processes"
    liftIO $ sleepSeconds 10

startLeader :: String -> String -> String -> String -> Int -> IO ()
startLeader host port followerAddress message numTimes = do
  localNode  <- startLocalNode host port
  let nodeId = nodeFromEndpoint followerAddress
  leaderProcess localNode nodeId message numTimes

startFollower :: String -> String -> IO ()
startFollower host port = do
  localNode  <- startLocalNode host port
  print (localNodeId localNode)
  sleepSeconds 10

rtable :: RemoteTable
rtable = __remoteTable initRemoteTable

main = do
  args <- getArgs
  case args of
    ["leader", host, port, followerAddress, message, numTimes] ->
      startLeader host port followerAddress message (read numTimes)
    ["follower", host, port] ->
      startFollower host port