@@ -70,6 +70,7 @@ import Control.Concurrent.MVar ( MVar
7070 , putMVar
7171 , newEmptyMVar
7272 , withMVar
73+ , tryPutMVar
7374 )
7475import Control.Category ((>>>))
7576import Control.Applicative ((<$>))
@@ -691,33 +692,61 @@ requestConnectionTo ourEndPoint theirAddress = go
691692-- RemoteEndPointInvalid.
692693setupRemoteEndPoint :: EndPointPair -> IO ()
693694setupRemoteEndPoint (ourEndPoint, theirEndPoint) = do
694− result <- socketToEndPoint (localAddress ourEndPoint) (remoteAddress theirEndPoint)
695− case result of
696− Right (sock, ConnectionRequestAccepted) -> do
697− let vst = ValidRemoteEndPointState
698− { remoteSocket = sock
699− , _remoteOutgoing = 0
700− , _remoteIncoming = IntSet.empty
701− , sendOn = sendMany sock
702− }
703− putMVar theirState (RemoteEndPointValid vst)
704− handleIncomingMessages (ourEndPoint, theirEndPoint)
705− Right (sock, ConnectionRequestEndPointInvalid) -> do
706− -- We remove the endpoint from our local state again because the next
707− -- call to 'connect' might give a different result. Threads that were
708− -- waiting on the result of this call to connect will get the
709− -- RemoteEndPointInvalid; subsequent threads will initiate a new
710− -- connection requests.
711− removeRemoteEndPoint (ourEndPoint, theirEndPoint)
712− putMVar theirState (RemoteEndPointInvalid (invalidAddress "Invalid endpoint"))
713− N.sClose sock
714− Right (sock, ConnectionRequestCrossed) -> do
715− N.sClose sock
716− Left err -> do
717− -- See comment above
718− removeRemoteEndPoint (ourEndPoint, theirEndPoint)
719− putMVar theirState (RemoteEndPointInvalid err)
695+ didAccept <- bracketOnError (socketToEndPoint (localAddress ourEndPoint) (remoteAddress theirEndPoint))
696+ onError $ \result -> do
697+ case result of
698+ Right (sock, ConnectionRequestAccepted) -> do
699+ let vst = ValidRemoteEndPointState
700+ { remoteSocket = sock
701+ , _remoteOutgoing = 0
702+ , _remoteIncoming = IntSet.empty
703+ , sendOn = sendMany sock
704+ }
705+ putMVar theirState (RemoteEndPointValid vst)
706+ return True
707+ Right (sock, ConnectionRequestEndPointInvalid) -> do
708+ -- We remove the endpoint from our local state again because the next
709+ -- call to 'connect' might give a different result. Threads that were
710+ -- waiting on the result of this call to connect will get the
711+ -- RemoteEndPointInvalid; subsequent threads will initiate a new
712+ -- connection requests.
713+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
714+ putMVar theirState (RemoteEndPointInvalid (invalidAddress "Invalid endpoint"))
715+ N.sClose sock
716+ return False
717+ Right (sock, ConnectionRequestCrossed) -> do
718+ N.sClose sock
719+ return False
720+ Left err -> do
721+ removeRemoteEndPoint (ourEndPoint, theirEndPoint) -- See comment above
722+ putMVar theirState (RemoteEndPointInvalid err)
723+ return False
724+
725+ -- If we get to this point without an exception, then
726+ -- * if didAccept is False the socket has already been closed
727+ -- * if didAccept is True, the socket has been stored as part of the remote
728+ -- state so we no longer need to worry about closing it when an
729+ -- asynchronous exception occurs
730+ when didAccept $ handleIncomingMessages (ourEndPoint, theirEndPoint)
720731 where
732+ -- If an asynchronous exception occurs while we set up the remote endpoint
733+ -- we need to make sure to close the socket. It is also useful to
734+ -- initialize the remote state (to "invalid") so that concurrent threads
735+ -- that are blocked on reading the remote state are unblocked. It is
736+ -- possible, however, that the exception occurred after we already
737+ -- initialized the remote state, which is why we use tryPutMVar here.
738+ onError :: Either (TransportError ConnectErrorCode) (N.Socket, ConnectionRequestResponse) -> IO ()
739+ onError result = do
740+ removeRemoteEndPoint (ourEndPoint, theirEndPoint)
741+ case result of
742+ Left err -> do
743+ tryPutMVar theirState (RemoteEndPointInvalid (TransportError ConnectFailed (show err)))
744+ return ()
745+ Right (sock, _) -> do
746+ tryPutMVar theirState (RemoteEndPointInvalid (TransportError ConnectFailed "setupRemoteEndPoint failed"))
747+ tryIO $ N.sClose sock
748+ return ()
749+
721750 theirState = remoteState theirEndPoint
722751 invalidAddress = TransportError ConnectNotFound
723752