@@ -162,20 +162,20 @@ data LocalEndPoint = LocalEndPoint { _localAddress :: EndPointAddress
162162data LocalState = LocalState { _nextConnectionId :: ConnectionId
163163 , _pendingCtrlRequests :: IntMap (MVar [ByteString])
164164 , _nextCtrlRequestId :: ControlRequestId
165− , _localConnections :: Map EndPointAddress (Maybe Closing, MVar RemoteState)
165+ , _localConnections :: Map EndPointAddress (IsClosing, MVar RemoteState)
166166 }
167167data RemoteEndPoint = RemoteEndPoint { _remoteAddress :: EndPointAddress
168168 , _remoteSocket :: N.Socket
169169 , sendTo :: [ByteString] -> IO ()
170170 }
171−data RemoteState = InvalidEndPoint (FailedWith ConnectErrorCode)
172− | ValidEndPoint !RefCount RemoteEndPoint
173− | DeletedEndPoint
171+data RemoteState = EndPointInvalid (FailedWith ConnectErrorCode)
172+ | EndPointValid !RefCount RemoteEndPoint
173+ | EndPointClosed
174174
175175type EndPointId = Int32
176176type ControlRequestId = Int32
177177type RefCount = Int32
178−type Closing = MVar [ByteString]
178+type IsClosing = Maybe RemoteEndPoint
179179type Conn = (LocalEndPoint, RemoteEndPoint)
180180
181181-- | Control headers
@@ -189,7 +189,7 @@ data ControlHeader =
189189-- Response sent by /B/ to /A/ when /A/ tries to connect
190190data ConnectionRequestResponse =
191191 ConnectionRequestAccepted -- ^ /B/ accepts the connection
192− | ConnectionRequestInvalidEndPoint -- ^ /A/ requested an invalid endpoint
192+ | ConnectionRequestEndPointInvalid -- ^ /A/ requested an invalid endpoint
193193 | ConnectionRequestCrossed -- ^ /A/s request crossed with a request from /B/ (see protocols)
194194 deriving Enum
195195
@@ -297,14 +297,14 @@ findRemoteEndPoint ourEndPoint theirAddress = ErrorT $ go
297297 return (connectionTo theirAddress ^= Just (Nothing, theirMVar) $ st, theirMVar)
298298 theirState <- readMVar endPoint
299299 case theirState of
300− InvalidEndPoint err ->
300+ EndPointInvalid err ->
301301 -- We remove the endpoint from the state again because the next call
302302 -- to 'connect' might give a different result
303303 modifyMVar state $ \st ->
304304 return (connectionTo theirAddress ^= Nothing $ st, Left err)
305− ValidEndPoint _ ep ->
305+ EndPointValid _ ep ->
306306 return $ Right ep
307− DeletedEndPoint ->
307+ EndPointClosed ->
308308 go
309309
310310 state :: MVar LocalState
@@ -316,16 +316,16 @@ connectToRemoteEndPoint ourEndPoint theirMVar theirAddress = do
316316 case result of
317317 Right (ConnectionRequestAccepted, sock) -> do
318318 theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock
319− putMVar theirMVar (ValidEndPoint 0 theirEndPoint)
320− Right (ConnectionRequestInvalidEndPoint, sock) -> do
321− putMVar theirMVar (InvalidEndPoint (invalidAddress "Invalid endpoint"))
319+ putMVar theirMVar (EndPointValid 0 theirEndPoint)
320+ Right (ConnectionRequestEndPointInvalid, sock) -> do
321+ putMVar theirMVar (EndPointInvalid (invalidAddress "Invalid endpoint"))
322322 N.sClose sock
323323 Right (ConnectionRequestCrossed, sock) ->
324324 N.sClose sock
325325 Left err ->
326326 -- TODO: this isn't quite right; socket *might* have been opened here,
327327 -- in which case we should close it
328− putMVar theirMVar (InvalidEndPoint err)
328+ putMVar theirMVar (EndPointInvalid err)
329329 where
330330 go :: ErrorT (FailedWith ConnectErrorCode) IO (ConnectionRequestResponse, N.Socket)
331331 go = do
@@ -391,16 +391,18 @@ increaseRemoteRefCount :: Conn -> IO ()
391391increaseRemoteRefCount (ourEndPoint, theirEndPoint) = do
392392 let theirAddr = theirEndPoint ^. remoteAddress
393393 state = ourEndPoint ^. localState
394− -- First check if this remote endpoint is awaiting closure
395− (mClosing, theirMVar) <- modifyMVar state $ \st -> do
396− let Just (closure, theirMVar) = st ^. connectionTo theirAddr
397− return (connectionTo theirAddr ^= Just (Nothing, theirMVar) $ st, (closure, theirMVar))
398− case mClosing of
399− Just mvar -> putMVar mvar []
400− _ -> return ()
394+ -- If this endpoint was closing, cancel that
395+ theirMVar <- modifyMVar state $ \st -> do
396+ let Just (isClosing, theirMVar) = st ^. connectionTo theirAddr
397+ case isClosing of
398+ Nothing ->
399+ return (st, theirMVar)
400+ Just ep -> do
401+ putMVar theirMVar (EndPointValid 0 ep)
402+ return (connectionTo theirAddr ^= Just (Nothing, theirMVar) $ st, theirMVar)
401403 -- Increase the refcount
402− modifyMVar_ theirMVar $ \(ValidEndPoint rc ep) ->
403− return (ValidEndPoint (rc + 1) ep)
404+ modifyMVar_ theirMVar $ \(EndPointValid rc ep) ->
405+ return (EndPointValid (rc + 1) ep)
404406
405407decreaseRemoteRefCount :: Conn -> IO ()
406408decreaseRemoteRefCount (ourEndPoint, theirEndPoint) = do
@@ -410,19 +412,15 @@ decreaseRemoteRefCount (ourEndPoint, theirEndPoint) = do
410412 -- We don't use modifyMVar here because we want to leave the MVar
411413 -- empty when refcount reaches 0
412414 -- TODO: deal with exceptions
413− ValidEndPoint rc ep <- takeMVar theirMVar
415+ EndPointValid rc ep <- takeMVar theirMVar
414416 if rc > 1
415417 then do
416− putMVar theirMVar (ValidEndPoint (rc - 1) ep)
418+ putMVar theirMVar (EndPointValid (rc - 1) ep)
417419 else do
418− reply <- sendRemoteRequest (ourEndPoint, theirEndPoint) CloseSocket
420+ sendTo theirEndPoint [encodeInt32 CloseSocket]
419421 modifyMVar_ (ourEndPoint ^. localState) $ \st -> do
420422 let Just (_, mvar) = st ^. connectionTo theirAddr
421− return (connectionTo theirAddr ^= Just (Just reply, mvar) $ st)
422− void . forkIO $ do
423− -- Wait for the reply from the remote endpoint
424− canCloseSocket <- takeMVar reply
425− putStrLn $ "canCloseSocket = " ++ show canCloseSocket
423+ return (connectionTo theirAddr ^= Just (Just ep, mvar) $ st)
426424
427425--------------------------------------------------------------------------------
428426-- Incoming requests --
@@ -437,7 +435,7 @@ handleConnectionRequest transport sock = do
437435 return (theirAddress, ourEndPoint)
438436 case request of
439437 Right (_, Nothing) -> do
440− sendMany sock [encodeInt32 ConnectionRequestInvalidEndPoint]
438+ sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]
441439 N.sClose sock
442440 Right (theirAddress, Just ourEndPoint) -> do
443441 (crossed, mvar) <- modifyMVar (ourEndPoint ^. localState) $ \st ->
@@ -453,7 +451,7 @@ handleConnectionRequest transport sock = do
453451 N.sClose sock
454452 else do
455453 theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock
456− putMVar mvar (ValidEndPoint 0 theirEndPoint)
454+ putMVar mvar (EndPointValid 0 theirEndPoint)
457455 sendMany sock [encodeInt32 ConnectionRequestAccepted]
458456 Left _ ->
459457 -- Invalid request
@@ -494,12 +492,11 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
494492 if connId >= firstNonReservedConnectionId
495493 then readMessage connId
496494 else do
497− reqId <- recvInt32 sock
498495 case toEnum (fromIntegral connId) of
499− RequestConnectionId -> createNewConnection reqId
500− ControlResponse -> readControlResponse reqId
501− CloseConnection -> closeConnection reqId
502− CloseSocket -> closeSocket reqId
496+ RequestConnectionId -> recvInt32 sock >>= createNewConnection
497+ ControlResponse -> recvInt32 sock >>= readControlResponse
498+ CloseConnection -> recvInt32 sock >>= closeConnection
499+ CloseSocket -> closeSocket
503500
504501 -- Create a new connection
505502 createNewConnection :: ControlRequestId -> IO ()
@@ -529,7 +526,7 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
529526 putMVar mvar response
530527 go
531528
532− -- Close a connection (in this case the reqId is abused for the connection ID)
529+ -- Close a connection
533530 closeConnection :: ConnectionId -> IO ()
534531 closeConnection cid = do
535532 writeChan ourChannel (ConnectionClosed cid)
@@ -538,22 +535,22 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do
538535 go
539536
540537 -- Close the socket (if we don't have any outgoing connections)
541− closeSocket :: ControlRequestId -> IO ()
542− closeSocket cid = do
543− let ok = encodeInt32 ControlResponse : encodeInt32 cid : prependLength [encodeInt32 (1 :: Int32)]
538+ closeSocket :: IO ()
539+ closeSocket = do
544540 didClose <- modifyMVar (ourEndPoint ^. localState) $ \st -> do
545541 case st ^. connectionTo theirAddr of
546− Just (Just _, _) -> do
547− sendTo theirEndPoint ok
542+ Just (Just _, mvar) -> do
543+ sendTo theirEndPoint [encodeInt32 CloseSocket]
548544 N.sClose sock
549− return (st, True)
545+ putMVar mvar EndPointClosed
546+ return (connectionTo theirAddr ^= Nothing $ st, True)
550547 Just (Nothing, theirMVar) -> do
551− ValidEndPoint rc _ <- readMVar theirMVar
548+ EndPointValid rc _ <- readMVar theirMVar
552549 if rc == 0
553550 then do
554− sendTo theirEndPoint ok
551+ sendTo theirEndPoint [encodeInt32 CloseSocket]
555552 N.sClose sock
556− return (st, True)
553+ return (connectionTo theirAddr ^= Nothing $ st, True)
557554 else do
558555 return (st, False)
559556 Nothing -> do
@@ -622,7 +619,7 @@ nextCtrlRequestId = lens _nextCtrlRequestId (\cid st -> st { _nextCtrlRequestId
622619nextConnectionId :: Lens LocalState ConnectionId
623620nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = cix })
624621
625−localConnections :: Lens LocalState (Map EndPointAddress (Maybe Closing, MVar RemoteState))
622+localConnections :: Lens LocalState (Map EndPointAddress (IsClosing, MVar RemoteState))
626623localConnections = lens _localConnections (\cs st -> st { _localConnections = cs })
627624
628625remoteAddress :: Lens RemoteEndPoint EndPointAddress
@@ -637,5 +634,5 @@ localEndPointAt ix = localEndPoints >>> intMapLens (fromIntegral ix)
637634pendingCtrlRequestsAt :: ControlRequestId -> Lens LocalState (Maybe (MVar [ByteString]))
638635pendingCtrlRequestsAt ix = pendingCtrlRequests >>> intMapLens (fromIntegral ix)
639636
640−connectionTo :: EndPointAddress -> Lens LocalState (Maybe (Maybe Closing, MVar RemoteState))
637+connectionTo :: EndPointAddress -> Lens LocalState (Maybe (IsClosing, MVar RemoteState))
641638connectionTo addr = localConnections >>> mapLens addr