GitLab

Rodrigo Mesquita pushed to branch wip/romes/async-rpc at Glasgow Haskell Compiler / GHC

Commits:

11710a3d

by Rodrigo Mesquita at 2026-10-12T18:12:00+01:00

Asynchronous RPC with external interpreter

There's a single pipe to connect the host compiler to the external
interpreter. On this pipe, we have historically done synchronous
request/reply communication to invoke methods on the external
interpreter.

However, with the advent of multi-threaded debugging, it has become
necessary to do asynchronous calls on the external interpreter, where we
may block on the host waiting for a reply, but the external interpreter
should handle the request asynchronously and reply when ready. For
instance, we want to start the evaluation of the debuggee but not block
the interpreter waiting for it to finish or hit a breakpoint. Or e.g. we
want to resume two threads of the debuggee simulatenously instead of
blocking waiting for the resume result of each one sequentially.

And it turns out, we were already doing some ad-hoc asynchronous
communication to evaluate Template Haskell on the external interpreter.
It used to go as an ad-hoc "send request without expecting reply",
followed by handling on the host a set of 'THMessage's sent by the
interpreter, followed by "expect a response to arrive now". It worked
because there was still a fixed order of messages, but it fits an async
model in a non-ad-hoc way.

Asynchronous communication that unblocks the external interpreter main
thread will also make it possible to resolve issues such as #18852
(exceptions leaving the external interpreter in an unrecoverrable state)

The new design
--------------

The new invariant: every single request should receive a matching reply.

The pipe is bidirectional (the ext. interpreter can receive requests,
but it can also send requests for TH evaluation; the host can send
requests, but it can also receive TH requests from the interpreter).

When a request is sent in whichever direction, it always contains a
correlation-id header. The Sender places an empty MVar on a
pending-reply map keyed by id and waits on it for a response.
See 'remoteCall' and 'remoteTHCall'.

The pipe is associated to a thread which drains the read end. Because
it's bidirectional, two things may be read off the pipe:
- A Reply, with a correlation-id header, which we match against the
  pending requests map to unblock the corresponding pending-reply MVar
- A Request, which is queued to the pending-request queue.

To handle a message, we take a Request from the pending-request queue
and run a handler that produces a Reply. Currently, the handler is run
synchronously in the thread, but it is trivial to make the handler run
in a forked thread for certain ("interruptible") messages in an upcoming
patch (when it is necessary for the multi-threaded debugger).
See 'handleRemoteCall' and 'handleRemoteTHCall'.

The old interp lock now just protect the pipe reading and writing, it
doesn't need nor should lock the full remote procedure call.

11 changed files:

compiler/GHC/Runtime/Interpreter.hs
compiler/GHC/Runtime/Interpreter/JS.hs
compiler/GHC/Runtime/Interpreter/Process.hs
compiler/GHC/Runtime/Interpreter/Types.hs
compiler/GHC/Runtime/Interpreter/Wasm.hs
compiler/GHC/Tc/Gen/Splice.hs
libraries/ghci/GHCi/Message.hs
libraries/ghci/GHCi/Server.hs
libraries/ghci/GHCi/TH.hs
libraries/ghci/ghci.cabal.in
testsuite/tests/ghci/custom-external-interpreter-commands/Main.hs

Changes:

compiler/GHC/Runtime/Interpreter.hs

@@ -626,11 +626,9 @@ spawnIServ conf = do
                                            []
                                            (iservConfOpts    conf)
    interpPipe <- mkPipeFromHandles rh wh
 -  lock <- newMVar ()
    let process = InterpProcess
                    { interpHandle = ph
                    , interpPipe
 -                  , interpLock   = lock
+                   }
    pending_frees <- newMVar []

compiler/GHC/Runtime/Interpreter/JS.hs

@@ -127,11 +127,9 @@ startTHRunnerProcess interp_js settings = do
    std_in <- readIORef interp_in
    interpPipe <- mkPipeFromHandles rh wh
 -  lock <- newMVar ()
    let proc = InterpProcess
                { interpHandle = hdl
                , interpPipe
 -              , interpLock   = lock
+               }
    pure (std_in, proc)
@@ -382,7 +380,7 @@ jsLinkPlan logger tmpfs tmp_dir ar_cache link_cfg cfg inst link_plan = do
    --   Haskell server.
    server_started <- jsServerStarted <$> readMVar (instJSState (instExtra inst))
    if server_started
 -    then sendMessageNoResponse inst $ LoadObj all_js
 +    then sendMessage inst $ LoadObj all_js
      else jsLoadFile            inst all_js
    ----------------------------------------------------------------

compiler/GHC/Runtime/Interpreter/Process.hs

@@ -2,16 +2,9 @@ module GHC.Runtime.Interpreter.Process
+   (
    -- * Message API
      Message(..)
 -  , DelayedResponse (..)
 -  -- * Top-level message API (these acquire/release a lock)
    , sendMessage
 -  , sendMessageNoResponse
 -  , sendMessageDelayedResponse
 -  , receiveDelayedResponse
 -  -- * Nested message API (these require the interpreter to already be locked)
 -  , sendAnyValue
 -  , receiveAnyValue
 -  , receiveTHMessage
 +  , sendTHComputation
 +  , RemoteCallFailed(..)
+   )
  where
