9679355 · distributed-process

git.iko.soy/distributed-process

on master help sign in

Implement and test closeEndPoint

9679355Edsko de Vries14 years ago

network-transport-2/src/Network/Transport/TCP.hs+18−19

compare with master

unifiedsplit

@@ -58,13 +58,11 @@ import Control.Concurrent.MVar ( MVar

58 , putMVar

59 , newEmptyMVar

60 , withMVar

61− , tryTakeMVar

62 )

61 )

63import Control.Category ((>>>))

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

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

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

65−import Control.Monad (forM_, void, when, mzero, unless)

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

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

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

67−import Control.Monad.Trans.Maybe (MaybeT, runMaybeT)

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

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

69import Control.Exception (IOException, SomeException)

67import Control.Exception (IOException, SomeException)

70import Data.IORef (IORef, newIORef, writeIORef, readIORef)

68import Data.IORef (IORef, newIORef, writeIORef, readIORef)

@@ -75,7 +73,7 @@ import Data.Int (Int32)

75import Data.IntMap (IntMap)

73import Data.IntMap (IntMap)

76import qualified Data.IntMap as IntMap (empty)

74import qualified Data.IntMap as IntMap (empty)

77import Data.IntSet (IntSet)

75import Data.IntSet (IntSet)

78−import qualified Data.IntSet as IntSet (empty, insert, elems, singleton, null, delete)

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

79import Data.Map (Map)

77import Data.Map (Map)

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

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

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

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

@@ -396,28 +394,30 @@ apiCloseEndPoint transport ourEndPoint = do

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

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

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

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

398 ourState <- takeMVar (localState ourEndPoint)

396 ourState <- takeMVar (localState ourEndPoint)

399− --incoming <- IntSet.elems <$> takeMVar (localIncoming ourEndPoint)

397+ -- Close all endpoints and kill all threads

400− let incoming = undefined

398+ incoming <- IntSet.elems . IntSet.unions <$> (forM (Map.elems $ ourState ^. localConnections) $ tryCloseRemoteSocket)

401− -- Kill all threads associated with this endpoint and close all sockets

402 forM_ (ourState ^. internalThreads) killThread

399 forM_ (ourState ^. internalThreads) killThread

403− forM_ (Map.elems $ ourState ^. localConnections) $ \theirEndPoint ->

404− runMaybeT $ remoteSocketOf theirEndPoint >>= liftIO . N.sClose

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

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

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

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

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

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

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

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

409 writeChan (localChannel ourEndPoint) (ErrorEvent $ ErrorEventEndPointClosed incoming)

404 writeChan (localChannel ourEndPoint) (ErrorEvent $ ErrorEventEndPointClosed incoming)

410 where

405 where

411− remoteSocketOf :: RemoteEndPoint -> MaybeT IO N.Socket

406+ -- Close the remote socket and return the set of all incoming connections

407+ tryCloseRemoteSocket :: RemoteEndPoint -> IO IntSet

412− remoteSocketOf theirEndPoint = do

408+ tryCloseRemoteSocket theirEndPoint = do

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

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

414− mRemoteState <- liftIO $ tryTakeMVar (remoteState theirEndPoint)

410+ st <- takeMVar (remoteState theirEndPoint)

411+ -- We make an attempt to close the connection nicely (by sending a CloseSocket first)

415− case mRemoteState of

412+ case st of

416− Nothing -> mzero

417− Just (RemoteEndPointInvalid _) -> mzero

413+ RemoteEndPointInvalid _ -> do return IntSet.empty

418− Just (RemoteEndPointValid conn) -> return $ remoteSocket conn

414+ RemoteEndPointValid conn -> do tryIO $ do

415+ sendOn conn [encodeInt32 CloseSocket]

416+ N.sClose (remoteSocket conn)

417+ return (conn ^. remoteIncoming)

419− Just (RemoteEndPointClosing _ conn) -> return $ remoteSocket conn

418+ RemoteEndPointClosing _ conn -> do tryIO $ N.sClose (remoteSocket conn)

419+ return IntSet.empty

420− Just RemoteEndPointClosed -> mzero

420+ RemoteEndPointClosed -> do return IntSet.empty

421

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

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

@@ -906,7 +906,6 @@ handleIncomingMessages (ourEndPoint, theirEndPoint) = do

906 Just unclosedConnections -> writeChan ourChannel . ErrorEvent $

907 ErrorEventConnectionLost (remoteAddress theirEndPoint) (IntSet.elems unclosedConnections)

908

909−

910-- | Get the next connection ID

909-- | Get the next connection ID

911getNextConnectionId :: LocalEndPoint -> IO ConnectionId

910getNextConnectionId :: LocalEndPoint -> IO ConnectionId

