de72d52 · distributed-process

5555 , takeMVar

5656 , putMVar

5757 , newEmptyMVar

58+ , withMVar

5859 )

60+import Data.IORef (IORef, newIORef, writeIORef, readIORef, modifyIORef)

5961import Control.Category ((>>>))

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

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

63+import Control.Monad (forM_, void, when)

6264import Control.Monad.Error (ErrorT(ErrorT), runErrorT)

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

64−import Control.Exception (IOException, bracket_)

66+import Control.Exception (IOException)

6567import Data.ByteString (ByteString)

6668import qualified Data.ByteString as BS (concat)

6769import qualified Data.ByteString.Char8 as BSC (pack, unpack)

7274import qualified Data.IntSet as IntSet (empty, insert, delete, elems)

7375import Data.Map (Map)

7476import qualified Data.Map as Map (empty)

75−import Data.IORef (newIORef, modifyIORef, readIORef)

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

7778import Text.Regex.Applicative (RE, few, anySym, sym, many, match)

7879

9192-- > | |==========================| |

9293-- > +-------+ +-------+

9394--

94−-- Ignoring the complications detailed below, the TCP connection is created is

95−-- created the first lightweight connection is created (in either direction),

96−-- and torn down when the last lightweight connection (in either direction) is

95+-- Ignoring the complications detailed below, the TCP connection is set up is

96+-- when the first lightweight connection is created (in either direction), and

97+-- torn down when the last lightweight connection (in either direction) is

9798-- closed.

9899--

99100-- [Connecting]

103104-- sends a request to /B/. As part of the request /A/ sends its own endpoint

104105-- address to /B/ (so that /B/ can reuse the connection in the other direction).

105106--

106−-- When /B/ receives the connection request it first checks if it had not

107−-- already initiated a connection request to /A/. If it had not, then it will

108−-- acknowledge the connection request by sending 'ConnectionRequestAccepted' to

109−-- /A/ and record that it has a TCP connection to /A/.

107+-- When /B/ receives the connection request it first checks if it did not

108+-- already initiate a connection request to /A/. If not it will acknowledge the

109+-- connection request by sending 'ConnectionRequestAccepted' to /A/ and record

110+-- that it has a TCP connection to /A/.

110111--

111112-- The tricky case arises when /A/ sends a connection request to /B/ and /B/

112113-- finds that it had already sent a connection request to /A/. In this case /B/

170171-- Internal datatypes --

171172--------------------------------------------------------------------------------

172173

