git.iko.soy/distributed-process
on bc4312e help sign in
distributed-process/benchmarks/JustPingCacheHeader.hs
110 lines 3.2 kB history raw
1{-# LANGUAGE CPP, BangPatterns #-}
3module Main where
5import Control.Monad
7import Data.Int
8import qualified Data.Serialize as Ser
9import Data.Word (Word8)
10import Network.Socket
11 ( AddrInfoFlag (AI_PASSIVE), HostName, ServiceName, Socket
12 , SocketType (Stream), SocketOption (ReuseAddr)
13 , accept, addrAddress, addrFlags, addrFamily, bindSocket, defaultProtocol
14 , defaultHints
15 , getAddrInfo, listen, setSocketOption, socket, sClose, withSocketsDo )
16import System.Environment (getArgs, withArgs)
17import Data.Time (getCurrentTime, diffUTCTime, NominalDiffTime)
18import System.IO (withFile, IOMode(..), hPutStrLn, Handle, stderr)
19import Control.Concurrent (forkIO)
20import Control.Concurrent.MVar (newEmptyMVar, takeMVar, putMVar)
21import qualified Network.Socket as N
22import Debug.Trace
23import Data.ByteString (ByteString)
24import Data.ByteString.Char8 (pack, unpack)
25import qualified Data.ByteString as BS
26import qualified Network.Socket.ByteString as NBS
27import Data.Time (getCurrentTime, diffUTCTime, NominalDiffTime)
29main :: IO ()
30main = do
31 [pingsStr] <- getArgs
32 serverReady <- newEmptyMVar
34 -- Start the server
35 forkIO $ do
36 putStrLn "server: creating TCP connection"
37 serverAddrs <- getAddrInfo
38 (Just (defaultHints { addrFlags = [AI_PASSIVE] } ))
39 Nothing
40 (Just "8080")
41 let serverAddr = head serverAddrs
42 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
43 setSocketOption sock ReuseAddr 1
44 bindSocket sock (addrAddress serverAddr)
46 putStrLn "server: awaiting client connection"
47 putMVar serverReady ()
48 listen sock 1
49 (clientSock, clientAddr) <- accept sock
51 putStrLn "server: listening for pings"
52 pong clientSock
54 -- Client
55 takeMVar serverReady
56 let pings = read pingsStr
57 serverAddrs <- getAddrInfo
58 Nothing
59 (Just "127.0.0.1")
60 (Just "8080")
61 let serverAddr = head serverAddrs
62 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
64 N.connect sock (addrAddress serverAddr)
66 ping sock pings
68pingMessage :: ByteString
69pingMessage = pack "ping123"
71ping :: Socket -> Int -> IO ()
72ping sock pings = go pings
73 where
74 go :: Int -> IO ()
75 go 0 = do
76 putStrLn $ "client did " ++ show pings ++ " pings"
77 go !i = do
78 before <- getCurrentTime
79 send sock pingMessage
80 bs <- recv sock 8
81 after <- getCurrentTime
82 -- putStrLn $ "client received " ++ unpack bs
83 let latency = (1e6 :: Double) * realToFrac (diffUTCTime after before)
84 hPutStrLn stderr $ show i ++ " " ++ show latency
85 go (i - 1)
87pong :: Socket -> IO ()
88pong sock = do
89 bs <- recv sock 8
90 -- putStrLn $ "server received " ++ unpack bs
91 when (BS.length bs > 0) $ do
92 send sock bs
93 pong sock
95-- | Wrapper around NBS.recv (for profiling)
96recv :: Socket -> Int -> IO ByteString
97recv sock i = do
98 (header, payload) <- BS.splitAt 4 `fmap` NBS.recv sock (4 + i)
99 case Ser.decode header :: Either String Int32 of
100 Left _ -> error "Could not decode header"
101 Right _ -> return payload
103-- | Cached header
104header :: ByteString
105header = Ser.encode (8 :: Int32)
107-- | Wrapper around NBS.send (for profiling)
108send :: Socket -> ByteString -> IO ()
109send sock bs = do
110 NBS.sendMany sock [header, bs]