@@ -27,97 +20,67 @@ import GHC.Utils.Exception as Ex
  import Data.Binary
  import System.Exit
  import System.Process
 -import Control.Concurrent.MVar (MVar, withMVar, takeMVar, putMVar, isEmptyMVar)
+-
 -data DelayedResponse a = DelayedResponse
 +import Control.Concurrent
 +import Data.IORef
  -- -----------------------------------------------------------------------------
  -- Top-level Message API
 --- | Send a message to the interpreter process that doesn't expect a response
 ---   (locks the interpreter while sending)
 -sendMessageNoResponse :: ExtInterpInstance d -> Message () -> IO ()
 -sendMessageNoResponse i m =
 -  withLock i $ writeInterpProcess (instProcess i) (putMessage m)
+-
 --- | Send a message to the interpreter that expects a response
 ---   (locks the interpreter while until the response is received)
 +-- | Send a message to the interpreter and expect a response, blocking until it arrives.
 +--
 +-- This issues an asynchronous RPC on the external interpreter which does not
 +-- block the external interpreter thread. We block here on the host because we
 +-- wait for the asynchronous response to arrive.
 +--
 +-- Throws 'RemoteCallFailed' if the remote call threw an exception in the
 +-- external interpreter
  sendMessage :: Binary a => ExtInterpInstance d -> Message a -> IO a
 -sendMessage i m = withLock i $ callInterpProcess (instProcess i) m
+-
 --- | Send a message to the interpreter process whose response is expected later
 +sendMessage i msg = unwrapRight =<<
 +  remoteCall (interpPipe proc) msg
 +    `catchException` \(e :: SomeException) -> handleInterpProcessFailure proc e
 +  where
 +    proc = instProcess i
 +    unwrapRight :: Either SomeException b -> IO b
 +    unwrapRight (Right a) = pure a
 +    unwrapRight (Left er) = do
 +      throwIO (RemoteCallFailed er)
++
 +-- | 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 is useful to avoid forgetting to receive the value and to ensure that
 --- the type of the response isn't lost. Use receiveDelayedResponse to read it.
 --- (locks the interpreter until the response is received using
 --- `receiveDelayedResponse`)
 -sendMessageDelayedResponse :: ExtInterpInstance d -> Message a -> IO (DelayedResponse a)
 -sendMessageDelayedResponse i m = do
 -  lock i
 -  writeInterpProcess (instProcess i) (putMessage m)
 -  pure DelayedResponse
+-
 --- | Expect a delayed result to be received now
 -receiveDelayedResponse :: Binary a => ExtInterpInstance d -> DelayedResponse a -> IO a
 -receiveDelayedResponse i DelayedResponse = do
 -  ensureLocked i
 -  r <- readInterpProcess (instProcess i) get
 -  unlock i
 -  pure r
+-
 --- -----------------------------------------------------------------------------
 --- Nested Message API
+-
 --- | Send any value (requires locked interpreter)
 -sendAnyValue :: Binary a => ExtInterpInstance d -> a -> IO ()
 -sendAnyValue i m = ensureLocked i >> writeInterpProcess (instProcess i) (put m)
+-
 --- | Expect a value to be received (requires locked interpreter)
 -receiveAnyValue :: ExtInterpInstance d -> Get a -> IO a
 -receiveAnyValue i get = ensureLocked i >> readInterpProcess (instProcess i) get
+-
 --- | Wait for a Template Haskell message (requires locked interpreter)
 -receiveTHMessage :: ExtInterpInstance d -> IO THMsg
 -receiveTHMessage i = ensureLocked i >> receiveAnyValue i getTHMessage
+-
 --- -----------------------------------------------------------------------------
+-
 -getLock :: ExtInterpInstance d -> MVar ()
 -getLock = interpLock . instProcess
+-
 -withLock :: ExtInterpInstance d -> IO a -> IO a
 -withLock i f = withMVar (getLock i) (const f)
+-
 -lock :: ExtInterpInstance d -> IO ()
 -lock i = takeMVar (getLock i)
+-
 -unlock :: ExtInterpInstance d -> IO ()
 -unlock i = putMVar (getLock i) ()
+-
 -ensureLocked :: ExtInterpInstance d -> IO ()
 -ensureLocked i =
 -  isEmptyMVar (getLock i) >>= \case
 -    False -> panic "ensureLocked: external interpreter not locked"
 -    _     -> pure ()
+-
+-
 --- | Send a 'Message' and receive the response from the interpreter process
 -callInterpProcess :: Binary a => InterpProcess -> Message a -> IO a
 -callInterpProcess i msg =
 -  remoteCall (interpPipe i) msg
 -    `catchException` \(e :: SomeException) -> handleInterpProcessFailure i e
+-
 --- | Read a value from the interpreter process
 -readInterpProcess :: InterpProcess -> Get a -> IO a
 -readInterpProcess i get =
 -  readPipe (interpPipe i) get
 -    `catchException` \(e :: SomeException) -> handleInterpProcessFailure i e
