@@ -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