@@ -46,7 +46,7 @@ import qualified Network.Socket as N ( HostName
4646 , connect
4747 )
4848import Network.Socket.ByteString (sendMany)
49−import Control.Concurrent (forkIO)
49+import Control.Concurrent (forkIO, ThreadId)
5050import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
5151import Control.Concurrent.MVar ( MVar
5252 , newMVar
@@ -77,6 +77,7 @@ import Data.Map (Map)
7777import qualified Data.Map as Map (empty)
7878import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=))
7979import Text.Regex.Applicative (RE, few, anySym, sym, many, match)
80+import System.IO (hPutStrLn, stderr)
8081
8182-- $design
8283--
@@ -275,7 +276,7 @@ createTransport host port = do
275276 , transportPort = port
276277 }
277278 tryIO $ do
278− forkServer host port $ handleConnectionRequest transport
279+ forkServer host port $ void . handleConnectionRequest transport
279280 return Transport { newEndPoint = apiNewEndPoint transport }
280281
281282--------------------------------------------------------------------------------
@@ -340,14 +341,12 @@ apiSend :: RemoteEndPoint -> ConnectionId -> IORef Bool -> [ByteString] -> IO (E
340341apiSend theirEndPoint connId connAlive payload = do
341342 withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) -> do
342343 alive <- readIORef connAlive
343− if alive
344− then do
345− result <- tryIO $ sendOn remoteConn (encodeInt32 connId : prependLength payload)
346− case result of
347− Left err -> return . Left $ FailedWith SendFailed (show err)
348− Right _ -> return . Right $ ()
349− else do
350− return . Left $ FailedWith SendConnectionClosed "Connection closed"
344+ runErrorT $ if alive
345+ then
346+ failWithIO (FailedWith SendFailed . show) $
347+ sendOn remoteConn (encodeInt32 connId : prependLength payload)
348+ else
349+ throwError $ FailedWith SendConnectionClosed "Connection closed"
351350
352351-- | Special case of 'apiConnect': connect an endpoint to itself
353352connectToSelf :: LocalEndPoint -> IO (Either (FailedWith ConnectErrorCode) Connection)
@@ -449,12 +448,6 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go
449448 -- matter.
450449 case endPointStateSnapshot of
451450 RemoteEndPointInvalid err -> do
452− -- We remove the endpoint from the state again because the next call
453− -- to 'connect' might give a different result. (We definitely don't
454− -- want to do this while holding the endpoint lock, because holding
455− -- multiple locks introduces all kinds of trouble.)
456− modifyMVar ourState $ \st ->
457− return (connectionTo theirAddress ^= Nothing $ st, Left err)
458451 return . Left $ err
459452
460453 RemoteEndPointClosing _ resolved -> do
@@ -475,8 +468,8 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go
475468 -- endpoint replies (but note that we don't hold any locks at this
476469 -- point)
477470 mcid <- tryIO $ do
478− reply <- sendRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId
479− decodeInt32 . BS.concat <$> takeMVar reply
471+ reply <- doRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId
472+ return (decodeInt32 . BS.concat $ reply)
480473
481474 -- On a failure we decrement the refcount again and return an error.
482475 -- The only state the remote endpoint can be in at this point is
@@ -499,7 +492,9 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go
499492-- RemoteEndPointInvalid.
500493connectToRemoteEndPoint :: EndPointPair -> IO ()
501494connectToRemoteEndPoint (ourEndPoint, theirEndPoint) = do
502− let theirState = remoteState theirEndPoint
495+ let ourState = localState ourEndPoint
496+ theirState = remoteState theirEndPoint
497+ theirAddress = remoteAddress theirEndPoint
503498 result <- runErrorT go
504499 case result of
505500 Right (ConnectionRequestAccepted, sock) -> do
@@ -510,14 +505,22 @@ connectToRemoteEndPoint (ourEndPoint, theirEndPoint) = do
510505 putMVar theirState (RemoteEndPointValid remoteConn)
511506 void . forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)
512507 Right (ConnectionRequestEndPointInvalid, sock) -> do
508+ -- We remove the endpoint from our local state again because the next
509+ -- call to 'connect' might give a different result. Threads that were
510+ -- waiting on the result of this call to connect will get the
511+ -- RemoteEndPointInvalid; subsequent threads will initiate a new
512+ -- connection requests.
513+ modifyMVar_ ourState $ return . (connectionTo theirAddress ^= Nothing)
513514 putMVar theirState (RemoteEndPointInvalid (invalidAddress "Invalid endpoint"))
514515 N.sClose sock
515516 Right (ConnectionRequestCrossed, sock) -> do
516517 N.sClose sock
517518 Left err -> do
519+ -- See comment above
520+ modifyMVar_ ourState $ return . (connectionTo theirAddress ^= Nothing)
521+ putMVar theirState (RemoteEndPointInvalid err)
518522 -- TODO: this isn't quite right; socket *might* have been opened here,
519523 -- in which case we should close it
520− putMVar theirState (RemoteEndPointInvalid err)
521524 where
522525 go :: ErrorT (FailedWith ConnectErrorCode) IO (ConnectionRequestResponse, N.Socket)
523526 go = do
@@ -553,20 +556,19 @@ decodeEndPointAddress (EndPointAddress bs) = match endPointAddressRE (BSC.unpack
553556 endPointAddressRE :: RE Char (N.HostName, N.ServiceName, EndPointId)
554557 endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))
555558
556−-- | Send a remote request and return the MVar that will contain the response
559+-- | Do a (blocking) remote request
557560--
558561-- RELY: Remote endpoint must be in valid state.
559562-- GUARANTEE: Will not change the state of the remote endpoint.
560−sendRemoteRequest :: EndPointPair -> ControlHeader -> IO (MVar [ByteString])
561−sendRemoteRequest (ourEndPoint, theirEndPoint) header = do
563+doRemoteRequest :: EndPointPair -> ControlHeader -> IO ([ByteString])
564+doRemoteRequest (ourEndPoint, theirEndPoint) header = do
562565 reply <- newEmptyMVar
563566 reqId <- modifyMVar (localState ourEndPoint) $ \st -> do
564567 let reqId = st ^. nextCtrlRequestId
565568 return ((nextCtrlRequestId ^+= 1) . (pendingCtrlRequestsAt reqId ^= Just reply) $ st, reqId)
566569 withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) ->
567570 sendOn remoteConn [encodeInt32 header, encodeInt32 reqId]
568− return reply
569−
571+ takeMVar reply
570572
571573-- | Decrease the remote refcount of a remote endpoint. When the refcount
572574-- reaches zero, initiate socket closure.
@@ -591,7 +593,9 @@ decreaseRemoteRefCount theirEndPoint = do
591593
592594-- | Handle a connection request (that is, a remote endpoint that is trying to
593595-- establish a TCP connection with us)
594−handleConnectionRequest :: TCPTransport -> N.Socket -> IO ()
596+--
597+-- Returns the thread ID of the thread that is handling incoming messages.
598+handleConnectionRequest :: TCPTransport -> N.Socket -> IO (Maybe ThreadId)
595599handleConnectionRequest transport sock = do
596600 request <- tryIO $ do
597601 ourEndPointId <- recvInt32 sock
@@ -602,6 +606,7 @@ handleConnectionRequest transport sock = do
602606 Right (_, Nothing) -> do
603607 sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]
604608 N.sClose sock
609+ return Nothing
605610 Right (theirAddress, Just ourEndPoint) -> do
606611 (crossed, theirEndPoint) <- modifyMVar (localState ourEndPoint) $ \st ->
607612 case st ^. connectionTo theirAddress of
@@ -617,23 +622,24 @@ handleConnectionRequest transport sock = do
617622 then do
618623 sendMany sock [encodeInt32 ConnectionRequestCrossed]
619624 N.sClose sock
625+ return Nothing
620626 else do
621627 let remoteConn = RemoteConnection { remoteSocket = sock
622628 , _remoteRefCount = 0
623629 , sendOn = sendMany sock
624630 }
625631 putMVar (remoteState theirEndPoint) (RemoteEndPointValid remoteConn)
626− forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)
627632 sendMany sock [encodeInt32 ConnectionRequestAccepted]
628− Left _ ->
633+ tid <- forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)
634+ return (Just tid)
635+ Left _ -> do
629636 -- Invalid request
630637 N.sClose sock
638+ return Nothing
631639
632640-- | Handle requests from a remote endpoint.
633641--
634642-- Returns only if the remote party closes the socket or if an error occurs.
635−-- The return value is the list of connections that are still open (normally,
636−-- this should be an empty list).
637643--
638644-- RELY: The remote endpoint must be in RemoteEndPointValid or
639645-- RemoteEndPointClosing state. If the latter, then the 'resolved' MVar
@@ -684,7 +690,7 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
684690 go sock openConnections
685691 Nothing ->
686692 -- Invalid control request, exit
687− putStrLn "Warning: invalid control request"
693+ hPutStrLn stderr "Warning: invalid control request"
688694
689695 -- Create a new connection
690696 createNewConnection :: IORef IntSet -> ControlRequestId -> IO ()
@@ -722,7 +728,8 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
722728 mmvar <- modifyMVar ourState $ \st ->
723729 return (pendingCtrlRequestsAt reqId ^= Nothing $ st, st ^. pendingCtrlRequestsAt reqId)
724730 case mmvar of
725− Nothing ->
731+ Nothing -> do
732+ hPutStrLn stderr $ "Warning: Invalid request ID"
726733 return () -- Invalid request ID. TODO: We just ignore it?
727734 Just mvar ->
728735 putMVar mvar response
@@ -744,6 +751,9 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
744751 N.sClose sock
745752 return (RemoteEndPointClosed, True)
746753 RemoteEndPointClosing _ resolved -> do
754+ -- If the socket is already in closing state, then we have already
755+ -- sent a CloseSocket request; hence, we don't need to do it again
756+ -- at this point
747757 putMVar resolved ()
748758 N.sClose sock
749759 return (RemoteEndPointClosed, True)