401984b · distributed-process

git.iko.soy/distributed-process

on master help sign in

Better bookkeeping

401984bEdsko de Vries14 years ago

network-transport-2/src/Network/Transport.hs+7−6

compare with master

unifiedsplit

@@ -16,6 +16,7 @@ module Network.Transport ( -- * Types

16 , NewMulticastGroupErrorCode(..)

17 , ResolveMulticastGroupErrorCode(..)

18 , SendErrorCode(..)

19+ , EventErrorCode(..)

19 ) where

20 ) where

21import Data.ByteString (ByteString)

22import Data.ByteString (ByteString)

@@ -69,7 +70,7 @@ data Event =

69 -- | Received multicast

70 -- | Received multicast

70 | ReceivedMulticast MulticastAddress [ByteString]

71 | ReceivedMulticast MulticastAddress [ByteString]

71 -- | Error that caused a bunch of connections to close (possibly none)

72 -- | Error that caused a bunch of connections to close (possibly none)

72− | ErrorEvent ErrorEventErrorCode [ConnectionId]

73+ | ErrorEvent EventErrorCode [ConnectionId]

73 deriving Show

74 deriving Show

75-- | Connection IDs enable receivers to distinguish one connection from another.

76-- | Connection IDs enable receivers to distinguish one connection from another.

@@ -142,12 +143,12 @@ data ResolveMulticastGroupErrorCode =

142

143

143-- | Failure during sending a message

144-- | Failure during sending a message

144data SendErrorCode =

145data SendErrorCode =

145− SendConnectionClosed -- ^ Connection was closed manually or because of an error

146+ SendConnectionClosed -- ^ Connection was closed manually or because of an error

146− | SendFailed -- ^ Send failed for some other reason

147+ | SendFailed -- ^ Send failed for some other reason

147 deriving Show

148 deriving Show

148

149

149-- | Error codes used when reporting errors to endpoints (through receive)

150-- | Error codes used when reporting errors to endpoints (through receive)

150−data ErrorEventErrorCode =

151+data EventErrorCode =

151− ErrorEventEndPointClosed -- ^ Endpoint was closed (manually or because of an error)

152+ EventErrorEndPointClosed -- ^ Endpoint was closed (manually or because of an error)

152− | TransportFatalError -- ^ Transport-wide fatal error

153+ | EventErrorTransportFailure -- ^ Transport-wide fatal error

153 deriving Show

154 deriving Show

network-transport-2/src/Network/Transport/Internal/TCP.hs+1−1

@@ -26,7 +26,7 @@ import qualified Network.Socket as N ( HostName

26import qualified Network.Socket.ByteString as NBS (recv)

27import Control.Concurrent (forkIO, ThreadId)

28import Control.Monad (mzero, MonadPlus, liftM, forever)

29−import Control.Exception (SomeException, try, handle)

29+import Control.Exception (SomeException, handle)

30import Data.ByteString (ByteString)

31import qualified Data.ByteString as BS (length, concat, null)

32import Data.Int (Int32)

network-transport-2/src/Network/Transport/TCP.hs+190−111

@@ -46,7 +46,7 @@ import qualified Network.Socket as N ( HostName

46 , connect

47 )

48import Network.Socket.ByteString (sendMany)

49−import Control.Concurrent (forkIO, ThreadId, killThread)

49+import Control.Concurrent (forkIO, ThreadId, killThread, myThreadId)

50import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)

51import Control.Concurrent.MVar ( MVar

52 , newMVar

@@ -57,25 +57,27 @@ import Control.Concurrent.MVar ( MVar

57 , putMVar

58 , newEmptyMVar

59 , withMVar

60+ , tryTakeMVar

60 )

61 )

61import Control.Category ((>>>))

62import Control.Category ((>>>))

62import Control.Applicative ((<*>), (*>), (<$>))

63import Control.Applicative ((<*>), (*>), (<$>))

63−import Control.Monad (forM_, void, when)

64+import Control.Monad (forM_, void, when, mzero, unless)

64import Control.Monad.Error (ErrorT(ErrorT), runErrorT, throwError)

65import Control.Monad.Error (ErrorT(ErrorT), runErrorT, throwError)

66+import Control.Monad.Trans.Maybe (MaybeT, runMaybeT)

65import Control.Monad.IO.Class (liftIO)

67import Control.Monad.IO.Class (liftIO)

66import Control.Exception (IOException, SomeException)

68import Control.Exception (IOException, SomeException)

67−import Data.IORef (IORef, newIORef, writeIORef, readIORef, modifyIORef)

69+import Data.IORef (IORef, newIORef, writeIORef, readIORef)

68import Data.ByteString (ByteString)

70import Data.ByteString (ByteString)

69import qualified Data.ByteString as BS (concat)

71import qualified Data.ByteString as BS (concat)

70import qualified Data.ByteString.Char8 as BSC (pack, unpack)

72import qualified Data.ByteString.Char8 as BSC (pack, unpack)

71import Data.Int (Int32)

73import Data.Int (Int32)

72import Data.IntMap (IntMap)

74import Data.IntMap (IntMap)

73−import qualified Data.IntMap as IntMap (empty, size)

75+import qualified Data.IntMap as IntMap (empty)

74import Data.IntSet (IntSet)

76import Data.IntSet (IntSet)

75−import qualified Data.IntSet as IntSet (empty, insert, delete, elems)

77+import qualified Data.IntSet as IntSet (empty, insert, elems, singleton, null, delete)

76import Data.Map (Map)

78import Data.Map (Map)

77−import qualified Data.Map as Map (empty)

79+import qualified Data.Map as Map (empty, elems, size)

78−import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=))

80+import Data.Lens.Lazy (Lens, lens, intMapLens, mapLens, (^.), (^=), (^+=), (^%=))

79import Text.Regex.Applicative (RE, few, anySym, sym, many, match)

81import Text.Regex.Applicative (RE, few, anySym, sym, many, match)

80import System.IO (hPutStrLn, stderr)

82import System.IO (hPutStrLn, stderr)

@@ -183,18 +185,18 @@ data LocalEndPoint = LocalEndPoint { localAddress :: EndPointAddress

183 }

185 }

184data RemoteEndPoint = RemoteEndPoint { remoteAddress :: EndPointAddress

186data RemoteEndPoint = RemoteEndPoint { remoteAddress :: EndPointAddress

185 , remoteState :: MVar RemoteState

187 , remoteState :: MVar RemoteState

186− , remoteThread :: ThreadId

187 }

188 }

188

189

189-- We use underscores for fields that we might update (using lenses).

190-- We use underscores for fields that we might update (using lenses).

190

191

191−data TransportState = TransportState { _localEndPoints :: IntMap LocalEndPoint

192+data TransportState = TransportState { _localEndPoints :: Map EndPointAddress LocalEndPoint

192 , _acceptNewConnections :: Bool

193 , _acceptNewConnections :: Bool

193 }

194 }

194data LocalState = LocalState { _nextConnectionId :: ConnectionId

195data LocalState = LocalState { _nextConnectionId :: ConnectionId

195 , _pendingCtrlRequests :: IntMap (MVar [ByteString])

196 , _pendingCtrlRequests :: IntMap (MVar [ByteString])

196 , _nextCtrlRequestId :: ControlRequestId

197 , _nextCtrlRequestId :: ControlRequestId

197 , _localConnections :: Map EndPointAddress RemoteEndPoint

198 , _localConnections :: Map EndPointAddress RemoteEndPoint

199+ , _internalThreads :: [ThreadId]

198 }

200 }