+-
 --- | Send a value to the interpreter process
 -writeInterpProcess :: InterpProcess -> Put -> IO ()
 -writeInterpProcess i put =
 -  writePipe (interpPipe i) put
 -    `catchException` \(e :: SomeException) -> handleInterpProcessFailure i e
 +-- 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.
 +sendTHComputation
 +  :: Binary a
 +  => ExtInterpInstance d
 +  -> Message a
 +  -> (forall b. Binary b => THMessage b -> IO (Bool, b))
 +  -- ^ Handle one message. Return True to end the loop.
 +  -> IO a
 +sendTHComputation i msg reqHandler = mask $ \restore -> do
 +  -- inline "withAsync/wait"
 +  var <- newEmptyMVar
 +  ht  <- forkFinally (restore handlerLoop) (putMVar var)
 +  restore (sendMessage i msg <* takeMVar var)
 +    `finally` cancel ht var
 +  where
 +    proc = instProcess i
 +    handleOne doneRef = do
 +      handleRemoteTHCall (interpPipe proc) $ \req -> do
 +        (stop, res) <- reqHandler req
 +        writeIORef doneRef stop
 +        return res
 +    handlerLoop = do
 +      doneRef <- newIORef False
 +      let loop = do
 +            handleOne doneRef
 +            isDone <- readIORef doneRef
 +            if isDone then return ()
 +                      else loop
 +       in loop
 +    cancel ht var = uninterruptibleMask_ (throwTo ht CancelTHHandler <* takeMVar var)
  handleInterpProcessFailure :: InterpProcess -> SomeException -> IO a
  handleInterpProcessFailure i e = do
@@ -130,3 +93,17 @@ handleInterpProcessFailure i e = do
        terminateProcess hdl
        _ <- waitForProcess hdl
        throw e
++
 +-- | The external interpreter server threw an exception while executing this
 +-- remote call. The ext. interp is still live and ready to take further
 +-- commands (such as 'Shutdown').
 +data RemoteCallFailed = RemoteCallFailed SomeException deriving Show
 +instance Exception RemoteCallFailed where
 +  displayException (RemoteCallFailed e)
 +    = "A remote call in the external interpreter threw an exception:\n"
 +    ++ displayException e
++
 +data CancelTHHandler = CancelTHHandler deriving Show
 +instance Exception CancelTHHandler where
 +  fromException = asyncExceptionFromException
 +  toException = asyncExceptionToException

compiler/GHC/Runtime/Interpreter/Types.hs

@@ -111,7 +111,6 @@ type WasmInterp = ExtInterpState WasmInterpConfig ()
  data InterpProcess = InterpProcess
    { interpPipe   :: !Pipe           -- ^ Pipe to communicate with the server
    , interpHandle :: !ProcessHandle  -- ^ Process handle of the server
 -  , interpLock   :: !(MVar ())      -- ^ Lock to prevent concurrent access to the stream
+   }
  -- | Status of an external interpreter

compiler/GHC/Runtime/Interpreter/Wasm.hs

