@@ -3,7 +3,8 @@
33-- TODOs:
44-- * Various errors are still left as "undefined"
55-- * Many possible errors are not yet dealt with
6−-- * Connections are not yet lightweight
6+-- * Multiple connections to the same endpoint (which reuse the same channel)
7+-- may be scrambed together
78module Network.Transport.TCP ( -- * Main API
89 createTransport
910 , -- TCP specific functionality
@@ -14,7 +15,6 @@ module Network.Transport.TCP ( -- * Main API
1415import Network.Transport
1516import Network.Transport.Internal.TCP ( forkServer
1617 , connectTo
17− , recvExact
1818 , recvWithLength
1919 , sendWithLength
2020 , recvInt16
@@ -22,6 +22,7 @@ import Network.Transport.Internal.TCP ( forkServer
2222 , sendInt32
2323 )
2424import Network.Transport.Internal ( encodeInt16
25+ , encodeInt32
2526 , failWith
2627 , failWithT
2728 )
@@ -33,15 +34,15 @@ import qualified Network.Socket as N ( HostName
3334 )
3435import Control.Concurrent (forkIO)
3536import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
36−import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, modifyMVar_, readMVar)
37+import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, modifyMVar_, readMVar, takeMVar, putMVar)
3738import Control.Category ((>>>))
3839import Control.Applicative ((<*>), (*>), (<$>))
39−import Control.Monad (forever, forM_)
40+import Control.Monad (forever, forM_, void)
4041import Control.Monad.Error (ErrorT, liftIO, runErrorT)
4142import Control.Monad.Trans.Maybe (MaybeT(MaybeT), runMaybeT)
4243import Control.Monad.Trans.Writer (WriterT, execWriterT)
4344import Control.Monad.Writer.Class (tell)
44−import Data.ByteString (ByteString)
45+import Control.Monad.IO.Class (MonadIO)
4546import qualified Data.ByteString as BS (concat)
4647import qualified Data.ByteString.Char8 as BSC (pack, unpack)
4748import Data.Int (Int16)
@@ -52,26 +53,28 @@ import qualified Data.Map as Map (empty)
5253import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=))
5354import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
5455
55−data TransportState = TransportState { _endPoints :: IntMap (Chan Event, MVar EndPointState)
56− , _sockets :: Map EndPointAddress N.Socket
56+data TCPTransport = TCPTransport { _transportHost :: N.HostName
57+ , _transportPort :: N.ServiceName
58+ , _transportState :: MVar TransportState
59+ }
60+data TransportState = TransportState { _localEndPoints :: IntMap LocalEndPoint
61+ , _remoteEndPoints :: Map EndPointAddress RemoteEndPoint
62+ }
63+data LocalEndPoint = LocalEndPoint { _localAddress :: EndPointAddress
64+ , _localChannel :: Chan Event
65+ , _localState :: MVar EndPointState
66+ }
67+data RemoteEndPoint = RemoteEndPoint { _remoteAddress :: EndPointAddress
68+ , _remoteSocket :: N.Socket
69+ , _remoteConnectionIds :: Chan ConnectionId
70+ , _remoteLock :: MVar ()
71+ }
72+data EndPointState = EndPointState { _nextConnectionId :: ConnectionId
5773 }
58−data EndPointState = EndPointState { _nextConnectionId :: ConnectionId }
5974type EndPointId = Int16
6075
61−endPoints :: Lens TransportState (IntMap (Chan Event, MVar EndPointState))
62−endPoints = lens _endPoints (\es st -> st { _endPoints = es })
63−
64−sockets :: Lens TransportState (Map EndPointAddress N.Socket)
65−sockets = lens _sockets (\ss st -> st { _sockets = ss })
66−
67−endPointAt :: EndPointId -> Lens TransportState (Maybe (Chan Event, MVar EndPointState))
68−endPointAt ix = endPoints >>> intMapLens (fromIntegral ix)
69−
70−socketAt :: EndPointAddress -> Lens TransportState (Maybe N.Socket)
71−socketAt addr = sockets >>> mapLens addr
72−
73−nextConnectionId :: Lens EndPointState ConnectionId
74−nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = cix })
76+#define REQUEST_CONNECTION_ID 0
77+#define RESPOND_CONNECTION_ID 1
7578
7679-- | Create a TCP transport
7780--
@@ -85,47 +88,69 @@ nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = c
8588-- * Deal with all exceptions that may occur
8689createTransport :: N.HostName -> N.ServiceName -> IO Transport
8790createTransport host port = do
88− state <- newMVar TransportState { _endPoints = IntMap.empty
89− , _sockets = Map.empty
91+ state <- newMVar TransportState { _localEndPoints = IntMap.empty
92+ , _remoteEndPoints = Map.empty
9093 }
91− forkServer host port (transportServer state)
92− return Transport { newEndPoint = tcpNewEndPoint state host port }
94+ let transport = TCPTransport { _transportState = state
95+ , _transportHost = host
96+ , _transportPort = port
97+ }
98+ forkServer host port (transportServer transport)
99+ return Transport { newEndPoint = tcpNewEndPoint transport }
93100
94101-- | The transport server handles incoming connections for all endpoints
95102--
96103-- TODO: lots of error checking
97−transportServer :: MVar TransportState -> N.Socket -> IO ()
98−transportServer state sock = forever $ do
104+transportServer :: TCPTransport -> N.Socket -> IO ()
105+transportServer transport sock = forever $ do
99106 (clientSock, _) <- N.accept sock
100− forkIO $ handleClient state clientSock >> N.sClose clientSock
101−
102−-- | Handle a new incoming connection
103−--
104−-- Returns when the client closes the socket and on invalid input.
105−handleClient :: MVar TransportState -> N.Socket -> IO ()
106−handleClient state sock = do
107− mendpoint <- runMaybeT $ do
107+ handleConnectionRequest transport clientSock
108+
109+handleConnectionRequest :: TCPTransport -> N.Socket -> IO ()
110+handleConnectionRequest transport sock = do
111+ request <- runMaybeT $ do
108112 ourEndPointId <- recvInt16 sock
109− theirAddress <- BS.concat <$> recvWithLength sock
110− ourEndPoint <- MaybeT $ (^. endPointAt ourEndPointId) <$> readMVar state
111− return (ourEndPoint, EndPointAddress theirAddress)
112− case mendpoint of
113− Just ((endPointCh, endPointSt), theirAddress) -> do
114− unclosedConnections <- endpointServer sock theirAddress endPointCh endPointSt
115− forM_ unclosedConnections $ \cix -> writeChan endPointCh (ConnectionClosed cix)
116− Nothing ->
117− -- Invalid request
118− return ()
119−
120−-- | Handle requests for one endpoint. Returns only if an error occurs,
121−-- in which case it returns the connections that are still open (normally,
113+ theirAddress <- EndPointAddress . BS.concat <$> recvWithLength sock
114+ ourEndPoint <- MaybeT $ (^. localEndPointAt ourEndPointId) <$> readMVar (transport ^. transportState)
115+ return (ourEndPoint, theirAddress)
116+ case request of
117+ Just (ourEndPoint, theirAddress) ->
118+ void $ forkRemoteEndPoint transport ourEndPoint theirAddress sock
119+ Nothing ->
120+ return ()
121+
122+forkRemoteEndPoint :: TCPTransport
123+ -> LocalEndPoint
124+ -> EndPointAddress
125+ -> N.Socket
126+ -> IO RemoteEndPoint
127+forkRemoteEndPoint transport ourEndPoint theirAddress theirSock = do
128+ theirConnectionIds <- newChan
129+ theirLock <- newMVar ()
130+ let theirEndPoint = RemoteEndPoint { _remoteAddress = theirAddress
131+ , _remoteSocket = theirSock
132+ , _remoteConnectionIds = theirConnectionIds
133+ , _remoteLock = theirLock
134+ }
135+ modifyMVar_ (transport ^. transportState) $ return . (remoteEndPointAt theirAddress ^= Just theirEndPoint)
136+ forkIO $ do
137+ unclosedConnections <- handleIncomingMessages ourEndPoint theirEndPoint
138+ forM_ unclosedConnections $ \cix -> writeChan (ourEndPoint ^. localChannel) (ConnectionClosed cix)
139+ -- TODO: should we close the socket here? (forkRemoteEndPoint gets called from multiple places)
140+ return theirEndPoint
141+
142+-- | Handle requests from a remote endpoint.
143+--
144+-- Returns only if the remote party closes the socket or if an error occurs.
145+-- The return value is the list of connections that are still open (normally,
122146-- this should be an empty list).
123−endpointServer :: N.Socket -> EndPointAddress -> Chan Event -> MVar EndPointState -> IO [ConnectionId]
124−endpointServer sock theirAddress chan state = execWriterT . runMaybeT . forever $ do
147+handleIncomingMessages :: LocalEndPoint -> RemoteEndPoint -> IO [ConnectionId]
148+handleIncomingMessages ourEndPoint theirEndPoint = execWriterT . runMaybeT . forever $ do
125149 connId <- recvInt32 sock
126150 case connId of
127− 0 -> createNewConnection
128− _ -> readMessage connId
151+ REQUEST_CONNECTION_ID -> createNewConnection
152+ RESPOND_CONNECTION_ID -> readNewConnectionId
153+ _ -> readMessage connId
129154 where
130155 -- A brief note on the types:
131156 -- These functions may fail (connection error): hence MaybeT
@@ -134,28 +159,51 @@ endpointServer sock theirAddress chan state = execWriterT . runMaybeT . forever
134159 -- Create a new connection
135160 createNewConnection :: MaybeT (WriterT [ConnectionId] IO) ()
136161 createNewConnection = do
137− newId <- liftIO $ modifyMVar state $ \st -> return (nextConnectionId ^+= 1 $ st, st ^. nextConnectionId)
138− sendInt32 sock newId
139− liftIO $ writeChan chan (ConnectionOpened newId ReliableOrdered theirAddress)
162+ newId <- liftIO $ getNextConnectionId ourEndPoint
163+ withRemoteLock theirEndPoint $ do
164+ sendInt32 sock RESPOND_CONNECTION_ID
165+ sendInt32 sock newId
166+ liftIO $ writeChan ourChannel (ConnectionOpened newId ReliableOrdered (theirEndPoint ^. remoteAddress))
140167 -- We add the new connection ID to the list of open connections only once the
141168 -- endpoint has been notified of the new connection (sendInt32 may fail)
142169 tell [newId]
143170
144171 -- Read a message and output it on the endPoint's channel
145172 readMessage :: ConnectionId -> MaybeT (WriterT [ConnectionId] IO) ()
146− readMessage connId = recvWithLength sock >>= liftIO . writeChan chan . Received connId
173+ readMessage connId = recvWithLength sock >>= liftIO . writeChan ourChannel . Received connId
174+
175+ -- Read a new connection ID
176+ readNewConnectionId :: MaybeT (WriterT [ConnectionId] IO) ()
177+ readNewConnectionId = recvInt32 sock >>= liftIO . writeChan (theirEndPoint ^. remoteConnectionIds)
178+
179+ -- Breakdown of the arguments
180+ ourChannel = ourEndPoint ^. localChannel
181+ sock = theirEndPoint ^. remoteSocket
182+
183+-- | Get the next connection ID
184+getNextConnectionId :: LocalEndPoint -> IO ConnectionId
185+getNextConnectionId endpoint =
186+ modifyMVar (endpoint ^. localState) $ \st ->
187+ return (nextConnectionId ^+= 1 $ st, st ^. nextConnectionId)
147188
148189-- | Create a new endpoint
149−tcpNewEndPoint :: MVar TransportState -> N.HostName -> N.ServiceName -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
150−tcpNewEndPoint transportState host port = do
151− chan <- newChan
152− endPointState <- newMVar EndPointState { _nextConnectionId = firstNonReservedConnectionId }
153− endPointId <- modifyMVar transportState $ \st -> do let ix = fromIntegral $ IntMap.size (st ^. endPoints)
154− return ((endPointAt ix ^= Just (chan, endPointState)) $ st, ix)
155− let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show endPointId
156− return . Right $ EndPoint { receive = readChan chan
157− , address = addr
158− , connect = tcpConnect transportState addr
190+tcpNewEndPoint :: TCPTransport -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
191+tcpNewEndPoint transport = do
192+ endPoint <- modifyMVar (transport ^. transportState) $ \st -> do
193+ chan <- newChan
194+ state <- newMVar EndPointState { _nextConnectionId = firstNonReservedConnectionId }
195+ let ix = fromIntegral $ IntMap.size (st ^. localEndPoints)
196+ let host = transport ^. transportHost
197+ let port = transport ^. transportPort
198+ let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show ix
199+ let localEndPoint = LocalEndPoint { _localAddress = addr
200+ , _localChannel = chan
201+ , _localState = state
202+ }
203+ return ((localEndPointAt ix ^= Just localEndPoint) $ st, localEndPoint)
204+ return . Right $ EndPoint { receive = readChan (endPoint ^. localChannel)
205+ , address = endPoint ^. localAddress
206+ , connect = tcpConnect transport endPoint
159207 , newMulticastGroup = return . Left $ newMulticastGroupError
160208 , resolveMulticastGroup = \_ -> return . Left $ resolveMulticastGroupError
161209 }
@@ -168,49 +216,62 @@ firstNonReservedConnectionId :: ConnectionId
168216firstNonReservedConnectionId = 1024
169217
170218-- | Connnect to an endpoint
171−tcpConnect :: MVar TransportState -- ^ Transport state
172− -> EndPointAddress -- ^ Our address
173− -> EndPointAddress -- ^ Their address
174− -> Reliability -- ^ Connection reliability (ignored)
219+tcpConnect :: TCPTransport -- ^ Transport
220+ -> LocalEndPoint -- ^ Local end point
221+ -> EndPointAddress -- ^ Remote address
222+ -> Reliability -- ^ Reliability (ignored)
175223 -> IO (Either (FailedWith ConnectErrorCode) Connection)
176−tcpConnect state ourAddress theirAddress _ = runErrorT $ do
177− sock <- socketForEndPoint state ourAddress theirAddress
178− connId <- failWithT undefined $ requestNewConnection sock
224+tcpConnect transport ourEndPoint theirAddress _ = runErrorT $ do
225+ theirEndPoint <- remoteEndPoint transport ourEndPoint theirAddress
226+ connId <- failWithT undefined $ requestNewConnection theirEndPoint
227+ connBs <- encodeInt32 connId
228+ let sock = theirEndPoint ^. remoteSocket
179229 return Connection { send = runErrorT .
180230 failWithT (FailedWith SendFailed "Send failed") .
181− sendWithLength sock (Just connId)
231+ withRemoteLock theirEndPoint .
232+ sendWithLength sock (Just connBs)
182233 , close = N.sClose sock
183234 }
184235
185236-- | Request a new connection
186−-- Returns the new connection ID in serialized form
187−requestNewConnection :: N.Socket -> MaybeT IO ByteString
188−requestNewConnection sock = sendInt32 sock 0 >> BS.concat <$> recvExact sock 4
237+requestNewConnection :: RemoteEndPoint -> MaybeT IO ConnectionId
238+requestNewConnection endPoint = do
239+ withRemoteLock endPoint $ sendInt32 (endPoint ^. remoteSocket) 0
240+ liftIO $ readChan (endPoint ^. remoteConnectionIds)
189241
190−-- | Find a socket to the specified endpoint
242+-- | Find a remote endpoint
191243--
192−-- TODOs:
193−-- * Reuse connections
194−-- * Hints?
195−socketForEndPoint :: MVar TransportState -- ^ Transport state (for reusing sockets)
196− -> EndPointAddress -- ^ Our address
197− -> EndPointAddress -- ^ Their address
198− -> ErrorT (FailedWith ConnectErrorCode) IO N.Socket
199−socketForEndPoint state (EndPointAddress ourAddress) theirAddress = do
200− msock <- liftIO $ (^. socketAt theirAddress) <$> readMVar state
201− case msock of
202− Just sock ->
203− return sock
244+-- TODO: hints?
245+remoteEndPoint :: TCPTransport
246+ -> LocalEndPoint
247+ -> EndPointAddress
248+ -> ErrorT (FailedWith ConnectErrorCode) IO RemoteEndPoint
249+remoteEndPoint transport ourEndPoint theirAddress = do
250+ mendpoint <- liftIO $ (^. remoteEndPointAt theirAddress) <$> readMVar (transport ^. transportState)
251+ case mendpoint of
252+ Just endpoint ->
253+ return endpoint
204254 Nothing -> do
205− liftIO $ putStrLn $ "Creating socket from " ++ show ourAddress ++ " to " ++ show theirAddress
255+ -- liftIO $ putStrLn $ "Creating socket from " ++ show ourAddress ++ " to " ++ show theirAddress
206256 (host, port, endPointId) <- failWith invalidAddress (decodeEndPointAddress theirAddress)
207257 sock <- failWithT undefined (connectTo host port)
208258 endPointBS <- encodeInt16 endPointId
209259 failWithT undefined $ sendWithLength sock (Just endPointBS) [ourAddress]
210− liftIO $ modifyMVar_ state $ return . (socketAt theirAddress ^= Just sock)
211− return sock
260+ liftIO $ forkRemoteEndPoint transport ourEndPoint theirAddress sock
212261 where
213262 invalidAddress = FailedWith ConnectInvalidAddress "Invalid address"
263+ EndPointAddress ourAddress = ourEndPoint ^. localAddress
264+
265+-- | We must avoid scrambling concurrent sends to the same endpoint
266+--
267+-- TODO: deal with exceptions
268+withRemoteLock :: MonadIO m => RemoteEndPoint -> m a -> m a
269+withRemoteLock theirEndPoint p = do
270+ let lock = theirEndPoint ^. remoteLock
271+ liftIO $ takeMVar lock
272+ x <- p
273+ liftIO $ putMVar lock ()
274+ return x
214275
215276-- | Decode end point address
216277-- TODO: This uses regular expression parsing, which is nice, but unnecessary
@@ -219,3 +280,52 @@ decodeEndPointAddress (EndPointAddress bs) = match endPointAddressRE (BSC.unpack
219280 where
220281 endPointAddressRE :: RE Char (N.HostName, N.ServiceName, EndPointId)
221282 endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))
283+
284+--------------------------------------------------------------------------------
285+-- Lens definitions
286+--------------------------------------------------------------------------------
287+
288+transportState :: Lens TCPTransport (MVar TransportState)
289+transportState = lens _transportState (\st trans -> trans { _transportState = st })
290+
291+transportHost :: Lens TCPTransport N.HostName
292+transportHost = lens _transportHost (\host trans -> trans { _transportHost = host })
293+
294+transportPort :: Lens TCPTransport N.ServiceName
295+transportPort = lens _transportPort (\port trans -> trans { _transportPort = port })
296+
297+localEndPoints :: Lens TransportState (IntMap LocalEndPoint)
298+localEndPoints = lens _localEndPoints (\es st -> st { _localEndPoints = es })
299+
300+remoteEndPoints :: Lens TransportState (Map EndPointAddress RemoteEndPoint)
301+remoteEndPoints = lens _remoteEndPoints (\es st -> st { _remoteEndPoints = es })
302+
303+localAddress :: Lens LocalEndPoint EndPointAddress
304+localAddress = lens _localAddress (\addr ep -> ep { _localAddress = addr })
305+
306+localChannel :: Lens LocalEndPoint (Chan Event)
307+localChannel = lens _localChannel (\ch ep -> ep { _localChannel = ch })
308+
309+localState :: Lens LocalEndPoint (MVar EndPointState)
310+localState = lens _localState (\st ep -> ep { _localState = st })
311+
312+remoteAddress :: Lens RemoteEndPoint EndPointAddress
313+remoteAddress = lens _remoteAddress (\addr ep -> ep { _remoteAddress = addr })
314+
315+remoteSocket :: Lens RemoteEndPoint N.Socket
316+remoteSocket = lens _remoteSocket (\sock ep -> ep { _remoteSocket = sock })
317+
318+remoteConnectionIds :: Lens RemoteEndPoint (Chan ConnectionId)
319+remoteConnectionIds = lens _remoteConnectionIds (\ch ep -> ep { _remoteConnectionIds = ch })
320+
321+remoteLock :: Lens RemoteEndPoint (MVar ())
322+remoteLock = lens _remoteLock (\lock ep -> ep { _remoteLock = lock })
323+
324+nextConnectionId :: Lens EndPointState ConnectionId
325+nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = cix })
326+
327+localEndPointAt :: EndPointId -> Lens TransportState (Maybe LocalEndPoint)
328+localEndPointAt ix = localEndPoints >>> intMapLens (fromIntegral ix)
329+
330+remoteEndPointAt :: EndPointAddress -> Lens TransportState (Maybe RemoteEndPoint)
331+remoteEndPointAt addr = remoteEndPoints >>> mapLens addr