199

201

200-- A remote endpoint has incoming and outgoing connections, and when the total

202-- A remote endpoint has incoming and outgoing connections, and when the total

@@ -238,9 +240,10 @@ data LocalState = LocalState { _nextConnectionId :: ConnectionId

238

240

239data RemoteState = RemoteEndPointInvalid (FailedWith ConnectErrorCode)

241data RemoteState = RemoteEndPointInvalid (FailedWith ConnectErrorCode)

240 | RemoteEndPointValid RemoteConnection

242 | RemoteEndPointValid RemoteConnection

241− | RemoteEndPointClosing RemoteConnection (MVar ())

243+ | RemoteEndPointClosing (MVar ()) RemoteConnection

242 | RemoteEndPointClosed

244 | RemoteEndPointClosed

243−data RemoteConnection = RemoteConnection { _remoteRefCount :: !Int

245+data RemoteConnection = RemoteConnection { _remoteOutgoing :: !Int

246+ , _remoteIncoming :: IntSet

244 , remoteSocket :: N.Socket

247 , remoteSocket :: N.Socket

245 , sendOn :: [ByteString] -> IO ()

248 , sendOn :: [ByteString] -> IO ()

246 }

249 }

@@ -273,7 +276,7 @@ data ConnectionRequestResponse =

273-- TODOs: deal with hints

276-- TODOs: deal with hints

274createTransport :: N.HostName -> N.ServiceName -> IO (Either IOException Transport)

277createTransport :: N.HostName -> N.ServiceName -> IO (Either IOException Transport)

275createTransport host port = do

278createTransport host port = do

276− state <- newMVar TransportState { _localEndPoints = IntMap.empty

279+ state <- newMVar TransportState { _localEndPoints = Map.empty

277 , _acceptNewConnections = True

280 , _acceptNewConnections = True

278 }

281 }

279 let transport = TCPTransport { transportState = state

282 let transport = TCPTransport { transportState = state

@@ -308,9 +311,10 @@ apiCloseTransport transportThread = do

308apiNewEndPoint :: TCPTransport -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)

311apiNewEndPoint :: TCPTransport -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)

309apiNewEndPoint transport = do

312apiNewEndPoint transport = do

310 ourEndPoint <- createLocalEndPoint transport

313 ourEndPoint <- createLocalEndPoint transport

311− return . Right $ EndPoint { receive = readChan (localChannel ourEndPoint)

314+ return . Right $ EndPoint { receive = readChan (localChannel ourEndPoint)

312− , address = localAddress ourEndPoint

315+ , address = localAddress ourEndPoint

313− , connect = apiConnect ourEndPoint

316+ , connect = apiConnect ourEndPoint

317+ , closeEndPoint = apiCloseEndPoint transport ourEndPoint

314 , newMulticastGroup = return . Left $ newMulticastGroupError

318 , newMulticastGroup = return . Left $ newMulticastGroupError

315 , resolveMulticastGroup = return . Left . const resolveMulticastGroupError

319 , resolveMulticastGroup = return . Left . const resolveMulticastGroupError

316 }

320 }

@@ -345,14 +349,16 @@ apiConnect ourEndPoint theirAddress _ = runErrorT $ do

345-- TODO: We ignore errors during a close. Is that right?

349-- TODO: We ignore errors during a close. Is that right?

346apiClose :: RemoteEndPoint -> ConnectionId -> IORef Bool -> IO ()

350apiClose :: RemoteEndPoint -> ConnectionId -> IORef Bool -> IO ()

347apiClose theirEndPoint connId connAlive = void . tryIO $ do

351apiClose theirEndPoint connId connAlive = void . tryIO $ do

348− wasAlive <- withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) -> do

352+ modifyMVar_ (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) -> do

349 alive <- readIORef connAlive

353 alive <- readIORef connAlive

350− when alive $ do

354+ if alive

355+ then do

351− sendOn remoteConn [encodeInt32 CloseConnection, encodeInt32 connId]

356+ sendOn remoteConn [encodeInt32 CloseConnection, encodeInt32 connId]

352− writeIORef connAlive False

357+ writeIORef connAlive False

353− return alive

358+ return (RemoteEndPointValid . (remoteOutgoing ^+= (-1)) $ remoteConn)

354− -- We don't want to call decreaseRemoteRefCount while we still hold the lock

359+ else

355− when wasAlive $ decreaseRemoteRefCount theirEndPoint

360+ return (RemoteEndPointValid remoteConn)

361+ closeIfUnused theirEndPoint

356

362

357-- | Send data across a connection

363-- | Send data across a connection

358--

364--

@@ -369,6 +375,37 @@ apiSend theirEndPoint connId connAlive payload = do

369 else

375 else

370 throwError $ FailedWith SendConnectionClosed "Connection closed"

376 throwError $ FailedWith SendConnectionClosed "Connection closed"

371

377

378+-- | Force-close the endpoint

379+apiCloseEndPoint :: TCPTransport -> LocalEndPoint -> IO ()

380+apiCloseEndPoint transport ourEndPoint = do

381+ -- Remove the reference from the transport state

382+ modifyMVar_ (transportState transport) $ return . (localEndPointAt (localAddress ourEndPoint) ^= Nothing)

383+ -- We take the endpoint state and don't put it back

384+ -- (this may block some threads, but we don't care: we are killing them anyway)

385+ ourState <- takeMVar (localState ourEndPoint)

386+ --incoming <- IntSet.elems <$> takeMVar (localIncoming ourEndPoint)

387+ let incoming = undefined

388+ -- Kill all threads associated with this endpoint and close all sockets

389+ forM_ (ourState ^. internalThreads) killThread

390+ forM_ (Map.elems $ ourState ^. localConnections) $ \theirEndPoint ->

391+ runMaybeT $ remoteSocketOf theirEndPoint >>= liftIO . N.sClose

392+ -- We send a single message to the endpoint that it is closed. Subsequent

393+ -- calls will block. We could change this so that all subsequent calls to

394+ -- receive return an error, but this would mean checking for some state on

395+ -- every call to receive, which is an unnecessary overhead.

396+ writeChan (localChannel ourEndPoint) (ErrorEvent EventErrorEndPointClosed incoming)

397+ where

398+ remoteSocketOf :: RemoteEndPoint -> MaybeT IO N.Socket

399+ remoteSocketOf theirEndPoint = do

400+ -- Like above, we take but don't put back

401+ mRemoteState <- liftIO $ tryTakeMVar (remoteState theirEndPoint)

402+ case mRemoteState of

403+ Nothing -> mzero

404+ Just (RemoteEndPointInvalid _) -> mzero

405+ Just (RemoteEndPointValid conn) -> return $ remoteSocket conn

406+ Just (RemoteEndPointClosing _ conn) -> return $ remoteSocket conn

407+ Just RemoteEndPointClosed -> mzero

408+

372-- | Special case of 'apiConnect': connect an endpoint to itself

409-- | Special case of 'apiConnect': connect an endpoint to itself

373connectToSelf :: LocalEndPoint -> IO (Either (FailedWith ConnectErrorCode) Connection)

410connectToSelf :: LocalEndPoint -> IO (Either (FailedWith ConnectErrorCode) Connection)

374connectToSelf ourEndPoint = do

411connectToSelf ourEndPoint = do

@@ -411,17 +448,16 @@ createLocalEndPoint transport = do

411 , _pendingCtrlRequests = IntMap.empty

448 , _pendingCtrlRequests = IntMap.empty

412 , _nextCtrlRequestId = 0

449 , _nextCtrlRequestId = 0

413 , _localConnections = Map.empty

450 , _localConnections = Map.empty

451+ , _internalThreads = []

414 }

