git.iko.soy/distributed-process
on c7f72ce help sign in
distributed-process/benchmarks/JustPingThroughChan.hs
139 lines 4.2 kB history raw
1{-# LANGUAGE CPP, BangPatterns #-}
3module Main where
5import Control.Monad
7import Data.Int
8import Network.Socket
9 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
10 , SocketType (Stream), SocketOption (ReuseAddr)
11 , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
12 , defaultHints
13 , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
14import System.Environment (getArgs, withArgs)
15import Data.Time (getCurrentTime, diffUTCTime, NominalDiffTime)
16import System.IO (withFile, IOMode(..), hPutStrLn, Handle, stderr)
17import Control.Concurrent (forkIO)
18import Control.Concurrent.MVar (newEmptyMVar, takeMVar, putMVar)
19import qualified Network.Socket as N
20import Debug.Trace
21import Data.ByteString (ByteString)
22import Data.ByteString.Char8 (pack, unpack)
23import qualified Data.ByteString as BS
24import qualified Network.Socket.ByteString as NBS
25import Data.Time (getCurrentTime, diffUTCTime, NominalDiffTime)
26import Data.ByteString.Internal as BSI
27import Foreign.Storable (pokeByteOff, peekByteOff)
28import Foreign.C (CInt(..))
29import Foreign.ForeignPtr (withForeignPtr)
30import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
32foreign import ccall unsafe "htonl" htonl :: CInt -> CInt
33foreign import ccall unsafe "ntohl" ntohl :: CInt -> CInt
35main :: IO ()
36main = do
37 [pingsStr] <- getArgs
38 serverReady <- newEmptyMVar
39 clientDone <- newEmptyMVar
41 -- Start the server
42 forkIO $ do
43 putStrLn "server: creating TCP connection"
44 serverAddrs <- getAddrInfo
45 (Just (defaultHints { addrFlags = [AI_PASSIVE] } ))
46 Nothing
47 (Just "8080")
48 let serverAddr = head serverAddrs
49 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
50 setSocketOption sock ReuseAddr 1
51 bindSocket sock (addrAddress serverAddr)
53 putStrLn "server: awaiting client connection"
54 putMVar serverReady ()
55 listen sock 1
56 (clientSock, clientAddr) <- accept sock
58 -- Set up channel that we will listen on and echo to the client
59 chan <- newChan
60 forkIO $ forever $ do
61 msg <- readChan chan
62 send clientSock msg
64 -- Read from the client and output on the channel
65 putStrLn "server: listening for pings"
66 pong clientSock chan
68 -- Start the client
69 forkIO $ do
70 takeMVar serverReady
71 let pings = read pingsStr
72 serverAddrs <- getAddrInfo
73 Nothing
74 (Just "127.0.0.1")
75 (Just "8080")
76 let serverAddr = head serverAddrs
77 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
79 N.connect sock (addrAddress serverAddr)
81 ping sock pings
82 putMVar clientDone ()
84 -- Wait for the client to finish
85 takeMVar clientDone
87pingMessage :: ByteString
88pingMessage = pack "ping123"
90ping :: Socket -> Int -> IO ()
91ping sock pings = go pings
92 where
93 go :: Int -> IO ()
94 go 0 = do
95 putStrLn $ "client did " ++ show pings ++ " pings"
96 go !i = do
97 before <- getCurrentTime
98 send sock pingMessage
99 bs <- recv sock 8
100 after <- getCurrentTime
101 -- putStrLn $ "client received " ++ unpack bs
102 let latency = (1e6 :: Double) * realToFrac (diffUTCTime after before)
103 hPutStrLn stderr $ show i ++ " " ++ show latency
104 go (i - 1)
106pong :: Socket -> Chan ByteString -> IO ()
107pong sock chan = do
108 bs <- recv sock 8
109 -- putStrLn $ "server received " ++ unpack bs
110 when (BS.length bs > 0) $ do
111 writeChan chan bs
112 pong sock chan
114-- | Wrapper around NBS.recv (for profiling)
115recv :: Socket -> Int -> IO ByteString
116recv sock _ = do
117 header <- NBS.recv sock 4
118 length <- decodeLength header
119 NBS.recv sock (fromIntegral (length :: Int32))
121-- | Wrapper around NBS.send (for profiling)
122send :: Socket -> ByteString -> IO ()
123send sock bs = do
124 length <- encodeLength (fromIntegral (BS.length bs))
125 NBS.sendMany sock [length, bs]
127-- | Encode length (manual for now)
128encodeLength :: Int32 -> IO ByteString
129encodeLength i32 =
130 BSI.create 4 $ \p ->
131 pokeByteOff p 0 (htonl (fromIntegral i32))
133-- | Decode length (manual for now)
134decodeLength :: ByteString -> IO Int32
135decodeLength bs =
136 let (fp, _, _) = BSI.toForeignPtr bs in
137 withForeignPtr fp $ \p -> do
138 w32 <- peekByteOff p 0
139 return (fromIntegral (ntohl w32))