76d810b · distributed-process

git.iko.soy/distributed-process

on master help sign in

Support self-connections in the TCP transport

76d810bEdsko de Vries14 years ago

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

compare with master

unifiedsplit

@@ -23,7 +23,7 @@ import Data.ByteString (ByteString)

23import qualified Data.ByteString as BS (length)

24import qualified Data.ByteString.Internal as BSI (unsafeCreate, toForeignPtr, inlinePerformIO)

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

26−import Control.Monad (MonadPlus, guard)

26+import Control.Monad (MonadPlus, guard, when)

27import Control.Monad.IO.Class (MonadIO, liftIO)

28import Control.Monad.Error (MonadError, throwError)

29import Control.Exception (IOException, catch, try)

@@ -90,10 +90,10 @@ tryIO :: IO a -> IO (Either IOException a)

90tryIO = try

92-- | Logging (for debugging)

93−tlog :: String -> IO ()

93+tlog :: MonadIO m => String -> m ()

94tlog _ = return ()

95{-

96−tlog msg = do

96+tlog msg = liftIO $ do

97 tid <- myThreadId

98 putStrLn $ show tid ++ ": " ++ msg

99-}

@@ -108,4 +108,4 @@ whileM cond body = go

108 where

109 go = do

110 b <- cond

111− if b then body >> go else return ()

111+ when b $ body >> go

network-transport-2/src/Network/Transport/TCP.hs+40−9

@@ -254,6 +254,37 @@ apiConnect :: LocalEndPoint -- ^ Local end point

254 -> EndPointAddress -- ^ Remote address

255 -> Reliability -- ^ Reliability (ignored)

256 -> IO (Either (FailedWith ConnectErrorCode) Connection)

257+-- Special case for self-connection

258+apiConnect ourEndPoint theirAddress _ | ourEndPoint ^. localAddress == theirAddress = do

259+ connAlive <- newMVar True

260+ connId <- modifyMVar (ourEndPoint ^. localState) $ \st -> do

261+ let connId = st ^. nextConnectionId

262+ return (nextConnectionId ^= connId + 1 $ st, connId)

263+

264+ let ourChan :: Chan Event

265+ ourChan = ourEndPoint ^. localChannel

266+

267+ selfSend :: [ByteString] -> IO (Either (FailedWith SendErrorCode) ())

268+ selfSend msg = do

269+ modifyMVar connAlive $ \alive ->

270+ if alive

271+ then do

272+ writeChan ourChan (Received connId msg)

273+ return (alive, Right ())

274+ else

275+ return (alive, Left (FailedWith SendConnectionClosed "Connection closed"))

276+

277+ selfClose :: IO ()

278+ selfClose = do

279+ modifyMVar_ connAlive $ \alive -> do

280+ when alive $ writeChan ourChan (ConnectionClosed connId)

281+ return False

282+

283+ writeChan ourChan (ConnectionOpened connId ReliableOrdered theirAddress)

284+ return . Right $ Connection { send = selfSend

285+ , close = selfClose

286+ }

287+-- Regular case

257apiConnect ourEndPoint theirAddress _ = runErrorT $ do

288apiConnect ourEndPoint theirAddress _ = runErrorT $ do

258 theirEndPoint <- findRemoteEndPoint ourEndPoint theirAddress

289 theirEndPoint <- findRemoteEndPoint ourEndPoint theirAddress

259 connId <- requestNewConnection (ourEndPoint, theirEndPoint)

290 connId <- requestNewConnection (ourEndPoint, theirEndPoint)

@@ -325,14 +356,14 @@ findRemoteEndPoint ourEndPoint theirAddress = ErrorT $ go

325 return (connectionTo theirAddress ^= Just (Nothing, theirMVar) $ st, theirMVar)

356 return (connectionTo theirAddress ^= Just (Nothing, theirMVar) $ st, theirMVar)

326 theirState <- readMVar endPoint

357 theirState <- readMVar endPoint

327 case theirState of

358 case theirState of

328− EndPointInvalid err ->

359+ EndPointInvalid err -> do

329 -- We remove the endpoint from the state again because the next call

360 -- We remove the endpoint from the state again because the next call

330 -- to 'connect' might give a different result

361 -- to 'connect' might give a different result

331 modifyMVar state $ \st ->

362 modifyMVar state $ \st ->

332 return (connectionTo theirAddress ^= Nothing $ st, Left err)

363 return (connectionTo theirAddress ^= Nothing $ st, Left err)

333− EndPointValid _ ep ->

364+ EndPointValid _ ep -> do

334 return $ Right ep

365 return $ Right ep

335− EndPointClosed ->

366+ EndPointClosed -> do

336 go

367 go

337

368

338 state :: MVar LocalState

369 state :: MVar LocalState

@@ -348,9 +379,9 @@ connectToRemoteEndPoint ourEndPoint theirMVar theirAddress = do

348 Right (ConnectionRequestEndPointInvalid, sock) -> do

379 Right (ConnectionRequestEndPointInvalid, sock) -> do

349 putMVar theirMVar (EndPointInvalid (invalidAddress "Invalid endpoint"))

380 putMVar theirMVar (EndPointInvalid (invalidAddress "Invalid endpoint"))

350 N.sClose sock

381 N.sClose sock

351− Right (ConnectionRequestCrossed, sock) ->

382+ Right (ConnectionRequestCrossed, sock) -> do

352 N.sClose sock

383 N.sClose sock

353− Left err ->

384+ Left err -> do

354 -- TODO: this isn't quite right; socket *might* have been opened here,