452 }

415− let host = transportHost transport

416− let port = transportPort transport

417 modifyMVar (transportState transport) $ \st -> do

453 modifyMVar (transportState transport) $ \st -> do

418− let ix = fromIntegral $ IntMap.size (st ^. localEndPoints)

454+ let ix = fromIntegral $ Map.size (st ^. localEndPoints)

419− let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show ix

455+ let addr = encodeEndPointAddress transport ix

420− let localEndPoint = LocalEndPoint { localAddress = addr

456+ let localEndPoint = LocalEndPoint { localAddress = addr

421− , localChannel = chan

457+ , localChannel = chan

422− , localState = state

458+ , localState = state

423 }

459 }

424− return ((localEndPointAt ix ^= Just localEndPoint) $ st, localEndPoint)

460+ return ((localEndPointAt addr ^= Just localEndPoint) $ st, localEndPoint)

425

461

426-- | Request a connection to a remote endpoint

462-- | Request a connection to a remote endpoint

427--

463--

@@ -435,17 +471,22 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go

435 where

471 where

436 go = do

472 go = do

437 let ourState = localState ourEndPoint

473 let ourState = localState ourEndPoint

438− theirEndPoint <- modifyMVar ourState $ \st ->

474+ (theirEndPoint, isNew) <- modifyMVar ourState $ \st ->

439 case st ^. connectionTo theirAddress of

475 case st ^. connectionTo theirAddress of

440 Just theirEndPoint ->

476 Just theirEndPoint ->

441− return (st, theirEndPoint)

477+ return (st, (theirEndPoint, False))

442 Nothing -> do

478 Nothing -> do

443 theirState <- newEmptyMVar

479 theirState <- newEmptyMVar

444 let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress

480 let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress

445 , remoteState = theirState

481 , remoteState = theirState

446 }

482 }

447− forkIO $ connectToRemoteEndPoint (ourEndPoint, theirEndPoint)

448− return (connectionTo theirAddress ^= Just theirEndPoint $ st, theirEndPoint)

483+ return (connectionTo theirAddress ^= Just theirEndPoint $ st, (theirEndPoint, True))

484+

485+ -- If this is a new endpoint, fork a thread to listen for incoming

486+ -- connections. We don't want to do this while we hold the lock, because

487+ -- forkEndPointThread modifies the local state too (to record the thread ID)

488+ when isNew . forkEndPointThread ourEndPoint $

489+ connectToRemoteEndPoint (ourEndPoint, theirEndPoint)

449

490

450 -- Before we initiate the new connection request we want to make sure

491 -- Before we initiate the new connection request we want to make sure

451 -- that refcount on the endpoint is incremented so that a concurrent

492 -- that refcount on the endpoint is incremented so that a concurrent

@@ -456,7 +497,7 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go

456 endPointStateSnapshot <- modifyMVar theirState $ \st ->

497 endPointStateSnapshot <- modifyMVar theirState $ \st ->

457 case st of

498 case st of

458 RemoteEndPointValid ep ->

499 RemoteEndPointValid ep ->

459− return (RemoteEndPointValid . (remoteRefCount ^+= 1) $ ep, st)

500+ return (RemoteEndPointValid . (remoteOutgoing ^+= 1) $ ep, st)

460 _ ->

501 _ ->

461 return (st, st)

502 return (st, st)

462

503

@@ -471,7 +512,7 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go

471 RemoteEndPointInvalid err -> do

512 RemoteEndPointInvalid err -> do

472 return . Left $ err

513 return . Left $ err

473

514

474− RemoteEndPointClosing _ resolved -> do

515+ RemoteEndPointClosing resolved _ -> do

475 -- If the remote endpoint is closing, then we need to block until

516 -- If the remote endpoint is closing, then we need to block until

476 -- this is resolved and we then try again

517 -- this is resolved and we then try again

477 readMVar resolved

518 readMVar resolved

@@ -498,12 +539,13 @@ requestConnectionTo ourEndPoint theirAddress = ErrorT go

498 -- valid state at this point.

539 -- valid state at this point.

499 let failed err = do

540 let failed err = do

500 modifyMVar_ theirState $ \(RemoteEndPointValid ep) ->

541 modifyMVar_ theirState $ \(RemoteEndPointValid ep) ->

501− return (RemoteEndPointValid . (remoteRefCount ^+= (-1)) $ ep)

542+ return (RemoteEndPointValid . (remoteOutgoing ^+= (-1)) $ ep)

543+ -- TODO: should we call closeIfUnused here?

502 return . Left $ FailedWith ConnectFailed err

544 return . Left $ FailedWith ConnectFailed err

503

545

504 case mcid of

546 case mcid of

505 Left err -> failed (show err)

547 Left err -> failed (show err)

506− Right Nothing -> failed ("Invalid integer")

548+ Right Nothing -> failed "Invalid integer"

507 Right (Just cid) -> return . Right $ (theirEndPoint, cid)

549 Right (Just cid) -> return . Right $ (theirEndPoint, cid)

508

550

509-- | Establish a connection to a remote endpoint

551-- | Establish a connection to a remote endpoint

@@ -520,11 +562,12 @@ connectToRemoteEndPoint (ourEndPoint, theirEndPoint) = do

520 case result of

562 case result of

521 Right (ConnectionRequestAccepted, sock) -> do

563 Right (ConnectionRequestAccepted, sock) -> do

522 let remoteConn = RemoteConnection { remoteSocket = sock

564 let remoteConn = RemoteConnection { remoteSocket = sock

523− , _remoteRefCount = 0

565+ , _remoteOutgoing = 0

566+ , _remoteIncoming = IntSet.empty

524 , sendOn = sendMany sock

567 , sendOn = sendMany sock

525 }

568 }

526 putMVar theirState (RemoteEndPointValid remoteConn)

569 putMVar theirState (RemoteEndPointValid remoteConn)

527− void . forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)

570+ handleIncomingMessages (ourEndPoint, theirEndPoint)

528 Right (ConnectionRequestEndPointInvalid, sock) -> do

571 Right (ConnectionRequestEndPointInvalid, sock) -> do

529 -- We remove the endpoint from our local state again because the next

572 -- We remove the endpoint from our local state again because the next

530 -- call to 'connect' might give a different result. Threads that were

573 -- call to 'connect' might give a different result. Threads that were

@@ -569,6 +612,11 @@ connectToRemoteEndPoint (ourEndPoint, theirEndPoint) = do

569 ourAddress :: ByteString

612 ourAddress :: ByteString

570 EndPointAddress ourAddress = localAddress ourEndPoint

613 EndPointAddress ourAddress = localAddress ourEndPoint

571

614

615+-- | Encode end point address

616+encodeEndPointAddress :: TCPTransport -> EndPointId -> EndPointAddress

617+encodeEndPointAddress transport ix = EndPointAddress . BSC.pack $

618+ transportHost transport ++ ":" ++ transportPort transport ++ ":" ++ show ix

619+

572-- | Decode end point address

620-- | Decode end point address

573-- TODO: This uses regular expression parsing, which is nice, but unnecessary

621-- TODO: This uses regular expression parsing, which is nice, but unnecessary

