GitLab

Andrea Vezzosi pushed to branch wip/romes/async-rpc at Glasgow Haskell Compiler / GHC

Commits:

3e76e7c9
by Andrea Vezzosi at 2026-10-30T13:04:38+01:00
```
interp: single-threaded sendTHComputation
```

3 changed files:

compiler/GHC/Runtime/Interpreter/Process.hs
libraries/ghci/GHCi/Message.hs
libraries/ghci/GHCi/Message/Async.hs

Changes:

compiler/GHC/Runtime/Interpreter/Process.hs

@@ -4,6 +4,7 @@ module GHC.Runtime.Interpreter.Process
      Message(..)
    , sendMessage
    , sendTHComputation
 +  , sendMessageWithResponse
+   )
  where
@@ -16,9 +17,7 @@ import GHC.Utils.Panic
  import GHC.Utils.Exception as Ex
  import Data.Binary
 -import Control.Concurrent
  import Data.IORef
 -import Control.Concurrent.STM
  -- -----------------------------------------------------------------------------
  -- Top-level Message API
@@ -34,29 +33,27 @@ sendMessage i msg =
    where
      proc = instProcess i
 +sendMessageWithResponse :: Binary a => ExtInterpInstance c -> Message a -> (IO a -> IO b) -> IO b
 +sendMessageWithResponse i msg k =
 +  remoteCallWithResponse (interpPipe proc) msg k
 +  where
 +    proc = instProcess i
++
  -- | Send a message with a TH computation to the ext. interpreter (like
  -- 'RunTH', or 'RunModFinalizers') and handle 'THMessage' requests from the
 --- ext. interpreter to us in the meantime.
 ---
 --- This function blocks until the TH computation returns a value, but spawns in
 --- the background a thread to handle the 'THMessage's until the TH computation
 --- is fully done.
 +-- ext. interpreter to us until done, then wait for the response.
  sendTHComputation
    :: Binary a
    => ExtInterpInstance d
    -> (InboxId -> Message a)
    -> (forall b. Binary b => THMessage b -> IO (Bool, b))
    -- ^ Handle one message. Return True to end the loop.
 -  -- This handler must never uninterrutibly mask an expression which may get
 -  -- stuck. If it does, then 'sendTHComputation' will also get uninterruptibly stuck.
    -> IO a
 -sendTHComputation i msg reqHandler = mask $ \restore -> do
 -  -- inline "withAsync/wait"
 -  var  <- newEmptyTMVarIO
 -  mbox <- newMailbox (pipeAsyncMngr (interpPipe proc))
 -  ht   <- forkFinally (restore (handlerLoop mbox)) (atomically . putTMVar var)
 -  restore (sendMessage i (msg mbox) <* atomically (readTMVar var))
 -    `onException` uninterruptibleMask_ (throwTo ht CancelTHHandler)
 +sendTHComputation i msg reqHandler =
 +  withNewMailbox (pipeAsyncMngr (interpPipe proc)) $ \ mbox ->
 +    sendMessageWithResponse i (msg mbox) $ \ waitRes -> do
 +      handlerLoop mbox
 +      waitRes
    where
      proc = instProcess i
      handleOne mbox doneRef = do

libraries/ghci/GHCi/Message.hs

@@ -25,13 +25,12 @@ module GHCi.Message
    , getMessage, putMessage, getTHMessage, putTHMessage
    , Pipe(..), mkPipeFromHandles, mkPipeFromContinuations, closePipe
 -  , remoteCall, remoteTHCall
 +  , remoteCall, remoteCallWithResponse, remoteTHCall
    , handleRemoteCall, handleRemoteTHCall, writeReply
    , RemoteProcessException(..), PipeClosedException(..)
    , InboxId(..)
 -  , newMailbox
 -  , closeMailbox
 +  , withNewMailbox
    , BreakModule
    , BreakUnitId
