@@ -46,7 +46,7 @@ import qualified Network.Socket as N ( HostName
4646 , connect
4747 )
4848import Network.Socket.ByteString (sendMany)
49−import Control.Concurrent (forkIO, ThreadId, killThread)
49+import Control.Concurrent (forkIO, ThreadId, killThread, myThreadId)
5050import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
5151import Control.Concurrent.MVar ( MVar
5252 , newMVar
@@ -57,25 +57,27 @@ import Control.Concurrent.MVar ( MVar
5757 , putMVar
5858 , newEmptyMVar
5959 , withMVar
60+ , tryTakeMVar
6061 )
6162import Control.Category ((>>>))
6263import Control.Applicative ((<*>), (*>), (<$>))
63−import Control.Monad (forM_, void, when)
64+import Control.Monad (forM_, void, when, mzero, unless)
6465import Control.Monad.Error (ErrorT(ErrorT), runErrorT, throwError)
66+import Control.Monad.Trans.Maybe (MaybeT, runMaybeT)
6567import Control.Monad.IO.Class (liftIO)
6668import Control.Exception (IOException, SomeException)
67−import Data.IORef (IORef, newIORef, writeIORef, readIORef, modifyIORef)
69+import Data.IORef (IORef, newIORef, writeIORef, readIORef)
6870import Data.ByteString (ByteString)
6971import qualified Data.ByteString as BS (concat)
7072import qualified Data.ByteString.Char8 as BSC (pack, unpack)
7173import Data.Int (Int32)
7274import Data.IntMap (IntMap)
73−import qualified Data.IntMap as IntMap (empty, size)
75+import qualified Data.IntMap as IntMap (empty)
7476import Data.IntSet (IntSet)
75−import qualified Data.IntSet as IntSet (empty, insert, delete, elems)
77+import qualified Data.IntSet as IntSet (empty, insert, elems, singleton, null, delete)
7678import Data.Map (Map)
77−import qualified Data.Map as Map (empty)
78−import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=))
79+import qualified Data.Map as Map (empty, elems, size)
80+import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=), (^%=))
7981import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
8082import System.IO (hPutStrLn, stderr)
8183
@@ -183,18 +185,18 @@ data LocalEndPoint = LocalEndPoint { localAddress :: EndPointAddress
183185 }
184186data RemoteEndPoint = RemoteEndPoint { remoteAddress :: EndPointAddress
185187 , remoteState :: MVar RemoteState
186− , remoteThread :: ThreadId
187188 }
188189
189190-- We use underscores for fields that we might update (using lenses).
190191
191−data TransportState = TransportState { _localEndPoints :: IntMap LocalEndPoint
192+data TransportState = TransportState { _localEndPoints :: Map EndPointAddress LocalEndPoint
192193 , _acceptNewConnections :: Bool
193194 }
194195data LocalState = LocalState { _nextConnectionId :: ConnectionId
195196 , _pendingCtrlRequests :: IntMap (MVar [ByteString])
196197 , _nextCtrlRequestId :: ControlRequestId
197198 , _localConnections :: Map EndPointAddress RemoteEndPoint
199+ , _internalThreads :: [ThreadId]
198200 }
199201
200202-- A remote endpoint has incoming and outgoing connections, and when the total
@@ -238,9 +240,10 @@ data LocalState = LocalState { _nextConnectionId :: ConnectionId
238240
239241data RemoteState = RemoteEndPointInvalid (FailedWith ConnectErrorCode)
240242 | RemoteEndPointValid RemoteConnection
241− | RemoteEndPointClosing RemoteConnection (MVar ())
243+ | RemoteEndPointClosing (MVar ()) RemoteConnection
242244 | RemoteEndPointClosed
243−data RemoteConnection = RemoteConnection { _remoteRefCount :: !Int
245+data RemoteConnection = RemoteConnection { _remoteOutgoing :: !Int
246+ , _remoteIncoming :: IntSet
244247 , remoteSocket :: N.Socket
245248 , sendOn :: [ByteString] -> IO ()
246249 }
@@ -273,7 +276,7 @@ data ConnectionRequestResponse =
273276-- TODOs: deal with hints
274277createTransport :: N.HostName -> N.ServiceName -> IO (Either IOException Transport)
275278createTransport host port = do
276− state <- newMVar TransportState { _localEndPoints = IntMap.empty
279+ state <- newMVar TransportState { _localEndPoints = Map.empty
277280 , _acceptNewConnections = True
278281 }
279282 let transport = TCPTransport { transportState = state
@@ -308,9 +311,10 @@ apiCloseTransport transportThread = do
308311apiNewEndPoint :: TCPTransport -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)
309312apiNewEndPoint transport = do
310313 ourEndPoint <- createLocalEndPoint transport
311− return . Right $ EndPoint { receive = readChan (localChannel ourEndPoint)
312− , address = localAddress ourEndPoint
313− , connect = apiConnect ourEndPoint
314+ return . Right $ EndPoint { receive = readChan (localChannel ourEndPoint)
315+ , address = localAddress ourEndPoint
316+ , connect = apiConnect ourEndPoint
317+ , closeEndPoint = apiCloseEndPoint transport ourEndPoint
314318 , newMulticastGroup = return . Left $ newMulticastGroupError
315319 , resolveMulticastGroup = return . Left . const resolveMulticastGroupError
316320 }
@@ -345,14 +349,16 @@ apiConnect ourEndPoint theirAddress _ = runErrorT $ do
345349-- TODO: We ignore errors during a close. Is that right?
346350apiClose :: RemoteEndPoint -> ConnectionId -> IORef Bool -> IO ()
347351apiClose theirEndPoint connId connAlive = void . tryIO $ do
348− wasAlive <- withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) -> do
352+ modifyMVar_ (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) -> do
349353 alive <- readIORef connAlive
350− when alive $ do
351− sendOn remoteConn [encodeInt32 CloseConnection, encodeInt32 connId]
352− writeIORef connAlive False
353− return alive
354− -- We don't want to call decreaseRemoteRefCount while we still hold the lock
355− when wasAlive $ decreaseRemoteRefCount theirEndPoint
354+ if alive
355+ then do
356+ sendOn remoteConn [encodeInt32 CloseConnection, encodeInt32 connId]
357+ writeIORef connAlive False
358+ return (RemoteEndPointValid . (remoteOutgoing ^+= (-1)) $ remoteConn)
359+ else
360+ return (RemoteEndPointValid remoteConn)
361+ closeIfUnused theirEndPoint
356362
357363-- | Send data across a connection
358364--
@@ -369,6 +375,37 @@ apiSend theirEndPoint connId connAlive payload = do
369375 else
370376 throwError $ FailedWith SendConnectionClosed "Connection closed"
371377
378+-- | Force-close the endpoint
379+apiCloseEndPoint :: TCPTransport -> LocalEndPoint -> IO ()
380+apiCloseEndPoint transport ourEndPoint = do
381+ -- Remove the reference from the transport state
382+ modifyMVar_ (transportState transport) $ return . (localEndPointAt (localAddress ourEndPoint) ^= Nothing)
383+ -- We take the endpoint state and don't put it back
384+ -- (this may block some threads, but we don't care: we are killing them anyway)
385+ ourState <- takeMVar (localState ourEndPoint)
386+ --incoming <- IntSet.elems <$> takeMVar (localIncoming ourEndPoint)
387+ let incoming = undefined
388+ -- Kill all threads associated with this endpoint and close all sockets
389+ forM_ (ourState ^. internalThreads) killThread
390+ forM_ (Map.elems $ ourState ^. localConnections) $ \theirEndPoint ->
391+ runMaybeT $ remoteSocketOf theirEndPoint >>= liftIO . N.sClose
392+ -- We send a single message to the endpoint that it is closed. Subsequent
393+ -- calls will block. We could change this so that all subsequent calls to
394+ -- receive return an error, but this would mean checking for some state on
395+ -- every call to receive, which is an unnecessary overhead.
396+ writeChan (localChannel ourEndPoint) (ErrorEvent EventErrorEndPointClosed incoming)
397+ where
398+ remoteSocketOf :: RemoteEndPoint -> MaybeT IO N.Socket
399+ remoteSocketOf theirEndPoint = do
400+ -- Like above, we take but don't put back
401+ mRemoteState <- liftIO $ tryTakeMVar (remoteState theirEndPoint)
402+ case mRemoteState of
403+ Nothing -> mzero
404+ Just (RemoteEndPointInvalid _) -> mzero
405+ Just (RemoteEndPointValid conn) -> return $ remoteSocket conn
406+ Just (RemoteEndPointClosing _ conn) -> return $ remoteSocket conn
407+ Just RemoteEndPointClosed -> mzero
408+
372409-- | Special case of 'apiConnect': connect an endpoint to itself
373410connectToSelf :: LocalEndPoint -> IO (Either (FailedWith ConnectErrorCode) Connection)
374411connectToSelf ourEndPoint = do
@@ -411,17 +448,16 @@ createLocalEndPoint transport = do
411448 , _pendingCtrlRequests = IntMap.empty
412449 , _nextCtrlRequestId = 0
413450 , _localConnections = Map.empty
451+ , _internalThreads = []
414452 }
415− let host = transportHost transport
416− let port = transportPort transport
417453 modifyMVar (transportState transport) $ \st -> do
418− let ix = fromIntegral $ IntMap.size (st ^. localEndPoints)
419− let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show ix
420− let localEndPoint = LocalEndPoint { localAddress = addr
421− , localChannel = chan
422− , localState = state
454+ let ix = fromIntegral $ Map.size (st ^. localEndPoints)
455+ let addr = encodeEndPointAddress transport ix
456+ let localEndPoint = LocalEndPoint { localAddress = addr
457+ , localChannel = chan
458+ , localState = state
423459 }
424− return ((localEndPointAt ix ^= Just localEndPoint) $ st, localEndPoint)
460+ return ((localEndPointAt addr ^= Just localEndPoint) $ st, localEndPoint)
425461
426462-- | Request a connection to a remote endpoint
427463--
@@ -435,17 +471,22 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go
435471 where
436472 go = do
437473 let ourState = localState ourEndPoint
438− theirEndPoint <- modifyMVar ourState $ \st ->
474+ (theirEndPoint, isNew) <- modifyMVar ourState $ \st ->
439475 case st ^. connectionTo theirAddress of
440476 Just theirEndPoint ->
441− return (st, theirEndPoint)
477+ return (st, (theirEndPoint, False))
442478 Nothing -> do
443479 theirState <- newEmptyMVar
444480 let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress
445481 , remoteState = theirState
446482 }
447− forkIO $ connectToRemoteEndPoint (ourEndPoint, theirEndPoint)
448− return (connectionTo theirAddress ^= Just theirEndPoint $ st, theirEndPoint)
483+ return (connectionTo theirAddress ^= Just theirEndPoint $ st, (theirEndPoint, True))
484+
485+ -- If this is a new endpoint, fork a thread to listen for incoming
486+ -- connections. We don't want to do this while we hold the lock, because
487+ -- forkEndPointThread modifies the local state too (to record the thread ID)
488+ when isNew . forkEndPointThread ourEndPoint $
489+ connectToRemoteEndPoint (ourEndPoint, theirEndPoint)
449490
450491 -- Before we initiate the new connection request we want to make sure
451492 -- that refcount on the endpoint is incremented so that a concurrent
@@ -456,7 +497,7 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go
456497 endPointStateSnapshot <- modifyMVar theirState $ \st ->
457498 case st of
458499 RemoteEndPointValid ep ->
459− return (RemoteEndPointValid . (remoteRefCount ^+= 1) $ ep, st)
500+ return (RemoteEndPointValid . (remoteOutgoing ^+= 1) $ ep, st)
460501 _ ->
461502 return (st, st)
462503
@@ -471,7 +512,7 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go
471512 RemoteEndPointInvalid err -> do
472513 return . Left $ err
473514
474− RemoteEndPointClosing _ resolved -> do
515+ RemoteEndPointClosing resolved _ -> do
475516 -- If the remote endpoint is closing, then we need to block until
476517 -- this is resolved and we then try again
477518 readMVar resolved
@@ -498,12 +539,13 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go
498539 -- valid state at this point.
499540 let failed err = do
500541 modifyMVar_ theirState $ \(RemoteEndPointValid ep) ->
501− return (RemoteEndPointValid . (remoteRefCount ^+= (-1)) $ ep)
542+ return (RemoteEndPointValid . (remoteOutgoing ^+= (-1)) $ ep)
543+ -- TODO: should we call closeIfUnused here?
502544 return . Left $ FailedWith ConnectFailed err
503545
504546 case mcid of
505547 Left err -> failed (show err)
506− Right Nothing -> failed ("Invalid integer")
548+ Right Nothing -> failed "Invalid integer"
507549 Right (Just cid) -> return . Right $ (theirEndPoint, cid)
508550
509551-- | Establish a connection to a remote endpoint
@@ -520,11 +562,12 @@ connectToRemoteEndPoint (ourEndPoint, theirEndPoint) = do
520562 case result of
521563 Right (ConnectionRequestAccepted, sock) -> do
522564 let remoteConn = RemoteConnection { remoteSocket = sock
523− , _remoteRefCount = 0
565+ , _remoteOutgoing = 0
566+ , _remoteIncoming = IntSet.empty
524567 , sendOn = sendMany sock
525568 }
526569 putMVar theirState (RemoteEndPointValid remoteConn)
527− void . forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)
570+ handleIncomingMessages (ourEndPoint, theirEndPoint)
528571 Right (ConnectionRequestEndPointInvalid, sock) -> do
529572 -- We remove the endpoint from our local state again because the next
530573 -- call to 'connect' might give a different result. Threads that were
@@ -569,6 +612,11 @@ connectToRemoteEndPoint (ourEndPoint, theirEndPoint) = do
569612 ourAddress :: ByteString
570613 EndPointAddress ourAddress = localAddress ourEndPoint
571614
615+-- | Encode end point address
616+encodeEndPointAddress :: TCPTransport -> EndPointId -> EndPointAddress
617+encodeEndPointAddress transport ix = EndPointAddress . BSC.pack $
618+ transportHost transport ++ ":" ++ transportPort transport ++ ":" ++ show ix
619+
572620-- | Decode end point address
573621-- TODO: This uses regular expression parsing, which is nice, but unnecessary
574622decodeEndPointAddress :: EndPointAddress -> Maybe (N.HostName, N.ServiceName, EndPointId)
@@ -581,7 +629,7 @@ decodeEndPointAddress (EndPointAddress bs) = match endPointAddressRE (BSC.unpack
581629--
582630-- RELY: Remote endpoint must be in valid state.
583631-- GUARANTEE: Will not change the state of the remote endpoint.
584−doRemoteRequest :: EndPointPair -> ControlHeader -> IO ([ByteString])
632+doRemoteRequest :: EndPointPair -> ControlHeader -> IO [ByteString]
585633doRemoteRequest (ourEndPoint, theirEndPoint) header = do
586634 reply <- newEmptyMVar
587635 reqId <- modifyMVar (localState ourEndPoint) $ \st -> do
@@ -591,22 +639,32 @@ doRemoteRequest (ourEndPoint, theirEndPoint) header = do
591639 sendOn remoteConn [encodeInt32 header, encodeInt32 reqId]
592640 takeMVar reply
593641
594−-- | Decrease the remote refcount of a remote endpoint. When the refcount
595−-- reaches zero, initiate socket closure.
596−--
597−-- RELY: Remote endpoint must be in valid state.
598−-- GUARANTEE: Will leave the endpoint in its current state or change it to
599−-- RemoteEndPointClosing.
600−decreaseRemoteRefCount :: RemoteEndPoint -> IO ()
601−decreaseRemoteRefCount theirEndPoint = do
602− modifyMVar_ (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) ->
603− if remoteConn ^. remoteRefCount == 1
604− then do
605− resolved <- newEmptyMVar
606− sendOn remoteConn [encodeInt32 CloseSocket]
607− return (RemoteEndPointClosing remoteConn resolved)
608− else do
609− return (RemoteEndPointValid . (remoteRefCount ^+= (-1)) $ remoteConn)
642+-- | Check if the remote endpoint is unused, and if so, send a CloseSocket request
643+closeIfUnused :: RemoteEndPoint -> IO ()
644+closeIfUnused theirEndPoint = modifyMVar_ (remoteState theirEndPoint) $ \st ->
645+ case st of
646+ RemoteEndPointValid remoteConn -> do
647+ if remoteConn ^. remoteOutgoing == 0 && IntSet.null (remoteConn ^. remoteIncoming)
648+ then do
649+ resolved <- newEmptyMVar
650+ sendOn remoteConn [encodeInt32 CloseSocket]
651+ return (RemoteEndPointClosing resolved remoteConn)
652+ else
653+ return st
654+ _ -> do
655+ return st
656+
657+-- | Fork a new thread and store its ID as part of the transport state
658+forkEndPointThread :: LocalEndPoint -> IO () -> IO ()
659+forkEndPointThread ourEndPoint p = do
660+ modifyMVar_ (localState ourEndPoint) $ \st -> do
661+ tid <- forkIO (p >> removeThread)
662+ return (internalThreads ^%= (tid :) $ st)
663+ where
664+ removeThread :: IO ()
665+ removeThread = do
666+ tid <- myThreadId
667+ modifyMVar_ (localState ourEndPoint) $ return . (internalThreads ^%= filter (/= tid))
610668
611669--------------------------------------------------------------------------------
612670-- Incoming requests --
@@ -616,17 +674,17 @@ decreaseRemoteRefCount theirEndPoint = do
616674-- establish a TCP connection with us)
617675handleConnectionRequest :: TCPTransport -> N.Socket -> IO ()
618676handleConnectionRequest transport sock = do
619− let tState = transportState transport
620677 request <- tryIO $ do
621678 ourEndPointId <- recvInt32 sock
622679 theirAddress <- EndPointAddress . BS.concat <$> recvWithLength sock
623− ourEndPoint <- (^. localEndPointAt ourEndPointId) <$> readMVar tState
680+ let ourAddress = encodeEndPointAddress transport ourEndPointId
681+ ourEndPoint <- (^. localEndPointAt ourAddress) <$> readMVar (transportState transport)
624682 return (theirAddress, ourEndPoint)
625683 case request of
626684 Right (_, Nothing) -> do
627685 sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]
628686 N.sClose sock
629− Right (theirAddress, Just ourEndPoint) -> do
687+ Right (theirAddress, Just ourEndPoint) -> forkEndPointThread ourEndPoint $ do
630688 (crossed, theirEndPoint) <- modifyMVar (localState ourEndPoint) $ \st ->
631689 case st ^. connectionTo theirAddress of
632690 Nothing -> do
@@ -643,14 +701,14 @@ handleConnectionRequest transport sock = do
643701 N.sClose sock
644702 else do
645703 let remoteConn = RemoteConnection { remoteSocket = sock
646− , _remoteRefCount = 0
704+ , _remoteOutgoing = 0
705+ , _remoteIncoming = IntSet.empty
647706 , sendOn = sendMany sock
648707 }
708+ -- TODO: this putMVar might block
649709 putMVar (remoteState theirEndPoint) (RemoteEndPointValid remoteConn)
650710 sendMany sock [encodeInt32 ConnectionRequestAccepted]
651− -- TODO: we should store the tid as part of the endpoint state
652− tid <- forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)
653− return ()
711+ handleIncomingMessages (ourEndPoint, theirEndPoint)
654712 Left _ -> do
655713 -- Invalid request
656714 N.sClose sock
@@ -665,59 +723,72 @@ handleConnectionRequest transport sock = do
665723-- GUARANTEE: May change the remote endpoint to RemoteEndPointClosed state.
666724handleIncomingMessages :: EndPointPair -> IO ()
667725handleIncomingMessages (ourEndPoint, theirEndPoint) = do
668− -- We use an IORef rather than the state monad because when an exception occurs we
669− -- want to know the current list of open connections
670− openConnections <- newIORef IntSet.empty
671−
672726 -- For efficiency sake we get the socket once and for all
673727 sock <- withMVar theirState $ \st ->
674728 case st of
675729 RemoteEndPointValid ep ->
676730 return (remoteSocket ep)
677− RemoteEndPointClosing ep _ ->
731+ RemoteEndPointClosing _ ep ->
678732 return (remoteSocket ep)
679733 _ ->
680734 error "handleIncomingMessages RELY violation"
681735
682− tryIO $ go sock openConnections
683− unclosedConnections <- readIORef openConnections
684− forM_ (map fromIntegral $ IntSet.elems unclosedConnections) $
685− writeChan ourChannel . ConnectionClosed
736+ tryIO $ go sock
737+
738+ -- Normally, there would be no unclosed connections, but if the client
739+ -- exits prematurely there might be
740+ unclosedConnections <- modifyMVar theirState $ \st ->
741+ case st of
742+ RemoteEndPointInvalid _ ->
743+ error "handleIncomingMessages RELY violation"
744+ RemoteEndPointValid remoteConn ->
745+ return (RemoteEndPointValid . (remoteIncoming ^= IntSet.empty) $ remoteConn, remoteConn ^. remoteIncoming)
746+ RemoteEndPointClosing resolved remoteConn ->
747+ return (RemoteEndPointClosing resolved . (remoteIncoming ^= IntSet.empty) $ remoteConn, remoteConn ^. remoteIncoming)
748+ RemoteEndPointClosed ->
749+ return (st, IntSet.empty)
750+
751+ -- TODO: should we give a single ErrorEvent instead?
752+ forM_ (map fromIntegral $ IntSet.elems unclosedConnections) $ \cid -> do
753+ writeChan ourChannel $ ConnectionClosed cid
754+
755+ -- TODO: if the endpoint was not in closing state, should we make it so?
686756 where
687757 -- Dispatch
688− go :: N.Socket -> IORef IntSet -> IO ()
689− go sock openConnections = do
758+ go :: N.Socket -> IO ()
759+ go sock = do
690760 connId <- recvInt32 sock
691761 if connId >= firstNonReservedConnectionId
692762 then do
693763 readMessage sock connId
694− go sock openConnections
764+ go sock
695765 else do
696766 case tryToEnum (fromIntegral connId) of
697767 Just RequestConnectionId -> do
698− recvInt32 sock >>= createNewConnection openConnections
699− go sock openConnections
768+ recvInt32 sock >>= createNewConnection
769+ go sock
700770 Just ControlResponse -> do
701771 recvInt32 sock >>= readControlResponse sock
702− go sock openConnections
772+ go sock
703773 Just CloseConnection -> do
704− recvInt32 sock >>= closeConnection openConnections
705− go sock openConnections
774+ recvInt32 sock >>= closeConnection
775+ go sock
706776 Just CloseSocket -> do
707− closeSocket sock
708− go sock openConnections
777+ didClose <- closeSocket sock
778+ unless didClose $ go sock
709779 Nothing ->
710780 -- Invalid control request, exit
711781 hPutStrLn stderr "Warning: invalid control request"
712782
713783 -- Create a new connection
714− createNewConnection :: IORef IntSet -> ControlRequestId -> IO ()
715− createNewConnection openConnections reqId = do
716− modifyMVar_ theirState $ \st ->
717− case st of
718− RemoteEndPointValid ep ->
719− return (RemoteEndPointValid . (remoteRefCount ^+= 1) $ ep)
720− RemoteEndPointClosing ep resolved -> do
784+ createNewConnection :: ControlRequestId -> IO ()
785+ createNewConnection reqId = do
786+ newId <- getNextConnectionId ourEndPoint
787+ modifyMVar_ theirState $ \st -> do
788+ remoteConn <- case st of
789+ RemoteEndPointValid remoteConn ->
790+ return (remoteIncoming ^%= IntSet.insert (fromIntegral newId) $ remoteConn)
791+ RemoteEndPointClosing resolved remoteConn -> do
721792 -- If the endpoint is in closing state that means we send a
722793 -- CloseSocket request to the remote endpoint. If the remote
723794 -- endpoint replies with the request to create a new connection, it
@@ -725,19 +796,17 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
725796 -- ours. Either way, at this point we simply restore the endpoint
726797 -- to RemoteEndPointValid
727798 putMVar resolved ()
728− return (RemoteEndPointValid . (remoteRefCount ^= 1) $ ep)
799+ return (remoteIncoming ^= IntSet.singleton newId $ remoteConn)
729800 _ ->
730801 error "handleIncomingMessages RELY violation"
731− newId <- getNextConnectionId ourEndPoint
732− withMVar theirState $ \(RemoteEndPointValid remoteConn) ->
733802 sendOn remoteConn ( encodeInt32 ControlResponse
734803 : encodeInt32 reqId
735804 : prependLength [encodeInt32 newId]
736805 )
806+ -- We add the new connection ID to the list of open connections only once the
807+ -- endpoint has been notified of the new connection (sendOn may fail)
808+ return (RemoteEndPointValid remoteConn)
737809 writeChan ourChannel (ConnectionOpened newId ReliableOrdered theirAddr)
738− -- We add the new connection ID to the list of open connections only once the
739− -- endpoint has been notified of the new connection (sendOn may fail)
740− modifyIORef openConnections (IntSet.insert (fromIntegral newId))
741810
742811 -- Read a control response
743812 readControlResponse :: N.Socket -> ControlRequestId -> IO ()
@@ -753,22 +822,24 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
753822 putMVar mvar response
754823
755824 -- Close a connection
756− closeConnection :: IORef IntSet -> ConnectionId -> IO ()
757− closeConnection openConnections cid = do
825+ closeConnection :: ConnectionId -> IO ()
826+ closeConnection cid = do
827+ -- TODO: we should check that this connection is in fact open
758828 writeChan ourChannel (ConnectionClosed cid)
759− decreaseRemoteRefCount theirEndPoint
760− modifyIORef openConnections (IntSet.delete (fromIntegral cid))
829+ modifyMVar_ theirState $ \(RemoteEndPointValid remoteConn) ->
830+ return (RemoteEndPointValid . (remoteIncoming ^%= IntSet.delete cid) $ remoteConn)
831+ closeIfUnused theirEndPoint
761832
762833 -- Close the socket (if we don't have any outgoing connections)
763834 closeSocket :: N.Socket -> IO Bool
764835 closeSocket sock = do
765836 didClose <- modifyMVar theirState $ \st ->
766837 case st of
767− RemoteEndPointValid remoteConn | remoteConn ^. remoteRefCount == 0 -> do
838+ RemoteEndPointValid remoteConn | remoteConn ^. remoteOutgoing == 0 -> do
768839 sendOn remoteConn [encodeInt32 CloseSocket]
769840 N.sClose sock
770841 return (RemoteEndPointClosed, True)
771− RemoteEndPointClosing _ resolved -> do
842+ RemoteEndPointClosing resolved _ -> do
772843 -- If the socket is already in closing state, then we have already
773844 -- sent a CloseSocket request; hence, we don't need to do it again
774845 -- at this point
@@ -785,10 +856,12 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
785856 readMessage sock connId = recvWithLength sock >>= writeChan ourChannel . Received connId
786857
787858 -- Arguments
788− ourChannel = localChannel ourEndPoint
789− ourState = localState ourEndPoint
790− theirState = remoteState theirEndPoint
791− theirAddr = remoteAddress theirEndPoint
859+ ourChannel = localChannel ourEndPoint
860+ ourState = localState ourEndPoint
861+ theirState = remoteState theirEndPoint
862+ theirAddr = remoteAddress theirEndPoint
863+
864+
792865
793866-- | Get the next connection ID
794867getNextConnectionId :: LocalEndPoint -> IO ConnectionId
@@ -809,7 +882,7 @@ firstNonReservedConnectionId = 1024
809882-- Lens definitions --
810883--------------------------------------------------------------------------------
811884
812−localEndPoints :: Lens TransportState (IntMap LocalEndPoint)
885+localEndPoints :: Lens TransportState (Map EndPointAddress LocalEndPoint)
813886localEndPoints = lens _localEndPoints (\es st -> st { _localEndPoints = es })
814887
815888acceptNewConnections :: Lens TransportState Bool
@@ -827,8 +900,11 @@ nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = c
827900localConnections :: Lens LocalState (Map EndPointAddress RemoteEndPoint)
828901localConnections = lens _localConnections (\es st -> st { _localConnections = es })
829902
830−localEndPointAt :: EndPointId -> Lens TransportState (Maybe LocalEndPoint)
831−localEndPointAt ix = localEndPoints >>> intMapLens (fromIntegral ix)
903+internalThreads :: Lens LocalState [ThreadId]
904+internalThreads = lens _internalThreads (\ts st -> st { _internalThreads = ts })
905+
906+localEndPointAt :: EndPointAddress -> Lens TransportState (Maybe LocalEndPoint)
907+localEndPointAt addr = localEndPoints >>> mapLens addr
832908
833909pendingCtrlRequestsAt :: ControlRequestId -> Lens LocalState (Maybe (MVar [ByteString]))
834910pendingCtrlRequestsAt ix = pendingCtrlRequests >>> intMapLens (fromIntegral ix)
@@ -836,5 +912,8 @@ pendingCtrlRequestsAt ix = pendingCtrlRequests >>> intMapLens (fromIntegral ix)
836912connectionTo :: EndPointAddress -> Lens LocalState (Maybe RemoteEndPoint)
837913connectionTo addr = localConnections >>> mapLens addr
838914
839−remoteRefCount :: Lens RemoteConnection Int
840−remoteRefCount = lens _remoteRefCount (\rc conn -> conn { _remoteRefCount = rc })
915+remoteOutgoing :: Lens RemoteConnection Int
916+remoteOutgoing = lens _remoteOutgoing (\cs conn -> conn { _remoteOutgoing = cs })
917+
918+remoteIncoming :: Lens RemoteConnection IntSet
919+remoteIncoming = lens _remoteIncoming (\cs conn -> conn { _remoteIncoming = cs })