GitLab

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

Commits:

73b952d2

by Andrea Vezzosi at 2026-10-27T19:56:51+01:00

interp: cancel requests if async exceptions while waiting for reply.

d748ef6a

by Andrea Vezzosi at 2026-10-27T19:56:51+01:00

interp: do not terminate ext. interp. process on exception

`sendMessage` throwing an exception does not warrant terminating the external
interpreter process anymore, because communication can keep going no matter what.

`stopInterp` now resets the whole `Interp` state, so that a new
ExtInterpInstance can be later created and work correctly.

`withExtInterp` now calls `stopInterp` on a `PipeClosedException` and rethrows.

Fixes #18852

8f45d12a

by Andrea Vezzosi at 2026-10-27T21:59:17+01:00

fixup! interp: do not terminate ext. interp. process on exception

7 changed files:

compiler/GHC/Runtime/Interpreter.hs
compiler/GHC/Runtime/Interpreter/Process.hs
compiler/GHC/Tc/Gen/Splice.hs
libraries/ghci/GHCi/Message.hs
libraries/ghci/GHCi/Message/Async.hs
libraries/ghci/GHCi/Server.hs
libraries/ghci/GHCi/Server/Threads.hs

Changes:

compiler/GHC/Runtime/Interpreter.hs

@@ -105,7 +105,7 @@ import GHCi.Run
  import Control.Concurrent
  import Control.Monad
  import Control.Monad.IO.Class
 -import Control.Monad.Catch as MC (mask)
 +import Control.Monad.Catch as MC (mask, handle, MonadThrow (..))
  import Data.Binary
  import Data.ByteString (ByteString)
  import qualified Data.ByteString.Short as SBS
@@ -113,6 +113,7 @@ import qualified GHC.Exts.Heap as Heap
  import GHC.Stack.CCS (CostCentre,CostCentreStack)
  import System.Directory
  import System.Process
 +import System.Timeout (timeout)
  import qualified GHC.InfoProv as InfoProv
  import GHC.Builtin.Modules( gHC_PRIM, gHC_PRIMOPWRAPPERS )
@@ -196,14 +197,16 @@ interpCmd interp msg = case interpInstance interp of
  #if defined(HAVE_INTERNAL_INTERPRETER)
    InternalInterp     -> run msg -- Just run it directly
  #endif
 -  ExternalInterp ext -> withExtInterp ext $ \inst ->
 +  ExternalInterp ext -> withExtInterp interp ext $ \inst ->
      sendMessage inst msg
 -withExtInterp :: ExceptionMonad m => ExtInterp -> (forall d. ExtInterpInstance d -> m a) -> m a
 -withExtInterp ext action = case ext of
 +withExtInterp :: ExceptionMonad m => Interp -> ExtInterp -> (forall d. ExtInterpInstance d -> m a) -> m a
 +withExtInterp interp ext action = stopInterpOnPipeClosed $ case ext of
    ExtJS    i -> withJSInterp i action
    ExtWasm  i -> withWasmInterp i action
    ExtIServ i -> withIServ    i action
 +  where
 +    stopInterpOnPipeClosed = MC.handle $ \ (e :: PipeClosedException) -> liftIO (stopInterp interp) >> MC.throwM e
  withExtInterpStatus :: ExtInterp -> (forall d. ExtInterpStatusVar d -> m a) -> m a
  withExtInterpStatus ext action = case ext of
