@@ -85,7 +85,7 @@ import qualified Data.IntMap as IntMap (empty)
8585import Data.IntSet (IntSet)
8686import qualified Data.IntSet as IntSet (empty, insert, elems, singleton, null, delete)
8787import Data.Map (Map)
88−import qualified Data.Map as Map (empty, elems, size)
88+import qualified Data.Map as Map (empty, elems)
8989import Data.Accessor (Accessor, accessor, (^.), (^=), (^:))
9090import qualified Data.Accessor.Container as DAC (mapMaybe, intMapMaybe)
9191import System.IO (hPutStrLn, stderr)
@@ -209,6 +209,7 @@ data TransportState =
209209
210210data ValidTransportState = ValidTransportState
211211 { _localEndPoints :: Map EndPointAddress LocalEndPoint
212+ , _nextEndPointId :: EndPointId
212213 }
213214
214215data LocalEndPoint = LocalEndPoint
@@ -330,7 +331,10 @@ createTransport host port = do
330331-- | You should probably not use this function (used for unit testing only)
331332createTransportExposeInternals :: N.HostName -> N.ServiceName -> IO (Either IOException (Transport, ThreadId))
332333createTransportExposeInternals host port = do
333− state <- newMVar . TransportValid $ ValidTransportState { _localEndPoints = Map.empty }
334+ state <- newMVar . TransportValid $ ValidTransportState
335+ { _localEndPoints = Map.empty
336+ , _nextEndPointId = 0
337+ }
334338 let transport = TCPTransport { transportState = state
335339 , transportHost = host
336340 , transportPort = port
@@ -378,7 +382,6 @@ apiCloseTransport transport mTransportThread evs = do
378382 -- transport thread died (in which case we don't need to kill it)
379383 return ()
380384
381−
382385-- | Create a new endpoint
383386apiNewEndPoint :: TCPTransport -> IO (Either (TransportError NewEndPointErrorCode) EndPoint)
384387apiNewEndPoint transport = try $ do
@@ -418,8 +421,6 @@ apiConnect ourEndPoint theirAddress _reliability _hints = try $ do
418421-- GUARANTEE: If the connection is alive on entry then the remote endpoint will
419422-- either be RemoteEndPointValid or RemoteEndPointClosing. Otherwise, the state
420423-- of the remote endpoint will not be changed.
421−--
422−-- TODO: We ignore errors during a close. Is that right?
423424apiClose :: RemoteEndPoint -> ConnectionId -> IORef Bool -> IO ()
424425apiClose theirEndPoint connId connAlive = void . tryIO $ do
425426 modifyMVar_ (remoteState theirEndPoint) $ \st -> case st of
@@ -428,6 +429,7 @@ apiClose theirEndPoint connId connAlive = void . tryIO $ do
428429 if alive
429430 then do
430431 writeIORef connAlive False
432+ -- Possible exception caught top-level
431433 sendOn vst [encodeInt32 CloseConnection, encodeInt32 connId]
432434 return (RemoteEndPointValid . (remoteOutgoing ^: (\x -> x - 1)) $ vst)
433435 else
@@ -454,7 +456,7 @@ apiSend theirEndPoint connId connAlive payload = do
454456 then mapExceptionIO sendFailed $ sendOn vst (encodeInt32 connId : prependLength payload)
455457 else throw $ TransportError SendClosed "Connection closed"
456458 RemoteEndPointClosed -> do
457− return (Left $ TransportError SendClosed "Connection lost")
459+ return . Left $ TransportError SendClosed "Connection lost"
458460 RemoteEndPointClosing _ _ ->
459461 -- The only way for the endpoint to be in closing state, while a
460462 -- connection is still active, is for the application to call 'close'
@@ -463,7 +465,7 @@ apiSend theirEndPoint connId connAlive payload = do
463465 -- yet. Even if the remote endpoint comes back with "please don't
464466 -- close", this still means that our *outgoing* connection has been
465467 -- closed
466− return (Left $ TransportError SendClosed "Connection lost")
468+ return . Left $ TransportError SendClosed "Connection lost"
467469 RemoteEndPointInvalid _ ->
468470 error "apiSend RELY violation"
469471 where
@@ -520,12 +522,15 @@ connectToSelf :: LocalEndPoint -> IO (Either (TransportError ConnectErrorCode) C
520522connectToSelf ourEndPoint = do
521523 -- Here connAlive must an MVar because it is not protected by another lock
522524 connAlive <- newMVar True
523− -- TODO: catch exception
524− connId <- getNextConnectionId ourEndPoint
525− writeChan ourChan (ConnectionOpened connId ReliableOrdered (localAddress ourEndPoint))
526− return . Right $ Connection { send = selfSend connAlive connId
527− , close = selfClose connAlive connId
528− }
525+ mConnId <- tryIO (getNextConnectionId ourEndPoint)
526+ case mConnId of
527+ Left err ->
528+ return . Left $ TransportError ConnectNotFound (show err)
529+ Right connId -> do
530+ writeChan ourChan (ConnectionOpened connId ReliableOrdered (localAddress ourEndPoint))
531+ return . Right $ Connection { send = selfSend connAlive connId
532+ , close = selfClose connAlive connId
533+ }
529534 where
530535 ourChan :: Chan Event
531536 ourChan = localChannel ourEndPoint
@@ -565,13 +570,13 @@ createLocalEndPoint transport = do
565570 }
566571 modifyMVar (transportState transport) $ \st -> case st of
567572 TransportValid vst -> do
568− let ix = fromIntegral $ Map.size (vst ^. localEndPoints)
573+ let ix = vst ^. nextEndPointId
569574 let addr = encodeEndPointAddress (transportHost transport) (transportPort transport) ix
570575 let localEndPoint = LocalEndPoint { localAddress = addr
571576 , localChannel = chan
572577 , localState = state
573578 }
574− return (TransportValid . (localEndPointAt addr ^= Just localEndPoint) $ vst, localEndPoint)
579+ return (TransportValid . (localEndPointAt addr ^= Just localEndPoint) . (nextEndPointId ^= ix + 1) $ vst, localEndPoint)
575580 TransportClosed ->
576581 throw (TransportError NewEndPointFailed "Transport closed")
577582
@@ -644,8 +649,10 @@ requestConnectionTo ourEndPoint theirAddress = go
644649 -- point)
645650 reply <- handle failureHandler $ doRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId
646651 case decodeInt32 . BS.concat $ reply of
647− Nothing -> failureHandler (userError "Invalid integer")
648− Just cid -> return (theirEndPoint, cid)
652+ Nothing ->
653+ failureHandler (userError "Invalid integer")
654+ Just cid -> do
655+ return (theirEndPoint, cid)
649656
650657 -- If this is a new endpoint, fork a thread to listen for incoming
651658 -- connections. We don't want to do this while we hold the lock, because
@@ -979,6 +986,8 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
979986 -- Create a new connection
980987 createNewConnection :: ControlRequestId -> IO ()
981988 createNewConnection reqId = do
989+ -- getNextConnectionId throws an exception if ourEndPoint is closed; but
990+ -- if this endpoint is closed, this thread will soon die anyway
982991 newId <- getNextConnectionId ourEndPoint
983992 modifyMVar_ theirState $ \st -> do
984993 vst <- case st of
@@ -1135,6 +1144,9 @@ firstNonReservedConnectionId = 1024
11351144localEndPoints :: Accessor ValidTransportState (Map EndPointAddress LocalEndPoint)
11361145localEndPoints = accessor _localEndPoints (\es st -> st { _localEndPoints = es })
11371146
1147+nextEndPointId :: Accessor ValidTransportState EndPointId
1148+nextEndPointId = accessor _nextEndPointId (\eid st -> st { _nextEndPointId = eid })
1149+
11381150pendingCtrlRequests :: Accessor ValidLocalEndPointState (IntMap (MVar [ByteString]))
11391151pendingCtrlRequests = accessor _pendingCtrlRequests (\rep st -> st { _pendingCtrlRequests = rep })
11401152