@@ -0,0 +1,131 @@
1+{-# LANGUAGE CPP, BangPatterns #-}
2+
3+module Main where
4+
5+import Control.Monad
6+
7+import Data.Int
8+import 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 )
14+import System.Environment (getArgs, withArgs)
15+import Data.Time (getCurrentTime, diffUTCTime, NominalDiffTime)
16+import System.IO (withFile, IOMode(..), hPutStrLn, Handle, stderr)
17+import Control.Concurrent (forkIO)
18+import Control.Concurrent.MVar (newEmptyMVar, takeMVar, putMVar)
19+import qualified Network.Socket as N
20+import Debug.Trace
21+import Data.ByteString (ByteString)
22+import Data.ByteString.Char8 (pack, unpack)
23+import qualified Data.ByteString as BS
24+import qualified Network.Socket.ByteString as NBS
25+import Data.Time (getCurrentTime, diffUTCTime, NominalDiffTime)
26+import Data.ByteString.Internal as BSI
27+import Foreign.Storable (pokeByteOff, peekByteOff)
28+import Foreign.C (CInt(..))
29+import Foreign.ForeignPtr (withForeignPtr)
30+
31+foreign import ccall unsafe "htonl" htonl :: CInt -> CInt
32+foreign import ccall unsafe "ntohl" ntohl :: CInt -> CInt
33+
34+main :: IO ()
35+main = do
36+ [pingsStr] <- getArgs
37+ serverReady <- newEmptyMVar
38+ clientDone <- newEmptyMVar
39+
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)
51+
52+ putStrLn "server: awaiting client connection"
53+ putMVar serverReady ()
54+ listen sock 1
55+ (clientSock, clientAddr) <- accept sock
56+
57+ putStrLn "server: listening for pings"
58+ pong clientSock
59+
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
70+
71+ N.connect sock (addrAddress serverAddr)
72+
73+ ping sock pings
74+ putMVar clientDone ()
75+
76+ -- Wait for the client to finish
77+ takeMVar clientDone
78+
79+pingMessage :: ByteString
80+pingMessage = pack "ping123"
81+
82+ping :: Socket -> Int -> IO ()
83+ping 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)
97+
98+pong :: Socket -> IO ()
99+pong 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
105+
106+-- | Wrapper around NBS.recv (for profiling)
107+recv :: Socket -> Int -> IO ByteString
108+recv sock _ = do
109+ header <- NBS.recv sock 4
110+ length <- decodeLength header
111+ NBS.recv sock (fromIntegral (length :: Int32))
112+
113+-- | Wrapper around NBS.send (for profiling)
114+send :: Socket -> ByteString -> IO ()
115+send sock bs = do
116+ length <- encodeLength (fromIntegral (BS.length bs))
117+ NBS.sendMany sock [length, bs]
118+
119+-- | Encode length (manual for now)
120+encodeLength :: Int32 -> IO ByteString
121+encodeLength i32 =
122+ BSI.create 4 $ \p ->
123+ pokeByteOff p 0 (htonl (fromIntegral i32))
124+
125+-- | Decode length (manual for now)
126+decodeLength :: ByteString -> IO Int32
127+decodeLength bs =
128+ let (fp, _, _) = BSI.toForeignPtr bs in
129+ withForeignPtr fp $ \p -> do
130+ w32 <- peekByteOff p 0
131+ return (fromIntegral (ntohl w32))