@@ -17,6 +17,7 @@ module Network.Transport.TCP ( -- * Main API
1717 createTransport
1818 -- * Internals (exposed for unit tests)
1919 , createTransportExposeInternals
20+ , TransportInternals(..)
2021 , EndPointId
2122 , encodeEndPointAddress
2223 , decodeEndPointAddress
@@ -74,8 +75,18 @@ import Control.Concurrent.MVar ( MVar
7475 )
7576import Control.Category ((>>>))
7677import Control.Applicative ((<$>))
77−import Control.Monad (forM_, when, unless)
78−import Control.Exception (IOException, SomeException, handle, throw, try, bracketOnError, mask_)
78+import Control.Monad (when, unless, join)
79+import Control.Monad.Error (ErrorT(..), runErrorT)
80+import Control.Exception ( IOException
81+ , SomeException
82+ , handle
83+ , throw
84+ , try
85+ , bracketOnError
86+ , mask
87+ , mask_
88+ , onException
89+ )
7990import Data.IORef (IORef, newIORef, writeIORef, readIORef)
8091import Data.ByteString (ByteString)
8192import qualified Data.ByteString as BS (concat)
@@ -89,6 +100,7 @@ import Data.Map (Map)
89100import qualified Data.Map as Map (empty, elems)
90101import Data.Accessor (Accessor, accessor, (^.), (^=), (^:))
91102import qualified Data.Accessor.Container as DAC (mapMaybe, intMapMaybe)
103+import Data.Foldable (forM_)
92104import System.IO (hPutStrLn, stderr)
93105
94106-- $design
@@ -229,6 +241,7 @@ data ValidLocalEndPointState = ValidLocalEndPointState
229241 , _nextCtrlRequestId :: ControlRequestId
230242 , _localConnections :: Map EndPointAddress RemoteEndPoint
231243 , _internalThreads :: [ThreadId]
244+ , _nextRemoteId :: Int
232245 }
233246
234247-- A remote endpoint has incoming and outgoing connections, and when the total
@@ -270,15 +283,31 @@ data ValidLocalEndPointState = ValidLocalEndPointState
270283-- Since we don't do concurrent reads from the same socket we don't need to
271284-- take the lock when reading from the socket.
272285--
273−-- Moreover, we insist on the invariant (INV-CLOSE) that whenever we put an
274−-- endpoint in closed state we remove that endpoint from localConnections
275−-- first, so that if a concurrent thread reads the mvar, finds EndPointClosed,
276−-- and then looks up the endpoint in localConnections it is guaranteed to
277−-- either find a different remote endpoint, or else none at all.
286+-- Invariants:
287+--
288+-- INV-CLOSE: Whenever we put an endpoint in closed state we remove that
289+-- endpoint from localConnections first, so that if a concurrent thread reads
290+-- the mvar, finds EndPointClosed, and then looks up the endpoint in
291+-- localConnections it is guaranteed to either find a different remote
292+-- endpoint, or else none at all.
293+-- INV-RESOLVE: Whenever we move a endpoint from Closing to Closed state, we
294+-- signal on the corresponding MVar only *after* the endpoint has been put in
295+-- Closed state. This way when to threads try to resolve they don't both
296+-- attempt to write to the "resolved" MVar. TODO: Make sure that this
297+-- invariant is adhered too.
298+-- INV-LOST: If a send or recv fails, or a socket is closed unexpectedly, we
299+-- first put the remote endpoint in Closing or Closed state, and then send a
300+-- EventConnectionLost event. This guarantees that we only send this event
301+-- once.
302+-- INV-CLOSING: An endpoint in closing state is for all intents and purposes
303+-- closed; that is, we shouldn't do any 'send's on it (although 'recv' is
304+-- acceptable, of course -- as we are waiting for the remote endpoint to
305+-- confirm or deny the request).
278306
279307data RemoteEndPoint = RemoteEndPoint
280308 { remoteAddress :: EndPointAddress
281309 , remoteState :: MVar RemoteState
310+ , remoteId :: Int
282311 }
283312
284313data RemoteState =
@@ -313,6 +342,14 @@ data ConnectionRequestResponse =
313342 | ConnectionRequestCrossed -- ^ /A/s request crossed with a request from /B/ (see protocols)
314343 deriving (Enum, Bounded, Show)
315344
345+-- Internal functionality we expose for unit testing
346+data TransportInternals = TransportInternals
347+ { -- | The ID of the thread that listens for new incoming connections
348+ transportThread :: ThreadId
349+ -- | Find the socket between a local and a remote endpoint
350+ , socketBetween :: EndPointAddress -> EndPointAddress -> IO (Either String N.Socket)
351+ }
352+
316353--------------------------------------------------------------------------------
317354-- Top-level functionality --
318355--------------------------------------------------------------------------------
@@ -330,7 +367,7 @@ createTransport host port = do
330367 return $ Right transport
331368
332369-- | You should probably not use this function (used for unit testing only)
333−createTransportExposeInternals :: N.HostName -> N.ServiceName -> IO (Either IOException (Transport, ThreadId))
370+createTransportExposeInternals :: N.HostName -> N.ServiceName -> IO (Either IOException (Transport, TransportInternals))
334371createTransportExposeInternals host port = do
335372 state <- newMVar . TransportValid $ ValidTransportState
336373 { _localEndPoints = Map.empty
@@ -344,11 +381,13 @@ createTransportExposeInternals host port = do
344381 -- For a discussion of the use of N.sOMAXCONN, see
345382 -- http://tangentsoft.net/wskfaq/advanced.html#backlog
346383 -- http://www.linuxjournal.com/files/linuxjournal.com/linuxjournal/articles/023/2333/2333s2.html
347− transportThread <- forkServer host port N.sOMAXCONN (terminationHandler transport) (handleConnectionRequest transport)
384+ tid <- forkServer host port N.sOMAXCONN (terminationHandler transport) (handleConnectionRequest transport)
348385 return ( Transport { newEndPoint = apiNewEndPoint transport
349− , closeTransport = apiCloseTransport transport (Just transportThread) [EndPointClosed]
386+ , closeTransport = apiCloseTransport transport (Just tid) [EndPointClosed]
350387 }
351− , transportThread
388+ , TransportInternals { transportThread = tid
389+ , socketBetween = internalSocketBetween transport
390+ }
352391 )
353392 where
354393 terminationHandler :: TCPTransport -> SomeException -> IO ()
@@ -366,22 +405,13 @@ apiCloseTransport transport mTransportThread evs = do
366405 mTSt <- modifyMVar (transportState transport) $ \st -> case st of
367406 TransportValid vst -> return (TransportClosed, Just vst)
368407 TransportClosed -> return (TransportClosed, Nothing)
369− case mTSt of
370− Nothing ->
371− -- Transport already closed
372− return ()
373− Just tSt -> do
374− mapM_ (apiCloseEndPoint transport evs) (Map.elems $ tSt ^. localEndPoints)
375− case mTransportThread of
376− Just transportThread ->
377− -- This will invoke the termination handler, which in turn will call
378− -- apiCloseTransport again, but at that point the transport is
379− -- already in closed state so at that point we return immediately
380− killThread transportThread
381− Nothing ->
382− -- This will happen is apiCloseTransport was invoked *because* the
383− -- transport thread died (in which case we don't need to kill it)
384− return ()
408+ forM_ mTSt $ \tSt -> do
409+ mapM_ (apiCloseEndPoint transport evs) (Map.elems $ tSt ^. localEndPoints)
410+ forM_ mTransportThread $ \tid -> do
411+ -- This will invoke the termination handler, which in turn will call
412+ -- apiCloseTransport again, but at that point the transport is
413+ -- already in closed state so at that point we return immediately
414+ killThread tid
385415
386416-- | Create a new endpoint
387417apiNewEndPoint :: TCPTransport -> IO (Either (TransportError NewEndPointErrorCode) EndPoint)
@@ -412,8 +442,8 @@ apiConnect ourEndPoint theirAddress _reliability _hints = try $ do
412442 -- connAlive can be an IORef rather than an MVar because it is protected by
413443 -- the remoteState MVar. We don't need the overhead of locking twice.
414444 connAlive <- newIORef True
415− return $ Connection { send = apiSend theirEndPoint connId connAlive
416− , close = apiClose theirEndPoint connId connAlive
445+ return $ Connection { send = apiSend (ourEndPoint, theirEndPoint) connId connAlive
446+ , close = apiClose (ourEndPoint, theirEndPoint) connId connAlive
417447 }
418448
419449-- | Close a connection
@@ -422,53 +452,49 @@ apiConnect ourEndPoint theirAddress _reliability _hints = try $ do
422452-- GUARANTEE: If the connection is alive on entry then the remote endpoint will
423453-- either be RemoteEndPointValid or RemoteEndPointClosing. Otherwise, the state
424454-- of the remote endpoint will not be changed.
425−apiClose :: RemoteEndPoint -> ConnectionId -> IORef Bool -> IO ()
426−apiClose theirEndPoint connId connAlive = void . tryIO $ do
427− modifyMVar_ (remoteState theirEndPoint) $ \st -> case st of
428− RemoteEndPointValid vst -> do
455+apiClose :: EndPointPair -> ConnectionId -> IORef Bool -> IO ()
456+apiClose (ourEndPoint, theirEndPoint) connId connAlive = void . tryIO $ do
457+ modifyRemoteState_ (ourEndPoint, theirEndPoint)
458+ -- RemoteEndPointInvalid
459+ (\_ -> fail "apiClose RELY violation")
460+ -- RemoteEndPointValid
461+ (\vst -> do
429462 alive <- readIORef connAlive
430463 if alive
431464 then do
432465 writeIORef connAlive False
433− -- Possible exception caught top-level
434466 sendOn vst [encodeInt32 CloseConnection, encodeInt32 connId]
435467 return (RemoteEndPointValid . (remoteOutgoing ^: (\x -> x - 1)) $ vst)
436468 else
437− return (RemoteEndPointValid vst)
438− RemoteEndPointClosed -> do
439− return st
440− RemoteEndPointClosing _ _ ->
441− -- See discussion in apiSend why we must allow for this case
442− return st
443− RemoteEndPointInvalid _ ->
444− fail "apiClose RELY violation"
445− closeIfUnused theirEndPoint
469+ return (RemoteEndPointValid vst))
470+ -- RemoteEndPointClosing
471+ (\(resolved, vst) -> return $ RemoteEndPointClosing resolved vst)
472+ -- RemoteEndPoinClosed
473+ (return $ RemoteEndPointClosed)
474+ closeIfUnused (ourEndPoint, theirEndPoint)
446475
447476-- | Send data across a connection
448477--
449478-- RELY: The endpoint must not be in invalid state.
450479-- GUARANTEE: The state of the remote endpoint will not be changed.
451−apiSend :: RemoteEndPoint -> ConnectionId -> IORef Bool -> [ByteString] -> IO (Either (TransportError SendErrorCode) ())
452−apiSend theirEndPoint connId connAlive payload = do
453− withMVar (remoteState theirEndPoint) $ \st -> case st of
454− RemoteEndPointValid vst -> do
480+apiSend :: EndPointPair -> ConnectionId -> IORef Bool -> [ByteString] -> IO (Either (TransportError SendErrorCode) ())
481+apiSend (ourEndPoint, theirEndPoint) connId connAlive payload = do
482+ join <$> (try . mapExceptionIO sendFailed $ modifyRemoteState (ourEndPoint, theirEndPoint)
483+ -- RemoteEndPointInvalid
484+ (\_ -> fail "apiSend RELY violation")
485+ -- RemoteEndPointValid
486+ (\vst -> do
455487 alive <- readIORef connAlive
456− try $ if alive
457− then mapExceptionIO sendFailed $ sendOn vst (encodeInt32 connId : prependLength payload)
458− else throw $ TransportError SendClosed "Connection closed"
459− RemoteEndPointClosed -> do
460− return . Left $ TransportError SendClosed "Connection lost"
461− RemoteEndPointClosing _ _ ->
462− -- The only way for the endpoint to be in closing state, while a
463− -- connection is still active, is for the application to call 'close'
464− -- on the connection triggering "garbage collection" of the TCP channel
465− -- (CloseSocket request to the remote endpoint) which hasn't completed
466− -- yet. Even if the remote endpoint comes back with "please don't
467− -- close", this still means that our *outgoing* connection has been
468− -- closed
469− return . Left $ TransportError SendClosed "Connection lost"
470− RemoteEndPointInvalid _ ->
471− error "apiSend RELY violation"
488+ if alive
489+ then do
490+ sendOn vst (encodeInt32 connId : prependLength payload)
491+ return (RemoteEndPointValid vst, Right $ ())
492+ else
493+ return (RemoteEndPointValid vst, Left $ TransportError SendClosed "Connection closed"))
494+ -- RemoteEndPointClosing
495+ (\(resolved, vst) -> return (RemoteEndPointClosing resolved vst, Left $ TransportError SendClosed "Connection lost"))
496+ -- RemoteEndPointClosed
497+ (return (RemoteEndPointClosed, Left $ TransportError SendClosed "Connection lost")))
472498 where
473499 sendFailed :: IOException -> TransportError SendErrorCode
474500 sendFailed = TransportError SendFailed . show
@@ -488,14 +514,11 @@ apiCloseEndPoint transport evs ourEndPoint = do
488514 return (LocalEndPointClosed, Just vst)
489515 LocalEndPointClosed ->
490516 return (LocalEndPointClosed, Nothing)
491− case mOurState of
492− Nothing -> -- Already closed
493− return ()
494− Just vst -> do
495− -- Close all endpoints and kill all threads
496− forM_ (Map.elems $ vst ^. localConnections) $ tryCloseRemoteSocket
497− forM_ (vst ^. internalThreads) killThread
498− forM_ evs $ writeChan (localChannel ourEndPoint)
517+ forM_ mOurState $ \vst -> do
518+ -- Close all endpoints and kill all threads
519+ forM_ (Map.elems $ vst ^. localConnections) $ tryCloseRemoteSocket
520+ forM_ (vst ^. internalThreads) killThread
521+ forM_ evs $ writeChan (localChannel ourEndPoint)
499522 where
500523 -- Close the remote socket and return the set of all incoming connections
501524 tryCloseRemoteSocket :: RemoteEndPoint -> IO ()
@@ -552,6 +575,43 @@ connectToSelf ourEndPoint = do
552575 when alive $ writeChan ourChan (ConnectionClosed connId)
553576 return False
554577
578+--------------------------------------------------------------------------------
579+-- Functions from TransportInternals --
580+--------------------------------------------------------------------------------
581+
582+-- Find a socket between two endpoints
583+internalSocketBetween :: TCPTransport -- ^ Transport
584+ -> EndPointAddress -- ^ Local endpoint
585+ -> EndPointAddress -- ^ Remote endpoint
586+ -> IO (Either String N.Socket)
587+internalSocketBetween transport ourAddress theirAddress = runErrorT $ do
588+ ourEndPoint <- ErrorT $ do
589+ withMVar (transportState transport) $ \st -> case st of
590+ TransportClosed ->
591+ return . Left $ "Transport closed"
592+ TransportValid vst -> do
593+ case vst ^. localEndPointAt ourAddress of
594+ Nothing -> return . Left $ "Local endpoint not found"
595+ Just ep -> return . Right $ ep
596+ theirEndPoint <- ErrorT $ do
597+ withMVar (localState ourEndPoint) $ \st -> case st of
598+ LocalEndPointClosed ->
599+ return . Left $ "Local endpoint closed"
600+ LocalEndPointValid vst -> do
601+ case vst ^. localConnectionTo theirAddress of
602+ Nothing -> return . Left $ "Remote endpoint not found"
603+ Just ep -> return . Right $ ep
604+ ErrorT $ do
605+ withMVar (remoteState theirEndPoint) $ \st -> case st of
606+ RemoteEndPointValid vst ->
607+ return . Right $ remoteSocket vst
608+ RemoteEndPointClosing _ _ ->
609+ return . Left $ "Remote endpoint closing"
610+ RemoteEndPointClosed ->
611+ return . Left $ "Remote endpoint closed"
612+ RemoteEndPointInvalid _ ->
613+ return . Left $ "Remote endpoint invalid"
614+
555615--------------------------------------------------------------------------------
556616-- Lower level functionality --
557617--------------------------------------------------------------------------------
@@ -568,6 +628,7 @@ createLocalEndPoint transport = do
568628 , _nextCtrlRequestId = 0
569629 , _localConnections = Map.empty
570630 , _internalThreads = []
631+ , _nextRemoteId = 0
571632 }
572633 modifyMVar (transportState transport) $ \st -> case st of
573634 TransportValid vst -> do
@@ -635,23 +696,15 @@ requestConnectionTo ourEndPoint theirAddress = go
635696 go
636697
637698 RemoteEndPointValid _ -> do
638− -- On a failure we decrement the refcount again and return an error.
639− -- The only state the remote endpoint can be in at this point is
640− -- valid. As mentioned above, we can rely on the endpoint being in
641− -- valid state at this point.
642− let failureHandler :: IOException -> IO b
643− failureHandler err = do
644− modifyMVar_ theirState $ \(RemoteEndPointValid ep) ->
645− return (RemoteEndPointValid . (remoteOutgoing ^: (\x -> x - 1)) $ ep)
646− throw $ TransportError ConnectFailed (show err)
647−
648699 -- Do the actual connection request. This blocks until the remote
649700 -- endpoint replies (but note that we don't hold any locks at this
650− -- point)
651− reply <- handle failureHandler $ doRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId
701+ -- point). Note that doRemoteRequest may throw an error if the send
702+ -- fails, and if it does, it will have put the remote endpoint in
703+ -- closed state.
704+ reply <- mapExceptionIO connectFailed $ doRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId
652705 case decodeInt32 . BS.concat $ reply of
653706 Nothing ->
654− failureHandler (userError "Invalid integer")
707+ throw (connectFailed $ userError "Invalid integer")
655708 Just cid -> do
656709 return (theirEndPoint, cid)
657710
@@ -672,8 +725,9 @@ requestConnectionTo ourEndPoint theirAddress = go
672725 theirState <- newEmptyMVar
673726 let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress
674727 , remoteState = theirState
728+ , remoteId = vst ^. nextRemoteId
675729 }
676− return (LocalEndPointValid . (localConnectionTo theirAddress ^= Just theirEndPoint) $ vst, (theirEndPoint, True))
730+ return (LocalEndPointValid . (localConnectionTo theirAddress ^= Just theirEndPoint) . (nextRemoteId ^: (+ 1)) $ vst, (theirEndPoint, True))
677731
678732 -- The only way for forkEndPointThread to fail is if the local endpoint
679733 -- gets closed. This error will be caught elsewhere, so we ignore it
@@ -685,6 +739,9 @@ requestConnectionTo ourEndPoint theirAddress = go
685739 ourState :: MVar LocalEndPointState
686740 ourState = localState ourEndPoint
687741
742+ connectFailed :: IOException -> TransportError ConnectErrorCode
743+ connectFailed = TransportError ConnectFailed . show
744+
688745-- | Set up a remote endpoint
689746--
690747-- RELY: The state of the remote endpoint must be uninitialized.
@@ -700,7 +757,7 @@ setupRemoteEndPoint (ourEndPoint, theirEndPoint) = do
700757 { remoteSocket = sock
701758 , _remoteOutgoing = 0
702759 , _remoteIncoming = IntSet.empty
703− , sendOn = sendMany sock
760+ , sendOn = \msg -> sendMany sock msg `onException` (tryIO $ N.sClose sock)
704761 }
705762 putMVar theirState (RemoteEndPointValid vst)
706763 return True
@@ -783,11 +840,20 @@ socketToEndPoint (EndPointAddress ourAddress) theirAddress = try $ do
783840-- If the local endpoint is closed, do nothing
784841removeRemoteEndPoint :: EndPointPair -> IO ()
785842removeRemoteEndPoint (ourEndPoint, theirEndPoint) = do
786− modifyMVar_ (localState ourEndPoint) $ \st -> case st of
787− LocalEndPointValid vst ->
788− return (LocalEndPointValid . (localConnectionTo (remoteAddress theirEndPoint) ^= Nothing) $ vst)
789− LocalEndPointClosed ->
790− return LocalEndPointClosed
843+ modifyMVar_ ourState $ \st -> case st of
844+ LocalEndPointValid vst ->
845+ case vst ^. localConnectionTo theirAddress of
846+ Nothing ->
847+ return st
848+ Just remoteEndPoint' ->
849+ if remoteId remoteEndPoint' == remoteId theirEndPoint
850+ then return (LocalEndPointValid . (localConnectionTo (remoteAddress theirEndPoint) ^= Nothing) $ vst)
851+ else return st
852+ LocalEndPointClosed ->
853+ return LocalEndPointClosed
854+ where
855+ ourState = localState ourEndPoint
856+ theirAddress = remoteAddress theirEndPoint
791857
792858-- | Remove reference to a local endpoint from the transport state
793859--
@@ -817,11 +883,11 @@ decodeEndPointAddress (EndPointAddress bs) = case map BSC.unpack $ BSC.split ':'
817883
818884-- | Do a (blocking) remote request
819885--
820−-- RELY: Remote endpoint must be in valid state.
886+-- RELY: Remote endpoint must be in valid or closed state.
821887-- GUARANTEE: Will not change the state of the remote endpoint.
822888--
823−-- May throw IO (user) exception if the local endpoint is closed or if the send
824−-- fails.
889+-- May throw IO (user) exception if the local or the remote endpoint is closed,
890+-- or if the send fails.
825891doRemoteRequest :: EndPointPair -> ControlHeader -> IO [ByteString]
826892doRemoteRequest (ourEndPoint, theirEndPoint) header = do
827893 reply <- newEmptyMVar
@@ -831,24 +897,35 @@ doRemoteRequest (ourEndPoint, theirEndPoint) header = do
831897 return (LocalEndPointValid . (nextCtrlRequestId ^: (+ 1)) . (pendingCtrlRequestsAt reqId ^= Just reply) $ vst, reqId)
832898 LocalEndPointClosed ->
833899 throw (userError "Local endpoint closed")
834− withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid vst) ->
835− sendOn vst [encodeInt32 header, encodeInt32 reqId]
900+ withRemoteState (ourEndPoint, theirEndPoint)
901+ -- RemoteEndPointInvalid
902+ (\_ -> fail "doRemoteRequest RELY violation")
903+ -- RemoteEndPointValid
904+ (\vst -> sendOn vst [encodeInt32 header, encodeInt32 reqId])
905+ -- RemoteEndPointClosing
906+ (\_ -> fail "doRemoteRequest RELY violation")
907+ -- RemoteEndPointClosed
908+ (throw (userError "Remote endpoint closed"))
836909 takeMVar reply
837910
838911-- | Check if the remote endpoint is unused, and if so, send a CloseSocket request
839−closeIfUnused :: RemoteEndPoint -> IO ()
840−closeIfUnused theirEndPoint = modifyMVar_ (remoteState theirEndPoint) $ \st ->
841− case st of
842− RemoteEndPointValid vst -> do
843− if vst ^. remoteOutgoing == 0 && IntSet.null (vst ^. remoteIncoming)
844− then do
845− resolved <- newEmptyMVar
846− sendOn vst [encodeInt32 CloseSocket]
847− return (RemoteEndPointClosing resolved vst)
848− else
849− return st
850− _ -> do
851− return st
912+closeIfUnused :: EndPointPair -> IO ()
913+closeIfUnused (ourEndPoint, theirEndPoint) =
914+ modifyRemoteState_ (ourEndPoint, theirEndPoint)
915+ -- RemoteEndPointInvalid
916+ (return . RemoteEndPointInvalid)
917+ -- RemoteEndPointValid
918+ (\vst -> if vst ^. remoteOutgoing == 0 && IntSet.null (vst ^. remoteIncoming)
919+ then do
920+ sendOn vst [encodeInt32 CloseSocket]
921+ resolved <- newEmptyMVar
922+ return $ RemoteEndPointClosing resolved vst
923+ else
924+ return $ RemoteEndPointValid vst)
925+ -- RemoteEndPointClosing
926+ (return . uncurry RemoteEndPointClosing)
927+ -- RemoteEndPointClosed
928+ (return RemoteEndPointClosed)
852929
853930-- | Fork a new thread and store its ID as part of the transport state
854931--
@@ -886,6 +963,109 @@ forkEndPointThread ourEndPoint p =
886963 ourState :: MVar LocalEndPointState
887964 ourState = localState ourEndPoint
888965
966+-- | Like modifyMVar, but if an exception occurs don't restore the remote
967+-- endpoint to its original value but close it instead
968+modifyRemoteState :: EndPointPair -- ^ Local and remote endpoint
969+ -> (TransportError ConnectErrorCode -> IO (RemoteState, a)) -- ^ Case for RemoteEndPointInvalid
970+ -> (ValidRemoteEndPointState -> IO (RemoteState, a)) -- ^ Case for RemoteEndPointValid
971+ -> ((MVar (), ValidRemoteEndPointState) -> IO (RemoteState, a)) -- ^ Case for RemoteEndPointClosing
972+ -> IO (RemoteState, a) -- ^ Case for RemoteEndPointClosed
973+ -> IO a
974+modifyRemoteState (ourEndPoint, theirEndPoint) caseInvalid caseValid caseClosing caseClosed = do
975+ mask $ \restore -> do
976+ st <- takeMVar theirState
977+ case st of
978+ RemoteEndPointValid vst -> do
979+ mResult <- try $ restore (caseValid vst)
980+ case mResult of
981+ Right (st', a) -> do
982+ putMVar theirState st'
983+ return a
984+ Left ex -> do
985+ handleException ex vst
986+ -- The other cases are less interesting, because unless the endpoint is
987+ -- in Valid state we're not supposed to do any IO on it
988+ RemoteEndPointClosing resolved vst -> do
989+ (st', a) <- restore (caseClosing (resolved, vst)) `onException` putMVar theirState st
990+ putMVar theirState st'
991+ return a
992+ RemoteEndPointInvalid err -> do
993+ (st', a) <- restore (caseInvalid err) `onException` putMVar theirState st
994+ putMVar theirState st'
995+ return a
996+ RemoteEndPointClosed -> do
997+ (st', a) <- restore caseClosed `onException` putMVar theirState st
998+ putMVar theirState st'
999+ return a
1000+ where
1001+ theirState :: MVar RemoteState
1002+ theirState = remoteState theirEndPoint
1003+
1004+ handleException :: SomeException -> ValidRemoteEndPointState -> IO a
1005+ handleException ex vst = do
1006+ -- We need to remove the remote endpoint from the local endpoint before
1007+ -- putting it in Closed state (by INV-CLOSE), but we don't want to hold
1008+ -- two locks at the same time. So we put it in closing state first ..
1009+ resolved <- newEmptyMVar
1010+ putMVar theirState (RemoteEndPointClosing resolved vst)
1011+ -- .. then remove it from the local endpoint ..
1012+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
1013+ -- .. close it and notify the local endpoint we lost the connection ..
1014+ mIncoming <- closeRemoteEndPoint
1015+ forM_ mIncoming $ \incoming -> do
1016+ let err = TransportError (EventConnectionLost (remoteAddress theirEndPoint) incoming) (show ex)
1017+ writeChan (localChannel ourEndPoint) $ ErrorEvent err
1018+ -- .. and finally rethrow the exception
1019+ throw ex
1020+
1021+ -- Returns the set of incoming connections if we closed the connection, or
1022+ -- 'Nothing' if the endpoint was already closed
1023+ closeRemoteEndPoint :: IO (Maybe [ConnectionId])
1024+ closeRemoteEndPoint = do
1025+ st <- takeMVar theirState
1026+ case st of
1027+ RemoteEndPointClosing resolved vst -> do
1028+ putMVar theirState RemoteEndPointClosed
1029+ putMVar resolved ()
1030+ return . Just . IntSet.elems $ vst ^. remoteIncoming
1031+ RemoteEndPointClosed -> do
1032+ putMVar theirState RemoteEndPointClosed
1033+ return Nothing
1034+ _ ->
1035+ fail "the impossible happened"
1036+
1037+-- | Like 'modifyRemoteState' but without a return value
1038+modifyRemoteState_ :: EndPointPair -- ^ Local and remote endpoint
1039+ -> (TransportError ConnectErrorCode -> IO RemoteState) -- ^ Case for RemoteEndPointInvalid
1040+ -> (ValidRemoteEndPointState -> IO RemoteState) -- ^ Case for RemoteEndPointValid
1041+ -> ((MVar (), ValidRemoteEndPointState) -> IO RemoteState) -- ^ Case for RemoteEndPointClosing
1042+ -> IO RemoteState -- ^ Case for RemoteEndPointClosed
1043+ -> IO ()
1044+modifyRemoteState_ (ourEndPoint, theirEndPoint) caseInvalid caseValid caseClosing caseClosed =
1045+ modifyRemoteState (ourEndPoint, theirEndPoint)
1046+ (u . caseInvalid)
1047+ (u . caseValid)
1048+ (u . caseClosing)
1049+ (u caseClosed)
1050+ where
1051+ u :: IO a -> IO (a, ())
1052+ u p = p >>= \a -> return (a, ())
1053+
1054+-- | Like 'modifyRemoteState' but without the ability to change the state
1055+withRemoteState :: EndPointPair -- ^ Local and remote endpoint
1056+ -> (TransportError ConnectErrorCode -> IO a) -- ^ Case for RemoteEndPointInvalid
1057+ -> (ValidRemoteEndPointState -> IO a) -- ^ Case for RemoteEndPointValid
1058+ -> ((MVar (), ValidRemoteEndPointState) -> IO a) -- ^ Case for RemoteEndPointClosing
1059+ -> IO a -- ^ Case for RemoteEndPointClosed
1060+ -> IO a
1061+withRemoteState (ourEndPoint, theirEndPoint) caseInvalid caseValid caseClosing caseClosed = do
1062+ modifyRemoteState (ourEndPoint, theirEndPoint)
1063+ (\err -> (,) (RemoteEndPointInvalid err) <$> caseInvalid err)
1064+ (\vst -> (,) (RemoteEndPointValid vst) <$> caseValid vst)
1065+ (\(resolved, vst) -> (,) (RemoteEndPointClosing resolved vst) <$> caseClosing (resolved, vst))
1066+ ((,) RemoteEndPointClosed <$> caseClosed)
1067+
1068+
8891069--------------------------------------------------------------------------------
8901070-- Incoming requests --
8911071--------------------------------------------------------------------------------
@@ -928,8 +1108,9 @@ handleConnectionRequest transport sock = handle tryCloseSocket $ do
9281108 theirState <- newEmptyMVar
9291109 let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress
9301110 , remoteState = theirState
1111+ , remoteId = vst ^. nextRemoteId
9311112 }
932− return (LocalEndPointValid . (localConnectionTo theirAddress ^= Just theirEndPoint) $ vst, (False, theirEndPoint))
1113+ return (LocalEndPointValid . (localConnectionTo theirAddress ^= Just theirEndPoint) . (nextRemoteId ^: (+ 1)) $ vst, (False, theirEndPoint))
9331114 Just theirEndPoint ->
9341115 return (st, (localAddress ourEndPoint < theirAddress, theirEndPoint))
9351116 LocalEndPointClosed ->
@@ -958,11 +1139,7 @@ handleConnectionRequest transport sock = handle tryCloseSocket $ do
9581139 -- been recorded as part of the remote endpoint. Either way, we no longer
9591140 -- have to worry about closing the socket on receiving an asynchronous
9601141 -- exception from this point forward.
961− case mEndPoint of
962− Nothing ->
963− return ()
964− Just theirEndPoint ->
965− handleIncomingMessages (ourEndPoint, theirEndPoint)
1142+ forM_ mEndPoint $ handleIncomingMessages . (,) ourEndPoint
9661143
9671144 tryCloseSocket :: IOException -> IO ()
9681145 tryCloseSocket _ = void . tryIO $ N.sClose sock
@@ -1076,7 +1253,7 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
10761253 writeChan ourChannel (ConnectionClosed cid)
10771254 modifyMVar_ theirState $ \(RemoteEndPointValid vst) ->
10781255 return (RemoteEndPointValid . (remoteIncoming ^: IntSet.delete cid) $ vst)
1079− closeIfUnused theirEndPoint
1256+ closeIfUnused (ourEndPoint, theirEndPoint)
10801257
10811258 -- Close the socket (if we don't have any outgoing connections)
10821259 closeSocket :: N.Socket -> IO Bool
@@ -1138,24 +1315,18 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
11381315 prematureExit sock err = do
11391316 N.sClose sock
11401317 removeRemoteEndPoint (ourEndPoint, theirEndPoint)
1141− mUnclosedConnections <- modifyMVar theirState $ \st ->
1318+ modifyMVar_ theirState $ \st ->
11421319 case st of
11431320 RemoteEndPointInvalid _ ->
11441321 error "handleIncomingMessages RELY violation"
1145− RemoteEndPointValid vst ->
1146− return (RemoteEndPointClosed, Just $ vst ^. remoteIncoming)
1322+ RemoteEndPointValid vst -> do
1323+ writeChan ourChannel . ErrorEvent $
1324+ TransportError (EventConnectionLost (remoteAddress theirEndPoint) (IntSet.elems $ vst ^. remoteIncoming)) (show err)
1325+ return RemoteEndPointClosed
11471326 RemoteEndPointClosing _ _ ->
1148− return (RemoteEndPointClosed, Nothing)
1327+ return RemoteEndPointClosed
11491328 RemoteEndPointClosed ->
1150− return (st, Nothing)
1151−
1152− -- We send the connection lost message even if unclosedConnections is the
1153− -- empty set, because *outgoing* connections will be broken too now
1154− case mUnclosedConnections of
1155− Nothing ->
1156− return ()
1157− Just unclosedConnections -> writeChan ourChannel . ErrorEvent $
1158− TransportError (EventConnectionLost (remoteAddress theirEndPoint) (IntSet.elems unclosedConnections)) (show err)
1329+ return RemoteEndPointClosed
11591330
11601331-- | Get the next connection ID
11611332--
@@ -1202,6 +1373,9 @@ localConnections = accessor _localConnections (\es st -> st { _localConnections
12021373internalThreads :: Accessor ValidLocalEndPointState [ThreadId]
12031374internalThreads = accessor _internalThreads (\ts st -> st { _internalThreads = ts })
12041375
1376+nextRemoteId :: Accessor ValidLocalEndPointState Int
1377+nextRemoteId = accessor _nextRemoteId (\rid st -> st { _nextRemoteId = rid })
1378+
12051379localEndPointAt :: EndPointAddress -> Accessor ValidTransportState (Maybe LocalEndPoint)
12061380localEndPointAt addr = localEndPoints >>> DAC.mapMaybe addr
12071381