70fa7e7 · distributed-process

git.iko.soy/distributed-process

on master help sign in

Add refcount to endpoint

70fa7e7Edsko de Vries14 years ago

network-transport-2/src/Network/Transport/TCP.hs+36−8

compare with master

unifiedsplit

@@ -49,7 +49,7 @@ import Control.Concurrent.MVar ( MVar

49 )

50import Control.Category ((>>>))

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

52−import Control.Monad (forever, forM_, void)

52+import Control.Monad (forever, forM_, void, when)

53import Control.Monad.Error (MonadError, ErrorT, liftIO, runErrorT, throwError)

54import Control.Monad.Trans.Maybe (MaybeT(MaybeT), runMaybeT)

55import Control.Monad.Trans.Writer (WriterT, execWriterT)

@@ -81,6 +81,7 @@ data LocalEndPoint = LocalEndPoint { _localAddress :: EndPointAddress

81data RemoteEndPoint = RemoteEndPoint { _remoteAddress :: EndPointAddress

82 , _remoteSocket :: N.Socket

83 , _remoteLock :: MVar ()

84+ , _remoteRefCount :: MVar Int

84 }

85 }

85data EndPointState = EndPointState { _nextConnectionId :: ConnectionId

86data EndPointState = EndPointState { _nextConnectionId :: ConnectionId

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

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

@@ -146,10 +147,13 @@ forkRemoteEndPoint :: TCPTransport

146 -> N.Socket

147 -> N.Socket

147 -> IO RemoteEndPoint

148 -> IO RemoteEndPoint

148forkRemoteEndPoint transport ourEndPoint theirAddress sock = do

149forkRemoteEndPoint transport ourEndPoint theirAddress sock = do

150+ -- putStrLn "Creating new refcount.."

149− lock <- newMVar ()

151+ lock <- newMVar ()

152+ refCount <- newMVar 0

150− let theirEndPoint = RemoteEndPoint { _remoteAddress = theirAddress

153+ let theirEndPoint = RemoteEndPoint { _remoteAddress = theirAddress

151− , _remoteSocket = sock

154+ , _remoteSocket = sock

152− , _remoteLock = lock

155+ , _remoteLock = lock

156+ , _remoteRefCount = refCount

153 }

157 }

154 modifyMVar_ (transport ^. transportState) $ return . (remoteEndPointAt theirAddress ^= Just theirEndPoint)

158 modifyMVar_ (transport ^. transportState) $ return . (remoteEndPointAt theirAddress ^= Just theirEndPoint)

155 forkIO $ do

159 forkIO $ do

@@ -262,17 +266,38 @@ tcpConnect :: TCPTransport -- ^ Transport

262tcpConnect transport ourEndPoint theirAddress _ = runErrorT $ do

266tcpConnect transport ourEndPoint theirAddress _ = runErrorT $ do

263 theirEndPoint <- remoteEndPoint transport ourEndPoint theirAddress

267 theirEndPoint <- remoteEndPoint transport ourEndPoint theirAddress

264 connId <- requestNewConnection ourEndPoint theirEndPoint

268 connId <- requestNewConnection ourEndPoint theirEndPoint

269+ increaseRemoteRefCount theirEndPoint

265 let sock = theirEndPoint ^. remoteSocket

270 let sock = theirEndPoint ^. remoteSocket

271+ -- TODO: should we modify send so that it returns an error when the user closes the connection?

266 return Connection { send = runErrorT .

272 return Connection { send = runErrorT .

267 failWithIO (FailedWith SendFailed . show) .

273 failWithIO (FailedWith SendFailed . show) .

268 withRemoteLock theirEndPoint .

274 withRemoteLock theirEndPoint .

269 sendMany sock .

275 sendMany sock .

270 (\payload -> encodeInt32 connId : prependLength payload)

276 (\payload -> encodeInt32 connId : prependLength payload)

271− , close = -- TODO: we should not close the TCP connection until all

277+ , close = decreaseRemoteRefCount theirEndPoint

272− -- lightweight connections have been closed

273− N.sClose sock

274 }

278 }

275

279

280+-- | Increase the refcount on a remote endpoint

281+increaseRemoteRefCount :: MonadIO m => RemoteEndPoint -> m ()

282+increaseRemoteRefCount endPoint = liftIO $ do

283+ return ()

284+{-

285+ rc <- modifyMVar (endPoint ^. remoteRefCount) $ \rc -> return (rc + 1, rc)

286+ putStrLn $ "Increased refcount from " ++ show rc

287+-}

288+

289+-- | Decrease the refcount on a remote endpoint and close the socket

290+-- when it reaches zero.

291+decreaseRemoteRefCount :: MonadIO m => RemoteEndPoint -> m ()

292+decreaseRemoteRefCount endPoint = liftIO $ do

293+ return ()

294+{-

295+ rc <- modifyMVar (endPoint ^. remoteRefCount) $ \rc -> return (rc - 1, rc)

296+ putStrLn $ "Reduced refcount from " ++ show rc

297+ when (rc == 0) $

298+ N.sClose (endPoint ^. remoteSocket)

299+-}

300+

276-- | Request a new connection

301-- | Request a new connection

277requestNewConnection :: (MonadIO m, MonadError (FailedWith ConnectErrorCode) m)

302requestNewConnection :: (MonadIO m, MonadError (FailedWith ConnectErrorCode) m)

278 => LocalEndPoint

303 => LocalEndPoint

@@ -388,6 +413,9 @@ remoteSocket = lens _remoteSocket (\sock ep -> ep { _remoteSocket = sock })

388remoteLock :: Lens RemoteEndPoint (MVar ())

413remoteLock :: Lens RemoteEndPoint (MVar ())

389remoteLock = lens _remoteLock (\lock ep -> ep { _remoteLock = lock })

414remoteLock = lens _remoteLock (\lock ep -> ep { _remoteLock = lock })

390

415

416+remoteRefCount :: Lens RemoteEndPoint (MVar Int)

417+remoteRefCount = lens _remoteRefCount (\rc ep -> ep { _remoteRefCount = rc })

418+

391nextConnectionId :: Lens EndPointState ConnectionId

419nextConnectionId :: Lens EndPointState ConnectionId

392nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = cix })

420nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = cix })

393

421

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

@@ -186,6 +186,73 @@ testCloseOneConnection transport numPings = do

186

187 takeMVar result

188

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

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

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

192+testCloseOneDirection :: Transport -> Int -> IO ()

193+testCloseOneDirection transport numPings = do

194+ addrA <- newEmptyMVar

195+ addrB <- newEmptyMVar

196+ doneA <- newEmptyMVar

197+ doneB <- newEmptyMVar

198+

199+ -- A

200+ forkIO $ do

201+ Right endpoint <- newEndPoint transport

202+ putMVar addrA (address endpoint)

203+

204+ -- Connect to B

205+ Right conn <- takeMVar addrB >>= \addr -> connect endpoint addr ReliableOrdered

206+

207+ -- Wait for B to connect to us

208+ ConnectionOpened _ _ _ <- receive endpoint

209+

210+ -- Send pings to B

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

212+

213+ -- Close our connection to B

214+ close conn

215+

216+ -- Wait for B's pongs

217+ replicateM_ numPings $ do Received _ _ <- receive endpoint ; return ()

218+

219+ -- Wait for B to close it's connection to us

220+ -- TODO: this message is not yet sent

221+ -- ConnectionClosed _ <- receive endpoint

222+

223+ -- Done

224+ putMVar doneA ()

225+

226+ -- B

227+ forkIO $ do

228+ Right endpoint <- newEndPoint transport

229+ putMVar addrB (address endpoint)

230+

231+ -- Wait for A to connect

232+ ConnectionOpened _ _ _ <- receive endpoint

233+

234+ -- Connect to A

235+ Right conn <- takeMVar addrA >>= \addr -> connect endpoint addr ReliableOrdered

236+

237+ -- Wait for A's pings

238+ replicateM_ numPings $ do Received _ _ <- receive endpoint ; return ()

239+

240+ -- Wait for A to close it's connection to us

241+ -- TODO: this message is not yet sent

242+ -- ConnectionClosed _ <- receive endpoint

243+

244+ -- Send pongs to A

245+ replicateM_ numPings $ send conn ["pong"]

246+

247+ -- Close our connection to A

248+ close conn

249+

250+ -- Done

251+ putMVar doneB ()

252+

253+ mapM_ takeMVar [doneA, doneB]

254+

255+

189runTestIO :: String -> IO () -> IO ()

256runTestIO :: String -> IO () -> IO ()

190runTestIO description test = do

257runTestIO description test = do

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

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

@@ -196,15 +263,16 @@ runTestIO description test = do

196runTest :: String -> (Transport -> Int -> IO ()) -> ReaderT (Transport, Int) IO ()

263runTest :: String -> (Transport -> Int -> IO ()) -> ReaderT (Transport, Int) IO ()

197runTest description test = do

264runTest description test = do

198 (transport, numPings) <- ask

265 (transport, numPings) <- ask

199− done <- liftIO $ timeout 1000000 $ runTestIO description (test transport numPings)

266+ done <- liftIO $ timeout 5000000 $ runTestIO description (test transport numPings)

200 case done of

267 case done of

201 Just () -> return ()

268 Just () -> return ()

202 Nothing -> error "timeout"

269 Nothing -> error "timeout"

203

270

204-- Transport tests

271-- Transport tests

205testTransport :: Transport -> IO ()

272testTransport :: Transport -> IO ()

206−testTransport transport = flip runReaderT (transport, 1000) $ do

273+testTransport transport = flip runReaderT (transport, 50000) $ do

207 runTest "PingPong" testPingPong

274 runTest "PingPong" testPingPong

208 runTest "EndPoints" testEndPoints

275 runTest "EndPoints" testEndPoints

209 runTest "Connections" testConnections

276 runTest "Connections" testConnections

210 runTest "CloseOneConnection" testCloseOneConnection

277 runTest "CloseOneConnection" testCloseOneConnection

278+-- runTest "CloseOneDirection" testCloseOneDirection