@@ -33,7 +33,7 @@ import qualified Network.Socket as N ( HostName
3333 )
3434import Control.Concurrent (forkIO)
3535import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
36−import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, readMVar)
36+import Control.Concurrent.MVar (MVar, newMVar, modifyMVar, modifyMVar_, readMVar)
3737import Control.Category ((>>>))
3838import Control.Applicative ((<*>), (*>), (<$>))
3939import Control.Monad (forever, forM_)
@@ -47,22 +47,32 @@ import qualified Data.ByteString.Char8 as BSC (pack, unpack)
4747import Data.Int (Int16)
4848import Data.IntMap (IntMap)
4949import qualified Data.IntMap as IntMap (empty, size)
50−import Data.Lens.Lazy (Lens, lens, intMapLens, (^.), (^=), (^+=))
50+import Data.Map (Map)
51+import qualified Data.Map as Map (empty)
52+import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=))
5153import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
5254
53−data TransportState = TransportState { _endPoints :: IntMap (Chan Event, MVar EndPointState) }
55+data TransportState = TransportState { _endPoints :: IntMap (Chan Event, MVar EndPointState)
56+ , _sockets :: Map EndPointAddress N.Socket
57+ }
5458data EndPointState = EndPointState { _nextConnectionId :: ConnectionId }
5559type EndPointId = Int16
5660
5761endPoints :: Lens TransportState (IntMap (Chan Event, MVar EndPointState))
5862endPoints = lens _endPoints (\es st -> st { _endPoints = es })
5963
60−nextConnectionId :: Lens EndPointState ConnectionId
61−nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = cix })
64+sockets :: Lens TransportState (Map EndPointAddress N.Socket)
65+sockets = lens _sockets (\ss st -> st { _sockets = ss })
6266
6367endPointAt :: EndPointId -> Lens TransportState (Maybe (Chan Event, MVar EndPointState))
6468endPointAt ix = endPoints >>> intMapLens (fromIntegral ix)
6569
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 })
75+
6676-- | Create a TCP transport
6777--
6878-- Hostname and service name should be "canonical" for the transport layer to
@@ -75,7 +85,9 @@ endPointAt ix = endPoints >>> intMapLens (fromIntegral ix)
7585-- * Deal with all exceptions that may occur
7686createTransport :: N.HostName -> N.ServiceName -> IO Transport
7787createTransport host port = do
78− state <- newMVar TransportState { _endPoints = IntMap.empty }
88+ state <- newMVar TransportState { _endPoints = IntMap.empty
89+ , _sockets = Map.empty
90+ }
7991 forkServer host port (transportServer state)
8092 return Transport { newEndPoint = tcpNewEndPoint state host port }
8193
@@ -143,7 +155,7 @@ tcpNewEndPoint transportState host port = do
143155 let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show endPointId
144156 return . Right $ EndPoint { receive = readChan chan
145157 , address = addr
146− , connect = tcpConnect addr
158+ , connect = tcpConnect transportState addr
147159 , newMulticastGroup = return . Left $ newMulticastGroupError
148160 , resolveMulticastGroup = \_ -> return . Left $ resolveMulticastGroupError
149161 }
@@ -156,9 +168,13 @@ firstNonReservedConnectionId :: ConnectionId
156168firstNonReservedConnectionId = 1024
157169
158170-- | Connnect to an endpoint
159−tcpConnect :: EndPointAddress -> EndPointAddress -> Reliability -> IO (Either (FailedWith ConnectErrorCode) Connection)
160−tcpConnect myAddress theirAddress _ = runErrorT $ do
161− sock <- socketForEndPoint myAddress theirAddress
171+tcpConnect :: MVar TransportState -- ^ Transport state
172+ -> EndPointAddress -- ^ Our address
173+ -> EndPointAddress -- ^ Their address
174+ -> Reliability -- ^ Connection reliability (ignored)
175+ -> IO (Either (FailedWith ConnectErrorCode) Connection)
176+tcpConnect state ourAddress theirAddress _ = runErrorT $ do
177+ sock <- socketForEndPoint state ourAddress theirAddress
162178 connId <- failWithT undefined $ requestNewConnection sock
163179 return Connection { send = runErrorT .
164180 failWithT (FailedWith SendFailed "Send failed") .
@@ -176,15 +192,23 @@ requestNewConnection sock = sendInt32 sock 0 >> BS.concat <$> recvExact sock 4
176192-- TODOs:
177193-- * Reuse connections
178194-- * Hints?
179−socketForEndPoint :: EndPointAddress -- ^ Our address
180− -> EndPointAddress -- ^ Their address
195+socketForEndPoint :: MVar TransportState -- ^ Transport state (for reusing sockets)
196+ -> EndPointAddress -- ^ Our address
197+ -> EndPointAddress -- ^ Their address
181198 -> ErrorT (FailedWith ConnectErrorCode) IO N.Socket
182−socketForEndPoint (EndPointAddress myAddress) theirAddress = do
183− (host, port, endPointId) <- failWith invalidAddress (decodeEndPointAddress theirAddress)
184− sock <- failWithT undefined (connectTo host port)
185− endPointBS <- encodeInt16 endPointId
186− failWithT undefined $ sendWithLength sock (Just endPointBS) [myAddress]
187− return sock
199+socketForEndPoint state (EndPointAddress ourAddress) theirAddress = do
200+ msock <- liftIO $ (^. socketAt theirAddress) <$> readMVar state
201+ case msock of
202+ Just sock ->
203+ return sock
204+ Nothing -> do
205+ liftIO $ putStrLn $ "Creating socket from " ++ show ourAddress ++ " to " ++ show theirAddress
206+ (host, port, endPointId) <- failWith invalidAddress (decodeEndPointAddress theirAddress)
207+ sock <- failWithT undefined (connectTo host port)
208+ endPointBS <- encodeInt16 endPointId
209+ failWithT undefined $ sendWithLength sock (Just endPointBS) [ourAddress]
210+ liftIO $ modifyMVar_ state $ return . (socketAt theirAddress ^= Just sock)
211+ return sock
188212 where
189213 invalidAddress = FailedWith ConnectInvalidAddress "Invalid address"
190214