@@ -660,8 +659,12 @@ serializeBCOs rbcos = parMap doChunk (chunkList 100 rbcos)
  -- | Send a 'Message' request to and wait for a reply from the interpreter
  -- (blocks on the calling side, doesn't necessarily block the interpreter)
  remoteCall :: Binary a => Pipe -> Message a -> IO a
 -remoteCall pipe msg = either throwIO (pure . runGet get)
 -                  =<< remoteAsyncReq pipe (putMessage msg) Nothing
 +remoteCall pipe msg = remoteCallWithResponse pipe msg id
++
 +remoteCallWithResponse :: Binary a => Pipe -> Message a -> (IO a -> IO b) -> IO b
 +remoteCallWithResponse pipe msg k =
 +  remoteAsyncReqWithResponse pipe (putMessage msg) Nothing $ \ m ->
 +    k $ either throwIO (pure . runGet get) =<< m
  -- | Handle a single 'Message' request on the interpreter and reply to the host
  handleRemoteCall :: Pipe

libraries/ghci/GHCi/Message/Async.hs

@@ -11,12 +11,14 @@ module GHCi.Message.Async
      -- * Interface for async bidirectional RPC
    , remoteAsyncReq, handleAsyncReq
 +  , remoteAsyncReqWithResponse
    , RemoteProcessException(..), DecodeAsyncReqException(..)
      -- ** Mailboxes
    , InboxId(..)
    , newMailbox
    , closeMailbox
 +  , withNewMailbox
      -- ** Async matching internals
    , AsyncManager(..), CorrelationId
@@ -247,24 +249,44 @@ newPendingReply pipe k = do
  --
  -- See Note [Exceptions during serialization] about exceptions the 'Put' argument might raise.
  remoteAsyncReq :: Pipe -> Put -> (Maybe InboxId) -> IO (Either SomeException LB.ByteString)
 -remoteAsyncReq pipe msg ii = newPendingReply pipe $ \ uq wv ->
 -  -- We might send a Cancel even if the request never got sent, but it will just
 -  -- get ignored in that case.
 -  handle (\e -> sendCancelMaybe uq e >> throwIO e) $ do
 +remoteAsyncReq pipe msg ii = remoteAsyncReqWithResponse pipe msg ii id
 -  -- Write request
 -  writePipe pipe (Request uq ii (runPut msg))
 -  -- Block waiting for reply
 -  atomically (takeTMVar wv)
 +remoteAsyncReqWithResponse :: Pipe
 +                           -> Put
 +                           -> Maybe InboxId
 +                           -> (IO (Either SomeException LB.ByteString) -> IO b)
 +                           -> IO b
 +remoteAsyncReqWithResponse pipe msg ii k = newPendingReply pipe $ \ uq wv -> do
 +  -- We use an MVar to only Cancel the request if the response did not get read by `k`.
 +  readRes <- newEmptyMVar
 +  -- We might send a Cancel even if the request never got sent, but it will just
 +  -- get ignored in that case.
 +  --
 +  -- Waiting for confirmation that the request got sent before "arming" the
 +  -- cancelation requires more collaboration from `fillPipe`, because
 +  -- writePipe's takeMVar is interruptible. Could be done if we generalize the
 +  -- type of `writeQueue` to `TQueue (IddMsg, (Either SomeException ()) -> STM ())`.
 +  cancelOnErrorMaybe uq readRes $ do
 +    -- Write request
 +    writePipe pipe (Request uq ii (runPut msg))
++
 +    k $ do
 +      -- Block waiting for reply
 +      atomically (takeTMVar wv) <* putMVar readRes ()
    where
 -    sendCancelMaybe uq e
 -      | Just PipeWriteClosedException <- fromException e = pure ()
 -      | otherwise = sendCancel uq
 -    sendCancel uq = void . atomically $ do
 -      -- we don't wait for Cancel to be written on the pipe,
 -      -- and won't get a reply: nowhere to report it anyway.
 -      writePipeAsync pipe (Cancel uq)
 +    cancelOnErrorMaybe :: CorrelationId -> MVar () -> IO a -> IO a
 +    cancelOnErrorMaybe uq readRes = handle (\ e -> sendCancelMaybe e >> throwIO e) where
 +      sendCancelMaybe = \case
 +        e | Just PipeWriteClosedException <- fromException e ->
 +          pure ()
 +          | otherwise ->
 +          tryTakeMVar readRes >>= \case
 +            Just () -> pure ()
 +            Nothing -> void . atomically $ do
 +              -- we don't wait for Cancel to be written on the pipe,
 +              -- and won't get a reply: nowhere to report it anyway.
 +              writePipeAsync pipe (Cancel uq)
  -- | Receive a single asynchronous request and send a matching reply
  -- asynchronously (the handler may fork a computation that eventually replies
@@ -352,6 +374,9 @@ closeMailbox :: AsyncManager -> InboxId -> IO ()
  closeMailbox AsyncManager{..} (InboxId ii) =
    atomically (modifyTVar mailboxQueues (IntMap.delete ii))
 +withNewMailbox :: AsyncManager -> (InboxId -> IO a) -> IO a
 +withNewMailbox am = bracket (newMailbox am) (closeMailbox am)
++
  -- | Get or create a new mailbox for this 'InboxId'
  getMailbox :: AsyncManager -> InboxId -> STM (TQueue (CorrelationId, LB.ByteString))
  getMailbox am (InboxId ii) = do