574decodeEndPointAddress :: EndPointAddress -> Maybe (N.HostName, N.ServiceName, EndPointId)

622decodeEndPointAddress :: EndPointAddress -> Maybe (N.HostName, N.ServiceName, EndPointId)

@@ -581,7 +629,7 @@ decodeEndPointAddress (EndPointAddress bs) = match endPointAddressRE (BSC.unpack

581--

629--

582-- RELY: Remote endpoint must be in valid state.

630-- RELY: Remote endpoint must be in valid state.

583-- GUARANTEE: Will not change the state of the remote endpoint.

631-- GUARANTEE: Will not change the state of the remote endpoint.

584−doRemoteRequest :: EndPointPair -> ControlHeader -> IO ([ByteString])

632+doRemoteRequest :: EndPointPair -> ControlHeader -> IO [ByteString]

585doRemoteRequest (ourEndPoint, theirEndPoint) header = do

633doRemoteRequest (ourEndPoint, theirEndPoint) header = do

586 reply <- newEmptyMVar

634 reply <- newEmptyMVar

587 reqId <- modifyMVar (localState ourEndPoint) $ \st -> do

635 reqId <- modifyMVar (localState ourEndPoint) $ \st -> do

@@ -591,22 +639,32 @@ doRemoteRequest (ourEndPoint, theirEndPoint) header = do

591 sendOn remoteConn [encodeInt32 header, encodeInt32 reqId]

639 sendOn remoteConn [encodeInt32 header, encodeInt32 reqId]

592 takeMVar reply

640 takeMVar reply

593

641

594−-- | Decrease the remote refcount of a remote endpoint. When the refcount

642+-- | Check if the remote endpoint is unused, and if so, send a CloseSocket request

595−-- reaches zero, initiate socket closure.

596−--

597−-- RELY: Remote endpoint must be in valid state.

598−-- GUARANTEE: Will leave the endpoint in its current state or change it to

599−-- RemoteEndPointClosing.

600−decreaseRemoteRefCount :: RemoteEndPoint -> IO ()

643+closeIfUnused :: RemoteEndPoint -> IO ()

601−decreaseRemoteRefCount theirEndPoint = do

644+closeIfUnused theirEndPoint = modifyMVar_ (remoteState theirEndPoint) $ \st ->

645+ case st of

602− modifyMVar_ (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) ->

646+ RemoteEndPointValid remoteConn -> do

603− if remoteConn ^. remoteRefCount == 1

647+ if remoteConn ^. remoteOutgoing == 0 && IntSet.null (remoteConn ^. remoteIncoming)

604− then do

648+ then do

605− resolved <- newEmptyMVar

649+ resolved <- newEmptyMVar

606− sendOn remoteConn [encodeInt32 CloseSocket]

650+ sendOn remoteConn [encodeInt32 CloseSocket]

607− return (RemoteEndPointClosing remoteConn resolved)

651+ return (RemoteEndPointClosing resolved remoteConn)

608− else do

652+ else

609− return (RemoteEndPointValid . (remoteRefCount ^+= (-1)) $ remoteConn)

653+ return st

654+ _ -> do

655+ return st

656+

657+-- | Fork a new thread and store its ID as part of the transport state

658+forkEndPointThread :: LocalEndPoint -> IO () -> IO ()

659+forkEndPointThread ourEndPoint p = do

660+ modifyMVar_ (localState ourEndPoint) $ \st -> do

661+ tid <- forkIO (p >> removeThread)

662+ return (internalThreads ^%= (tid :) $ st)

663+ where

664+ removeThread :: IO ()

665+ removeThread = do

666+ tid <- myThreadId

667+ modifyMVar_ (localState ourEndPoint) $ return . (internalThreads ^%= filter (/= tid))

610

668

611--------------------------------------------------------------------------------

669--------------------------------------------------------------------------------

612-- Incoming requests --

670-- Incoming requests --

@@ -616,17 +674,17 @@ decreaseRemoteRefCount theirEndPoint = do

616-- establish a TCP connection with us)

674-- establish a TCP connection with us)

617handleConnectionRequest :: TCPTransport -> N.Socket -> IO ()

675handleConnectionRequest :: TCPTransport -> N.Socket -> IO ()

618handleConnectionRequest transport sock = do

676handleConnectionRequest transport sock = do

619− let tState = transportState transport

620 request <- tryIO $ do

677 request <- tryIO $ do

621 ourEndPointId <- recvInt32 sock

678 ourEndPointId <- recvInt32 sock

622 theirAddress <- EndPointAddress . BS.concat <$> recvWithLength sock

679 theirAddress <- EndPointAddress . BS.concat <$> recvWithLength sock

680+ let ourAddress = encodeEndPointAddress transport ourEndPointId

623− ourEndPoint <- (^. localEndPointAt ourEndPointId) <$> readMVar tState

681+ ourEndPoint <- (^. localEndPointAt ourAddress) <$> readMVar (transportState transport)

624 return (theirAddress, ourEndPoint)

682 return (theirAddress, ourEndPoint)

625 case request of

683 case request of

626 Right (_, Nothing) -> do

684 Right (_, Nothing) -> do

627 sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]

685 sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]

628 N.sClose sock

686 N.sClose sock

629− Right (theirAddress, Just ourEndPoint) -> do

687+ Right (theirAddress, Just ourEndPoint) -> forkEndPointThread ourEndPoint $ do

630 (crossed, theirEndPoint) <- modifyMVar (localState ourEndPoint) $ \st ->

688 (crossed, theirEndPoint) <- modifyMVar (localState ourEndPoint) $ \st ->

631 case st ^. connectionTo theirAddress of

689 case st ^. connectionTo theirAddress of

632 Nothing -> do

690 Nothing -> do

@@ -643,14 +701,14 @@ handleConnectionRequest transport sock = do

643 N.sClose sock

701 N.sClose sock

644 else do

702 else do

645 let remoteConn = RemoteConnection { remoteSocket = sock

703 let remoteConn = RemoteConnection { remoteSocket = sock

646− , _remoteRefCount = 0

704+ , _remoteOutgoing = 0

705+ , _remoteIncoming = IntSet.empty

647 , sendOn = sendMany sock

706 , sendOn = sendMany sock

648 }

707 }

708+ -- TODO: this putMVar might block

649 putMVar (remoteState theirEndPoint) (RemoteEndPointValid remoteConn)

709 putMVar (remoteState theirEndPoint) (RemoteEndPointValid remoteConn)

650 sendMany sock [encodeInt32 ConnectionRequestAccepted]

710 sendMany sock [encodeInt32 ConnectionRequestAccepted]

651− -- TODO: we should store the tid as part of the endpoint state

652− tid <- forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)

711+ handleIncomingMessages (ourEndPoint, theirEndPoint)

653− return ()

654 Left _ -> do

712 Left _ -> do

655 -- Invalid request

713 -- Invalid request

656 N.sClose sock

714 N.sClose sock

@@ -665,59 +723,72 @@ handleConnectionRequest transport sock = do

665-- GUARANTEE: May change the remote endpoint to RemoteEndPointClosed state.

723-- GUARANTEE: May change the remote endpoint to RemoteEndPointClosed state.

666handleIncomingMessages :: EndPointPair -> IO ()

724handleIncomingMessages :: EndPointPair -> IO ()

667handleIncomingMessages (ourEndPoint, theirEndPoint) = do

725handleIncomingMessages (ourEndPoint, theirEndPoint) = do

