git.iko.soy/distributed-process
on c7f72ce help sign in
distributed-process/benchmarks/JustPingThroughMVar.hs
138 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 (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)
31foreign import ccall unsafe "htonl" htonl :: CInt -> CInt
32foreign import ccall unsafe "ntohl" ntohl :: CInt -> CInt
34main :: IO ()
35main = do
36 [pingsStr] <- getArgs
37 serverReady <- newEmptyMVar
38 clientDone <- newEmptyMVar
40 -- Start the server
41 forkIO $ do
42 putStrLn "server: creating TCP connection"
43 serverAddrs <- getAddrInfo
44 (Just (defaultHints { addrFlags = [AI_PASSIVE] } ))
45 Nothing
46 (Just "8080")
47 let serverAddr = head serverAddrs
48 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
49 setSocketOption sock ReuseAddr 1
50 bindSocket sock (addrAddress serverAddr)
52 putStrLn "server: awaiting client connection"
53 putMVar serverReady ()
54 listen sock 1
55 (clientSock, clientAddr) <- accept sock
57 -- Set up an mvar that we will listen on and echo to the client
58 mvar <- newEmptyMVar
59 forkIO $ forever $ do
60 msg <- takeMVar mvar
61 send clientSock msg
63 -- Read from the client and output on the channel
64 putStrLn "server: listening for pings"
65 pong clientSock mvar
67 -- Start the client
68 forkIO $ do
69 takeMVar serverReady
70 let pings = read pingsStr
71 serverAddrs <- getAddrInfo
72 Nothing
73 (Just "127.0.0.1")
74 (Just "8080")
75 let serverAddr = head serverAddrs
76 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
78 N.connect sock (addrAddress serverAddr)
80 ping sock pings
81 putMVar clientDone ()
83 -- Wait for the client to finish
84 takeMVar clientDone
86pingMessage :: ByteString
87pingMessage = pack "ping123"
89ping :: Socket -> Int -> IO ()
90ping sock pings = go pings
91 where
92 go :: Int -> IO ()
93 go 0 = do
94 putStrLn $ "client did " ++ show pings ++ " pings"
95 go !i = do
96 before <- getCurrentTime
97 send sock pingMessage
98 bs <- recv sock 8
99 after <- getCurrentTime
100 -- putStrLn $ "client received " ++ unpack bs
101 let latency = (1e6 :: Double) * realToFrac (diffUTCTime after before)
102 hPutStrLn stderr $ show i ++ " " ++ show latency
103 go (i - 1)
105pong :: Socket -> MVar ByteString -> IO ()
106pong sock mvar = do
107 bs <- recv sock 8
108 -- putStrLn $ "server received " ++ unpack bs
109 when (BS.length bs > 0) $ do
110 putMVar mvar bs
111 pong sock mvar
113-- | Wrapper around NBS.recv (for profiling)
114recv :: Socket -> Int -> IO ByteString
115recv sock _ = do
116 header <- NBS.recv sock 4
117 length <- decodeLength header
118 NBS.recv sock (fromIntegral (length :: Int32))
120-- | Wrapper around NBS.send (for profiling)
121send :: Socket -> ByteString -> IO ()
122send sock bs = do
123 length <- encodeLength (fromIntegral (BS.length bs))
124 NBS.sendMany sock [length, bs]
126-- | Encode length (manual for now)
127encodeLength :: Int32 -> IO ByteString
128encodeLength i32 =
129 BSI.create 4 $ \p ->
130 pokeByteOff p 0 (htonl (fromIntegral i32))
132-- | Decode length (manual for now)
133decodeLength :: ByteString -> IO Int32
134decodeLength bs =
135 let (fp, _, _) = BSI.toForeignPtr bs in
136 withForeignPtr fp $ \p -> do
137 w32 <- peekByteOff p 0
138 return (fromIntegral (ntohl w32))