git.iko.soy/distributed-process
on master help sign in
Add refcount to endpoint
commit70fa7e7589c14514124451fb689dd8b356796e7f authorEdsko de Vries <edsko@edsko.net> dateFri 27 Apr 2012, 09:03 (14 years ago) parent7a0b324
2 files changed+106−10
network-transport-2/src/Network/Transport/TCP.hs+36−8 network-transport-2/tests/TestTransport.hs+70−2
network-transport-2/src/Network/Transport/TCP.hs+36−8
@@ -49,7 +49,7 @@ import Control.Concurrent.MVar ( MVar
4949 )
5050import Control.Category ((>>>))
5151import Control.Applicative ((<*>), (*>), (<$>))
52−import Control.Monad (forever, forM_, void)
52+import Control.Monad (forever, forM_, void, when)
5353import Control.Monad.Error (MonadError, ErrorT, liftIO, runErrorT, throwError)
5454import Control.Monad.Trans.Maybe (MaybeT(MaybeT), runMaybeT)
5555import Control.Monad.Trans.Writer (WriterT, execWriterT)
@@ -81,6 +81,7 @@ data LocalEndPoint = LocalEndPoint { _localAddress :: EndPointAddress
8181data RemoteEndPoint = RemoteEndPoint { _remoteAddress :: EndPointAddress
8282 , _remoteSocket :: N.Socket
8383 , _remoteLock :: MVar ()
84+ , _remoteRefCount :: MVar Int
8485 }
8586data EndPointState = EndPointState { _nextConnectionId :: ConnectionId
8687 , _pendingCtrlRequests :: IntMap (MVar [ByteString])
@@ -146,10 +147,13 @@ forkRemoteEndPoint :: TCPTransport
146147 -> N.Socket
147148 -> IO RemoteEndPoint
148149forkRemoteEndPoint transport ourEndPoint theirAddress sock = do
149− lock <- newMVar ()
150− let theirEndPoint = RemoteEndPoint { _remoteAddress = theirAddress
151− , _remoteSocket = sock
152− , _remoteLock = lock
150+ -- putStrLn "Creating new refcount.."
151+ lock <- newMVar ()
152+ refCount <- newMVar 0
153+ let theirEndPoint = RemoteEndPoint { _remoteAddress = theirAddress
154+ , _remoteSocket = sock
155+ , _remoteLock = lock
156+ , _remoteRefCount = refCount
153157 }
154158 modifyMVar_ (transport ^. transportState) $ return . (remoteEndPointAt theirAddress ^= Just theirEndPoint)
155159 forkIO $ do
@@ -262,17 +266,38 @@ tcpConnect :: TCPTransport -- ^ Transport
262266tcpConnect transport ourEndPoint theirAddress _ = runErrorT $ do
263267 theirEndPoint <- remoteEndPoint transport ourEndPoint theirAddress
264268 connId <- requestNewConnection ourEndPoint theirEndPoint
269+ increaseRemoteRefCount theirEndPoint
265270 let sock = theirEndPoint ^. remoteSocket
271+ -- TODO: should we modify send so that it returns an error when the user closes the connection?
266272 return Connection { send = runErrorT .
267273 failWithIO (FailedWith SendFailed . show) .
268274 withRemoteLock theirEndPoint .
269275 sendMany sock .
270276 (\payload -> encodeInt32 connId : prependLength payload)
271− , close = -- TODO: we should not close the TCP connection until all
272− -- lightweight connections have been closed
273− N.sClose sock
277+ , close = decreaseRemoteRefCount theirEndPoint
274278 }
275279
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+
276301-- | Request a new connection
277302requestNewConnection :: (MonadIO m, MonadError (FailedWith ConnectErrorCode) m)
278303 => LocalEndPoint
@@ -388,6 +413,9 @@ remoteSocket = lens _remoteSocket (\sock ep -> ep { _remoteSocket = sock })
388413remoteLock :: Lens RemoteEndPoint (MVar ())
389414remoteLock = lens _remoteLock (\lock ep -> ep { _remoteLock = lock })
390415
416+remoteRefCount :: Lens RemoteEndPoint (MVar Int)
417+remoteRefCount = lens _remoteRefCount (\rc ep -> ep { _remoteRefCount = rc })
418+
391419nextConnectionId :: Lens EndPointState ConnectionId
392420nextConnectionId = lens _nextConnectionId (\cix st -> st { _nextConnectionId = cix })
393421
network-transport-2/tests/TestTransport.hs+70−2
@@ -186,6 +186,73 @@ testCloseOneConnection transport numPings = do
186186
187187 takeMVar result
188188
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+
189256runTestIO :: String -> IO () -> IO ()
190257runTestIO description test = do
191258 putStr $ "Running " ++ show description ++ ": "
@@ -196,15 +263,16 @@ runTestIO description test = do
196263runTest :: String -> (Transport -> Int -> IO ()) -> ReaderT (Transport, Int) IO ()
197264runTest description test = do
198265 (transport, numPings) <- ask
199− done <- liftIO $ timeout 1000000 $ runTestIO description (test transport numPings)
266+ done <- liftIO $ timeout 5000000 $ runTestIO description (test transport numPings)
200267 case done of
201268 Just () -> return ()
202269 Nothing -> error "timeout"
203270
204271-- Transport tests
205272testTransport :: Transport -> IO ()
206−testTransport transport = flip runReaderT (transport, 1000) $ do
273+testTransport transport = flip runReaderT (transport, 50000) $ do
207274 runTest "PingPong" testPingPong
208275 runTest "EndPoints" testEndPoints
209276 runTest "Connections" testConnections
210277 runTest "CloseOneConnection" testCloseOneConnection
278+-- runTest "CloseOneDirection" testCloseOneDirection