d063618 · distributed-process

162data LocalState = LocalState { _nextConnectionId :: ConnectionId

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

164 , _nextCtrlRequestId :: ControlRequestId

165− , _localConnections :: Map EndPointAddress (Maybe Closing, MVar RemoteState)

165+ , _localConnections :: Map EndPointAddress (IsClosing, MVar RemoteState)

166 }

167data RemoteEndPoint = RemoteEndPoint { _remoteAddress :: EndPointAddress

168 , _remoteSocket :: N.Socket

169 , sendTo :: [ByteString] -> IO ()

170 }

171−data RemoteState = InvalidEndPoint (FailedWith ConnectErrorCode)

171+data RemoteState = EndPointInvalid (FailedWith ConnectErrorCode)

172− | ValidEndPoint !RefCount RemoteEndPoint

172+ | EndPointValid !RefCount RemoteEndPoint

173− | DeletedEndPoint

173+ | EndPointClosed

174

175type EndPointId = Int32

176type ControlRequestId = Int32

177type RefCount = Int32

178−type Closing = MVar [ByteString]

178+type IsClosing = Maybe RemoteEndPoint

179type Conn = (LocalEndPoint, RemoteEndPoint)

180

181-- | Control headers

189-- Response sent by /B/ to /A/ when /A/ tries to connect

190data ConnectionRequestResponse =

191 ConnectionRequestAccepted -- ^ /B/ accepts the connection

192− | ConnectionRequestInvalidEndPoint -- ^ /A/ requested an invalid endpoint

192+ | ConnectionRequestEndPointInvalid -- ^ /A/ requested an invalid endpoint

193 | ConnectionRequestCrossed -- ^ /A/s request crossed with a request from /B/ (see protocols)

194 deriving Enum

195

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

298 theirState <- readMVar endPoint

299 case theirState of

300− InvalidEndPoint err ->

300+ EndPointInvalid err ->

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

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

303 modifyMVar state $ \st ->

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

305− ValidEndPoint _ ep ->

305+ EndPointValid _ ep ->

306 return $ Right ep

307− DeletedEndPoint ->

307+ EndPointClosed ->

308 go

309

310 state :: MVar LocalState

316 case result of

317 Right (ConnectionRequestAccepted, sock) -> do

318 theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock

319− putMVar theirMVar (ValidEndPoint 0 theirEndPoint)

319+ putMVar theirMVar (EndPointValid 0 theirEndPoint)

320− Right (ConnectionRequestInvalidEndPoint, sock) -> do

320+ Right (ConnectionRequestEndPointInvalid, sock) -> do

321− putMVar theirMVar (InvalidEndPoint (invalidAddress "Invalid endpoint"))

321+ putMVar theirMVar (EndPointInvalid (invalidAddress "Invalid endpoint"))

322 N.sClose sock

323 Right (ConnectionRequestCrossed, sock) ->

324 N.sClose sock

325 Left err ->

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

327 -- in which case we should close it

328− putMVar theirMVar (InvalidEndPoint err)

328+ putMVar theirMVar (EndPointInvalid err)

329 where

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

331 go = do

391increaseRemoteRefCount (ourEndPoint, theirEndPoint) = do

392 let theirAddr = theirEndPoint ^. remoteAddress

393 state = ourEndPoint ^. localState

394− -- First check if this remote endpoint is awaiting closure

394+ -- If this endpoint was closing, cancel that

395− (mClosing, theirMVar) <- modifyMVar state $ \st -> do

395+ theirMVar <- modifyMVar state $ \st -> do

396− let Just (closure, theirMVar) = st ^. connectionTo theirAddr

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

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

398− case mClosing of

397+ case isClosing of

399− Just mvar -> putMVar mvar []

398+ Nothing ->

400− _ -> return ()

399+ return (st, theirMVar)

400+ Just ep -> do

401+ putMVar theirMVar (EndPointValid 0 ep)

402+ return (connectionTo theirAddr ^= Just (Nothing, theirMVar) $ st, theirMVar)

401 -- Increase the refcount

403 -- Increase the refcount

402− modifyMVar_ theirMVar $ \(ValidEndPoint rc ep) ->

404+ modifyMVar_ theirMVar $ \(EndPointValid rc ep) ->

403− return (ValidEndPoint (rc + 1) ep)

405+ return (EndPointValid (rc + 1) ep)

404

406

405decreaseRemoteRefCount :: Conn -> IO ()

407decreaseRemoteRefCount :: Conn -> IO ()

406decreaseRemoteRefCount (ourEndPoint, theirEndPoint) = do

408decreaseRemoteRefCount (ourEndPoint, theirEndPoint) = do

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

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

411 -- empty when refcount reaches 0

413 -- empty when refcount reaches 0

412 -- TODO: deal with exceptions

414 -- TODO: deal with exceptions

413− ValidEndPoint rc ep <- takeMVar theirMVar

415+ EndPointValid rc ep <- takeMVar theirMVar

414 if rc > 1

416 if rc > 1

415 then do

417 then do

416− putMVar theirMVar (ValidEndPoint (rc - 1) ep)

418+ putMVar theirMVar (EndPointValid (rc - 1) ep)

417 else do

419 else do

418− reply <- sendRemoteRequest (ourEndPoint, theirEndPoint) CloseSocket

420+ sendTo theirEndPoint [encodeInt32 CloseSocket]

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

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

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

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

421− return (connectionTo theirAddr ^= Just (Just reply, mvar) $ st)