668− -- We use an IORef rather than the state monad because when an exception occurs we

669− -- want to know the current list of open connections

670− openConnections <- newIORef IntSet.empty

671−

672 -- For efficiency sake we get the socket once and for all

726 -- For efficiency sake we get the socket once and for all

673 sock <- withMVar theirState $ \st ->

727 sock <- withMVar theirState $ \st ->

674 case st of

728 case st of

675 RemoteEndPointValid ep ->

729 RemoteEndPointValid ep ->

676 return (remoteSocket ep)

730 return (remoteSocket ep)

677− RemoteEndPointClosing ep _ ->

731+ RemoteEndPointClosing _ ep ->

678 return (remoteSocket ep)

732 return (remoteSocket ep)

679 _ ->

733 _ ->

680 error "handleIncomingMessages RELY violation"

734 error "handleIncomingMessages RELY violation"

681

735

682− tryIO $ go sock openConnections

736+ tryIO $ go sock

737+

738+ -- Normally, there would be no unclosed connections, but if the client

739+ -- exits prematurely there might be

683− unclosedConnections <- readIORef openConnections

740+ unclosedConnections <- modifyMVar theirState $ \st ->

741+ case st of

742+ RemoteEndPointInvalid _ ->

743+ error "handleIncomingMessages RELY violation"

744+ RemoteEndPointValid remoteConn ->

745+ return (RemoteEndPointValid . (remoteIncoming ^= IntSet.empty) $ remoteConn, remoteConn ^. remoteIncoming)

746+ RemoteEndPointClosing resolved remoteConn ->

747+ return (RemoteEndPointClosing resolved . (remoteIncoming ^= IntSet.empty) $ remoteConn, remoteConn ^. remoteIncoming)

748+ RemoteEndPointClosed ->

749+ return (st, IntSet.empty)

750+

751+ -- TODO: should we give a single ErrorEvent instead?

684− forM_ (map fromIntegral $ IntSet.elems unclosedConnections) $

752+ forM_ (map fromIntegral $ IntSet.elems unclosedConnections) $ \cid -> do

685− writeChan ourChannel . ConnectionClosed

753+ writeChan ourChannel $ ConnectionClosed cid

754+

755+ -- TODO: if the endpoint was not in closing state, should we make it so?

686 where

756 where

687 -- Dispatch

757 -- Dispatch

688− go :: N.Socket -> IORef IntSet -> IO ()

758+ go :: N.Socket -> IO ()

689− go sock openConnections = do

759+ go sock = do

690 connId <- recvInt32 sock

760 connId <- recvInt32 sock

691 if connId >= firstNonReservedConnectionId

761 if connId >= firstNonReservedConnectionId

692 then do

762 then do

693 readMessage sock connId

763 readMessage sock connId

694− go sock openConnections

764+ go sock

695 else do

765 else do

696 case tryToEnum (fromIntegral connId) of

766 case tryToEnum (fromIntegral connId) of

697 Just RequestConnectionId -> do

767 Just RequestConnectionId -> do

698− recvInt32 sock >>= createNewConnection openConnections

768+ recvInt32 sock >>= createNewConnection

699− go sock openConnections

769+ go sock

700 Just ControlResponse -> do

770 Just ControlResponse -> do

701 recvInt32 sock >>= readControlResponse sock

771 recvInt32 sock >>= readControlResponse sock

702− go sock openConnections

772+ go sock

703 Just CloseConnection -> do

773 Just CloseConnection -> do

704− recvInt32 sock >>= closeConnection openConnections

774+ recvInt32 sock >>= closeConnection

705− go sock openConnections

775+ go sock

706 Just CloseSocket -> do

776 Just CloseSocket -> do

707− closeSocket sock

777+ didClose <- closeSocket sock

708− go sock openConnections

778+ unless didClose $ go sock

709 Nothing ->

779 Nothing ->

710 -- Invalid control request, exit

780 -- Invalid control request, exit

711 hPutStrLn stderr "Warning: invalid control request"

781 hPutStrLn stderr "Warning: invalid control request"

712

782

713 -- Create a new connection

783 -- Create a new connection

714− createNewConnection :: IORef IntSet -> ControlRequestId -> IO ()

784+ createNewConnection :: ControlRequestId -> IO ()

715− createNewConnection openConnections reqId = do

785+ createNewConnection reqId = do

786+ newId <- getNextConnectionId ourEndPoint

716− modifyMVar_ theirState $ \st ->

787+ modifyMVar_ theirState $ \st -> do

717− case st of

788+ remoteConn <- case st of

718− RemoteEndPointValid ep ->

789+ RemoteEndPointValid remoteConn ->

719− return (RemoteEndPointValid . (remoteRefCount ^+= 1) $ ep)

790+ return (remoteIncoming ^%= IntSet.insert (fromIntegral newId) $ remoteConn)

720− RemoteEndPointClosing ep resolved -> do

791+ RemoteEndPointClosing resolved remoteConn -> do

721 -- If the endpoint is in closing state that means we send a

792 -- If the endpoint is in closing state that means we send a

722 -- CloseSocket request to the remote endpoint. If the remote

793 -- CloseSocket request to the remote endpoint. If the remote

723 -- endpoint replies with the request to create a new connection, it

794 -- endpoint replies with the request to create a new connection, it

@@ -725,19 +796,17 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do

725 -- ours. Either way, at this point we simply restore the endpoint

796 -- ours. Either way, at this point we simply restore the endpoint

726 -- to RemoteEndPointValid

797 -- to RemoteEndPointValid

727 putMVar resolved ()

798 putMVar resolved ()

728− return (RemoteEndPointValid . (remoteRefCount ^= 1) $ ep)

799+ return (remoteIncoming ^= IntSet.singleton newId $ remoteConn)

729 _ ->

800 _ ->

730 error "handleIncomingMessages RELY violation"

801 error "handleIncomingMessages RELY violation"

731− newId <- getNextConnectionId ourEndPoint

732− withMVar theirState $ \(RemoteEndPointValid remoteConn) ->

733 sendOn remoteConn ( encodeInt32 ControlResponse

802 sendOn remoteConn ( encodeInt32 ControlResponse

734 : encodeInt32 reqId

803 : encodeInt32 reqId

735 : prependLength [encodeInt32 newId]

804 : prependLength [encodeInt32 newId]

736 )

805 )

806+ -- We add the new connection ID to the list of open connections only once the

807+ -- endpoint has been notified of the new connection (sendOn may fail)

808+ return (RemoteEndPointValid remoteConn)

737 writeChan ourChannel (ConnectionOpened newId ReliableOrdered theirAddr)

809 writeChan ourChannel (ConnectionOpened newId ReliableOrdered theirAddr)

738− -- We add the new connection ID to the list of open connections only once the

739− -- endpoint has been notified of the new connection (sendOn may fail)

740− modifyIORef openConnections (IntSet.insert (fromIntegral newId))

741

810

742 -- Read a control response

811 -- Read a control response

743 readControlResponse :: N.Socket -> ControlRequestId -> IO ()

812 readControlResponse :: N.Socket -> ControlRequestId -> IO ()

@@ -753,22 +822,24 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do

753 putMVar mvar response

822 putMVar mvar response

754

823

755 -- Close a connection

824 -- Close a connection

756− closeConnection :: IORef IntSet -> ConnectionId -> IO ()

825+ closeConnection :: ConnectionId -> IO ()

757− closeConnection openConnections cid = do

826+ closeConnection cid = do

