@@ -1,12 +1,25 @@
1−-- | TCP transport
2−
3−module Network.Transport.TCP ( -- * Main API
1+-- | TCP implementation of the transport layer.
2+--
3+-- The TCP implementation guarantees that only a single TCP connection (socket)
4+-- will be used between endpoints, provided that the addresses specified are
5+-- canonical. If /A/ connects to /B/ and reports its address as
6+-- @192.168.0.1:8080@ and /B/ subsequently connects tries to connect to /A/ as
7+-- @client1.local:http-alt@ then the transport layer will not realize that the
8+-- TCP connection can be reused.
9+module Network.Transport.TCP {- ( -- * Main API
410 createTransport
5− , -- TCP specific functionality
11+ , -- * TCP specific functionality
612 EndPointId
713 , decodeEndPointAddress
814 , ControlHeader(..)
9− ) where
15+ -- * Design notes
16+ -- ** Overview
17+ -- $overview
18+ -- ** Connecting
19+ -- $connecting
20+ -- ** Disconnecting
21+ -- $disconnecting
22+ ) -} where
1023
1124import Network.Transport
1225import Network.Transport.Internal.TCP ( forkServer
@@ -19,13 +32,11 @@ import Network.Transport.Internal ( encodeInt32
1932 , failWith
2033 , failWithIO
2134 , tryIO
22− , tlog
2335 )
2436import qualified Network.Socket as N ( HostName
2537 , ServiceName
2638 , Socket
2739 , sClose
28− , accept
2940 , getAddrInfo
3041 , socket
3142 , addrFamily
@@ -36,7 +47,7 @@ import qualified Network.Socket as N ( HostName
3647 , SocketOption(ReuseAddr)
3748 , connect
3849 )
39−import Network.Socket.ByteString (sendMany)
50+import Network.Socket.ByteString (sendMany)
4051import Control.Concurrent (forkIO)
4152import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
4253import Control.Concurrent.MVar ( MVar
@@ -50,9 +61,9 @@ import Control.Concurrent.MVar ( MVar
5061 )
5162import Control.Category ((>>>))
5263import Control.Applicative ((<*>), (*>), (<$>))
53−import Control.Monad (forever, forM_, void)
54−import Control.Monad.Error (MonadError, ErrorT, liftIO, runErrorT, throwError)
55−import Control.Monad.IO.Class (MonadIO)
64+import Control.Monad (forM_, void, unless)
65+import Control.Monad.Error (MonadError, ErrorT(ErrorT), runErrorT, throwError)
66+import Control.Monad.IO.Class (liftIO)
5667import Control.Exception (IOException, bracket_)
5768import Data.ByteString (ByteString)
5869import qualified Data.ByteString as BS (concat)
@@ -65,163 +76,445 @@ import qualified Data.IntSet as IntSet (empty, insert, delete, elems)
6576import Data.Map (Map)
6677import qualified Data.Map as Map (empty)
6778import Data.IORef (newIORef, modifyIORef, readIORef)
68−import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=), (^%=), (^%%=))
79+import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=))
6980import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
7081
82+-- $overview
83+--
84+-- The TCP transport maps multiple logical connections between A and B
85+-- (in either direction) to a single TCP connection:
86+--
87+-- > +-------+ +-------+
88+-- > | A |==========================| B |
89+-- > | |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~\ |
90+-- > | Q |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~Q |
91+-- > | \~~~|~~~~~~~~~~~~~~~~~~~~~~~~~<| |
92+-- > | |==========================| |
93+-- > +-------+ +-------+
94+
95+-- $connecting
96+--
97+-- Let /A/, /B/ be two endpoints without any connections. When /A/ wants to
98+-- connect to /B/, it locally records that it is trying to connect to /B/ and
99+-- sends a request to /B/. As part of the request /A/ sends its own endpoint
100+-- address to /B/.
101+--
102+-- When /B/ receives the connection request it first checks if it had not
103+-- already initiated a connection request to /A/. If it had not, then it will
104+-- acknowledge the connection request and record that it has a TCP connection to
105+-- /A/.
106+--
107+-- The tricky case arises when /A/ sends a connection request to /B/ and /B/
108+-- finds that it had already sent a connection request to /A/. In this case /B/
109+-- will accept the connection request from /A/ if /A/s endpoint address is
110+-- smaller (lexicographically) than /B/s, and reject it otherwise.
111+--
112+-- Upon a rejection the /A/ thread that initiated the request just needs to wait
113+-- until the /A/ thread that is dealing with /B/s connection request completes.
114+
115+-- $disconnecting
116+--
117+-- The TCP connection is created as soon as the first logical connection
118+-- from A to B (or B to A) is established. At this point a thread (@#@)
119+-- is spawned that listens for incoming connections from B:
120+--
121+-- > +-------+ +-------+
122+-- > | A |==========================| B |
123+-- > | |>~~~~~~~~~~~~~~~~~~~~~~~~~|~~~\ |
124+-- > | | | Q |
125+-- > | #| | |
126+-- > | |==========================| |
127+-- > +-------+ +-------+
128+--
129+-- The question is when the TCP connection can be closed again.
130+-- Conceptually, we want to do reference counting: when there are no
131+-- logical connections left between A and B we want to close the socket
132+-- (possibly after some timeout).
133+--
134+-- However, A and B need to agree that the refcount has reached zero.
135+-- It might happen that B sends a connection request over the existing
136+-- socket at the same time that A closes its logical connection to B and
137+-- closes the socket. This will cause a failure in B (which will have to
138+-- retry) which is not caused by a network failure, which is
139+-- unfortunate. (Note that the connection request from B might succeed
140+-- even if A closes the socket.)
141+--
142+-- Instead, A will first ask B if it's okay to close the connection. If
143+-- B has already initiated a logical connection to A then it can simply
144+-- ignore this message, because any message other than 'OkayToClose'
145+-- will imply to A that the connection should not yet be closed. On the
146+-- other hand, B must guarantee that once it does send an 'OkayToClose'
147+-- message it will not reuse the socket anymore.
148+
149+--------------------------------------------------------------------------------
150+-- Internal datatypes --
151+--------------------------------------------------------------------------------
152+
71153data TCPTransport = TCPTransport { _transportHost :: N.HostName
72154 , _transportPort :: N.ServiceName
73− , _transportState :: MVar TransportState
74− }
75−data TransportState = TransportState { _localEndPoints :: IntMap LocalEndPoint
76− , _remoteEndPoints :: Map EndPointAddress (RemoteEndPoint, RefCount)
155+ , _transportState :: MVar TransportState
77156 }
157+data TransportState = TransportState { _localEndPoints :: IntMap LocalEndPoint }
78158data LocalEndPoint = LocalEndPoint { _localAddress :: EndPointAddress
79159 , _localChannel :: Chan Event
80− , _localState :: MVar EndPointState
160+ , _localState :: MVar LocalState
161+ }
162+data LocalState = LocalState { _nextConnectionId :: ConnectionId
163+ , _pendingCtrlRequests :: IntMap (MVar [ByteString])
164+ , _nextCtrlRequestId :: ControlRequestId
165+ , _localConnections :: Map EndPointAddress (Maybe Closing, MVar RemoteState)
81166 }
82167data RemoteEndPoint = RemoteEndPoint { _remoteAddress :: EndPointAddress
83168 , _remoteSocket :: N.Socket
84169 , sendTo :: [ByteString] -> IO ()
85170 }
86−data EndPointState = EndPointState { _nextConnectionId :: ConnectionId
87− , _pendingCtrlRequests :: IntMap (MVar [ByteString])
88− , _nextCtrlRequestId :: ControlRequestId
89− }
171+data RemoteState = InvalidEndPoint (FailedWith ConnectErrorCode)
172+ | ValidEndPoint !RefCount RemoteEndPoint
173+ | DeletedEndPoint
174+
90175type EndPointId = Int32
91176type ControlRequestId = Int32
92−data RefCount = RefCount !Int
177+type RefCount = Int32
178+type Closing = MVar [ByteString]
179+type Conn = (LocalEndPoint, RemoteEndPoint)
180+
181+-- | Control headers
182+data ControlHeader =
183+ RequestConnectionId -- ^ Request a new connection ID from the remote endpoint
184+ | CloseConnection -- ^ Tell the remote endpoint we will no longer be using a connection
185+ | ControlResponse -- ^ Respond to a control request _from_ the remote endpoint
186+ | CloseSocket -- ^ Request to close the connection (see module description)
187+ deriving Enum
188+
189+-- Response sent by /B/ to /A/ when /A/ tries to connect
190+data ConnectionRequestResponse =
191+ ConnectionRequestAccepted -- ^ /B/ accepts the connection
192+ | ConnectionRequestInvalidEndPoint -- ^ /A/ requested an invalid endpoint
193+ | ConnectionRequestCrossed -- ^ /A/s request crossed with a request from /B/ (see protocols)
194+ deriving Enum
195+
196+--------------------------------------------------------------------------------
197+-- Top-level functionality --
198+--------------------------------------------------------------------------------
93199
94200-- | Create a TCP transport
95201--
96−-- Hostname and service name should be "canonical" for the transport layer to
97−-- use TCP channels bidirectionally (otherwise when A connects to B and B
98−-- connects to A two socket pairs will be created that will both be used
99−-- unidirectionally, at degraded performance).
100−--
101202-- TODOs: deal with hints
102203createTransport :: N.HostName -> N.ServiceName -> IO (Either IOException Transport)
103204createTransport host port = do
104− state <- newMVar TransportState { _localEndPoints = IntMap.empty
105− , _remoteEndPoints = Map.empty
106− }
205+ state <- newMVar TransportState { _localEndPoints = IntMap.empty }
107206 let transport = TCPTransport { _transportState = state
108207 , _transportHost = host
109208 , _transportPort = port
110209 }
111210 tryIO $ do
112− forkServer host port (transportServer transport)
113− return Transport { newEndPoint = tcpNewEndPoint transport }
211+ forkServer host port $ handleConnectionRequest transport
212+ return Transport { newEndPoint = apiNewEndPoint transport }
114213
115−-- | The transport server handles incoming connections for all endpoints
214+--------------------------------------------------------------------------------
215+-- API functions --
216+--------------------------------------------------------------------------------
217+
218+-- | Create a new endpoint
219+apiNewEndPoint :: TCPTransport -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
220+apiNewEndPoint transport = do
221+ ourEndPoint <- createLocalEndPoint transport
222+ return . Right $ EndPoint { receive = readChan (ourEndPoint ^. localChannel)
223+ , address = ourEndPoint ^. localAddress
224+ , connect = apiConnect ourEndPoint
225+ , newMulticastGroup = return . Left $ newMulticastGroupError
226+ , resolveMulticastGroup = return . Left . const resolveMulticastGroupError
227+ }
228+ where
229+ newMulticastGroupError = FailedWith NewMulticastGroupUnsupported "TCP does not support multicast"
230+ resolveMulticastGroupError = FailedWith ResolveMulticastGroupUnsupported "TCP does not support multicast"
231+
232+-- | Connnect to an endpoint
233+apiConnect :: LocalEndPoint -- ^ Local end point
234+ -> EndPointAddress -- ^ Remote address
235+ -> Reliability -- ^ Reliability (ignored)
236+ -> IO (Either (FailedWith ConnectErrorCode) Connection)
237+apiConnect ourEndPoint theirAddress _ = runErrorT $ do
238+ theirEndPoint <- findRemoteEndPoint ourEndPoint theirAddress
239+ connId <- requestNewConnection (ourEndPoint, theirEndPoint)
240+ return $ Connection { send = apiSend theirEndPoint connId
241+ , close = apiClose (ourEndPoint, theirEndPoint) connId
242+ }
243+
244+-- | Close a connection
116245--
117−-- TODO: lots of error checking
118−transportServer :: TCPTransport -> N.Socket -> IO ()
119−transportServer transport sock = forever $ do
120− (clientSock, _) <- N.accept sock
121− handleConnectionRequest transport clientSock
246+-- TODO: We ignore errors during a close. Is that right?
247+apiClose :: Conn -> ConnectionId -> IO ()
248+apiClose (ourEndPoint, theirEndPoint) connId = void . tryIO $ do
249+ sendTo theirEndPoint [encodeInt32 CloseConnection, encodeInt32 connId]
250+ decreaseRemoteRefCount (ourEndPoint, theirEndPoint)
251+
252+-- | Send data across a connection
253+--
254+-- TODO: we should have some per-connection state so that the user cannot send
255+-- to the connection after they closed it, or close it more than once
256+apiSend :: RemoteEndPoint -> ConnectionId -> [ByteString] -> IO (Either (FailedWith SendErrorCode) ())
257+apiSend theirEndPoint connId payload = runErrorT $
258+ failWithIO (FailedWith SendFailed . show) $
259+ sendTo theirEndPoint (encodeInt32 connId : prependLength payload)
260+
261+--------------------------------------------------------------------------------
262+-- Lower level functionality --
263+--------------------------------------------------------------------------------
264+
265+createLocalEndPoint :: TCPTransport -> IO LocalEndPoint
266+createLocalEndPoint transport = do
267+ chan <- newChan
268+ state <- newMVar LocalState { _nextConnectionId = firstNonReservedConnectionId
269+ , _pendingCtrlRequests = IntMap.empty
270+ , _nextCtrlRequestId = 0
271+ , _localConnections = Map.empty
272+ }
273+ let host = transport ^. transportHost
274+ let port = transport ^. transportPort
275+ modifyMVar (transport ^. transportState) $ \st -> do
276+ let ix = fromIntegral $ IntMap.size (st ^. localEndPoints)
277+ let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show ix
278+ let localEndPoint = LocalEndPoint { _localAddress = addr
279+ , _localChannel = chan
280+ , _localState = state
281+ }
282+ return ((localEndPointAt ix ^= Just localEndPoint) $ st, localEndPoint)
283+
284+findRemoteEndPoint :: LocalEndPoint
285+ -> EndPointAddress
286+ -> ErrorT (FailedWith ConnectErrorCode) IO RemoteEndPoint
287+findRemoteEndPoint ourEndPoint theirAddress = ErrorT $ go
288+ where
289+ go = do
290+ endPoint <- modifyMVar state $ \st ->
291+ case st ^. connectionTo theirAddress of
292+ Just (_, mvar) ->
293+ return (st, mvar)
294+ Nothing -> do
295+ theirMVar <- newEmptyMVar
296+ forkIO $ connectToRemoteEndPoint ourEndPoint theirMVar theirAddress
297+ return (connectionTo theirAddress ^= Just (Nothing, theirMVar) $ st, theirMVar)
298+ theirState <- readMVar endPoint
299+ case theirState of
300+ InvalidEndPoint err ->
301+ -- We remove the endpoint from the state again because the next call
302+ -- to 'connect' might give a different result
303+ modifyMVar state $ \st ->
304+ return (connectionTo theirAddress ^= Nothing $ st, Left err)
305+ ValidEndPoint _ ep ->
306+ return $ Right ep
307+ DeletedEndPoint ->
308+ go
309+
310+ state :: MVar LocalState
311+ state = ourEndPoint ^. localState
312+
313+connectToRemoteEndPoint :: LocalEndPoint -> MVar RemoteState -> EndPointAddress -> IO ()
314+connectToRemoteEndPoint ourEndPoint theirMVar theirAddress = do
315+ result <- runErrorT go
316+ case result of
317+ Right (ConnectionRequestAccepted, sock) -> do
318+ theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock
319+ putMVar theirMVar (ValidEndPoint 0 theirEndPoint)
320+ Right (ConnectionRequestInvalidEndPoint, sock) -> do
321+ putMVar theirMVar (InvalidEndPoint (invalidAddress "Invalid endpoint"))
322+ N.sClose sock
323+ Right (ConnectionRequestCrossed, sock) ->
324+ N.sClose sock
325+ Left err ->
326+ -- TODO: this isn't quite right; socket *might* have been opened here,
327+ -- in which case we should close it
328+ putMVar theirMVar (InvalidEndPoint err)
329+ where
330+ go :: ErrorT (FailedWith ConnectErrorCode) IO (ConnectionRequestResponse, N.Socket)
331+ go = do
332+ (host, port, theirEndPointId) <- failWith (invalidAddress "Could not parse") $
333+ decodeEndPointAddress theirAddress
334+ addr:_ <- failWithIO (invalidAddress . show) $
335+ N.getAddrInfo Nothing (Just host) (Just port)
336+ sock <- failWithIO (insufficientResources . show) $
337+ N.socket (N.addrFamily addr) N.Stream N.defaultProtocol
338+ failWithIO (failed . show) $ do
339+ N.setSocketOption sock N.ReuseAddr 1
340+ N.connect sock (N.addrAddress addr)
341+ sendMany sock (encodeInt32 theirEndPointId : prependLength [ourAddress])
342+ -- TODO: what happens when the server sends an int outside the range of ConnectionRequestResponse?
343+ response <- recvInt32 sock
344+ return (response, sock)
345+
346+ invalidAddress, insufficientResources, failed :: String -> FailedWith ConnectErrorCode
347+ invalidAddress = FailedWith ConnectInvalidAddress
348+ insufficientResources = FailedWith ConnectInsufficientResources
349+ failed = FailedWith ConnectFailed
350+
351+ ourAddress :: ByteString
352+ EndPointAddress ourAddress = ourEndPoint ^. localAddress
353+
354+-- | Request a new connection
355+requestNewConnection :: Conn -> ErrorT (FailedWith ConnectErrorCode) IO ConnectionId
356+requestNewConnection (ourEndPoint, theirEndPoint) = do
357+ mcid <- failWithIO (failed . show) $ do
358+ reply <- sendRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId
359+ decodeInt32 . BS.concat <$> takeMVar reply
360+ case mcid of
361+ Nothing ->
362+ throwError (failed "Invalid integer")
363+ Just cid -> do
364+ liftIO $ increaseRemoteRefCount (ourEndPoint, theirEndPoint)
365+ return cid
366+ where
367+ failed = FailedWith ConnectFailed
368+
369+-- | Decode end point address
370+-- TODO: This uses regular expression parsing, which is nice, but unnecessary
371+decodeEndPointAddress :: EndPointAddress -> Maybe (N.HostName, N.ServiceName, EndPointId)
372+decodeEndPointAddress (EndPointAddress bs) = match endPointAddressRE (BSC.unpack bs)
373+ where
374+ endPointAddressRE :: RE Char (N.HostName, N.ServiceName, EndPointId)
375+ endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))
376+
377+sendRemoteRequest :: Conn -> ControlHeader -> IO (MVar [ByteString])
378+sendRemoteRequest (ourEndPoint, theirEndPoint) header = do
379+ reply <- newEmptyMVar
380+ reqId <- modifyMVar (ourEndPoint ^. localState) $ \st -> do
381+ let reqId = st ^. nextCtrlRequestId
382+ return ((nextCtrlRequestId ^+= 1) . (pendingCtrlRequestsAt reqId ^= Just reply) $ st, reqId)
383+ sendTo theirEndPoint [encodeInt32 header, encodeInt32 reqId]
384+ return reply
385+
386+--------------------------------------------------------------------------------
387+-- Reference counting --
388+--------------------------------------------------------------------------------
389+
390+increaseRemoteRefCount :: Conn -> IO ()
391+increaseRemoteRefCount (ourEndPoint, theirEndPoint) = do
392+ let theirAddr = theirEndPoint ^. remoteAddress
393+ state = ourEndPoint ^. localState
394+ -- First check if this remote endpoint is awaiting closure
395+ (mClosing, theirMVar) <- modifyMVar state $ \st -> do
396+ let Just (closure, theirMVar) = st ^. connectionTo theirAddr
397+ return (connectionTo theirAddr ^= Just (Nothing, theirMVar) $ st, (closure, theirMVar))
398+ case mClosing of
399+ Just mvar -> putMVar mvar []
400+ _ -> return ()
401+ -- Increase the refcount
402+ modifyMVar_ theirMVar $ \(ValidEndPoint rc ep) ->
403+ return (ValidEndPoint (rc + 1) ep)
404+
405+decreaseRemoteRefCount :: Conn -> IO ()
406+decreaseRemoteRefCount (ourEndPoint, theirEndPoint) = do
407+ let theirAddr = theirEndPoint ^. remoteAddress
408+ state <- readMVar (ourEndPoint ^. localState)
409+ let Just (Nothing, theirMVar) = state ^. connectionTo theirAddr
410+ -- We don't use modifyMVar here because we want to leave the MVar
411+ -- empty when refcount reaches 0
412+ -- TODO: deal with exceptions
413+ ValidEndPoint rc ep <- takeMVar theirMVar
414+ if rc > 1
415+ then do
416+ putMVar theirMVar (ValidEndPoint (rc - 1) ep)
417+ else do
418+ reply <- sendRemoteRequest (ourEndPoint, theirEndPoint) CloseSocket
419+ modifyMVar_ (ourEndPoint ^. localState) $ \st -> do
420+ let Just (_, mvar) = st ^. connectionTo theirAddr
421+ return (connectionTo theirAddr ^= Just (Just reply, mvar) $ st)
422+ void . forkIO $ do
423+ -- Wait for the reply from the remote endpoint
424+ canCloseSocket <- takeMVar reply
425+ putStrLn $ "canCloseSocket = " ++ show canCloseSocket
426+
427+--------------------------------------------------------------------------------
428+-- Incoming requests --
429+--------------------------------------------------------------------------------
122430
123−-- | Set up a new TCP connection (at the request of a remote endpoint)
124431handleConnectionRequest :: TCPTransport -> N.Socket -> IO ()
125432handleConnectionRequest transport sock = do
126433 request <- tryIO $ do
127434 ourEndPointId <- recvInt32 sock
128435 theirAddress <- EndPointAddress . BS.concat <$> recvWithLength sock
129− ourEndPoint <- lookupEndPoint transport ourEndPointId
130− return (ourEndPoint, theirAddress)
436+ ourEndPoint <- (^. localEndPointAt ourEndPointId) <$> readMVar state
437+ return (theirAddress, ourEndPoint)
131438 case request of
132− Right (ourEndPoint, theirAddress) -> do
133− -- Send acknowledgement
134− sendMany sock [encodeInt32 ValidEndPoint]
135− -- We ignore the return value of 'forkRemoteEndPoint' here because _they_ are
136− -- connecting to _us_. 'forkRemoteEndPoint' will have stored the remote endpoint
137− -- in the transport state, however, in case _we_ later connect to _them_.
138− void $ forkRemoteEndPoint transport ourEndPoint theirAddress sock
139− Left _ ->
439+ Right (_, Nothing) -> do
440+ sendMany sock [encodeInt32 ConnectionRequestInvalidEndPoint]
441+ N.sClose sock
442+ Right (theirAddress, Just ourEndPoint) -> do
443+ (crossed, mvar) <- modifyMVar (ourEndPoint ^. localState) $ \st ->
444+ case st ^. connectionTo theirAddress of
445+ Nothing -> do
446+ mvar <- newEmptyMVar
447+ return (connectionTo theirAddress ^= Just (Nothing, mvar) $ st, (False, mvar))
448+ Just (_, mvar) ->
449+ return (st, (ourEndPoint ^. localAddress <= theirAddress, mvar))
450+ if crossed
451+ then do
452+ sendMany sock [encodeInt32 ConnectionRequestCrossed]
453+ N.sClose sock
454+ else do
455+ theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock
456+ putMVar mvar (ValidEndPoint 0 theirEndPoint)
457+ sendMany sock [encodeInt32 ConnectionRequestAccepted]
458+ Left _ ->
140459 -- Invalid request
141460 N.sClose sock
142461
143−-- | Look up a local endpoint
144−--
145−-- Throws an I/O exception if the endpoint cannot be found (this makes it easy to integrate
146−-- this with I/O exceptions that arise from socket reads)
147−lookupEndPoint :: TCPTransport -> EndPointId -> IO LocalEndPoint
148−lookupEndPoint transport endPointId = do
149− mep <- (^. localEndPointAt endPointId) <$> readMVar (transport ^. transportState)
150− case mep of
151− Nothing -> fail "Endpoint not found"
152− Just ep -> return ep
153−
154−-- | Set up a new remote endpoint (either because we connected to a
155−-- new remote endpoint, or because a remote endpoint connected to us)
156−forkRemoteEndPoint :: TCPTransport
157− -> LocalEndPoint
158− -> EndPointAddress
159− -> N.Socket
160− -> IO RemoteEndPoint
161−forkRemoteEndPoint transport ourEndPoint theirAddress sock = do
462+ where
463+ state :: MVar TransportState
464+ state = transport ^. transportState
465+
466+createRemoteEndPoint :: LocalEndPoint -> EndPointAddress -> N.Socket -> IO RemoteEndPoint
467+createRemoteEndPoint ourEndPoint theirAddress sock = do
162468 lock <- newMVar ()
163− let theirEndPoint = RemoteEndPoint { _remoteAddress = theirAddress
164− , _remoteSocket = sock
165− , sendTo = bracket_ (takeMVar lock) (putMVar lock ())
166− . sendMany sock
469+ let theirEndPoint = RemoteEndPoint { _remoteAddress = theirAddress
470+ , _remoteSocket = sock
471+ , sendTo = bracket_ (takeMVar lock) (putMVar lock ())
472+ . sendMany sock
167473 }
168− modifyMVar_ (transport ^. transportState) $
169− return . (remoteEndPointAt theirAddress ^= Just (theirEndPoint, RefCount 0))
170474 forkIO $ do
171− unclosedConnections <- handleIncomingMessages transport ourEndPoint theirEndPoint
475+ unclosedConnections <- handleIncomingMessages (ourEndPoint, theirEndPoint)
172476 forM_ (map fromIntegral $ IntSet.elems unclosedConnections) $
173477 writeChan (ourEndPoint ^. localChannel) . ConnectionClosed
174− -- TODO: should we close the socket here? (forkRemoteEndPoint gets called from multiple places)
175− -- TODO: and should we remove it from the transaction state?
176478 return theirEndPoint
177479
178−-- | Control headers
179−data ControlHeader =
180− ValidEndPoint -- ^ Acknowledgement sent when setting up a new TCP connection
181− | RequestConnectionId -- ^ Request a new connection ID from the remote endpoint
182− | CloseConnection -- ^ Tell the remote endpoint we will no longer be using a connection
183− | ControlResponse -- ^ Respond to a control request _from_ the remote endpoint
184− deriving Enum
185−
186480-- | Handle requests from a remote endpoint.
187481--
188482-- Returns only if the remote party closes the socket or if an error occurs.
189483-- The return value is the list of connections that are still open (normally,
190484-- this should be an empty list).
191−handleIncomingMessages :: TCPTransport -> LocalEndPoint -> RemoteEndPoint -> IO IntSet
192−handleIncomingMessages transport ourEndPoint theirEndPoint = do
193− tlog $ "Handling incoming messages from " ++ show (theirEndPoint ^. remoteAddress)
485+handleIncomingMessages :: Conn -> IO IntSet
486+handleIncomingMessages (ourEndPoint, theirEndPoint) = do
194487 -- We use an IORef rather than the state monad because when an exception occurs we
195488 -- want to know the current list of open connections
196489 openConnections <- newIORef IntSet.empty
197490
198− let -- The main loop terminates only when an exception occurs
199− mainLoop :: IO ()
200− mainLoop = forever $ do
491+ let go :: IO ()
492+ go = do
201493 connId <- recvInt32 sock
202494 if connId >= firstNonReservedConnectionId
203495 then readMessage connId
204496 else do
205497 reqId <- recvInt32 sock
206498 case toEnum (fromIntegral connId) of
207− ValidEndPoint -> fail "Unexpected control header"
208499 RequestConnectionId -> createNewConnection reqId
209500 ControlResponse -> readControlResponse reqId
210501 CloseConnection -> closeConnection reqId
502+ CloseSocket -> closeSocket reqId
211503
212504 -- Create a new connection
213505 createNewConnection :: ControlRequestId -> IO ()
214506 createNewConnection reqId = do
215− increaseRemoteRefCount transport theirEndPoint
507+ increaseRemoteRefCount (ourEndPoint, theirEndPoint)
216508 newId <- getNextConnectionId ourEndPoint
217509 sendTo theirEndPoint ( encodeInt32 ControlResponse
218− : encodeInt32 reqId
219− : prependLength [encodeInt32 newId]
220− )
221− writeChan ourChannel (ConnectionOpened newId ReliableOrdered (theirEndPoint ^. remoteAddress))
510+ : encodeInt32 reqId
511+ : prependLength [encodeInt32 newId]
512+ )
513+ writeChan ourChannel (ConnectionOpened newId ReliableOrdered theirAddr)
222514 -- We add the new connection ID to the list of open connections only once the
223515 -- endpoint has been notified of the new connection (sendMany may fail)
224516 modifyIORef openConnections (IntSet.insert (fromIntegral newId))
517+ go
225518
226519 -- Read a control response
227520 readControlResponse :: ControlRequestId -> IO ()
@@ -229,27 +522,56 @@ handleIncomingMessages transport ourEndPoint theirEndPoint = do
229522 response <- recvWithLength sock
230523 mmvar <- (^. pendingCtrlRequestsAt reqId) <$> readMVar ourState
231524 case mmvar of
232− Nothing -> return () -- Invalid request ID. TODO: We just ignore it?
233− Just mvar -> putMVar mvar response
525+ Nothing ->
526+ return () -- Invalid request ID. TODO: We just ignore it?
527+ Just mvar -> do
528+ modifyMVar_ ourState $ return . (pendingCtrlRequestsAt reqId ^= Nothing)
529+ putMVar mvar response
530+ go
234531
235532 -- Close a connection (in this case the reqId is abused for the connection ID)
236533 closeConnection :: ConnectionId -> IO ()
237534 closeConnection cid = do
238535 writeChan ourChannel (ConnectionClosed cid)
239− decreaseRemoteRefCount transport theirEndPoint
536+ decreaseRemoteRefCount (ourEndPoint, theirEndPoint)
240537 modifyIORef openConnections (IntSet.delete (fromIntegral cid))
538+ go
539+
540+ -- Close the socket (if we don't have any outgoing connections)
541+ closeSocket :: ControlRequestId -> IO ()
542+ closeSocket cid = do
543+ let ok = encodeInt32 ControlResponse : encodeInt32 cid : prependLength [encodeInt32 (1 :: Int32)]
544+ didClose <- modifyMVar (ourEndPoint ^. localState) $ \st -> do
545+ case st ^. connectionTo theirAddr of
546+ Just (Just _, _) -> do
547+ sendTo theirEndPoint ok
548+ N.sClose sock
549+ return (st, True)
550+ Just (Nothing, theirMVar) -> do
551+ ValidEndPoint rc _ <- readMVar theirMVar
552+ if rc == 0
553+ then do
554+ sendTo theirEndPoint ok
555+ N.sClose sock
556+ return (st, True)
557+ else do
558+ return (st, False)
559+ Nothing -> do
560+ error "Internal error"
561+ unless didClose go
241562
242563 -- Read a message and output it on the endPoint's channel
243564 readMessage :: ConnectionId -> IO ()
244− readMessage connId = recvWithLength sock >>= writeChan ourChannel . Received connId
565+ readMessage connId = recvWithLength sock >>= writeChan ourChannel . Received connId >> go
245566
246567 -- Breakdown of the arguments
247568 ourChannel = ourEndPoint ^. localChannel
248569 ourState = ourEndPoint ^. localState
249570 sock = theirEndPoint ^. remoteSocket
571+ theirAddr = theirEndPoint ^. remoteAddress
250572
251573 -- We keep handling messages until an exception occurs
252− tryIO mainLoop
574+ tryIO go
253575 readIORef openConnections
254576
255577-- | Get the next connection ID
@@ -257,162 +579,17 @@ getNextConnectionId :: LocalEndPoint -> IO ConnectionId
257579getNextConnectionId endpoint =
258580 modifyMVar (endpoint ^. localState) $ \st ->
259581 return (nextConnectionId ^+= 1 $ st, st ^. nextConnectionId)
260−
261−-- | Create a new endpoint
262−tcpNewEndPoint :: TCPTransport -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
263−tcpNewEndPoint transport = do
264− endPoint <- modifyMVar (transport ^. transportState) $ \st -> do
265− chan <- newChan
266− state <- newMVar EndPointState { _nextConnectionId = firstNonReservedConnectionId
267− , _pendingCtrlRequests = IntMap.empty
268− , _nextCtrlRequestId = 0
269− }
270− let ix = fromIntegral $ IntMap.size (st ^. localEndPoints)
271− let host = transport ^. transportHost
272− let port = transport ^. transportPort
273− let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show ix
274− let localEndPoint = LocalEndPoint { _localAddress = addr
275− , _localChannel = chan
276− , _localState = state
277− }
278− return ((localEndPointAt ix ^= Just localEndPoint) $ st, localEndPoint)
279− return . Right $ EndPoint { receive = readChan (endPoint ^. localChannel)
280− , address = endPoint ^. localAddress
281− , connect = tcpConnect transport endPoint
282− , newMulticastGroup = return . Left $ newMulticastGroupError
283− , resolveMulticastGroup = \_ -> return . Left $ resolveMulticastGroupError
284− }
285− where
286− newMulticastGroupError = FailedWith NewMulticastGroupUnsupported "TCP does not support multicast"
287− resolveMulticastGroupError = FailedWith ResolveMulticastGroupUnsupported "TCP does not support multicast"
582+
583+--------------------------------------------------------------------------------
584+-- Constants --
585+--------------------------------------------------------------------------------
288586
289587-- | We reserve a bunch of connection IDs for control messages
290588firstNonReservedConnectionId :: ConnectionId
291589firstNonReservedConnectionId = 1024
292590
293−-- | Connnect to an endpoint
294−tcpConnect :: TCPTransport -- ^ Transport
295− -> LocalEndPoint -- ^ Local end point
296− -> EndPointAddress -- ^ Remote address
297− -> Reliability -- ^ Reliability (ignored)
298− -> IO (Either (FailedWith ConnectErrorCode) Connection)
299−tcpConnect transport ourEndPoint theirAddress _ = runErrorT $ do
300− theirEndPoint <- remoteEndPoint transport ourEndPoint theirAddress
301− increaseRemoteRefCount transport theirEndPoint
302− connId <- requestNewConnection ourEndPoint theirEndPoint
303− return Connection { send = tcpSend connId theirEndPoint
304− , close = tcpClose transport connId theirEndPoint
305− }
306−
307−-- | Send data across a connection
308−--
309−-- TODO: should we modify send so that it returns an error when the user closes the connection?
310−tcpSend :: ConnectionId -> RemoteEndPoint -> [ByteString] -> IO (Either (FailedWith SendErrorCode) ())
311−tcpSend connId theirEndPoint payload = runErrorT $
312− failWithIO (FailedWith SendFailed . show) $
313− sendTo theirEndPoint (encodeInt32 connId : prependLength payload)
314−
315−-- | Close a connection
316−--
317−-- We ignore errors during a close
318−tcpClose :: TCPTransport -> ConnectionId -> RemoteEndPoint -> IO ()
319−tcpClose transport connId theirEndPoint = void . runErrorT . failWithIO id $ do
320− sendTo theirEndPoint [encodeInt32 CloseConnection, encodeInt32 connId]
321− decreaseRemoteRefCount transport theirEndPoint
322−
323−-- | Increase the refcount on a remote endpoint
324−increaseRemoteRefCount :: MonadIO m => TCPTransport -> RemoteEndPoint -> m ()
325−increaseRemoteRefCount transport endPoint = do
326− liftIO $ modifyMVar_ (transport ^. transportState) $
327− return . (remoteEndPointAt (endPoint ^. remoteAddress) ^%= go)
328− where
329− go :: Maybe (RemoteEndPoint, RefCount) -> Maybe (RemoteEndPoint, RefCount)
330− go (Just (ep, RefCount rc)) = (Just (ep, RefCount $ rc + 1))
331− go Nothing = Nothing
332−
333−-- | Decrease the refcount on a remote endpoint and close the socket when it reaches zero.
334−decreaseRemoteRefCount :: MonadIO m => TCPTransport -> RemoteEndPoint -> m ()
335−decreaseRemoteRefCount transport endPoint = do
336− liftIO $ modifyMVar_ (transport ^. transportState) $
337− (remoteEndPointAt (endPoint ^. remoteAddress) ^%%= go)
338− where
339− go :: Maybe (RemoteEndPoint, RefCount) -> IO (Maybe (RemoteEndPoint, RefCount))
340− go (Just (_, RefCount 1)) = do tlog $ "Closing socket to " ++ show (endPoint ^. remoteAddress)
341− N.sClose (endPoint ^. remoteSocket)
342− return Nothing
343− go (Just (ep, RefCount rc)) = return (Just (ep, RefCount $ rc - 1))
344− go Nothing = return Nothing
345−
346−-- | Request a new connection
347−requestNewConnection :: (MonadIO m, MonadError (FailedWith ConnectErrorCode) m)
348− => LocalEndPoint
349− -> RemoteEndPoint
350− -> m ConnectionId
351−requestNewConnection ourEndPoint theirEndPoint = do
352− mcid <- failWithIO (failed . show) $ do
353− reply <- newEmptyMVar
354− reqId <- modifyMVar state $ \st -> do
355− let reqId = st ^. nextCtrlRequestId
356− return ((nextCtrlRequestId ^+= 1) . (pendingCtrlRequestsAt reqId ^= Just reply) $ st, reqId)
357− sendTo theirEndPoint [encodeInt32 RequestConnectionId, encodeInt32 reqId]
358− response <- decodeInt32 . BS.concat <$> takeMVar reply
359− modifyMVar_ state $ return . (pendingCtrlRequestsAt reqId ^= Nothing)
360− return response
361− case mcid of
362− Nothing -> throwError (failed "Invalid integer")
363− Just cid -> return cid
364− where
365− failed = FailedWith ConnectFailed
366− state = ourEndPoint ^. localState
367−
368−-- | Find a remote endpoint
369−--
370−-- TODO: hints?
371−remoteEndPoint :: TCPTransport
372− -> LocalEndPoint
373− -> EndPointAddress
374− -> ErrorT (FailedWith ConnectErrorCode) IO RemoteEndPoint
375−remoteEndPoint transport ourEndPoint theirAddress = do
376− mendpoint <- liftIO $ (^. remoteEndPointAt theirAddress) <$> readMVar (transport ^. transportState)
377− case mendpoint of
378− Just (endpoint, _) ->
379− return endpoint
380− Nothing -> do
381− -- Create a new connection
382− liftIO . tlog $ "Creating socket from " ++ show ourAddress ++ " to " ++ show theirAddress
383− (host, port, endPointId) <- failWith (invalidAddress "Could not parse") $
384− decodeEndPointAddress theirAddress
385− addr:_ <- failWithIO (invalidAddress . show) $
386− N.getAddrInfo Nothing (Just host) (Just port)
387− sock <- failWithIO (insufficientResources . show) $
388− N.socket (N.addrFamily addr) N.Stream N.defaultProtocol
389− liftIO $ N.setSocketOption sock N.ReuseAddr 1
390− failWithIO (failed . show) $ N.connect sock (N.addrAddress addr)
391− -- Send the endpoint ID we're interested in and our own address
392− failWithIO (failed . show) $ sendMany sock (encodeInt32 endPointId : prependLength [ourAddress])
393− -- For the remote endpoint to acknowledge
394− -- We fail with 'invalidAddress' here because the remote endpoint will close
395− -- the connection if we send an invalid endpoint identifier
396− ack <- failWithIO (invalidAddress . show) $ recvInt32 sock
397− case ack of
398− ValidEndPoint -> liftIO $ forkRemoteEndPoint transport ourEndPoint theirAddress sock
399− _ -> throwError (invalidAddress "Unexpected response")
400− where
401− invalidAddress = FailedWith ConnectInvalidAddress
402− insufficientResources = FailedWith ConnectInsufficientResources
403− failed = FailedWith ConnectFailed
404− EndPointAddress ourAddress = ourEndPoint ^. localAddress
405−
406−-- | Decode end point address
407−-- TODO: This uses regular expression parsing, which is nice, but unnecessary
408−decodeEndPointAddress :: EndPointAddress -> Maybe (N.HostName, N.ServiceName, EndPointId)
409−decodeEndPointAddress (EndPointAddress bs) = match endPointAddressRE (BSC.unpack bs)
410− where
411− endPointAddressRE :: RE Char (N.HostName, N.ServiceName, EndPointId)
412− endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))
413−
414591--------------------------------------------------------------------------------
415−-- Lens definitions
592+-- Lens definitions --
416593--------------------------------------------------------------------------------
417594
418595transportState :: Lens TCPTransport (MVar TransportState)
@@ -427,40 +604,38 @@ transportPort = lens _transportPort (\port trans -> trans { _transportPort = por
427604localEndPoints :: Lens TransportState (IntMap LocalEndPoint)
428605localEndPoints = lens _localEndPoints (\es st -> st { _localEndPoints = es })
429606
430−remoteEndPoints :: Lens TransportState (Map EndPointAddress (RemoteEndPoint, RefCount))
431−remoteEndPoints = lens _remoteEndPoints (\es st -> st { _remoteEndPoints = es })
432−
433607localAddress :: Lens LocalEndPoint EndPointAddress
434608localAddress = lens _localAddress (\addr ep -> ep { _localAddress = addr })
435609
436610localChannel :: Lens LocalEndPoint (Chan Event)
437611localChannel = lens _localChannel (\ch ep -> ep { _localChannel = ch })
438612
439−localState :: Lens LocalEndPoint (MVar EndPointState)
613+localState :: Lens LocalEndPoint (MVar LocalState)
440614localState = lens _localState (\st ep -> ep { _localState = st })
441615
442−remoteAddress :: Lens RemoteEndPoint EndPointAddress
443−remoteAddress = lens _remoteAddress (\addr ep -> ep { _remoteAddress = addr })
616+pendingCtrlRequests :: Lens LocalState (IntMap (MVar [ByteString]))
617+pendingCtrlRequests = lens _pendingCtrlRequests (\rep st -> st { _pendingCtrlRequests = rep })
444618
445−remoteSocket :: Lens RemoteEndPoint N.Socket
446−remoteSocket = lens _remoteSocket (\sock ep -> ep { _remoteSocket = sock })
619+nextCtrlRequestId :: Lens LocalState ControlRequestId
620+nextCtrlRequestId = lens _nextCtrlRequestId (\cid st -> st { _nextCtrlRequestId = cid })
447621
448−nextConnectionId :: Lens EndPointState ConnectionId
622+nextConnectionId :: Lens LocalState ConnectionId
449623nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = cix })
450624
451−pendingCtrlRequests :: Lens EndPointState (IntMap (MVar [ByteString]))
452−pendingCtrlRequests = lens _pendingCtrlRequests (\rep st -> st { _pendingCtrlRequests = rep })
625+localConnections :: Lens LocalState (Map EndPointAddress (Maybe Closing, MVar RemoteState))
626+localConnections = lens _localConnections (\cs st -> st { _localConnections = cs })
453627
454−nextCtrlRequestId :: Lens EndPointState ControlRequestId
455−nextCtrlRequestId = lens _nextCtrlRequestId (\cid st -> st { _nextCtrlRequestId = cid })
628+remoteAddress :: Lens RemoteEndPoint EndPointAddress
629+remoteAddress = lens _remoteAddress (\addr ep -> ep { _remoteAddress = addr })
630+
631+remoteSocket :: Lens RemoteEndPoint N.Socket
632+remoteSocket = lens _remoteSocket (\sock ep -> ep { _remoteSocket = sock })
456633
457634localEndPointAt :: EndPointId -> Lens TransportState (Maybe LocalEndPoint)
458635localEndPointAt ix = localEndPoints >>> intMapLens (fromIntegral ix)
459636
460−remoteEndPointAt :: EndPointAddress -> Lens TransportState (Maybe (RemoteEndPoint, RefCount))
461−remoteEndPointAt addr = remoteEndPoints >>> mapLens addr
462−
463−pendingCtrlRequestsAt :: ControlRequestId -> Lens EndPointState (Maybe (MVar [ByteString]))
637+pendingCtrlRequestsAt :: ControlRequestId -> Lens LocalState (Maybe (MVar [ByteString]))
464638pendingCtrlRequestsAt ix = pendingCtrlRequests >>> intMapLens (fromIntegral ix)
465639
466−
640+connectionTo :: EndPointAddress -> Lens LocalState (Maybe (Maybe Closing, MVar RemoteState))
641+connectionTo addr = localConnections >>> mapLens addr