@@ -0,0 +1,110 @@
1+{-# LANGUAGE CPP, BangPatterns #-}
2+
3+module Main where
4+
5+import Control.Monad
6+
7+import Data.Int
8+import qualified Data.Serialize as Ser
9+import Data.Word (Word8)
10+import 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 )
16+import System.Environment (getArgs, withArgs)
17+import Data.Time (getCurrentTime, diffUTCTime, NominalDiffTime)
18+import System.IO (withFile, IOMode(..), hPutStrLn, Handle, stderr)
19+import Control.Concurrent (forkIO)
20+import Control.Concurrent.MVar (newEmptyMVar, takeMVar, putMVar)
21+import qualified Network.Socket as N
22+import Debug.Trace
23+import Data.ByteString (ByteString)
24+import Data.ByteString.Char8 (pack, unpack)
25+import qualified Data.ByteString as BS
26+import qualified Network.Socket.ByteString as NBS
27+import Data.Time (getCurrentTime, diffUTCTime, NominalDiffTime)
28+
29+main :: IO ()
30+main = do
31+ [pingsStr] <- getArgs
32+ serverReady <- newEmptyMVar
33+
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)
45+
46+ putStrLn "server: awaiting client connection"
47+ putMVar serverReady ()
48+ listen sock 1
49+ (clientSock, clientAddr) <- accept sock
50+
51+ putStrLn "server: listening for pings"
52+ pong clientSock
53+
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
63+
64+ N.connect sock (addrAddress serverAddr)
65+
66+ ping sock pings
67+
68+pingMessage :: ByteString
69+pingMessage = pack "ping123"
70+
71+ping :: Socket -> Int -> IO ()
72+ping 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)
86+
87+pong :: Socket -> IO ()
88+pong 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
94+
95+-- | Wrapper around NBS.recv (for profiling)
96+recv :: Socket -> Int -> IO ByteString
97+recv 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
102+
103+-- | Cached header
104+header :: ByteString
105+header = Ser.encode (8 :: Int32)
106+
107+-- | Wrapper around NBS.send (for profiling)
108+send :: Socket -> ByteString -> IO ()
109+send sock bs = do
110+ NBS.sendMany sock [header, bs]