827+ -- TODO: we should check that this connection is in fact open

758 writeChan ourChannel (ConnectionClosed cid)

828 writeChan ourChannel (ConnectionClosed cid)

759− decreaseRemoteRefCount theirEndPoint

829+ modifyMVar_ theirState $ \(RemoteEndPointValid remoteConn) ->

760− modifyIORef openConnections (IntSet.delete (fromIntegral cid))

830+ return (RemoteEndPointValid . (remoteIncoming ^%= IntSet.delete cid) $ remoteConn)

831+ closeIfUnused theirEndPoint

761

832

762 -- Close the socket (if we don't have any outgoing connections)

833 -- Close the socket (if we don't have any outgoing connections)

763 closeSocket :: N.Socket -> IO Bool

834 closeSocket :: N.Socket -> IO Bool

764 closeSocket sock = do

835 closeSocket sock = do

765 didClose <- modifyMVar theirState $ \st ->

836 didClose <- modifyMVar theirState $ \st ->

766 case st of

837 case st of

767− RemoteEndPointValid remoteConn | remoteConn ^. remoteRefCount == 0 -> do

838+ RemoteEndPointValid remoteConn | remoteConn ^. remoteOutgoing == 0 -> do

768 sendOn remoteConn [encodeInt32 CloseSocket]

839 sendOn remoteConn [encodeInt32 CloseSocket]

769 N.sClose sock

840 N.sClose sock

770 return (RemoteEndPointClosed, True)

841 return (RemoteEndPointClosed, True)

771− RemoteEndPointClosing _ resolved -> do

842+ RemoteEndPointClosing resolved _ -> do

772 -- If the socket is already in closing state, then we have already

843 -- If the socket is already in closing state, then we have already

773 -- sent a CloseSocket request; hence, we don't need to do it again

844 -- sent a CloseSocket request; hence, we don't need to do it again

774 -- at this point

845 -- at this point

@@ -785,10 +856,12 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do

785 readMessage sock connId = recvWithLength sock >>= writeChan ourChannel . Received connId

856 readMessage sock connId = recvWithLength sock >>= writeChan ourChannel . Received connId

786

857

787 -- Arguments

858 -- Arguments

788− ourChannel = localChannel ourEndPoint

859+ ourChannel = localChannel ourEndPoint

789− ourState = localState ourEndPoint

860+ ourState = localState ourEndPoint

790− theirState = remoteState theirEndPoint

861+ theirState = remoteState theirEndPoint

791− theirAddr = remoteAddress theirEndPoint

862+ theirAddr = remoteAddress theirEndPoint

863+

864+

792

865

793-- | Get the next connection ID

866-- | Get the next connection ID

794getNextConnectionId :: LocalEndPoint -> IO ConnectionId

867getNextConnectionId :: LocalEndPoint -> IO ConnectionId

@@ -809,7 +882,7 @@ firstNonReservedConnectionId = 1024

809-- Lens definitions --

882-- Lens definitions --

810--------------------------------------------------------------------------------

883--------------------------------------------------------------------------------

811

884

812−localEndPoints :: Lens TransportState (IntMap LocalEndPoint)

885+localEndPoints :: Lens TransportState (Map EndPointAddress LocalEndPoint)

813localEndPoints = lens _localEndPoints (\es st -> st { _localEndPoints = es })

886localEndPoints = lens _localEndPoints (\es st -> st { _localEndPoints = es })

814

887

815acceptNewConnections :: Lens TransportState Bool

888acceptNewConnections :: Lens TransportState Bool

@@ -827,8 +900,11 @@ nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = c

827localConnections :: Lens LocalState (Map EndPointAddress RemoteEndPoint)

900localConnections :: Lens LocalState (Map EndPointAddress RemoteEndPoint)

828localConnections = lens _localConnections (\es st -> st { _localConnections = es })

901localConnections = lens _localConnections (\es st -> st { _localConnections = es })

829

902

903+internalThreads :: Lens LocalState [ThreadId]

904+internalThreads = lens _internalThreads (\ts st -> st { _internalThreads = ts })

905+

830−localEndPointAt :: EndPointId -> Lens TransportState (Maybe LocalEndPoint)

906+localEndPointAt :: EndPointAddress -> Lens TransportState (Maybe LocalEndPoint)

831−localEndPointAt ix = localEndPoints >>> intMapLens (fromIntegral ix)

907+localEndPointAt addr = localEndPoints >>> mapLens addr

832

908

833pendingCtrlRequestsAt :: ControlRequestId -> Lens LocalState (Maybe (MVar [ByteString]))

909pendingCtrlRequestsAt :: ControlRequestId -> Lens LocalState (Maybe (MVar [ByteString]))

834pendingCtrlRequestsAt ix = pendingCtrlRequests >>> intMapLens (fromIntegral ix)

910pendingCtrlRequestsAt ix = pendingCtrlRequests >>> intMapLens (fromIntegral ix)

@@ -836,5 +912,8 @@ pendingCtrlRequestsAt ix = pendingCtrlRequests >>> intMapLens (fromIntegral ix)

836connectionTo :: EndPointAddress -> Lens LocalState (Maybe RemoteEndPoint)

912connectionTo :: EndPointAddress -> Lens LocalState (Maybe RemoteEndPoint)

837connectionTo addr = localConnections >>> mapLens addr

913connectionTo addr = localConnections >>> mapLens addr

838

914

839−remoteRefCount :: Lens RemoteConnection Int

915+remoteOutgoing :: Lens RemoteConnection Int

840−remoteRefCount = lens _remoteRefCount (\rc conn -> conn { _remoteRefCount = rc })

916+remoteOutgoing = lens _remoteOutgoing (\cs conn -> conn { _remoteOutgoing = cs })

917+

918+remoteIncoming :: Lens RemoteConnection IntSet

919+remoteIncoming = lens _remoteIncoming (\cs conn -> conn { _remoteIncoming = cs })

network-transport-2/tests/TestAuxiliary.hs+13−4

@@ -1,7 +1,8 @@

1−module TestAuxiliary (runTest, forkTry) where

1+module TestAuxiliary (runTest, runTests, forkTry) where

3import Prelude hiding (catch)

4import Control.Concurrent (myThreadId, forkIO, ThreadId, throwTo)

5+import Control.Monad (liftM2, unless)

5import Control.Exception (SomeException, try, catch)

6import Control.Exception (SomeException, try, catch)

6import System.Timeout (timeout)

7import System.Timeout (timeout)

7import System.IO (stdout, hFlush)

8import System.IO (stdout, hFlush)

@@ -19,7 +20,7 @@ forkTry p = do

19 forkIO $ catch p (\e -> throwTo tid (e :: SomeException))

20 forkIO $ catch p (\e -> throwTo tid (e :: SomeException))

21-- | Run the given test, catching timeouts and exceptions

22-- | Run the given test, catching timeouts and exceptions

22−runTest :: String -> IO () -> IO ()

23+runTest :: String -> IO () -> IO Bool

23runTest description test = do

24runTest description test = do

24 putStr $ "Running " ++ show description ++ ": "

25 putStr $ "Running " ++ show description ++ ": "

25 hFlush stdout

26 hFlush stdout

@@ -29,15 +30,23 @@ runTest description test = do

29 Right Nothing -> failed $ "(timeout)"

30 Right Nothing -> failed $ "(timeout)"

30 Right (Just ()) -> ok

31 Right (Just ()) -> ok

31 where

32 where

32− failed :: String -> IO ()

33+ failed :: String -> IO Bool

33 failed err = do

34 failed err = do

34 setSGR [SetColor Foreground Vivid Red]

35 setSGR [SetColor Foreground Vivid Red]

35 putStr "failed "

36 putStr "failed "

36 setSGR [Reset]

37 setSGR [Reset]

37 putStrLn err

38 putStrLn err

39+ return False

39− ok :: IO ()

41+ ok :: IO Bool

40 ok = do

42 ok = do

41 setSGR [SetColor Foreground Vivid Green]

43 setSGR [SetColor Foreground Vivid Green]

42 putStrLn "ok"

44 putStrLn "ok"

43 setSGR [Reset]

45 setSGR [Reset]

46+ return True

47+

48+-- | Run a bunch of tests and throw an exception if any fails

49+runTests :: [(String, IO ())] -> IO ()

50+runTests tests = do

51+ success <- foldr (liftM2 (&&) . uncurry runTest) (return True) $ tests

52+ unless success $ fail "Some tests failed"

network-transport-2/tests/TestTCP.hs+10−9

@@ -1,8 +1,8 @@

1module Main where

3import Prelude hiding (catch)

4−import TestTransport

4+import TestTransport (testTransport)

5−import TestAuxiliary (forkTry, runTest)

5+import TestAuxiliary (forkTry, runTests)

6import Network.Transport

7import Network.Transport.TCP (createTransport)

8import Data.Int (Int32)

@@ -184,11 +184,11 @@ testIgnoreCloseSocket = do

184 tlog "Ignoring it, requesting another connection"

185 let reqId' = 1 :: Int32

186 sendMany sock [encodeInt32 RequestConnectionId, encodeInt32 reqId']

187− replicateM 4 $ recvInt32 sock :: IO [Int32]

187+ response' <- replicateM 4 $ recvInt32 sock :: IO [Int32]

188

189 -- Close it again

190 tlog "Closing connection"

191− sendMany sock [encodeInt32 CloseConnection, encodeInt32 (response !! 3)]

191+ sendMany sock [encodeInt32 CloseConnection, encodeInt32 (response' !! 3)]

192

193 -- We now get a CloseSocket again, and this time we heed it

194 tlog "Waiting for second CloseSocket request"

@@ -289,10 +289,11 @@ testBlockAfterCloseSocket = do

289

290main :: IO ()

291main = do

292− runTest "EarlyDisconnect" testEarlyDisconnect

292+ runTests [ ("EarlyDisconnect", testEarlyDisconnect)

293− runTest "InvalidAddress" testInvalidAddress

293+ , ("InvalidAddress", testInvalidAddress)

294− runTest "InvalidConnect" testInvalidConnect

294+ , ("InvalidConnect", testInvalidConnect)

295− runTest "IgnoreCloseSocket" testIgnoreCloseSocket

295+ , ("IgnoreCloseSocket", testIgnoreCloseSocket)

296− runTest "BlockAfterCloseSocket" testBlockAfterCloseSocket

296+ , ("BlockAfterCloseSocket", testBlockAfterCloseSocket)

297+ ]

297 Right transport <- createTransport "127.0.0.1" "8080"

298 Right transport <- createTransport "127.0.0.1" "8080"

298 testTransport transport

299 testTransport transport

network-transport-2/tests/TestTransport.hs+35−27

@@ -1,11 +1,9 @@

1module TestTransport where

3import Prelude hiding (catch)

4−import TestAuxiliary (forkTry, runTest)

4+import TestAuxiliary (forkTry, runTests)

5import Control.Concurrent.MVar (newEmptyMVar, takeMVar, putMVar, readMVar)

6import Control.Monad (replicateM, replicateM_, when, guard, forM_)

7−import Control.Monad.Reader (ReaderT, runReaderT, ask)

8−import Control.Monad.IO.Class (liftIO)

9import Control.Exception (evaluate, throw, catch, SomeException)

7import Control.Exception (evaluate, throw, catch, SomeException)

10import Control.Applicative ((<$>))

8import Control.Applicative ((<$>))

11import Network.Transport

9import Network.Transport

@@ -16,7 +14,7 @@ import Data.ByteString.Char8 (pack)

16import Data.Map (Map)

14import Data.Map (Map)

17import qualified Data.Map as Map (empty, insert, delete, findWithDefault, adjust, null, toList, map)

15import qualified Data.Map as Map (empty, insert, delete, findWithDefault, adjust, null, toList, map)

19−-- Server that echoes messages straight back to the origin endpoint.

17+-- | Server that echoes messages straight back to the origin endpoint.

20echoServer :: EndPoint -> IO ()

18echoServer :: EndPoint -> IO ()

21echoServer endpoint = do

19echoServer endpoint = do

22 go Map.empty

20 go Map.empty

@@ -26,18 +24,23 @@ echoServer endpoint = do

26 event <- receive endpoint

24 event <- receive endpoint

27 case event of

25 case event of

28 ConnectionOpened cid rel addr -> do

26 ConnectionOpened cid rel addr -> do

27+ tlog $ "Opened new conncetion " ++ show cid

29 Right conn <- connect endpoint addr rel

28 Right conn <- connect endpoint addr rel

30 go (Map.insert cid conn cs)

29 go (Map.insert cid conn cs)

31 Received cid payload -> do

30 Received cid payload -> do

32 send (Map.findWithDefault (error $ "Received: Invalid cid " ++ show cid) cid cs) payload

31 send (Map.findWithDefault (error $ "Received: Invalid cid " ++ show cid) cid cs) payload

33 go cs

32 go cs

34 ConnectionClosed cid -> do

33 ConnectionClosed cid -> do

34+ tlog $ "Close connection " ++ show cid

35 close (Map.findWithDefault (error $ "ConnectionClosed: Invalid cid " ++ show cid) cid cs)

36 go (Map.delete cid cs)

37 ReceivedMulticast _ _ ->

38 -- Ignore

39 go cs

40+ ErrorEvent _ _ ->

41+ fail (show event)

43+-- | Wait for an event, throw an exception if the given predicate returns false

41expect :: EndPoint -> (Event -> (Bool, a)) -> IO a

44expect :: EndPoint -> (Event -> (Bool, a)) -> IO a

42expect endpoint predicate = do

45expect endpoint predicate = do

43 event <- receive endpoint

46 event <- receive endpoint

@@ -49,10 +52,12 @@ expect endpoint predicate = do

49 fail "Unexpected event"

52 fail "Unexpected event"

50 Right (True, a) -> return a

53 Right (True, a) -> return a

55+-- | Like 'expect' but without a return value

52expect' :: EndPoint -> (Event -> Bool) -> IO ()

56expect' :: EndPoint -> (Event -> Bool) -> IO ()

53expect' endpoint predicate =

57expect' endpoint predicate =

54 expect endpoint (\event -> (predicate event, ()))

58 expect endpoint (\event -> (predicate event, ()))

60+-- | Ping client used in a few tests

56ping :: EndPoint -> EndPointAddress -> Int -> ByteString -> IO ()

61ping :: EndPoint -> EndPointAddress -> Int -> ByteString -> IO ()

57ping endpoint server numPings msg = do

62ping endpoint server numPings msg = do

58 -- Open connection to the server

63 -- Open connection to the server

@@ -80,7 +85,7 @@ ping endpoint server numPings msg = do

80 -- Done

85 -- Done

81 tlog "Ping client done"

86 tlog "Ping client done"

83−-- Basic ping test

88+-- | Basic ping test

84testPingPong :: Transport -> Int -> IO ()

89testPingPong :: Transport -> Int -> IO ()

85testPingPong transport numPings = do

90testPingPong transport numPings = do

86 tlog "Starting ping pong test"

91 tlog "Starting ping pong test"

@@ -96,7 +101,7 @@ testPingPong transport numPings = do

101

97 takeMVar result

102 takeMVar result

103

99−-- Test that endpoints don't get confused

104+-- | Test that endpoints don't get confused

100testEndPoints :: Transport -> Int -> IO ()

105testEndPoints :: Transport -> Int -> IO ()

101testEndPoints transport numPings = do

106testEndPoints transport numPings = do

102 server <- spawn transport echoServer

107 server <- spawn transport echoServer

@@ -150,7 +155,7 @@ testConnections transport numPings = do

150

155

151 takeMVar result

156 takeMVar result

152

157

153−-- Test that closing one connection does not close the other

158+-- | Test that closing one connection does not close the other

154testCloseOneConnection :: Transport -> Int -> IO ()

159testCloseOneConnection :: Transport -> Int -> IO ()

155testCloseOneConnection transport numPings = do

160testCloseOneConnection transport numPings = do

156 server <- spawn transport echoServer

161 server <- spawn transport echoServer

@@ -190,7 +195,7 @@ testCloseOneConnection transport numPings = do

190

195

191 takeMVar result

196 takeMVar result

192

197

193−-- Test that if A connects to B and B connects to A, B can still send to A after

198+-- | Test that if A connects to B and B connects to A, B can still send to A after

194-- A closes its connection to B (for instance, in the TCP transport, the socket pair

199-- A closes its connection to B (for instance, in the TCP transport, the socket pair

195-- connecting A and B should not yet be closed).

200-- connecting A and B should not yet be closed).

196testCloseOneDirection :: Transport -> Int -> IO ()

201testCloseOneDirection :: Transport -> Int -> IO ()

@@ -274,6 +279,7 @@ testCloseOneDirection transport numPings = do

274

279

275 mapM_ takeMVar [doneA, doneB]

280 mapM_ takeMVar [doneA, doneB]

276

281

282+-- | Collect a given number of events and order them by connection ID

277collect :: EndPoint -> Int -> IO [(ConnectionId, [[ByteString]])]

283collect :: EndPoint -> Int -> IO [(ConnectionId, [[ByteString]])]

278collect endPoint numEvents = go numEvents Map.empty Map.empty

284collect endPoint numEvents = go numEvents Map.empty Map.empty

279 where

285 where

@@ -292,8 +298,10 @@ collect endPoint numEvents = go numEvents Map.empty Map.empty

292 go (n - 1) (Map.adjust (msg :) cid open) closed

298 go (n - 1) (Map.adjust (msg :) cid open) closed

293 ReceivedMulticast _ _ ->

299 ReceivedMulticast _ _ ->

294 fail "Unexpected multicast"

300 fail "Unexpected multicast"

301+ ErrorEvent _ _ ->

302+ fail "Unexpected error"

295

303

296−-- Open connection, close it, then reopen it

304+-- | Open connection, close it, then reopen it

297-- (In the TCP transport this means the socket will be closed, then reopened)

305-- (In the TCP transport this means the socket will be closed, then reopened)

298--

306--

299-- Note that B cannot expect to receive all of A's messages on the first connection

307-- Note that B cannot expect to receive all of A's messages on the first connection

@@ -349,7 +357,7 @@ testCloseReopen transport numPings = do

349

357

350 takeMVar doneB

358 takeMVar doneB

351

359

352−-- Test lots of parallel connection attempts

360+-- | Test lots of parallel connection attempts

353testParallelConnects :: Transport -> Int -> IO ()

361testParallelConnects :: Transport -> Int -> IO ()

354testParallelConnects transport numPings = do

362testParallelConnects transport numPings = do

355 server <- spawn transport echoServer

363 server <- spawn transport echoServer

@@ -373,7 +381,7 @@ testParallelConnects transport numPings = do

373

381

374 takeMVar done

382 takeMVar done

375

383

376−-- Test that sending on a closed connection gives an error

384+-- | Test that sending on a closed connection gives an error

377testSendAfterClose :: Transport -> Int -> IO ()

385testSendAfterClose :: Transport -> Int -> IO ()

378testSendAfterClose transport _ = do

386testSendAfterClose transport _ = do

379 server <- spawn transport echoServer

387 server <- spawn transport echoServer

@@ -401,7 +409,7 @@ testSendAfterClose transport _ = do

401

409

402 takeMVar clientDone

410 takeMVar clientDone

403

411

404−-- Test that closing the same connection twice has no effect

412+-- | Test that closing the same connection twice has no effect

405testCloseTwice :: Transport -> Int -> IO ()

413testCloseTwice :: Transport -> Int -> IO ()

406testCloseTwice transport _ = do

414testCloseTwice transport _ = do

407 server <- spawn transport echoServer

415 server <- spawn transport echoServer

@@ -433,6 +441,7 @@ testCloseTwice transport _ = do

433

441

434 takeMVar clientDone

442 takeMVar clientDone

435

443

444+-- | Test that we can connect an endpoint to itself

436testConnectToSelf :: Transport -> Int -> IO ()

445testConnectToSelf :: Transport -> Int -> IO ()

437testConnectToSelf transport numPings = do

446testConnectToSelf transport numPings = do

438 done <- newEmptyMVar

447 done <- newEmptyMVar

@@ -471,21 +480,20 @@ testConnectToSelf transport numPings = do

471

480

472 takeMVar done

481 takeMVar done

473

482

474−runTransportTest :: String -> (Transport -> Int -> IO ()) -> ReaderT (Transport, Int) IO ()

475−runTransportTest description test = do

476− (transport, numPings) <- ask

477− liftIO $ runTest description (test transport numPings)

478

483

479-- Transport tests

484-- Transport tests

480testTransport :: Transport -> IO ()

485testTransport :: Transport -> IO ()

481−testTransport transport = flip runReaderT (transport, 10000) $ do

486+testTransport transport = runTests

482− runTransportTest "PingPong" testPingPong

487+ [ ("PingPong", testPingPong transport numPings)

483− runTransportTest "EndPoints" testEndPoints

488+ , ("EndPoints", testEndPoints transport numPings)

484− runTransportTest "Connections" testConnections

489+ , ("Connections", testConnections transport numPings)

485− runTransportTest "CloseOneConnection" testCloseOneConnection

490+ , ("CloseOneConnection", testCloseOneConnection transport numPings)

486− runTransportTest "CloseOneDirection" testCloseOneDirection

491+ , ("CloseOneDirection", testCloseOneDirection transport numPings)

487− runTransportTest "CloseReopen" testCloseReopen

492+ , ("CloseReopen", testCloseReopen transport numPings)

488− runTransportTest "ParallelConnects" testParallelConnects

493+ , ("ParallelConnects", testParallelConnects transport numPings)

489− runTransportTest "SendAfterClose" testSendAfterClose

494+ , ("SendAfterClose", testSendAfterClose transport numPings)

490− runTransportTest "CloseTwice" testCloseTwice

495+ , ("CloseTwice", testCloseTwice transport numPings)

491− runTransportTest "ConnectToSelf" testConnectToSelf

496+ , ("ConnectToSelf", testConnectToSelf transport numPings)

497+ ]

498+ where

499+ numPings = 10000