423+ return (connectionTo theirAddr ^= Just (Just ep, mvar) $ st)

422− void . forkIO $ do

423− -- Wait for the reply from the remote endpoint

424− canCloseSocket <- takeMVar reply

425− putStrLn $ "canCloseSocket = " ++ show canCloseSocket

426

424

427--------------------------------------------------------------------------------

425--------------------------------------------------------------------------------

428-- Incoming requests --

426-- Incoming requests --

437 return (theirAddress, ourEndPoint)

435 return (theirAddress, ourEndPoint)

438 case request of

436 case request of

439 Right (_, Nothing) -> do

437 Right (_, Nothing) -> do

440− sendMany sock [encodeInt32 ConnectionRequestInvalidEndPoint]

438+ sendMany sock [encodeInt32 ConnectionRequestEndPointInvalid]

441 N.sClose sock

439 N.sClose sock

442 Right (theirAddress, Just ourEndPoint) -> do

440 Right (theirAddress, Just ourEndPoint) -> do

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

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

453 N.sClose sock

451 N.sClose sock

454 else do

452 else do

455 theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock

453 theirEndPoint <- createRemoteEndPoint ourEndPoint theirAddress sock

456− putMVar mvar (ValidEndPoint 0 theirEndPoint)

454+ putMVar mvar (EndPointValid 0 theirEndPoint)

457 sendMany sock [encodeInt32 ConnectionRequestAccepted]

455 sendMany sock [encodeInt32 ConnectionRequestAccepted]

458 Left _ ->

456 Left _ ->

459 -- Invalid request

457 -- Invalid request

494 if connId >= firstNonReservedConnectionId

492 if connId >= firstNonReservedConnectionId

495 then readMessage connId

493 then readMessage connId

496 else do

494 else do

497− reqId <- recvInt32 sock

498 case toEnum (fromIntegral connId) of

495 case toEnum (fromIntegral connId) of

499− RequestConnectionId -> createNewConnection reqId

496+ RequestConnectionId -> recvInt32 sock >>= createNewConnection

500− ControlResponse -> readControlResponse reqId

497+ ControlResponse -> recvInt32 sock >>= readControlResponse

501− CloseConnection -> closeConnection reqId

498+ CloseConnection -> recvInt32 sock >>= closeConnection

502− CloseSocket -> closeSocket reqId

499+ CloseSocket -> closeSocket

503

500

504 -- Create a new connection

501 -- Create a new connection

505 createNewConnection :: ControlRequestId -> IO ()

502 createNewConnection :: ControlRequestId -> IO ()

529 putMVar mvar response

526 putMVar mvar response

530 go

527 go

531

528

532− -- Close a connection (in this case the reqId is abused for the connection ID)

529+ -- Close a connection

533 closeConnection :: ConnectionId -> IO ()

530 closeConnection :: ConnectionId -> IO ()

534 closeConnection cid = do

531 closeConnection cid = do

535 writeChan ourChannel (ConnectionClosed cid)

532 writeChan ourChannel (ConnectionClosed cid)

538 go

535 go

539

536

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

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

541− closeSocket :: ControlRequestId -> IO ()

538+ closeSocket :: IO ()

542− closeSocket cid = do

539+ closeSocket = do

543− let ok = encodeInt32 ControlResponse : encodeInt32 cid : prependLength [encodeInt32 (1 :: Int32)]

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

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

545 case st ^. connectionTo theirAddr of

541 case st ^. connectionTo theirAddr of

546− Just (Just _, _) -> do

542+ Just (Just _, mvar) -> do

547− sendTo theirEndPoint ok

543+ sendTo theirEndPoint [encodeInt32 CloseSocket]

548 N.sClose sock

544 N.sClose sock

545+ putMVar mvar EndPointClosed

549− return (st, True)

546+ return (connectionTo theirAddr ^= Nothing $ st, True)

550 Just (Nothing, theirMVar) -> do

547 Just (Nothing, theirMVar) -> do

551− ValidEndPoint rc _ <- readMVar theirMVar

548+ EndPointValid rc _ <- readMVar theirMVar

552 if rc == 0

549 if rc == 0

553 then do

550 then do

554− sendTo theirEndPoint ok

551+ sendTo theirEndPoint [encodeInt32 CloseSocket]

555 N.sClose sock

552 N.sClose sock

556− return (st, True)

553+ return (connectionTo theirAddr ^= Nothing $ st, True)

557 else do

554 else do

558 return (st, False)

555 return (st, False)

559 Nothing -> do

556 Nothing -> do

622nextConnectionId :: Lens LocalState ConnectionId

619nextConnectionId :: Lens LocalState ConnectionId

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

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

624

621

625−localConnections :: Lens LocalState (Map EndPointAddress (Maybe Closing, MVar RemoteState))

622+localConnections :: Lens LocalState (Map EndPointAddress (IsClosing, MVar RemoteState))

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

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

627

624

628remoteAddress :: Lens RemoteEndPoint EndPointAddress

625remoteAddress :: Lens RemoteEndPoint EndPointAddress

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

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

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

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

639

636

640−connectionTo :: EndPointAddress -> Lens LocalState (Maybe (Maybe Closing, MVar RemoteState))

637+connectionTo :: EndPointAddress -> Lens LocalState (Maybe (IsClosing, MVar RemoteState))

641connectionTo addr = localConnections >>> mapLens addr

638connectionTo addr = localConnections >>> mapLens addr