@@ -631,18 +634,41 @@ stopInterp interp = case interpInstance interp of
  #if defined(HAVE_INTERNAL_INTERPRETER)
      InternalInterp -> pure ()
  #endif
 -    ExternalInterp ext -> withExtInterpStatus ext $ \mstate -> do
 -      MC.mask $ \_restore -> modifyMVar_ mstate $ \state -> do
 +    ExternalInterp ext -> (maybe (pure ()) throwIO =<<) $ withExtInterpStatus ext $ \mstate -> do
 +      MC.mask $ \restore -> modifyMVar mstate $ \state -> do
          case state of
 -          InterpPending    -> pure state -- already stopped
 +          InterpPending    -> pure (state, Nothing) -- already stopped
            InterpRunning i  -> do
              ex <- getProcessExitCode (interpHandle (instProcess i))
 -            if isJust ex
 -               then pure ()
 -               else sendMessage i Shutdown
 -            -- We asked the server to shutdown, so we will also shutdown any communication.
 -            closePipe (interpPipe (instProcess i))
 -            pure InterpPending
 +            e <- if isJust ex
 +                 then gracefulShutdown i restore
 +                 else closePipe (interpPipe (instProcess i)) >> pure Nothing
 +            modifyMVar_ (loader_state (interpLoader interp)) (const $ pure Nothing)
 +            purgeInterpSymbolCache (interpSymbolCache interp)
 +            modifyMVar_ (interpStringCache interp) (const $ pure emptyFsEnv)
 +            pure (InterpPending,e)
 +  where
 +    gracefulShutdown i restore = do
 +      x <- try (timeLimited restore $ sendMessage i Shutdown `finally` closePipe (interpPipe (instProcess i)))
 +      let ph = interpHandle (instProcess i)
 +      e <- case x of
 +        -- happy path: ext. interp. should be shutting down of its own accord.
 +        Right (Just ()) -> pure Nothing
 +        -- timeout: ext. interp. couldn't handle Shutdown in time.
 +        Right Nothing   -> do
 +          terminateProcess ph
 +          pure Nothing
 +        -- something went wrong, probably Pipe already closed.
 +        Left e -> do
 +          terminateProcess ph
 +          returnUnexpectedException e
 +      void $ timeLimited restore $ void $ waitForProcess ph
 +      pure e
 +    timeLimited restore m = timeout 1_000_000 (restore m)
 +    returnUnexpectedException e
 +      | Just (_ :: PipeClosedException) <- fromException e = return Nothing
 +      | Just (_ :: IOException) <- fromException e = return Nothing
 +      | otherwise = pure (Just e)
  -- -----------------------------------------------------------------------------
  {- Note [External GHCi pointers]

compiler/GHC/Runtime/Interpreter/Process.hs

@@ -12,13 +12,10 @@ import GHC.Prelude
  import GHC.Runtime.Interpreter.Types
  import GHCi.Message
 -import GHC.IO (catchException)
  import GHC.Utils.Panic
  import GHC.Utils.Exception as Ex
  import Data.Binary
 -import System.Exit
 -import System.Process
  import Control.Concurrent
  import Data.IORef
  import Control.Concurrent.STM
@@ -34,7 +31,6 @@ import Control.Concurrent.STM
  sendMessage :: Binary a => ExtInterpInstance d -> Message a -> IO a
  sendMessage i msg =
    remoteCall (interpPipe proc) msg
 -    `catchException` \(e :: SomeException) -> handleInterpProcessFailure proc e
    where
      proc = instProcess i
@@ -64,7 +60,7 @@ sendTHComputation i msg reqHandler = mask $ \restore -> do
    where
      proc = instProcess i
      handleOne mbox doneRef = do
 -      handleRemoteTHCall (interpPipe proc) mbox $ \req -> do
 +      handleRemoteTHCall (interpPipe proc) mbox $ \_setCancel req -> do
          (stop, res) <- reqHandler req
          writeIORef doneRef stop
          return res
@@ -77,18 +73,6 @@ sendTHComputation i msg reqHandler = mask $ \restore -> do
                        else loop
         in loop
 -handleInterpProcessFailure :: InterpProcess -> SomeException -> IO a
 -handleInterpProcessFailure i e = do
 -  let hdl = interpHandle i
 -  ex <- getProcessExitCode hdl
 -  case ex of
 -    Just (ExitFailure n) ->
 -      throwIO (InstallationError ("External interpreter terminated (" ++ show n ++ ")"))
 -    _ -> do
 -      terminateProcess hdl
 -      _ <- waitForProcess hdl
 -      throw e
+-
  data CancelTHHandler = CancelTHHandler deriving Show
  instance Exception CancelTHHandler where
    fromException = asyncExceptionFromException

compiler/GHC/Tc/Gen/Splice.hs

@@ -1160,7 +1160,7 @@ runRemoteModFinalizers (ThModFinalizers finRefs) = do
        runQinTcM $ sequence_ qs
  #endif
 -    ExternalInterp ext -> withExtInterp ext $ \inst -> do
 +    ExternalInterp ext -> withExtInterp interp ext $ \inst -> do
        tcg <- getGblEnv
        th_state <- readTcRef (tcg_th_remote_state tcg)
        case th_state of
@@ -1817,7 +1817,7 @@ runTH ty fhv = do
        return r
  #endif
 -    ExternalInterp ext -> withExtInterp ext $ \inst -> do
 +    ExternalInterp ext -> withExtInterp interp ext $ \inst -> do
        -- Run it on the server.  For an overview of how TH works with
        -- Remote GHCi, see Note [Remote Template Haskell] in
        -- libraries/ghci/GHCi/TH.hs.

libraries/ghci/GHCi/Message.hs

@@ -27,7 +27,7 @@ module GHCi.Message
    , Pipe(..), mkPipeFromHandles, mkPipeFromContinuations, closePipe
    , remoteCall, remoteTHCall
    , handleRemoteCall, handleRemoteTHCall, writeReply
 -  , RemoteProcessException(..), PipeReadClosedException(..)
 +  , RemoteProcessException(..), PipeClosedException(..)
    , InboxId(..)
    , newMailbox
@@ -666,12 +666,12 @@ remoteCall pipe msg = either throwIO (pure . runGet get)
  -- | 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) => WriteReply a %1 -> Message a -> IO r)
 +                 -> (forall a. (Show a, Binary a) => WriteReply a %1 -> (IO () -> IO ()) -> Message a -> IO r)
                   -> IO r
  handleRemoteCall pipe hook k =
 -  handleAsyncReq pipe getMessage Nothing \(WriteReply writer) msg0 -> IO \s0 ->
 +  handleAsyncReq pipe getMessage Nothing \(WriteReply writer) sa msg0 -> IO \s0 ->
      case hook msg0 |#> s0 of -- unIO for linearity checking
 -      (# s1, Msg msg #) -> k (WriteReply \x -> writer (fmap put x)) msg |#> s1
 +      (# s1, Msg msg #) -> k (WriteReply \x -> writer (fmap put x)) sa msg |#> s1
  -- | Send a 'THMessage' request to the host and wait for a reply from the host
  -- to the interpreter.
@@ -687,9 +687,9 @@ remoteTHCall pipe ii msg = either throwIO (pure . runGet get)
  -- Symmterically to 'remoteTHCall', we always want to specify the 'InboxId' for
  -- RunTH "reverse requests"
  handleRemoteTHCall
 -  :: Pipe -> InboxId -> (forall a. (Show a, Binary a) => THMessage a -> IO a) -> IO ()
 +  :: Pipe -> InboxId -> (forall a. (Show a, Binary a) => (IO () -> IO ()) -> THMessage a -> IO a) -> IO ()
  handleRemoteTHCall pipe ii k =
 -  handleAsyncReq pipe getTHMessage (Just ii) \writer (THMsg msg) -> IO \s0 ->
 -    case k msg |#> s0 of -- unIO for linearity checking
 +  handleAsyncReq pipe getTHMessage (Just ii) \writer sa (THMsg msg) -> IO \s0 ->
 +    case k sa msg |#> s0 of -- unIO for linearity checking
        (# s1, a #) -> writeReply writer (put a) |#> s1

libraries/ghci/GHCi/Message/Async.hs

@@ -7,7 +7,7 @@ module GHCi.Message.Async
+   (
      -- * Pipes
      Pipe(..), mkPipeFromHandles, mkPipeFromContinuations
 -  , closePipe, PipeReadClosedException(..)
 +  , closePipe, PipeClosedException(..)
      -- * Interface for async bidirectional RPC
    , remoteAsyncReq, handleAsyncReq
@@ -145,13 +145,13 @@ closePipe p@Pipe{..} = do
        unless writeAlreadyClosed
          (throwTo (asyncFillTid pipeAsyncMngr) PipeWriteClosedException)
 --- | We've closed the pipe read end on this process
 -data PipeReadClosedException = PipeReadClosedException deriving Show
 -instance Exception PipeReadClosedException
 +data PipeClosedException
 +  = PipeReadClosedException -- ^ We've closed the pipe read end on this process
++
 +  | PipeWriteClosedException -- ^ We've closed the pipe write end on this process
 +  deriving Show
 +instance Exception PipeClosedException
 --- | We've closed the pipe write end on this process
 -data PipeWriteClosedException = PipeWriteClosedException deriving Show
 -instance Exception PipeWriteClosedException
  --------------------------------------------------------------------------------
  -- ** Pipe thread-safe read/write
@@ -160,14 +160,19 @@ instance Exception PipeWriteClosedException
  -- and identify whether they are requests or replies.
  writePipe :: Pipe -> IddMsg -> IO ()
 -writePipe Pipe{..} msg = do
 -  var <- newEmptyTMVarIO
 -  atomically $ do
 +writePipe p msg = do
 +  var <- atomically $ writePipeAsync p msg
 +  either throwIO pure =<< atomically (takeTMVar var)
++
 +writePipeAsync :: Pipe -> IddMsg -> STM (TMVar (Either SomeException ()))
 +writePipeAsync Pipe{..} msg = do
 +    var <- newEmptyTMVar
      readTVar pipeWriteClosed >>= \case
        True -> throwSTM PipeWriteClosedException
        False -> pure ()
      writeTQueue (writeQueue pipeAsyncMngr) (msg, var)
 -  either throwIO pure =<< atomically (takeTMVar var)
 +    pure var
++
  -- | Only the 'drainPipe' thread should read from the pipe directly. Other
  -- functions should wait for the matching reply to be delivered asynchronously.
@@ -239,14 +244,18 @@ remoteAsyncReq pipe msg ii = do
    -- Write request
    writePipe pipe (Request uq ii (runPut msg))
      `onException`
 -      -- the req wasn't done, or no one will be reading the reply:
 -      -- in either case, drop it from the pending map.
 -      atomically (modifyTVar pendingReplies (IntMap.delete uq))
 +        -- the req wasn't done, or no one will be reading the reply:
 +        -- in either case, drop it from the pending map.
 +        atomically (modifyTVar pendingReplies (IntMap.delete uq))
    -- Block waiting for reply
    atomically (takeTMVar wv)
      `onException`
 -      atomically (modifyTVar pendingReplies (IntMap.delete uq))
 +       atomically (do
 +         modifyTVar pendingReplies (IntMap.delete uq)
 +         -- 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
@@ -258,7 +267,7 @@ remoteAsyncReq pipe msg ii = do
  -- branch, and if the handler otherwise throws, it is handled by 'handleAsyncReq'.
  --
  -- See Note [Exceptions during serialization] about exceptions the 'Put' argument might raise.
 -handleAsyncReq :: Pipe -> Get a -> (Maybe InboxId) -> (WriteReply Put %1 -> a -> IO r) -> IO r
 +handleAsyncReq :: Pipe -> Get a -> (Maybe InboxId) -> (WriteReply Put %1 -> (IO () -> IO ()) -> a -> IO r) -> IO r
  handleAsyncReq pipe getit mii k = mask $ \restore -> do
    let am = pipeAsyncMngr pipe
    ack <- newEmptyMVar
@@ -274,7 +283,7 @@ handleAsyncReq pipe getit mii k = mask $ \restore -> do
        Nothing -> tryReadTQueue  $  requestsQueue am
        Just ii -> tryReadTQueue =<< getMailbox am ii
      case taken of
 -      Just req -> pure (Just req)
 +      Just req@(uq,_) -> fetchedRequest am uq >> pure (Just req)
        Nothing  | isClosed  -> pure Nothing -- empty and closed
                 | otherwise -> retry -- wait for msg
@@ -282,9 +291,9 @@ handleAsyncReq pipe getit mii k = mask $ \restore -> do
      Nothing -> throwIO PipeReadClosedException -- can't expect reqs
      Just (uq, bs) -> do
        let reply p = do
 -            writePipe pipe (Reply uq (bimap toSerializableException runPut p))
 +            b <- atomically $ doneHandling am uq
 +            when b $ writePipe pipe (Reply uq (bimap toSerializableException runPut p))
              putMVar ack ()
+-
        -- Run the continuation which must reply exactly once with 'writeReply'
        case runGetOrFail getit bs of
          Left (_, _, err) -> do
@@ -292,7 +301,7 @@ handleAsyncReq pipe getit mii k = mask $ \restore -> do
            reply (Left (toException exc))
            throwIO exc
          Right (_, _, it) ->
 -          restore (k (WriteReply reply) it)
 +          restore (k (WriteReply reply) (setActive am uq) it)
              -- If there's an exception before the reply is sent, reply with an error instead.
              `catch` \(e :: SomeException) -> do
@@ -306,6 +315,7 @@ handleAsyncReq pipe getit mii k = mask $ \restore -> do
                      reply (Left e)
                      throwIO e
++
  -- | 'handleSyncReq' failed to decode an incoming request payload using the given @'Get' a@
  newtype DecodeAsyncReqException = DecodeAsyncReqException String deriving Show
  instance Exception DecodeAsyncReqException
@@ -347,6 +357,28 @@ getMailbox am (InboxId ii) = do
  -- ** Async matching "internals" -----------------------------------------------
 +-- | Tracks handling status of a received 'Request'. The values are states of a
 +-- tiny state machine used to coordinate cancelations.
 +--
 +-- 'handleAsyncReq' will pop a request from the appropriate queue, setting the
 +-- status to `Fetched`. Then one of two things can happen first:
 +--  - we receive a 'Cancel', then status is set to `Canceling` so the cancel
 +--    action will run when set.
 +--  - handler sets a cancel action, then reception of a `Cancel` runs that
 +--    action and deletes the RequestStatus.
 +--
 +-- If RequestStatus is `Canceling` or absent, then a `Reply` won't be sent.
 +--
 +-- The two possible transition paths:
 +-- requestsQueue -> Fetched -> Active cancel -> gone
 +-- requestsQueue -> Fetched -> Canceling     -> gone
++
 +data RequestStatus
 +  = Fetched        -- ^ 'Fetched' from a request queue, to pass to handler.
 +  | Active (IO ()) -- ^ Handler set a cancel action, and has not sent a reply
 +                   -- yet. The cancel action must not block.
 +  | Canceling      -- ^ `Cancel` received while status is still `Fetched`.
++
  data AsyncManager = AsyncManager
    { pendingReply :: !(TVar (IntMap (TMVar (Either SomeException LB.ByteString))))
      -- ^ When a request is sent, keep track that we're expecting a reply. The
@@ -357,6 +389,9 @@ data AsyncManager = AsyncManager
      -- /without/ a specific 'InboxId', queue it here. A request handler will
      -- pop the next request from it. (Nothing except the 'drainPipe' thread
      -- should read from the pipe). Read with @'handleAsyncReq' p get Nothing@.
 +  , cancelStatus :: !(TVar (IntMap RequestStatus))
 +    -- ^ When a `Cancel uq` is received, we lookup `uq` in this map and call the
 +    -- callback if present.
    , mailboxQueues :: !(TVar (IntMap (TQueue (CorrelationId, LB.ByteString))))
      -- ^ Requests for a specific 'InboxId' that can be read using
      -- @'handleAsyncReq' p get (Just <mailboxid>)@.
@@ -389,15 +424,21 @@ type CorrelationId = Int
  -- exclusively.
  data IddMsg  = Request !CorrelationId !(Maybe InboxId) LB.ByteString
               | Reply   !CorrelationId (Either SerializableException LB.ByteString)
 +             | Cancel  !CorrelationId -- ^ Technically a Request, but no Reply
 +                                      -- is expected, the CorrelationId is the
 +                                      -- one of the Request being canceled.
  instance Binary IddMsg where
    put (Request i ii bs) = putWord8 0 >> put i >> put ii >> put bs
    put (Reply i bs)      = putWord8 1 >> put i >> put bs
 +  put (Cancel i)        = putWord8 2 >> put i
++
    get = do
      b <- getWord8
      case b of
 -> Request <$> get <*> get <*> get
 -> Reply   <$> get <*> get
 +      2 -> Cancel  <$> get
        _ -> fail $ "impossible: Binary IddMsg - " ++ show b
  newAsyncManager :: Pipe -> IO AsyncManager
@@ -405,6 +446,7 @@ newAsyncManager pipe = do
    asyncFreshId  <- newIORef 0
    pendingReply  <- newTVarIO IntMap.empty
    requestsQueue <- newTQueueIO
 +  cancelStatus  <- newTVarIO IntMap.empty
    mailboxQueues <- newTVarIO IntMap.empty
    writeQueue <- newTQueueIO
@@ -430,41 +472,85 @@ drainPipe :: Pipe -> (forall x. IO x -> IO x) -> IO ()
  drainPipe p unmask = catch (unmask $ loop `finally` pipeCloseRead p) replyToPendings  where
    loop = do
      readPipe p >>= \case
 -      -- Receive a reply to a pending notification
 -      Just (Reply uq bs) -> do
 -        atomically $ do
 -          m <- readTVar (pendingReply am)
 -          case IntMap.lookup uq m of
 -            Nothing ->
 -              -- the req was dropped from pending bc there was an exception, and no
 -              -- one's listening any more. the reply is ignored
 -              pure ()
 -            Just wv -> do
 -              writeTVar (pendingReply am) (IntMap.delete uq m)
 -              -- signal matching request
 -              putTMVar wv (bimap (toException . RemoteProcessException) id bs)
 -        loop
 -      -- Receive a request and queue it
 -      Just (Request uq mbi bs) -> do
 -        atomically $ case mbi of
 -          Nothing ->
 -            -- "Global" request
 -            writeTQueue (requestsQueue am) (uq,bs)
 -          Just ii -> do
 -            -- "Specific mailbox" request
 -            q <- getMailbox am ii
 -            writeTQueue q (uq,bs)
 +      Just msg -> do
 +        handleMsg msg
          loop
        -- EOF on read handle
        Nothing -> do
          -- clean up pipe then exit loop
          replyToPendings (toException PipeReadClosedException)
          pure ()
 +  handleMsg = \case
 +    -- Receive a reply to a pending notification
 +    Reply uq bs -> do
 +      atomically $ do
 +        m <- readTVar (pendingReply am)
 +        case IntMap.lookup uq m of
 +          Nothing ->
 +            -- the req was dropped from pending bc there was an exception, and no
 +            -- one's listening any more. the reply is ignored
 +            pure ()
 +          Just wv -> do
 +            writeTVar (pendingReply am) (IntMap.delete uq m)
 +            -- signal matching request
 +            putTMVar wv (bimap (toException . RemoteProcessException) id bs)
 +    -- Receive a request and queue it
 +    Request uq mbi bs -> do
 +      atomically $ case mbi of
 +        Nothing ->
 +          -- "Global" request
 +          writeTQueue (requestsQueue am) (uq,bs)
 +        Just ii -> do
 +          -- "Specific mailbox" request
 +          q <- getMailbox am ii
 +          writeTQueue q (uq,bs)
 +    Cancel uq -> cancelRequest am uq
    replyToPendings e = atomically $ failPendingReplies p e
    am = pipeAsyncMngr p
 +cancelRequest :: AsyncManager -> CorrelationId -> IO ()
 +cancelRequest am uq = join . atomically $ do
 +  cs <- readTVar (cancelStatus am)
 +  case IntMap.lookup uq cs of
 +    Nothing -> pure (pure ())
 +    Just Fetched -> writeTVar (cancelStatus am) (IntMap.insert uq Canceling cs) >> pure (pure ())
 +    Just (Active cancel) -> writeTVar (cancelStatus am) (IntMap.delete uq cs) >> pure cancel
 +    Just Canceling -> pure (pure ())
++
 +-- | Returns @True@ when the handler should write its 'Reply'.
 +doneHandling :: AsyncManager -> CorrelationId -> STM Bool
 +doneHandling am uq = do
 +  cs <- readTVar (cancelStatus am)
 +  case IntMap.lookup uq cs of
 +    Nothing -> pure False
 +    Just s -> do
 +      writeTVar (cancelStatus am) $ IntMap.delete uq cs
 +      case s of
 +        Canceling -> pure False
 +        _         -> pure True
++
 +setActive :: AsyncManager -> CorrelationId -> IO () -> IO ()
 +setActive am uq cancel = join . atomically $ do
 +  cs <- readTVar (cancelStatus am)
 +  case IntMap.lookup uq cs of
 +    Nothing -> pure (pure ()) -- reply already received.
 +    Just Fetched -> writeTVar (cancelStatus am) (IntMap.insert uq (Active cancel) cs) >> pure (pure ())
 +    Just (Active _) -> writeTVar (cancelStatus am) (IntMap.insert uq (Active cancel) cs) >> pure (pure ())
 +    Just Canceling -> pure cancel
++
 +fetchedRequest :: AsyncManager -> CorrelationId -> STM ()
 +fetchedRequest am uq = do
 +  cs <- readTVar (cancelStatus am)
 +  case IntMap.lookup uq cs of
 +    Nothing -> writeTVar (cancelStatus am) (IntMap.insert uq Fetched cs)
 +    Just Fetched -> error "impossible"
 +    Just (Active _) -> error "impossible"
 +    Just Canceling -> error "impossible"
++
++
++
  {-
  Note [Exceptions during serialization]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -553,7 +639,7 @@ runWriteReply (WriteReply r) = r
  (|#>) (IO fs) s = fs s
  (%>>=) :: (IO a) -> (a -> IO b) %1 -> IO b
 -(%>>=) = \ m f -> IO (\s0 ->
 +(%>>=) = \ m -> \ f -> IO (\s0 ->
          case m |#> s0 of -- unIO for linearity checking
            (# s1, x #) -> f x |#> s1)
  infixl 1 %>>=

libraries/ghci/GHCi/Server.hs

@@ -66,13 +66,13 @@ servWithCustom
    -> IO ()
  servWithCustom verbose hook pipe restore customHandler = do
    ref <- newIORef False
 -  withThreadManager $ \mngr ->
 +  withThreadManager $ \mngr -> do
      loop ref mngr `finally` closePipe pipe
   where
    loop doneRef mngr = do
      when verbose $ trace "reading pipe..."
 -    handleRemoteCall pipe hook $ \wr (msg :: Message a) ->
 +    handleRemoteCall pipe hook $ \wr setCancel (msg :: Message a) ->
        discardCtrlC                                 %>>= \_ ->
        when verbose (trace ("msg: " ++ (show msg))) %>>= \_ ->
@@ -82,9 +82,11 @@ servWithCustom verbose hook pipe restore customHandler = do
          Shutdown                  -> writeIORef doneRef True >>
                                       reply ()                 %>>= writeReply wr
          RunTH ii st q ty loc      -> spawnFinally mngr         (runWriteReply wr)
 -                                     (wrapRunTH ii (runTH pipe ii st q ty loc))
 +                                     (restore $ wrapRunTH ii (runTH pipe ii st q ty loc))
 +                                     (setCancel . void . forkIO . killThread)
          RunModFinalizers ii st qs -> spawnFinally mngr         (runWriteReply wr)
 -                                     (wrapRunTH ii (runModFinalizerRefs pipe ii st qs))
 +                                     (restore $ wrapRunTH ii (runModFinalizerRefs pipe ii st qs))
 +                                     (setCancel . void . forkIO . killThread)
          _other                    -> run msg >>= reply        %>>= writeReply wr
      isDone <- readIORef doneRef

libraries/ghci/GHCi/Server/Threads.hs

@@ -41,8 +41,8 @@ withThreadManager k = do
  -- | Linear and resource-safe because 'ThreadManager' is guaranteed to wait or
  -- cancel all threads.
 -spawnFinally :: ThreadManager -> (Either SomeException a -> IO ()) %1 -> IO a -> IO ()
 -spawnFinally ThreadManager{..} = unsafe_linear \fin -> \action -> void $
 +spawnFinally :: ThreadManager -> (Either SomeException a -> IO ()) %1 -> IO a -> (ThreadId -> IO b) -> IO b
 +spawnFinally ThreadManager{..} = unsafe_linear \fin -> \action -> \k ->
    forkFinally
      ( mask $ \unmask -> do
          tid <- myThreadId
@@ -57,6 +57,7 @@ spawnFinally ThreadManager{..} = unsafe_linear \fin -> \action -> void $
            else throwIO ThreadKilled
+     )
      ( \ea -> fin ea )
 +    >>= k
  unsafe_linear :: (a -> b) %1 -> (a %1 -> b)
  unsafe_linear = Unsafe.unsafeCoerce (Unsafe.unsafeCoerce :: (a -> b) -> (a %1 -> b))