@@ -73,14 +73,12 @@ spawnWasmInterp WasmInterpConfig {..} = do
    hSetBuffering rh NoBuffering
    interpPipe <- mkPipeFromHandles rh wh
    pending_frees <- newMVar []
 -  lock <- newMVar ()
    pure
      $ ExtInterpInstance
        { instProcess =
            InterpProcess
              { interpHandle = ph,
 -              interpPipe,
 -              interpLock = lock
 +              interpPipe
              },
          instPendingFrees = pending_frees,
          instExtra = ()

compiler/GHC/Tc/Gen/Splice.hs

@@ -170,6 +170,7 @@ import Data.IORef
  import GHC.Parser.HaddockLex (lexHsDoc)
  import GHC.Parser (parseIdentifier)
  import GHC.Rename.Doc (rnHsDoc)
 +import Data.Function
  {-
  Note [Template Haskell state diagram]
@@ -1165,11 +1166,12 @@ runRemoteModFinalizers (ThModFinalizers finRefs) = do
        case th_state of
          Nothing -> return () -- TH was not started, nothing to do
          Just fhv -> do
 -          r <- liftIO $ withForeignRef fhv $ \st ->
 -            withForeignRefs finRefs $ \qrefs ->
 -              sendMessageDelayedResponse inst (RunModFinalizers st qrefs)
 -          () <- runRemoteTH inst []
 -          qr <- liftIO $ receiveDelayedResponse inst r
 +          qr <- withRunInIO $ \runInIO -> do
 +            recovers <- newIORef []
 +            withForeignRef fhv $ \st ->
 +              withForeignRefs finRefs $ \qrefs ->
 +                sendTHComputation inst (RunModFinalizers st qrefs) $ \msg ->
 +                  runInIO (runRemoteTH recovers msg)
            checkQResult qr
  runQResult
@@ -1824,48 +1826,47 @@ runTH ty fhv = do
        rstate <- getTHState inst
        loc <- location
        -- run a remote TH request
 -      r <- liftIO $
 +      qr <- withRunInIO $ \runInIO -> do
 +        recovers <- newIORef []
          withForeignRef rstate $ \state_hv ->
 -        withForeignRef fhv $ \q_hv ->
 -          sendMessageDelayedResponse inst (RunTH state_hv q_hv ty (Just loc))
 -      -- respond to requests from the interpreter
 -      runRemoteTH inst []
 -      -- get the final result
 -      qr <- liftIO $ receiveDelayedResponse inst r
 +          withForeignRef fhv $ \q_hv ->
 +            sendTHComputation inst (RunTH state_hv q_hv ty (Just loc)) $ \msg ->
 +              runInIO (runRemoteTH recovers msg)
        bs <- checkQResult qr
        return $! runGet get (LB.fromStrict bs)
 --- | communicate with a remotely-running TH computation until it finishes.
 +-- | Communicate with a remotely-running TH computation until it finishes.
  -- See Note [Remote Template Haskell] in libraries/ghci/GHCi/TH.hs.
  runRemoteTH
 -  :: ExtInterpInstance d
 -  -> [Messages TcRnMessage]   --  saved from nested calls to qRecover
 -  -> TcM ()
 -runRemoteTH inst recovers = do
 -  THMsg msg <- liftIO $ receiveTHMessage inst
 +  :: Binary a
 +  => IORef [Messages TcRnMessage] -- ^ saved from nested calls to qRecover
 +  -> THMessage a
 +  -> TcM (Bool, a)
 +runRemoteTH recovers msg = do
    case msg of
 -    RunTHDone -> return ()
 +    RunTHDone -> return (True, ())
      StartRecover -> do -- Note [TH recover with -fexternal-interpreter]
        v <- getErrsVar
        msgs <- readTcRef v
        writeTcRef v emptyMessages
 -      runRemoteTH inst (msgs : recovers)
 +      modifyIORef recovers (msgs:) & liftIO
 +      return (False, ())
      EndRecover caught_error -> do
 -      let (prev_msgs, rest) = case recovers of
 -             [] -> panic "EndRecover"
 -             a : b -> (a,b)
 +      (prev_msgs, rest) <- liftIO $
 +        readIORef recovers >>= \case
 +          [] -> panic "EndRecover"
 +          a : b -> pure (a,b)
        v <- getErrsVar
        warn_msgs <- getWarningMessages <$> readTcRef v
        -- keep the warnings only if there were no errors
        writeTcRef v $ if caught_error
          then prev_msgs
          else mkMessages warn_msgs `unionMessages` prev_msgs
 -      runRemoteTH inst rest
 +      writeIORef recovers rest & liftIO
 +      return (False, ())
      _other -> do
 -      r <- handleTHMessage msg
 -      liftIO $ sendAnyValue inst r
 -      runRemoteTH inst recovers
 +      (False,) <$> handleTHMessage msg
  -- | Check a QResult
  checkQResult :: QResult a -> TcM a

libraries/ghci/GHCi/Message.hs

 -{-# LANGUAGE GADTs, DeriveGeneric, StandaloneDeriving, ScopedTypeVariables,
 -    GeneralizedNewtypeDeriving, ExistentialQuantification, RecordWildCards,
 -    CPP, NamedFieldPuns #-}
 +{-# LANGUAGE RecursiveDo, GADTs, DeriveGeneric, StandaloneDeriving, ScopedTypeVariables,
 +    GeneralizedNewtypeDeriving, ExistentialQuantification, RecordWildCards, CPP,
 +    NamedFieldPuns, LambdaCase, TypeAbstractions, TypeApplications,
 +    LinearTypes, RankNTypes, UnboxedTuples, BlockArguments, TupleSections, MagicHash #-}
  {-# OPTIONS_GHC -fno-warn-name-shadowing -fno-warn-orphans #-}
  -- |
@@ -22,10 +23,21 @@ module GHCi.Message
    , ResumeContext(..)
    , QState(..)
    , getMessage, putMessage, getTHMessage, putTHMessage
 -  , Pipe, mkPipeFromHandles, mkPipeFromContinuations, remoteCall, remoteTHCall, readPipe, writePipe
++
 +  , Pipe, mkPipeFromHandles, mkPipeFromContinuations
 +  , remoteCall, remoteTHCall
 +  , handleRemoteCall, handleRemoteTHCall
++
    , BreakModule
    , BreakUnitId
    , LoadedDLL
++
 +  -- * Low-level interface for async bidirectional RPC
 +  , withAsyncReq
 +  , withAsyncHandler
 +  -- ** Consume `reply` exactly once
 +  , WriteReply -- abstract
 +  , writeReply
    ) where
  import Prelude -- See note [Why do we import Prelude here?]
@@ -45,6 +57,7 @@ import qualified GHC.Exts.Heap as Heap
  import GHC.ForeignSrcLang
  import GHC.Fingerprint
  import GHC.Conc (pseq, par)
 +import GHC.IO (IO(..))
  import Control.Concurrent
  import Control.DeepSeq
  import Control.Exception
@@ -63,6 +76,8 @@ import Data.Dynamic
  import Data.Typeable (TypeRep)
  import Data.IORef
  import Data.Map (Map)
 +import Data.IntMap (IntMap)
 +import qualified Data.IntMap as IntMap
  import Foreign
  import GHC.Generics
  import GHC.Stack.CCS
@@ -71,6 +86,11 @@ import qualified GHC.Boot.TH.Monad         as TH
  import System.Exit
  import System.IO
  import System.IO.Error
 +import Control.Monad
 +import Control.Concurrent.STM
 +import Data.Bifunctor
 +import GHC.Base (State#)
 +import GHC.Exts (RealWorld)
  -- -----------------------------------------------------------------------------
  -- The RPC protocol between GHC and the interactive server
@@ -693,9 +713,11 @@ serializeBCOs rbcos = parMap doChunk (chunkList 100 rbcos)
  -- | An opaque pipe for bidirectional binary data transmission.
  data Pipe = Pipe
 -  { getSome :: !(IO ByteString)
 -  , putAll :: !(B.Builder -> IO ())
 -  , pipeLeftovers :: !(IORef (Maybe ByteString))
 +  { getSome         :: !(IO ByteString)
 +  , putAll          :: !(B.Builder -> IO ())
 +  , pipeLeftovers   :: !(IORef (Maybe ByteString))
 +  , pipeLock        :: !(MVar ()) -- ^ Lock to prevent concurrent access to the stream
 +  , pipeAsyncMngr   :: !AsyncManager
+   }
  -- | Make a 'Pipe' from a 'Handle' to read and a 'Handle' to write.
@@ -705,30 +727,193 @@ mkPipeFromHandles pipeRead pipeWrite = do
        putAll b = do
          B.hPutBuilder pipeWrite b
          hFlush pipeWrite
 -  pipeLeftovers <- newIORef Nothing
 -  pure $ Pipe { getSome, putAll, pipeLeftovers }
 +  mkPipeFromContinuations getSome putAll
  -- | Make a 'Pipe' from a reader function and a writer function.
  mkPipeFromContinuations :: IO ByteString -> (B.Builder -> IO ()) -> IO Pipe
 -mkPipeFromContinuations getSome putAll = do
 -  pipeLeftovers <- newIORef Nothing
 -  pure $ Pipe { getSome, putAll, pipeLeftovers }
 +mkPipeFromContinuations getSome putAll = mdo
 +  pipeLeftovers   <- newIORef Nothing
 +  pipeLock        <- newMVar ()
 +  mdo
 +    pipeAsyncMngr <- newAsyncManager p
 +    let p = Pipe { .. }
 +    pure p
++
 +-- | 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 (Either SomeException a)
 +remoteCall pipe msg = bimap fromSerializableException id . runGet get
 +                  <$> withAsyncReq pipe (putMessage msg)
++
 +-- | Handle a single 'Message' request on the interpreter and reply to the host
 +handleRemoteCall :: Pipe
 +                 -> (Msg -> IO Msg) {-^ Msg hook -}
 +                 -> (forall a. (Show a, Binary a) => Message a -> IO a)
 +                 -> IO ()
 +handleRemoteCall pipe hook k =
 +  withAsyncHandler pipe getMessage \writer msg0 -> IO \s0 ->
 +    case hook msg0 |#> s0 of -- unIO for linearity checking
 +      (# s1, Msg msg #) -> case k msg |#> s1 of
 +        (# s2, a #) -> writeReply writer (put a) |#> s2
++
 +-- | Send a 'THMessage' request to the host and wait for a reply from the host
 +-- to the interpreter
 +remoteTHCall :: Binary a => Pipe -> THMessage a -> IO (Either SomeException a)
 +remoteTHCall pipe msg = bimap fromSerializableException id . runGet get
 +                    <$> withAsyncReq pipe (putTHMessage msg)
++
 +-- | Handle a 'THMessage' request on the host and reply to the interpreter.
 +handleRemoteTHCall
 +  :: Pipe -> (forall a. (Show a, Binary a) => THMessage a -> IO a) -> IO ()
 +handleRemoteTHCall pipe k =
 +  withAsyncHandler pipe getTHMessage \writer (THMsg msg) -> IO \s0 ->
 +    case k msg |#> s0 of -- unIO for linearity checking
 +      (# s1, a #) -> writeReply writer (put a) |#> s1
++
 +--------------------------------------------------------------------------------
 +-- * Matching requests with async replies
 +--------------------------------------------------------------------------------
++
 +-- | Send an asynchronous request and block waiting for a matching reply
 +withAsyncReq :: Pipe -> Put -> IO LB.ByteString
 +withAsyncReq pipe msg = do
++
 +  -- Mark pending request
 +  uq <- freshReqId (pipeAsyncMngr pipe)
 +  wv <- newEmptyMVar
 +  atomicModifyIORef' (pendingSent (pipeAsyncMngr pipe)) $
 +    \m -> (IntMap.insert uq wv m, ())
++
 +  -- Write request
 +  writePipe pipe (Request uq (runPut msg))
++
 +  -- Block waiting for reply
 +  takeMVar wv
++
 +-- | Receive a single asynchronous request and send a matching reply
 +-- asynchronously (the handler may fork a computation that eventually replies
 +-- and can return before actually replying, as long as it eventually does).
 +--
 +-- The handler receives as an argument the function to use to reply to the request.
 +-- It must be used exactly once (all requests expect a reply, and you can't
 +-- reply more than once). Linear types enforce the handler replies in every
 +-- branch, and if the handler otherwise throws, it is handled by 'withAsyncHandler'.
 +withAsyncHandler :: Pipe -> Get a -> (WriteReply Put %1 -> a -> IO r) -> IO r
 +withAsyncHandler pipe getit k = do
 +  ack <- newEmptyMVar
++
 +  -- Block waiting for request
 +  (uq, bs) <- atomically (readTQueue (pendingRecv (pipeAsyncMngr pipe)))
++
 +  let reply p = do
 +        writePipe pipe (Reply uq (runPut $ either (put . toSerializableException) id p))
 +        putMVar ack ()
++
 +  -- Run the continuation which must reply exactly once with 'writeReply'
 +  k (WriteReply reply) (runGet getit bs)
++
 +    -- If there's an exception before the reply is sent, reply with an error instead.
 +    `catch` \(e :: SomeException) -> do
++
 +        tryTakeMVar ack >>= \case
 +          Just () ->
 +            -- Already replied
 +            throwIO e
 +          Nothing -> do
 +            -- Reply with an error
 +            reply (Left e)
 +            throwIO e
++
 +-- ** Async matching "internals" -----------------------------------------------
++
 +data AsyncManager = AsyncManager
 +  { pendingSent :: !(IORef (IntMap (MVar LB.ByteString)))
 +    -- ^ When a request is sent, keep track that we're expecting a reply. The
 +    -- MVar is signaled when the reply arrives.
 +  , pendingRecv :: !(TQueue (CorrelationId, LB.ByteString))
 +    -- ^ A queue of pending received requests. When a request is received,
 +    -- queue it here. A request handler will pop the next request from it.
 +    -- (Nothing except the 'drainPipe' thread should read from the pipe).
 +    --
 +    -- Seq contract: Dequeue from left, enqueue on right
 +  , asyncDrainTid :: !ThreadId
 +  , asyncNextReq  :: !(IORef Int)
 +  }
++
 +type CorrelationId = Int
++
 +-- | We have a single type for requests and replies because both the host and
 +-- interpreter may send requests *and* replies on the write end *and* on the read end.
 +--
 +-- For instance:
 +--  1. The host sends a RunTH request on its Host Write End (HWE)
 +--  2. The interpreter sends a NewName TH request on its Interpreter Write End (IWE)
 +--  3. The host sends a reply to the TH request on its HWE
 +--  4. The interpreter sends the reply to the RunTH request on its IWE
 +data IddMsg  = Request !CorrelationId LB.ByteString
 +             | Reply   !CorrelationId LB.ByteString
++
 +instance Binary IddMsg where
 +  put (Request i bs) = putWord8 0 >> put i >> put bs
 +  put (Reply i bs)   = putWord8 1 >> put i >> put bs
 +  get = do
 +    b <- getWord8
 +    case b of
 +      0 -> Request <$> get <*> get
 +      1 -> Reply   <$> get <*> get
 +      _ -> error $ "impossible: Binary IddMsg - " ++ show b
 -remoteCall :: Binary a => Pipe -> Message a -> IO a
 -remoteCall pipe msg = do
 -  writePipe pipe (putMessage msg)
 -  readPipe pipe get
 +newAsyncManager :: Pipe -> IO AsyncManager
 +newAsyncManager pipe = do
 +  asyncNextReq <- newIORef 0
 +  pendingSent  <- newIORef IntMap.empty
 +  pendingRecv  <- newTQueueIO
 -writePipe :: Pipe -> Put -> IO ()
 -writePipe Pipe{..} put = putAll $ execPut put
 +  mdo
 +    -- Fork thread to drain read end.
 +    asyncDrainTid  <- forkIO (drainPipe pipe am)
 +    -- todo: proper thread clean up and no leaks
 -remoteTHCall :: Binary a => Pipe -> THMessage a -> IO a
 -remoteTHCall pipe msg = do
 -  writePipe pipe (putTHMessage msg)
 -  readPipe pipe get
 +    let am = AsyncManager{..}
 +    pure am
 -readPipe :: Pipe -> Get a -> IO a
 -readPipe Pipe{..} get = do
 +-- | Drains the read end of the Pipe and match 'Reply'es against pending 'Request's.
 +--
 +-- In pratice, we'll have both a thread draining the read end of the host's
 +-- 'Pipe' and a thread draining the read end of the interpreter's 'Pipe', since
 +-- both processes construct a 'Pipe' and send and receive messages:
 +drainPipe :: Pipe -> AsyncManager -> IO ()
 +drainPipe p am = forever $ do
 +  readPipe p >>= \case
 +    -- Receive a reply to a pending notification
 +    Reply uq bs -> do
 +      wv <- atomicModifyIORef' (pendingSent am) $ \m -> do
 +        case IntMap.lookup uq m of
 +          Nothing -> error "drainPipe: Received reply to no pending request"
 +          Just wv -> (IntMap.delete uq m, wv)
 +      putMVar wv bs -- signal matching request
 +    -- Receive a request and queue it
 +    Request uq bs ->
 +      atomically $ writeTQueue (pendingRecv am) (uq,bs)
++
 +freshReqId :: AsyncManager -> IO CorrelationId
 +freshReqId AsyncManager{asyncNextReq} =
 +  atomicModifyIORef' asyncNextReq $ \i -> (i+1, i)
++
 +--------------------------------------------------------------------------------
 +-- * Pipe thread-safe read/write
 +--------------------------------------------------------------------------------
 +-- All messages on the wire are 'IddMsg', which contain a correlation header
 +-- and identify whether they are requests or replies.
++
 +writePipe :: Pipe -> IddMsg -> IO ()
 +writePipe p@Pipe{..} = withLock p . putAll . execPut . put
++
 +-- | Only the 'drainPipe' thread should read from the pipe directly. Other
 +-- functions should wait for the matching reply to be delivered asynchronously.
 +-- (see 'withAsyncReq'). It doesn't lock to read bc it assumes single reader is 'drainPipe'.
 +readPipe :: Pipe -> IO IddMsg
 +readPipe Pipe{..} = do
    leftovers <- readIORef pipeLeftovers
    m <- getBin getSome get leftovers
    case m of
@@ -758,3 +943,22 @@ getBin getsome get leftover = go leftover (runGetIncremental get)
          else go Nothing (fun (Just b))
     go _lft (Fail _rest _off str) =
       throwIO (ErrorCall ("getBin: " ++ str))
++
 +withLock :: Pipe -> IO c -> IO c
 +withLock Pipe{..} = bracket (takeMVar pipeLock) (putMVar pipeLock) . const
++
 +--------------------------------------------------------------------------------
 +-- * Consume `reply` exactly once
 +--------------------------------------------------------------------------------
++
 +-- | Every request should get exactly one reply.
 +data WriteReply a where
 +  WriteReply :: (Either SomeException a -> IO ()) %1 -> WriteReply a
++
 +-- | Write a reply exactly once (consuming the ability to write more replies)
 +writeReply :: WriteReply a %1 -> a -> IO ()
 +writeReply (WriteReply r) = \x -> r (Right x)
++
 +-- | Infix, linear, 'unIO'
 +(|#>) :: IO a %1 -> State# RealWorld -> (# State# RealWorld, a #)
 +(|#>) (IO fs) s = fs s

libraries/ghci/GHCi/Server.hs

 -{-# LANGUAGE CPP, RankNTypes, RecordWildCards, GADTs, ScopedTypeVariables #-}
 +{-# LANGUAGE LambdaCase, CPP, RankNTypes, RecordWildCards, GADTs, ScopedTypeVariables #-}
  module GHCi.Server
    ( MessageHook
    , CustomMessageHandler
@@ -15,6 +15,7 @@ import GHCi.Signals
  import GHCi.TH
  import GHCi.Message
  import Data.ByteString (ByteString)
 +import Data.IORef
  #if defined(wasm32_HOST_ARCH)
  import qualified Data.ByteString.Builder as B
  import qualified Data.ByteString.Internal as B
@@ -60,27 +61,26 @@ servWithCustom
    -> (forall a .IO a -> IO a)
    -> CustomMessageHandler
    -> IO ()
 -servWithCustom verbose hook pipe restore customHandler = loop
 +servWithCustom verbose hook pipe restore customHandler = newIORef False >>= loop
   where
 -  loop = do
 +  loop doneRef = do
      when verbose $ trace "reading pipe..."
 -    Msg msg <- readPipe pipe getMessage >>= hook
 +    handleRemoteCall pipe hook $ \(msg :: Message a) -> do
 -    discardCtrlC
 +      discardCtrlC
 -    when verbose $ trace ("msg: " ++ (show msg))
 -    case msg of
 -      CustomMessage tag payload -> handleCustom tag payload
 -      Shutdown -> return ()
 -      RunTH st q ty loc -> wrapRunTH $ runTH pipe st q ty loc
 -      RunModFinalizers st qrefs -> wrapRunTH $ runModFinalizerRefs pipe st qrefs
 -      _other -> run msg >>= reply
 +      when verbose $ trace ("msg: " ++ (show msg))
 +      case msg of
 +        CustomMessage tag payload -> handleCustom tag payload
 +        Shutdown -> writeIORef doneRef True >> reply ()
 +        -- TODO: These two should probably be done off the main thread
 +        RunTH st q ty loc -> wrapRunTH $ runTH pipe st q ty loc
 +        RunModFinalizers st qrefs -> wrapRunTH $ runModFinalizerRefs pipe st qrefs
 +        _other -> run msg >>= reply
 -  reply :: forall a. (Binary a, Show a) => a -> IO ()
 -  reply r = do
 -    when verbose $ trace ("writing pipe: " ++ show r)
 -    writePipe pipe (put r)
 -    loop
 +    isDone <- readIORef doneRef
 +    if isDone then return ()
 +              else loop doneRef
    handleCustom tag payload = do
      mresp <- customHandler tag payload
@@ -89,17 +89,24 @@ servWithCustom verbose hook pipe restore customHandler = loop
        Nothing ->
          error $ "GHCi.Server: unhandled CustomMessage with tag " ++ show tag
 +  reply :: forall a. Show a => a -> IO a
 +  reply r = do
 +    when verbose $ trace ("writing pipe: " ++ show r)
 +    return r -- the result of the handler gets sent
++
    -- Run some TH code, which may interact with GHC by sending
    -- THMessage requests, and then finally send RunTHDone followed by a
    -- QResult.  For an overview of how TH works with Remote GHCi, see
    -- Note [Remote Template Haskell] in libraries/ghci/GHCi/TH.hs.
 -  wrapRunTH :: forall a. (Binary a, Show a) => IO a -> IO ()
 +  wrapRunTH :: forall a. (Binary a, Show a) => IO a -> IO (QResult a)
    wrapRunTH io = do
      when verbose $ trace "wrapRunTH..."
      r <- try io
      when verbose $ trace "wrapRunTH done."
      when verbose $ trace "writing RunTHDone."
 -    writePipe pipe (putTHMessage RunTHDone)
 +    remoteTHCall pipe RunTHDone >>= \case
 +      Right () -> pure ()
 +      Left err -> throwIO (GHCiQUncaughtExceptionOnHost err)
      case r of
        Left e
          | Just (GHCiQException _ err) <- fromException e  -> do

libraries/ghci/GHCi/TH.hs

@@ -10,6 +10,7 @@ module GHCi.TH
    , runModFinalizerRefs
    , runTH
    , GHCiQException(..)
 +  , GHCiQUncaughtExceptionOnHost(..)
    ) where
  {- Note [Remote Template Haskell]
@@ -158,8 +159,13 @@ ghcCmd m = GHCiQ $ \sRef -> do
    s <- readIORef sRef
    r <- remoteTHCall (qsPipe s) m
    case r of
 -    THException str -> throwIO (GHCiQException s str)
 -    THComplete res -> return res
 +    Right (THException str) -> throwIO (GHCiQException s str)
 +    Right (THComplete res) -> return res
 +    Left e -> liftIO $ throwIO (GHCiQUncaughtExceptionOnHost e)
++
 +-- | An exception was thrown on the host while processing a 'remoteTHCall'
 +data GHCiQUncaughtExceptionOnHost = GHCiQUncaughtExceptionOnHost SomeException deriving Show
 +instance Exception GHCiQUncaughtExceptionOnHost
  instance MonadIO GHCiQ where
    liftIO m = GHCiQ $ \_ -> m
@@ -189,9 +195,9 @@ metaHandlersGHCiQ runInIO = TH.MetaHandlers {
    -- See Note [TH recover with -fexternal-interpreter] in GHC.Tc.Gen.Splice
    , mRecover = \h a -> runInIO $ GHCiQ $ \sRef -> mask $ \unmask -> do
        s <- readIORef sRef
 -      remoteTHCall (qsPipe s) StartRecover
 +      remoteTHCall (qsPipe s) StartRecover >>= either (throwIO . GHCiQUncaughtExceptionOnHost) pure
        e <- try $ unmask $ runGHCiQ (runQinGHCiQ a <* ghcCmd FailIfErrs) sRef
 -      remoteTHCall (qsPipe s) (EndRecover (isLeft e))
 +      remoteTHCall (qsPipe s) (EndRecover (isLeft e))  >>= either (throwIO . GHCiQUncaughtExceptionOnHost) pure
        case e of
          Left GHCiQException{} ->
            runGHCiQ (runQinGHCiQ h) sRef

libraries/ghci/ghci.cabal.in

@@ -95,7 +95,8 @@ library
          filepath         >= 1.4 && < 1.6,
          ghc-boot         == @ProjectVersionMunged@,
          ghc-heap         >= 9.10.1 && <=@ProjectVersionMunged@,
 -        transformers     >= 0.5 && < 0.7
 +        transformers     >= 0.5 && < 0.7,
 +        stm              >= 2.4 && < 2.6
      if flag(bootstrap)
        build-depends:

testsuite/tests/ghci/custom-external-interpreter-commands/Main.hs

@@ -144,7 +144,7 @@ customMessage :: (Show a, Bin.Binary a) => Pipe -> ClientCommand a -> IO a
  customMessage pipe c = do
    let payload = encodeLazy (Some @Bin.Binary c)
    putStrLn $ "Sending: " ++ show c
 -  respBytes <- remoteCall pipe (CustomMessage customTag payload)
 +  Right respBytes <- remoteCall pipe (CustomMessage customTag payload)
    case decodeLazy respBytes of
      Left err -> error $ "Decode error: " ++ err
      Right res -> pure res