git.iko.soy/distributed-process
on 706ba61 help sign in
distributed-process/benchmarks/JustPingCacheHeader.hs
134 lines 3.9 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)
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 putStrLn "server: listening for pings"
58 pong clientSock
60 -- Start the client
61 forkIO $ do
62 takeMVar serverReady
63 let pings = read pingsStr
64 serverAddrs <- getAddrInfo
65 Nothing
66 (Just "127.0.0.1")
67 (Just "8080")
68 let serverAddr = head serverAddrs
69 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
71 N.connect sock (addrAddress serverAddr)
73 ping sock pings
74 putMVar clientDone ()
76 -- Wait for the client to finish
77 takeMVar clientDone
79pingMessage :: ByteString
80pingMessage = pack "ping123"
82ping :: Socket -> Int -> IO ()
83ping sock pings = go pings
84 where
85 go :: Int -> IO ()
86 go 0 = do
87 putStrLn $ "client did " ++ show pings ++ " pings"
88 go !i = do
89 before <- getCurrentTime
90 send sock pingMessage
91 bs <- recv sock 8
92 after <- getCurrentTime
93 -- putStrLn $ "client received " ++ unpack bs
94 let latency = (1e6 :: Double) * realToFrac (diffUTCTime after before)
95 hPutStrLn stderr $ show i ++ " " ++ show latency
96 go (i - 1)
98pong :: Socket -> IO ()
99pong sock = do
100 bs <- recv sock 8
101 -- putStrLn $ "server received " ++ unpack bs
102 when (BS.length bs > 0) $ do
103 send sock bs
104 pong sock
106-- | Wrapper around NBS.recv (for profiling)
107recv :: Socket -> Int -> IO ByteString
108recv sock i = do
109 (header, payload) <- BS.splitAt 4 `fmap` NBS.recv sock (4 + i)
110 -- Ignore header
111 return payload
113-- | Cached header
114header :: ByteString
115header = pack "fake"
117-- | Wrapper around NBS.send (for profiling)
118send :: Socket -> ByteString -> IO ()
119send sock bs = do
120 NBS.sendMany sock [header, bs]
122-- | Encode length (manual for now)
123encodeLength :: Int32 -> IO ByteString
124encodeLength i32 =
125 BSI.create 4 $ \p ->
126 pokeByteOff p 0 (htonl (fromIntegral i32))
128-- | Decode length (manual for now)
129decodeLength :: ByteString -> IO Int32
130decodeLength bs =
131 let (fp, _, _) = BSI.toForeignPtr bs in
132 withForeignPtr fp $ \p -> do
133 w32 <- peekByteOff p 0
134 return (fromIntegral (ntohl w32))