912getNextConnectionId ourEndpoint =

911getNextConnectionId ourEndpoint =

network-transport-2/tests/TestTCP.hs−1

@@ -498,4 +498,3 @@ main = do

498 ]

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

500 testTransport transport

501− return ()

network-transport-2/tests/TestTransport.hs+78−1

@@ -512,6 +512,82 @@ testConnectToSelfTwice transport numPings = do

512

513 takeMVar done

514

515+-- | Test various aspects of 'closeEndPoint'

516+testCloseEndPoint :: Transport -> Int -> IO ()

517+testCloseEndPoint transport _ = do

518+ serverDone <- newEmptyMVar

519+ clientDone <- newEmptyMVar

520+ clientAddr1 <- newEmptyMVar

521+ clientAddr2 <- newEmptyMVar

522+ serverAddr <- newEmptyMVar

523+

524+ -- Server

525+ forkTry $ do

526+ Right endpoint <- newEndPoint transport

527+ putMVar serverAddr (address endpoint)

528+

529+ -- First test (see client)

530+ do

531+ theirAddr <- readMVar clientAddr1

532+ ConnectionOpened cid ReliableOrdered addr <- receive endpoint ; True <- return $ addr == theirAddr

533+ ConnectionClosed cid' <- receive endpoint ; True <- return $ cid == cid'

534+ return ()

535+

536+ -- Second test

537+ do

538+ theirAddr <- readMVar clientAddr2

539+

540+ ConnectionOpened cid ReliableOrdered addr <- receive endpoint ; True <- return $ addr == theirAddr

541+ Received cid' ["ping"] <- receive endpoint ; True <- return $ cid == cid'

542+

543+ Right conn <- connect endpoint theirAddr ReliableOrdered

544+ send conn ["pong"]

545+

546+ ConnectionClosed cid'' <- receive endpoint ; True <- return $ cid == cid''

547+ ErrorEvent (ErrorEventConnectionLost addr' []) <- receive endpoint ; True <- return $ addr' == theirAddr

548+

549+ Left (FailedWith SendConnectionClosed _) <- send conn ["pong2"]

550+

551+ return ()

552+

553+ putMVar serverDone ()

554+

555+ -- Client

556+ forkTry $ do

557+ theirAddr <- readMVar serverAddr

558+

559+ -- First test: close endpoint with one outgoing but no incoming connections

560+ do

561+ Right endpoint <- newEndPoint transport

562+ putMVar clientAddr1 (address endpoint)

563+

564+ -- Connect to the server, then close the endpoint without disconnecting explicitly

565+ Right _ <- connect endpoint theirAddr ReliableOrdered

566+ closeEndPoint endpoint

567+ ErrorEvent (ErrorEventEndPointClosed []) <- receive endpoint

568+ return ()

569+

570+ -- Second test: close endpoint with one outgoing and one incoming connection

571+ do

572+ Right endpoint <- newEndPoint transport

573+ putMVar clientAddr2 (address endpoint)

574+

575+ Right conn <- connect endpoint theirAddr ReliableOrdered

576+ send conn ["ping"]

577+

578+ -- Reply from the server

579+ ConnectionOpened cid ReliableOrdered addr <- receive endpoint ; True <- return $ addr == theirAddr

580+ Received cid' ["pong"] <- receive endpoint ; True <- return $ cid == cid'

581+

582+ -- Close the endpoint

583+ closeEndPoint endpoint

584+ ErrorEvent (ErrorEventEndPointClosed [cid'']) <- receive endpoint; True <- return $ cid == cid''

585+ return ()

586+

587+ putMVar clientDone ()

588+

589+ mapM_ takeMVar [serverDone, clientDone]

590+

515-- Transport tests

591-- Transport tests

516testTransport :: Transport -> IO ()

592testTransport :: Transport -> IO ()

517testTransport transport = runTests

593testTransport transport = runTests

@@ -525,7 +601,8 @@ testTransport transport = runTests

525 , ("SendAfterClose", testSendAfterClose transport numPings)

601 , ("SendAfterClose", testSendAfterClose transport numPings)

526 , ("CloseTwice", testCloseTwice transport numPings)

602 , ("CloseTwice", testCloseTwice transport numPings)

527 , ("ConnectToSelf", testConnectToSelf transport numPings)

603 , ("ConnectToSelf", testConnectToSelf transport numPings)

528− , ("ConnectToSelfTwice", testConnectToSelfTwice transport numPings)

604+ , ("ConnectToSelfTwice", testConnectToSelfTwice transport numPings)

605+ , ("CloseEndPoint", testCloseEndPoint transport numPings)

529 ]

606 ]

530 where

607 where

531 numPings = 10000 :: Int

608 numPings = 10000 :: Int