@@ -0,0 +1,138 @@
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 (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+ -- 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
62+
63+ -- Read from the client and output on the channel
64+ putStrLn "server: listening for pings"
65+ pong clientSock mvar
66+
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
77+
78+ N.connect sock (addrAddress serverAddr)
79+
80+ ping sock pings
81+ putMVar clientDone ()
82+
83+ -- Wait for the client to finish
84+ takeMVar clientDone
85+
86+pingMessage :: ByteString
87+pingMessage = pack "ping123"
88+
89+ping :: Socket -> Int -> IO ()
90+ping 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)
104+
105+pong :: Socket -> MVar ByteString -> IO ()
106+pong 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
112+
113+-- | Wrapper around NBS.recv (for profiling)
114+recv :: Socket -> Int -> IO ByteString
115+recv sock _ = do
116+ header <- NBS.recv sock 4
117+ length <- decodeLength header
118+ NBS.recv sock (fromIntegral (length :: Int32))
119+
120+-- | Wrapper around NBS.send (for profiling)
121+send :: Socket -> ByteString -> IO ()
122+send sock bs = do
123+ length <- encodeLength (fromIntegral (BS.length bs))
124+ NBS.sendMany sock [length, bs]
125+
126+-- | Encode length (manual for now)
127+encodeLength :: Int32 -> IO ByteString
128+encodeLength i32 =
129+ BSI.create 4 $ \p ->
130+ pokeByteOff p 0 (htonl (fromIntegral i32))
131+
132+-- | Decode length (manual for now)
133+decodeLength :: ByteString -> IO Int32
134+decodeLength bs =
135+ let (fp, _, _) = BSI.toForeignPtr bs in
136+ withForeignPtr fp $ \p -> do
137+ w32 <- peekByteOff p 0
138+ return (fromIntegral (ntohl w32))