385 -- TODO: this isn't quite right; socket *might* have been opened here,

355 -- in which case we should close it

386 -- in which case we should close it

356 putMVar theirMVar (EndPointInvalid err)

387 putMVar theirMVar (EndPointInvalid err)

@@ -421,12 +452,12 @@ sendRemoteRequest (ourEndPoint, theirEndPoint) header = do

421increaseRemoteRefCount :: Conn -> IO ()

452increaseRemoteRefCount :: Conn -> IO ()

422increaseRemoteRefCount (ourEndPoint, theirEndPoint) = do

453increaseRemoteRefCount (ourEndPoint, theirEndPoint) = do

423 let theirAddr = theirEndPoint ^. remoteAddress

454 let theirAddr = theirEndPoint ^. remoteAddress

424− state = ourEndPoint ^. localState

455+ ourState = ourEndPoint ^. localState

425 -- If this endpoint was closing, cancel that

456 -- If this endpoint was closing, cancel that

426− theirMVar <- modifyMVar state $ \st -> do

457+ theirMVar <- modifyMVar ourState $ \st -> do

427 let Just (isClosing, theirMVar) = st ^. connectionTo theirAddr

458 let Just (isClosing, theirMVar) = st ^. connectionTo theirAddr

428 case isClosing of

459 case isClosing of

429− Nothing ->

460+ Nothing -> do

430 return (st, theirMVar)

461 return (st, theirMVar)

431 Just ep -> do

462 Just ep -> do

432 putMVar theirMVar (EndPointValid 0 ep)

463 putMVar theirMVar (EndPointValid 0 ep)

@@ -475,7 +506,7 @@ handleConnectionRequest transport sock = do

475 mvar <- newEmptyMVar

506 mvar <- newEmptyMVar

476 return (connectionTo theirAddress ^= Just (Nothing, mvar) $ st, (False, mvar))

507 return (connectionTo theirAddress ^= Just (Nothing, mvar) $ st, (False, mvar))

477 Just (_, mvar) ->

508 Just (_, mvar) ->

478− return (st, (ourEndPoint ^. localAddress <= theirAddress, mvar))

509+ return (st, (ourEndPoint ^. localAddress < theirAddress, mvar))

479 if crossed

510 if crossed

480 then do

511 then do

481 sendMany sock [encodeInt32 ConnectionRequestCrossed]

512 sendMany sock [encodeInt32 ConnectionRequestCrossed]

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

@@ -230,8 +230,7 @@ testBlockAfterCloseSocket = do

230 -- At this point the server will have sent a CloseSocket request to the

231 -- client, and must block until the client responds

232 tlog "Server waiting to connect to the client.."

233− Right conn <- readMVar clientAddr >>= \addr -> connect endpoint addr ReliableOrdered

233+ Right _ <- readMVar clientAddr >>= \addr -> connect endpoint addr ReliableOrdered

234−

235

234

236 tlog "Server waiting.."

235 tlog "Server waiting.."

237

236

network-transport-2/tests/TestTransport.hs+41−2

@@ -274,8 +274,8 @@ testCloseOneDirection transport numPings = do

274

275 mapM_ takeMVar [doneA, doneB]

276

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

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

278−collect endPoint numEvents = go numEvents (Map.empty) (Map.empty)

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

279 where

280 go 0 !open !closed = if Map.null open

281 then return . Map.toList . Map.map reverse $ closed

@@ -433,6 +433,44 @@ testCloseTwice transport _ = do

433

434 takeMVar clientDone

435

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

437+testConnectToSelf transport numPings = do

438+ done <- newEmptyMVar

439+ Right endpoint <- newEndPoint transport

440+

441+ tlog "Creating self-connection"

442+ Right conn <- connect endpoint (address endpoint) ReliableOrdered

443+

444+ tlog "Talk to myself"

445+

446+ -- One thread to write to the endpoint

447+ forkTry $ do

448+ tlog $ "writing"

449+

450+ tlog $ "Sending ping"

451+ replicateM_ numPings $ send conn ["ping"]

452+

453+ tlog $ "Closing connection"

454+ close conn

455+

456+ -- And one thread to read

457+ forkTry $ do

458+ tlog $ "reading"

459+

460+ tlog "Waiting for ConnectionOpened"

461+ cid <- expect endpoint (\(ConnectionOpened cid _ addr) -> (addr == address endpoint, cid))

462+

463+ tlog "Waiting for Received"

464+ replicateM_ numPings $ expect' endpoint (\(Received cid' msg) -> cid == cid' && msg == ["ping"])

465+

466+ tlog "Waiting for ConnectionClosed"

467+ expect' endpoint (\(ConnectionClosed cid') -> cid == cid')

468+

469+ tlog "Done"

470+ putMVar done ()

471+

472+ takeMVar done

473+

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

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

437runTransportTest description test = do

475runTransportTest description test = do

438 (transport, numPings) <- ask

476 (transport, numPings) <- ask

@@ -450,3 +488,4 @@ testTransport transport = flip runReaderT (transport, 10000) $ do

450 runTransportTest "ParallelConnects" testParallelConnects

488 runTransportTest "ParallelConnects" testParallelConnects

451 runTransportTest "SendAfterClose" testSendAfterClose

489 runTransportTest "SendAfterClose" testSendAfterClose

452 runTransportTest "CloseTwice" testCloseTwice

490 runTransportTest "CloseTwice" testCloseTwice

491+ runTransportTest "ConnectToSelf" testConnectToSelf