173−data TCPTransport = TCPTransport { _transportHost :: N.HostName

174− , _transportPort :: N.ServiceName

175− , _transportState :: MVar TransportState

174+data TCPTransport = TCPTransport { transportHost :: N.HostName

175+ , transportPort :: N.ServiceName

176+ , transportState :: MVar TransportState

176177 }

177−data TransportState = TransportState { _localEndPoints :: IntMap LocalEndPoint }

178−data LocalEndPoint = LocalEndPoint { _localAddress :: EndPointAddress

179− , _localChannel :: Chan Event

180− , _localState :: MVar LocalState

178+data LocalEndPoint = LocalEndPoint { localAddress :: EndPointAddress

179+ , localChannel :: Chan Event

180+ , localState :: MVar LocalState

181181 }

182−data LocalState = LocalState { _nextConnectionId :: ConnectionId

183− , _pendingCtrlRequests :: IntMap (MVar [ByteString])

184− , _nextCtrlRequestId :: ControlRequestId

185− , _localConnections :: Map EndPointAddress (IsClosing, MVar RemoteState)

182+data RemoteEndPoint = RemoteEndPoint { remoteAddress :: EndPointAddress

183+ , remoteState :: MVar RemoteState

186184 }

187−data RemoteEndPoint = RemoteEndPoint { _remoteAddress :: EndPointAddress

188− , _remoteSocket :: N.Socket

189− , sendTo :: [ByteString] -> IO ()

190− }

191−data RemoteState = EndPointInvalid (FailedWith ConnectErrorCode)

192− | EndPointValid !RefCount RemoteEndPoint

193− | EndPointClosed

185+

186+-- We use underscores for fields that we might update (using lenses).

187+

188+data TransportState = TransportState { _localEndPoints :: IntMap LocalEndPoint }

189+data LocalState = LocalState { _nextConnectionId :: ConnectionId

190+ , _pendingCtrlRequests :: IntMap (MVar [ByteString])

191+ , _nextCtrlRequestId :: ControlRequestId

192+ , _localConnections :: Map EndPointAddress RemoteEndPoint

193+ }

194+

195+-- A remote endpoint has incoming and outgoing connections, and when the total

196+-- number of connections (that is, the 'remoteRefCount') drops to zero we want

197+-- to close the TCP connection to the endpoint.

198+--

199+-- What we need to avoid, however, is a situation with two concurrent threads

200+-- where one closes the last (incoming or outgoing) connection, initiating the

201+-- process of closing the connection, while another requests (but does not yet

202+-- have) a new connection.

203+--

204+-- We therefore insist that:

205+--

206+-- 1. All operations that change the state of the endpoint (ask for a new

207+-- connection, close a connection, close the endpoint completely) are

208+-- serialized (that is, they take the contents of the MVar containing the

209+-- endpoint state before starting and don't put the updated contents back

210+-- until they have completed).

211+-- 2. Writing to ('apiSend') or reading from (in 'handleIncomingMessages') must

212+-- maintain the invariant that the connection they are writing to/reading

213+-- from *must* be "included" in the 'remoteRefCount'.

214+-- 3. Since every endpoint is associated with a single socket, we regard writes

215+-- that endpoint a state change too (i.e., we take the MVar before the write

216+-- and put it back after). The reason is that we don't want to "scramble" the

217+-- output of multiple concurrent writes (either from an explicit 'send' or

218+-- the writes for control messages).

219+--

220+-- Of course, "serialize" does not mean that we want for the remote endpoint to

221+-- reply. "Send" takes the mvar, sends to the endpoint, and then puts the mvar

222+-- back, without waiting for the endpoint to receive the message. Similarly,

223+-- when requesting a new connection, we take the mvar, tentatively increment the

224+-- reference count, send the control request, and then put the mvar back. When

225+-- the remote host responds to the new connection request we might have to do

226+-- another state change (reduce the refcount) if the connection request was

227+-- refused but we don't want to increment the ref count only after the remote

228+-- host acknowledges the request because then a concurrent 'close' might

229+-- actually attempt to close the socket.

230+--

231+-- Since we don't do concurrent reads from the same socket we don't need to

232+-- take the lock when reading from the socket.

233+

234+data RemoteState = RemoteEndPointInvalid (FailedWith ConnectErrorCode)

235+ | RemoteEndPointValid RemoteConnection

236+ | RemoteEndPointClosing RemoteConnection (MVar ())

237+ | RemoteEndPointClosed

238+data RemoteConnection = RemoteConnection { _remoteRefCount :: !Int

239+ , remoteSocket :: N.Socket

240+ , sendOn :: [ByteString] -> IO ()

241+ }

194242

195243type EndPointId = Int32

196244type ControlRequestId = Int32

197−type RefCount = Int32

198−type IsClosing = Maybe RemoteEndPoint

199−type Conn = (LocalEndPoint, RemoteEndPoint)

245+type EndPointPair = (LocalEndPoint, RemoteEndPoint)

200246

201247-- | Control headers

202248data ControlHeader =

223269createTransport :: N.HostName -> N.ServiceName -> IO (Either IOException Transport)

224270createTransport host port = do

225271 state <- newMVar TransportState { _localEndPoints = IntMap.empty }

226− let transport = TCPTransport { _transportState = state

227− , _transportHost = host

228− , _transportPort = port

272+ let transport = TCPTransport { transportState = state

273+ , transportHost = host

274+ , transportPort = port

229275 }

230276 tryIO $ do

231277 forkServer host port $ handleConnectionRequest transport

239285apiNewEndPoint :: TCPTransport -> IO (Either (FailedWith NewEndPointErrorCode) EndPoint)

240286apiNewEndPoint transport = do

241287 ourEndPoint <- createLocalEndPoint transport

242− return . Right $ EndPoint { receive = readChan (ourEndPoint ^. localChannel)

243− , address = ourEndPoint ^. localAddress

288+ return . Right $ EndPoint { receive = readChan (localChannel ourEndPoint)

289+ , address = localAddress ourEndPoint

244290 , connect = apiConnect ourEndPoint

245291 , newMulticastGroup = return . Left $ newMulticastGroupError

246292 , resolveMulticastGroup = return . Left . const resolveMulticastGroupError

254300 -> EndPointAddress -- ^ Remote address

255301 -> Reliability -- ^ Reliability (ignored)

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

288−apiConnect ourEndPoint theirAddress _ = runErrorT $ do

289− theirEndPoint <- findRemoteEndPoint ourEndPoint theirAddress

290− connId <- requestNewConnection (ourEndPoint, theirEndPoint)

291− connAlive <- liftIO $ newMVar True

292− return $ Connection { send = apiSend theirEndPoint connId connAlive

293− , close = apiClose (ourEndPoint, theirEndPoint) connId connAlive

303+apiConnect ourEndPoint theirAddress _ | localAddress ourEndPoint == theirAddress =

304+ connectToSelf ourEndPoint

305+apiConnect ourEndPoint theirAddress _ = runErrorT $ do

306+ (theirEndPoint, connId) <- requestConnectionTo ourEndPoint theirAddress

307+ -- connAlive can be an IORef rather than an MVar because it is protected by

308+ -- the remoteState MVar. We don't need the overhead of locking twice.

309+ connAlive <- liftIO $ newIORef True

310+ return $ Connection { send = apiSend theirEndPoint connId connAlive

311+ , close = apiClose theirEndPoint connId connAlive

294312 }

295313

296314-- | Close a connection

297315--

316+-- RELY: If the connection is alive on entry then the remote endpoint must

317+-- remain in the state 'RemoteEndPointValid'.

318+-- GUARANTEE: If the connection is alive on entry then the remote endpoint will

319+-- either be RemoteEndPointValid or RemoteEndPointClosing. Otherwise, the state

320+-- of the remote endpoint will not be changed.

321+--

298322-- TODO: We ignore errors during a close. Is that right?

299−apiClose :: Conn -> ConnectionId -> MVar Bool -> IO ()

300−apiClose (ourEndPoint, theirEndPoint) connId connAlive = void . tryIO $ do

301− modifyMVar_ connAlive $ \alive -> do

302− when alive $ do

303− sendTo theirEndPoint [encodeInt32 CloseConnection, encodeInt32 connId]

304− decreaseRemoteRefCount (ourEndPoint, theirEndPoint)

305− return False

323+apiClose :: RemoteEndPoint -> ConnectionId -> IORef Bool -> IO ()

324+apiClose theirEndPoint connId connAlive = void . tryIO $ do

325+ wasAlive <- withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) -> do

326+ alive <- readIORef connAlive

327+ when alive $ do

328+ sendOn remoteConn [encodeInt32 CloseConnection, encodeInt32 connId]

329+ writeIORef connAlive False

330+ return alive

331+ -- We don't want to call decreaseRemoteRefCount while we still hold the lock

332+ when wasAlive $ decreaseRemoteRefCount theirEndPoint

306333

307334-- | Send data across a connection

308−apiSend :: RemoteEndPoint -> ConnectionId -> MVar Bool -> [ByteString] -> IO (Either (FailedWith SendErrorCode) ())

309−apiSend theirEndPoint connId connAlive payload =

310− modifyMVar connAlive $ \alive -> do

335+--

336+-- RELY: The remote endpoint must remain in state 'RemoteEndPointValid'.

337+-- GUARANTEE: The state of the remote endpoint will not be changed.

338+apiSend :: RemoteEndPoint -> ConnectionId -> IORef Bool -> [ByteString] -> IO (Either (FailedWith SendErrorCode) ())

339+apiSend theirEndPoint connId connAlive payload = do

340+ withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) -> do

341+ alive <- readIORef connAlive

311342 if alive

312343 then do

313− result <- tryIO $ sendTo theirEndPoint (encodeInt32 connId : prependLength payload)

344+ result <- tryIO $ sendOn remoteConn (encodeInt32 connId : prependLength payload)

314345 case result of

315− Left err -> return (alive, Left $ FailedWith SendFailed (show err))

316− Right _ -> return (alive, Right $ ())

346+ Left err -> return . Left $ FailedWith SendFailed (show err)

347+ Right _ -> return . Right $ ()

317348 else do

318− return (alive, Left $ FailedWith SendConnectionClosed "Connection closed")

349+ return . Left $ FailedWith SendConnectionClosed "Connection closed"

350+

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

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

353+connectToSelf ourEndPoint = do

354+ -- Here connAlive must an MVar because it is not protected by another lock

355+ connAlive <- newMVar True

356+ connId <- getNextConnectionId ourEndPoint

357+ writeChan ourChan (ConnectionOpened connId ReliableOrdered (localAddress ourEndPoint))

358+ return . Right $ Connection { send = selfSend connAlive connId

359+ , close = selfClose connAlive connId

360+ }

361+ where

362+ ourChan :: Chan Event

363+ ourChan = localChannel ourEndPoint

364+

365+ selfSend :: MVar Bool -> ConnectionId -> [ByteString] -> IO (Either (FailedWith SendErrorCode) ())

366+ selfSend connAlive connId msg = do

367+ modifyMVar connAlive $ \alive ->

368+ if alive

369+ then do

370+ writeChan ourChan (Received connId msg)

371+ return (alive, Right ())

372+ else

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

374+

375+ selfClose :: MVar Bool -> ConnectionId -> IO ()

376+ selfClose connAlive connId = do

377+ modifyMVar_ connAlive $ \alive -> do

378+ when alive $ writeChan ourChan (ConnectionClosed connId)

379+ return False

319380

320381--------------------------------------------------------------------------------

321382-- Lower level functionality --

322383--------------------------------------------------------------------------------

323384

385+-- | Create a new local endpoint

324386createLocalEndPoint :: TCPTransport -> IO LocalEndPoint

325387createLocalEndPoint transport = do

326388 chan <- newChan

329391 , _nextCtrlRequestId = 0

330392 , _localConnections = Map.empty

331393 }

332− let host = transport ^. transportHost

333− let port = transport ^. transportPort

334− modifyMVar (transport ^. transportState) $ \st -> do

394+ let host = transportHost transport

395+ let port = transportPort transport

396+ modifyMVar (transportState transport) $ \st -> do

335397 let ix = fromIntegral $ IntMap.size (st ^. localEndPoints)

336398 let addr = EndPointAddress . BSC.pack $ host ++ ":" ++ port ++ ":" ++ show ix

337− let localEndPoint = LocalEndPoint { _localAddress = addr

338− , _localChannel = chan

339− , _localState = state

399+ let localEndPoint = LocalEndPoint { localAddress = addr

400+ , localChannel = chan

401+ , localState = state

340402 }

341403 return ((localEndPointAt ix ^= Just localEndPoint) $ st, localEndPoint)

342404

343−findRemoteEndPoint :: LocalEndPoint

344− -> EndPointAddress

345− -> ErrorT (FailedWith ConnectErrorCode) IO RemoteEndPoint

346−findRemoteEndPoint ourEndPoint theirAddress = ErrorT $ go

405+-- | Request a connection to a remote endpoint

406+--

407+-- This will block until we get a connection ID from the remote endpoint; if

408+-- the remote endpoint was in 'RemoteEndPointClosing' state then we will

409+-- additionally block until that is resolved.

410+requestConnectionTo :: LocalEndPoint

411+ -> EndPointAddress

412+ -> ErrorT (FailedWith ConnectErrorCode) IO (RemoteEndPoint, ConnectionId)

413+requestConnectionTo ourEndPoint theirAddress = ErrorT go

347414 where

348415 go = do

349− endPoint <- modifyMVar state $ \st ->

350− case st ^. connectionTo theirAddress of

351− Just (_, mvar) ->

352− return (st, mvar)

353− Nothing -> do

354− theirMVar <- newEmptyMVar

355− forkIO $ connectToRemoteEndPoint ourEndPoint theirMVar theirAddress

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

357− theirState <- readMVar endPoint

358− case theirState of

359− EndPointInvalid err -> do

416+ let ourState = localState ourEndPoint

417+ theirEndPoint <- modifyMVar ourState $ \st ->

418+ case st ^. connectionTo theirAddress of

419+ Just theirEndPoint ->

420+ return (st, theirEndPoint)

421+ Nothing -> do

422+ theirState <- newEmptyMVar

423+ let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress

424+ , remoteState = theirState

425+ }

426+ forkIO $ connectToRemoteEndPoint (ourEndPoint, theirEndPoint)

427+ return (connectionTo theirAddress ^= Just theirEndPoint $ st, theirEndPoint)

428+

429+ -- Before we initiate the new connection request we want to make sure

430+ -- that refcount on the endpoint is incremented so that a concurrent

431+ -- thread will not close the connection. Note that if IF we return

432+ -- RemoteEndPointValid here then we can rely on the endpoint remaining in

433+ -- that state.

434+ let theirState = remoteState theirEndPoint

435+ endPointStateSnapshot <- modifyMVar theirState $ \st ->

436+ case st of

437+ RemoteEndPointValid ep ->

438+ return (RemoteEndPointValid . (remoteRefCount ^+= 1) $ ep, st)

439+ _ ->

440+ return (st, st)

441+

442+ -- From this point on we are guaranteed the refcount is positive,

443+ -- provided that the endpoint was valid. We still need to deal with the

444+ -- case where it was not valid, however, which we didn't want to do that

445+ -- while holding the endpoint lock.

446+ -- Although 'endPointStateSnapshot' here refers to a snapshot of the endpoint

447+ -- state, and might have changed in the meantime, these changes won't

448+ -- matter.

449+ case endPointStateSnapshot of

450+ RemoteEndPointInvalid err -> do

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

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

362− modifyMVar state $ \st ->

452+ -- to 'connect' might give a different result. (We definitely don't

453+ -- want to do this while holding the endpoint lock, because holding

454+ -- multiple locks introduces all kinds of trouble.)

455+ modifyMVar ourState $ \st ->

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

364− EndPointValid _ ep -> do

365− return $ Right ep

366− EndPointClosed -> do

457+ return . Left $ err

458+

459+ RemoteEndPointClosing _ resolved -> do

460+ -- If the remote endpoint is closing, then we need to block until

461+ -- this is resolved and we then try again

462+ readMVar resolved

463+ go

464+

465+ RemoteEndPointClosed -> do

466+ -- EndPointClosed indicates that a concurrent thread was in the

467+ -- process of closing the TCP connection to the remote endpoint when

468+ -- we obtained a reference to it. The remote endpoint will now have

469+ -- been removed from ourState, so we simply try again.

367470 go

471+

472+ RemoteEndPointValid _ -> do

473+ -- Do the actual connection request. This blocks until the remote

474+ -- endpoint replies (but note that we don't hold any locks at this

475+ -- point)

476+ mcid <- tryIO $ do

477+ reply <- sendRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId

478+ decodeInt32 . BS.concat <$> takeMVar reply

479+

480+ -- On a failure we decrement the refcount again and return an error.

481+ -- The only state the remote endpoint can be in at this point is

482+ -- valid. As mentioned above, we can rely on the endpoint being in

483+ -- valid state at this point.

484+ let failed err = do

485+ modifyMVar_ theirState $ \(RemoteEndPointValid ep) ->

486+ return (RemoteEndPointValid . (remoteRefCount ^+= (-1)) $ ep)

487+ return . Left $ FailedWith ConnectFailed err

488+

489+ case mcid of

490+ Left err -> failed (show err)

491+ Right Nothing -> failed ("Invalid integer")

492+ Right (Just cid) -> return . Right $ (theirEndPoint, cid)

368493

369− state :: MVar LocalState

370− state = ourEndPoint ^. localState

371−

372−connectToRemoteEndPoint :: LocalEndPoint -> MVar RemoteState -> EndPointAddress -> IO ()

373−connectToRemoteEndPoint ourEndPoint theirMVar theirAddress = do

494+-- | Establish a connection to a remote endpoint

495+--

496+-- RELY: The state of the remote endpoint must be uninitialized.

497+-- GUARANTEE: Will only change the state to RemoteEndPointValid or

498+-- RemoteEndPointInvalid.

499+connectToRemoteEndPoint :: EndPointPair -> IO ()

500+connectToRemoteEndPoint (ourEndPoint, theirEndPoint) = do

501+ let theirState = remoteState theirEndPoint

374502 result <- runErrorT go

375503 case result of

376504 Right (ConnectionRequestAccepted, sock) -> do

377− theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock

378− putMVar theirMVar (EndPointValid 0 theirEndPoint)

505+ let remoteConn = RemoteConnection { remoteSocket = sock

506+ , _remoteRefCount = 0

507+ , sendOn = sendMany sock

508+ }

509+ putMVar theirState (RemoteEndPointValid remoteConn)

510+ void . forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)

379511 Right (ConnectionRequestEndPointInvalid, sock) -> do

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

512+ putMVar theirState (RemoteEndPointInvalid (invalidAddress "Invalid endpoint"))

381513 N.sClose sock

382514 Right (ConnectionRequestCrossed, sock) -> do

383515 N.sClose sock

384516 Left err -> do

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

386518 -- in which case we should close it

387− putMVar theirMVar (EndPointInvalid err)

519+ putMVar theirState (RemoteEndPointInvalid err)

388520 where

389521 go :: ErrorT (FailedWith ConnectErrorCode) IO (ConnectionRequestResponse, N.Socket)

390522 go = do

523+ let theirAddress = remoteAddress theirEndPoint

391524 (host, port, theirEndPointId) <- failWith (invalidAddress "Could not parse") $

392525 decodeEndPointAddress theirAddress

393526 addr:_ <- failWithIO (invalidAddress . show) $

408541 failed = FailedWith ConnectFailed

409542

410543 ourAddress :: ByteString

411− EndPointAddress ourAddress = ourEndPoint ^. localAddress

412−

413−-- | Request a new connection

414−requestNewConnection :: Conn -> ErrorT (FailedWith ConnectErrorCode) IO ConnectionId

415−requestNewConnection (ourEndPoint, theirEndPoint) = ErrorT $ do

416− -- We increment the refcount before sending the request, because if a

417− -- CloseSocket request comes this way first we want to ignore it

418− increaseRemoteRefCount (ourEndPoint, theirEndPoint)

419− mcid <- tryIO $ do

420− reply <- sendRemoteRequest (ourEndPoint, theirEndPoint) RequestConnectionId

421− decodeInt32 . BS.concat <$> takeMVar reply

422− case mcid of

423− Left err -> do decreaseRemoteRefCount (ourEndPoint, theirEndPoint)

424− return . Left $ failed (show err)

425− Right Nothing -> do decreaseRemoteRefCount (ourEndPoint, theirEndPoint)

426− return . Left $ failed "Invalid integer"

427− Right (Just cid) -> do return . Right $ cid

428− where

429− failed = FailedWith ConnectFailed

544+ EndPointAddress ourAddress = localAddress ourEndPoint

430545

431546-- | Decode end point address

432547-- TODO: This uses regular expression parsing, which is nice, but unnecessary

436551 endPointAddressRE :: RE Char (N.HostName, N.ServiceName, EndPointId)

437552 endPointAddressRE = (,,) <$> few anySym <*> (sym ':' *> few anySym) <*> (sym ':' *> (read <$> many anySym))

438553

439−sendRemoteRequest :: Conn -> ControlHeader -> IO (MVar [ByteString])

554+-- | Send a remote request and return the MVar that will contain the response

555+--

556+-- RELY: Remote endpoint must be in valid state.

557+-- GUARANTEE: Will not change the state of the remote endpoint.

558+sendRemoteRequest :: EndPointPair -> ControlHeader -> IO (MVar [ByteString])

440559sendRemoteRequest (ourEndPoint, theirEndPoint) header = do

441560 reply <- newEmptyMVar

442− reqId <- modifyMVar (ourEndPoint ^. localState) $ \st -> do

561+ reqId <- modifyMVar (localState ourEndPoint) $ \st -> do

443562 let reqId = st ^. nextCtrlRequestId

444563 return ((nextCtrlRequestId ^+= 1) . (pendingCtrlRequestsAt reqId ^= Just reply) $ st, reqId)

445− sendTo theirEndPoint [encodeInt32 header, encodeInt32 reqId]

564+ withMVar (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) ->

565+ sendOn remoteConn [encodeInt32 header, encodeInt32 reqId]

446566 return reply

447567

448−--------------------------------------------------------------------------------

449−-- Reference counting --

450−--------------------------------------------------------------------------------

451568

452−increaseRemoteRefCount :: Conn -> IO ()

453−increaseRemoteRefCount (ourEndPoint, theirEndPoint) = do

454− let theirAddr = theirEndPoint ^. remoteAddress

455− ourState = ourEndPoint ^. localState

456− -- If this endpoint was closing, cancel that

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

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

459− case isClosing of

460− Nothing -> do

461− return (st, theirMVar)

462− Just ep -> do

463− putMVar theirMVar (EndPointValid 0 ep)

464− return (connectionTo theirAddr ^= Just (Nothing, theirMVar) $ st, theirMVar)

465− -- Increase the refcount

466− modifyMVar_ theirMVar $ \(EndPointValid rc ep) ->

467− return (EndPointValid (rc + 1) ep)

468−

469−decreaseRemoteRefCount :: Conn -> IO ()

470−decreaseRemoteRefCount (ourEndPoint, theirEndPoint) = do

471− let theirAddr = theirEndPoint ^. remoteAddress

472− state <- readMVar (ourEndPoint ^. localState)

473− let Just (Nothing, theirMVar) = state ^. connectionTo theirAddr

474− -- We don't use modifyMVar here because we want to leave the MVar

475− -- empty when refcount reaches 0

476− -- TODO: deal with exceptions

477− EndPointValid rc ep <- takeMVar theirMVar

478− if rc > 1

479− then do

480− putMVar theirMVar (EndPointValid (rc - 1) ep)

481− else do

482− sendTo theirEndPoint [encodeInt32 CloseSocket]

483− modifyMVar_ (ourEndPoint ^. localState) $ \st -> do

484− let Just (_, mvar) = st ^. connectionTo theirAddr

485− return (connectionTo theirAddr ^= Just (Just ep, mvar) $ st)

569+-- | Decrease the remote refcount of a remote endpoint. When the refcount

570+-- reaches zero, initiate socket closure.

571+--

572+-- RELY: Remote endpoint must be in valid state.

573+-- GUARANTEE: Will leave the endpoint in its current state or change it to

574+-- RemoteEndPointClosing.

575+decreaseRemoteRefCount :: RemoteEndPoint -> IO ()

576+decreaseRemoteRefCount theirEndPoint = do

577+ modifyMVar_ (remoteState theirEndPoint) $ \(RemoteEndPointValid remoteConn) ->

578+ if remoteConn ^. remoteRefCount == 1

579+ then do

580+ resolved <- newEmptyMVar

581+ sendOn remoteConn [encodeInt32 CloseSocket]

582+ return (RemoteEndPointClosing remoteConn resolved)

583+ else do

584+ return (RemoteEndPointValid . (remoteRefCount ^+= (-1)) $ remoteConn)

486585

487586--------------------------------------------------------------------------------

488587-- Incoming requests --

489588--------------------------------------------------------------------------------

490589

590+-- | Handle a connection request (that is, a remote endpoint that is trying to

591+-- establish a TCP connection with us)

491592handleConnectionRequest :: TCPTransport -> N.Socket -> IO ()

492593handleConnectionRequest transport sock = do

493594 request <- tryIO $ do

494595 ourEndPointId <- recvInt32 sock

495596 theirAddress <- EndPointAddress . BS.concat <$> recvWithLength sock

496− ourEndPoint <- (^. localEndPointAt ourEndPointId) <$> readMVar state

597+ ourEndPoint <- (^. localEndPointAt ourEndPointId) <$> readMVar (transportState transport)

497598 return (theirAddress, ourEndPoint)

498599 case request of

499600 Right (_, Nothing) -> do

500601 sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]

501602 N.sClose sock

502603 Right (theirAddress, Just ourEndPoint) -> do

503− (crossed, mvar) <- modifyMVar (ourEndPoint ^. localState) $ \st ->

604+ (crossed, theirEndPoint) <- modifyMVar (localState ourEndPoint) $ \st ->

504605 case st ^. connectionTo theirAddress of

505606 Nothing -> do

506− mvar <- newEmptyMVar

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

508− Just (_, mvar) ->

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

607+ theirState <- newEmptyMVar

608+ let theirEndPoint = RemoteEndPoint { remoteAddress = theirAddress

609+ , remoteState = theirState

610+ }

611+ return (connectionTo theirAddress ^= Just theirEndPoint $ st, (False, theirEndPoint))

612+ Just theirEndPoint ->

613+ return (st, (localAddress ourEndPoint < theirAddress, theirEndPoint))

510614 if crossed

511615 then do

512616 sendMany sock [encodeInt32 ConnectionRequestCrossed]

513617 N.sClose sock

514618 else do

515− theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock

516− putMVar mvar (EndPointValid 0 theirEndPoint)

619+ let remoteConn = RemoteConnection { remoteSocket = sock

620+ , _remoteRefCount = 0

621+ , sendOn = sendMany sock

622+ }

623+ putMVar (remoteState theirEndPoint) (RemoteEndPointValid remoteConn)

624+ forkIO $ handleIncomingMessages (ourEndPoint, theirEndPoint)

517625 sendMany sock [encodeInt32 ConnectionRequestAccepted]

518626 Left _ ->

519627 -- Invalid request

520628 N.sClose sock

521629

522− where

523− state :: MVar TransportState

524− state = transport ^. transportState

525−

526−createRemoteEndPoint :: LocalEndPoint -> EndPointAddress -> N.Socket -> IO RemoteEndPoint

527−createRemoteEndPoint ourEndPoint theirAddress sock = do

528− lock <- newMVar ()

529− let theirEndPoint = RemoteEndPoint { _remoteAddress = theirAddress

530− , _remoteSocket = sock

531− , sendTo = bracket_ (takeMVar lock) (putMVar lock ())

532− . sendMany sock

533− }

534− forkIO $ do

535− unclosedConnections <- handleIncomingMessages (ourEndPoint, theirEndPoint)

536− forM_ (map fromIntegral $ IntSet.elems unclosedConnections) $

537− writeChan (ourEndPoint ^. localChannel) . ConnectionClosed

538− return theirEndPoint

539−

540630-- | Handle requests from a remote endpoint.

541631--

542632-- Returns only if the remote party closes the socket or if an error occurs.

543633-- The return value is the list of connections that are still open (normally,

544634-- this should be an empty list).

545−handleIncomingMessages :: Conn -> IO IntSet

635+--

636+-- RELY: The remote endpoint must be in RemoteEndPointValid or

637+-- RemoteEndPointClosing state. If the latter, then the 'resolved' MVar

638+-- associated with the closing state must be empty.

639+-- GUARANTEE: May change the remote endpoint to RemoteEndPointClosed state.

640+handleIncomingMessages :: EndPointPair -> IO ()

546641handleIncomingMessages (ourEndPoint, theirEndPoint) = do

547− -- We use an IORef rather than the state monad because when an exception occurs we

548− -- want to know the current list of open connections

549− openConnections <- newIORef IntSet.empty

550−

551− let go :: IO ()

552− go = do

553− connId <- recvInt32 sock

554− if connId >= firstNonReservedConnectionId

555− then readMessage connId

556− else do

557− case toEnum (fromIntegral connId) of

558− RequestConnectionId -> recvInt32 sock >>= createNewConnection

559− ControlResponse -> recvInt32 sock >>= readControlResponse

560− CloseConnection -> recvInt32 sock >>= closeConnection

561− CloseSocket -> closeSocket

562−

563− -- Create a new connection

564− createNewConnection :: ControlRequestId -> IO ()

565− createNewConnection reqId = do

566− increaseRemoteRefCount (ourEndPoint, theirEndPoint)

567− newId <- getNextConnectionId ourEndPoint

568− sendTo theirEndPoint ( encodeInt32 ControlResponse

569− : encodeInt32 reqId

570− : prependLength [encodeInt32 newId]

571− )

572− writeChan ourChannel (ConnectionOpened newId ReliableOrdered theirAddr)

573− -- We add the new connection ID to the list of open connections only once the

574− -- endpoint has been notified of the new connection (sendMany may fail)

575− modifyIORef openConnections (IntSet.insert (fromIntegral newId))

576− go

577−

578− -- Read a control response

579− readControlResponse :: ControlRequestId -> IO ()

580− readControlResponse reqId = do

581− response <- recvWithLength sock

582− mmvar <- (^. pendingCtrlRequestsAt reqId) <$> readMVar ourState

583− case mmvar of

584− Nothing ->

585− return () -- Invalid request ID. TODO: We just ignore it?

586− Just mvar -> do

587− modifyMVar_ ourState $ return . (pendingCtrlRequestsAt reqId ^= Nothing)

588− putMVar mvar response

589− go

590−

591− -- Close a connection

592− closeConnection :: ConnectionId -> IO ()

593− closeConnection cid = do

594− writeChan ourChannel (ConnectionClosed cid)

595− decreaseRemoteRefCount (ourEndPoint, theirEndPoint)

596− modifyIORef openConnections (IntSet.delete (fromIntegral cid))

597− go

598−

599− -- Close the socket (if we don't have any outgoing connections)

600− closeSocket :: IO ()

601− closeSocket = do

602− didClose <- modifyMVar (ourEndPoint ^. localState) $ \st -> do

603− case st ^. connectionTo theirAddr of

604− Just (Just _, mvar) -> do

605− -- The connection is already in "closing" state, which means we

606− -- don't need to send another CloseSocket request (we already did)

607− N.sClose sock

608− putMVar mvar EndPointClosed

609− return (connectionTo theirAddr ^= Nothing $ st, True)

610− Just (Nothing, theirMVar) -> do

611− EndPointValid rc _ <- readMVar theirMVar

612− if rc == 0

613− then do

614− sendTo theirEndPoint [encodeInt32 CloseSocket]

615− N.sClose sock

616− return (connectionTo theirAddr ^= Nothing $ st, True)

617− else do

618− return (st, False)

619− Nothing -> do

620− error "Internal error"

621− unless didClose go

622−

623− -- Read a message and output it on the endPoint's channel

624− readMessage :: ConnectionId -> IO ()

625− readMessage connId = recvWithLength sock >>= writeChan ourChannel . Received connId >> go

626−

627− -- Breakdown of the arguments

628− ourChannel = ourEndPoint ^. localChannel

629− ourState = ourEndPoint ^. localState

630− sock = theirEndPoint ^. remoteSocket

631− theirAddr = theirEndPoint ^. remoteAddress

632−

633− -- We keep handling messages until an exception occurs

634− tryIO go

635− readIORef openConnections

642+ -- We use an IORef rather than the state monad because when an exception occurs we

643+ -- want to know the current list of open connections

644+ openConnections <- newIORef IntSet.empty

645+

646+ -- For efficiency sake we get the socket once and for all

647+ sock <- withMVar theirState $ \st ->

648+ case st of

649+ RemoteEndPointValid ep ->

650+ return (remoteSocket ep)

651+ RemoteEndPointClosing ep _ ->

652+ return (remoteSocket ep)

653+ _ ->

654+ error "handleIncomingMessages RELY violation"

655+

656+ tryIO $ go sock openConnections

657+ unclosedConnections <- readIORef openConnections

658+ forM_ (map fromIntegral $ IntSet.elems unclosedConnections) $

659+ writeChan ourChannel . ConnectionClosed

660+ where

661+ -- Dispatch

662+ go :: N.Socket -> IORef IntSet -> IO ()

663+ go sock openConnections = do

664+ connId <- recvInt32 sock

665+ if connId >= firstNonReservedConnectionId

666+ then do

667+ readMessage sock connId

668+ go sock openConnections

669+ else do

670+ case toEnum (fromIntegral connId) of

671+ RequestConnectionId -> do

672+ recvInt32 sock >>= createNewConnection openConnections

673+ go sock openConnections

674+ ControlResponse -> do

675+ recvInt32 sock >>= readControlResponse sock

676+ go sock openConnections

677+ CloseConnection -> do

678+ recvInt32 sock >>= closeConnection openConnections

679+ go sock openConnections

680+ CloseSocket -> do

681+ closeSocket sock

682+ go sock openConnections

683+

684+ -- Create a new connection

685+ createNewConnection :: IORef IntSet -> ControlRequestId -> IO ()

686+ createNewConnection openConnections reqId = do

687+ modifyMVar_ theirState $ \st ->

688+ case st of

689+ RemoteEndPointValid ep ->

690+ return (RemoteEndPointValid . (remoteRefCount ^+= 1) $ ep)

691+ RemoteEndPointClosing ep resolved -> do

692+ -- If the endpoint is in closing state that means we send a

693+ -- CloseSocket request to the remote endpoint. If the remote

694+ -- endpoint replies with the request to create a new connection, it

695+ -- either ignored our request or it sent the request before it got

696+ -- ours. Either way, at this point we simply restore the endpoint

697+ -- to RemoteEndPointValid

698+ putMVar resolved ()

699+ return (RemoteEndPointValid . (remoteRefCount ^= 1) $ ep)

700+ _ ->

701+ error "handleIncomingMessages RELY violation"

702+ newId <- getNextConnectionId ourEndPoint

703+ withMVar theirState $ \(RemoteEndPointValid remoteConn) ->

704+ sendOn remoteConn ( encodeInt32 ControlResponse

705+ : encodeInt32 reqId

706+ : prependLength [encodeInt32 newId]

707+ )

708+ writeChan ourChannel (ConnectionOpened newId ReliableOrdered theirAddr)

709+ -- We add the new connection ID to the list of open connections only once the

710+ -- endpoint has been notified of the new connection (sendOn may fail)

711+ modifyIORef openConnections (IntSet.insert (fromIntegral newId))

712+

713+ -- Read a control response

714+ readControlResponse :: N.Socket -> ControlRequestId -> IO ()

715+ readControlResponse sock reqId = do

716+ response <- recvWithLength sock

717+ mmvar <- modifyMVar ourState $ \st ->

718+ return (pendingCtrlRequestsAt reqId ^= Nothing $ st, st ^. pendingCtrlRequestsAt reqId)

719+ case mmvar of

720+ Nothing ->

721+ return () -- Invalid request ID. TODO: We just ignore it?

722+ Just mvar ->

723+ putMVar mvar response

724+

725+ -- Close a connection

726+ closeConnection :: IORef IntSet -> ConnectionId -> IO ()

727+ closeConnection openConnections cid = do

728+ writeChan ourChannel (ConnectionClosed cid)

729+ decreaseRemoteRefCount theirEndPoint

730+ modifyIORef openConnections (IntSet.delete (fromIntegral cid))

731+

732+ -- Close the socket (if we don't have any outgoing connections)

733+ closeSocket :: N.Socket -> IO Bool

734+ closeSocket sock = do

735+ didClose <- modifyMVar theirState $ \st ->

736+ case st of

737+ RemoteEndPointValid remoteConn | remoteConn ^. remoteRefCount == 0 -> do

738+ sendOn remoteConn [encodeInt32 CloseSocket]

739+ N.sClose sock

740+ return (RemoteEndPointClosed, True)

741+ RemoteEndPointClosing _ resolved -> do

742+ putMVar resolved ()

743+ N.sClose sock

744+ return (RemoteEndPointClosed, True)

745+ _ ->

746+ return (st, False)

747+ when didClose $ modifyMVar_ ourState $ return . (connectionTo theirAddr ^= Nothing)

748+ return didClose

749+

750+ -- Read a message and output it on the endPoint's channel

751+ readMessage :: N.Socket -> ConnectionId -> IO ()

752+ readMessage sock connId = recvWithLength sock >>= writeChan ourChannel . Received connId

753+

754+ -- Arguments

755+ ourChannel = localChannel ourEndPoint

756+ ourState = localState ourEndPoint

757+ theirState = remoteState theirEndPoint

758+ theirAddr = remoteAddress theirEndPoint

636759

637760-- | Get the next connection ID

638761getNextConnectionId :: LocalEndPoint -> IO ConnectionId

639−getNextConnectionId endpoint =

640− modifyMVar (endpoint ^. localState) $ \st ->

641− return (nextConnectionId ^+= 1 $ st, st ^. nextConnectionId)

762+getNextConnectionId ourEndpoint =

763+ modifyMVar (localState ourEndpoint) $ \st -> do

764+ let connId = st ^. nextConnectionId

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

642766

643767--------------------------------------------------------------------------------

644768-- Constants --

652776-- Lens definitions --

653777--------------------------------------------------------------------------------

654778

655−transportState :: Lens TCPTransport (MVar TransportState)

656−transportState = lens _transportState (\st trans -> trans { _transportState = st })

657−

658−transportHost :: Lens TCPTransport N.HostName

659−transportHost = lens _transportHost (\host trans -> trans { _transportHost = host })

660−

661−transportPort :: Lens TCPTransport N.ServiceName

662−transportPort = lens _transportPort (\port trans -> trans { _transportPort = port })

663−

664779localEndPoints :: Lens TransportState (IntMap LocalEndPoint)

665780localEndPoints = lens _localEndPoints (\es st -> st { _localEndPoints = es })

666781

667−localAddress :: Lens LocalEndPoint EndPointAddress

668−localAddress = lens _localAddress (\addr ep -> ep { _localAddress = addr })

669−

670−localChannel :: Lens LocalEndPoint (Chan Event)

671−localChannel = lens _localChannel (\ch ep -> ep { _localChannel = ch })

672−

673−localState :: Lens LocalEndPoint (MVar LocalState)

674−localState = lens _localState (\st ep -> ep { _localState = st })

675−

676782pendingCtrlRequests :: Lens LocalState (IntMap (MVar [ByteString]))

677783pendingCtrlRequests = lens _pendingCtrlRequests (\rep st -> st { _pendingCtrlRequests = rep })

678784

682788nextConnectionId :: Lens LocalState ConnectionId

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

684790

685−localConnections :: Lens LocalState (Map EndPointAddress (IsClosing, MVar RemoteState))

686−localConnections = lens _localConnections (\cs st -> st { _localConnections = cs })

687−

688−remoteAddress :: Lens RemoteEndPoint EndPointAddress

689−remoteAddress = lens _remoteAddress (\addr ep -> ep { _remoteAddress = addr })

690−

691−remoteSocket :: Lens RemoteEndPoint N.Socket

692−remoteSocket = lens _remoteSocket (\sock ep -> ep { _remoteSocket = sock })

791+localConnections :: Lens LocalState (Map EndPointAddress RemoteEndPoint)

792+localConnections = lens _localConnections (\es st -> st { _localConnections = es })

693793

694794localEndPointAt :: EndPointId -> Lens TransportState (Maybe LocalEndPoint)

695795localEndPointAt ix = localEndPoints >>> intMapLens (fromIntegral ix)

697797pendingCtrlRequestsAt :: ControlRequestId -> Lens LocalState (Maybe (MVar [ByteString]))

698798pendingCtrlRequestsAt ix = pendingCtrlRequests >>> intMapLens (fromIntegral ix)

699799

700−connectionTo :: EndPointAddress -> Lens LocalState (Maybe (IsClosing, MVar RemoteState))

800+connectionTo :: EndPointAddress -> Lens LocalState (Maybe RemoteEndPoint)

701801connectionTo addr = localConnections >>> mapLens addr

802+

803+remoteRefCount :: Lens RemoteConnection Int

804+remoteRefCount = lens _remoteRefCount (\rc conn -> conn { _remoteRefCount = rc })