From b4eab458faf2d3c50c56377c327982dbbe49deaf Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Mon, 2 Jun 2025 16:05:27 +0200
Subject: [PATCH 01/27] Pop the log hook after compilation to avoid memory
 leaks

Change-Id: Ieb3f3577749ee4569951d65b44cdaeba3ed96525
---
 internal/src/Internal/Session.hs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/internal/src/Internal/Session.hs b/internal/src/Internal/Session.hs
index 5922ca33..014aeb00 100644
--- a/internal/src/Internal/Session.hs
+++ b/internal/src/Internal/Session.hs
@@ -21,6 +21,7 @@ import GHC (
   gopt,
   parseDynamicFlags,
   parseTargetFiles,
+  popLogHookM,
   prettyPrintGhcErrors,
   pushLogHookM,
   setSessionDynFlags,
@@ -130,9 +131,10 @@ withGhcInSession env prog argv = do
   pushLogHookM (const (logToState log))
   cache <- liftIO $ readMVar env.cache
   (dflags0, logger, fileish_args, dynamicFlagWarnings) <- parseFlags (argv ++ map instrumentLocation (words cache.options.extraGhcOptions))
-  prettyPrintGhcErrors logger do
+  result <- prettyPrintGhcErrors logger do
     srcs <- initGhc dflags0 logger fileish_args dynamicFlagWarnings
     prog srcs
+  result <$ popLogHookM
 
 -- | Create a base session and store it in the cache.
 -- On subsequent calls, return the cached session, unless the cache is disabled or @reuse@ is true.

From 6d20b25a56e5101b155d02244a3483c8be0144c0 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 28 May 2025 17:52:14 +0200
Subject: [PATCH 02/27] Directly store DynFlags in the new unit env in the
 metadata step

This avoids initializing the session DynFlags before storing the unit
env in the cache.
It might not be necessary, but it's useful to be very deliberate about
this to be able to eliminate some problem origins.

Change-Id: I1911041ea5feac42a6ecef3ba0d659c4b36e680a
---
 ghc-worker/test/CompileHptTest.hs |  12 +--
 internal/src/Internal/Metadata.hs | 124 ++++++++++++++++++++++--------
 internal/src/Internal/Session.hs  |  46 ++++++++---
 3 files changed, 128 insertions(+), 54 deletions(-)

diff --git a/ghc-worker/test/CompileHptTest.hs b/ghc-worker/test/CompileHptTest.hs
index 07803d0b..a50fac19 100644
--- a/ghc-worker/test/CompileHptTest.hs
+++ b/ghc-worker/test/CompileHptTest.hs
@@ -27,7 +27,7 @@ import GHC.Utils.Panic (throwGhcExceptionIO)
 import Internal.Cache (Target (..), logMemStats)
 import Internal.CompileHpt (compileModuleWithDepsInHpt)
 import Internal.Log (dbg, dbgp, dbgs, newLog)
-import Internal.Metadata (computeMetadata, computeMetadataInSession)
+import Internal.Metadata (computeMetadata)
 import Internal.Session (Env (..), withGhcMhu)
 import Prelude hiding (log)
 import System.Directory (createDirectoryIfMissing, listDirectory, removeDirectoryRecursive)
@@ -85,13 +85,6 @@ initUnit specific = do
   setUnitDynFlags current dflags
   modifySession (hscSetActiveUnitId current)
 
--- | Approximate synthetic reproduction of what happens when the metadata step is performed by the worker.
-loadModuleGraph :: Env -> Unit -> [String] -> Ghc (Maybe Bool)
-loadModuleGraph env Unit {dir} specific = do
-  names <- liftIO $ listDirectory dir
-  let srcs = [dir </> name | name <- names, takeExtension name == ".hs"]
-  computeMetadataInSession (initUnit specific) env srcs
-
 stepMetadata :: Conf -> Unit -> [Unit] -> IO ()
 stepMetadata Conf {cache, tmp, args0} unit deps = do
   log <- newLog True
@@ -104,13 +97,12 @@ stepMetadata Conf {cache, tmp, args0} unit deps = do
   unless success do
     liftIO $ throwGhcExceptionIO (ProgramError "Metadata failed")
   where
-
     args srcs = args0 {
       ghcOptions = mkDependArgs ++ unitDepArgs ++ args0.ghcOptions ++ srcs,
       tempDir = Just sessionTmpDir
     }
 
-    unitDepArgs = concat [["-package-db", db, "-package-id", name]  | Unit {name, db} <- deps]
+    unitDepArgs = concat [["-package-id", name]  | Unit {name} <- deps]
 
     mkDependArgs = [
       "-i",
diff --git a/internal/src/Internal/Metadata.hs b/internal/src/Internal/Metadata.hs
index 8b9cb1c8..9a3d7a32 100644
--- a/internal/src/Internal/Metadata.hs
+++ b/internal/src/Internal/Metadata.hs
@@ -2,16 +2,19 @@ module Internal.Metadata where
 
 import Control.Concurrent (modifyMVar_, readMVar)
 import Control.Monad.IO.Class (liftIO)
-import Data.Maybe (fromMaybe)
-import GHC (DynFlags (..), Ghc, GhcMode (..))
-import GHC.Driver.Env (HscEnv (..), hscUpdateFlags)
-import GHC.Driver.Monad (modifySession, modifySessionM, withTempSession, withSession)
+import Data.Maybe (isJust)
+import GHC (DynFlags (..), Ghc, GhcMode (..), Logger, ModuleGraph)
+import GHC.Driver.Env (HscEnv (..), hscSetActiveUnitId, hscUpdateFlags, hscUpdateLoggerFlags)
+import GHC.Driver.Monad (modifySession, modifySessionM, withSession, withTempSession)
+import GHC.Driver.Session (updatePlatformConstants)
 import GHC.Platform.Ways (Way (WayDyn), addWay)
 import GHC.Runtime.Loader (initializeSessionPlugins)
-import GHC.Unit.Env (UnitEnv (..), unitEnv_union)
-import Internal.Cache (Cache (..), Target (..), insertUnitEnv, mergeHugs, newFinderCache, updateModuleGraph, logMemStats)
+import GHC.Unit (HomeUnit, UnitDatabase, UnitId, UnitState, initUnits)
+import GHC.Unit.Env (HomeUnitEnv (..), UnitEnv (..), unitEnv_insert, unitEnv_keys, unitEnv_union, updateHug)
+import GHC.Unit.Home.ModInfo (emptyHomePackageTable)
+import Internal.Cache (Cache (..), insertUnitEnv, mergeHugs, updateModuleGraph, logMemStats)
 import Internal.MakeFile (doMkDependHS)
-import Internal.Session (Env (..), runSession, withGhcInSession)
+import Internal.Session (Env (..), runSession, withDynFlags)
 
 -- | Copy the cached unit env and module graph to the given session.
 restoreEnv :: Cache -> HscEnv -> HscEnv
@@ -26,9 +29,80 @@ restoreEnv cache hsc_env = do
     current = hsc_env.hsc_unit_env.ue_home_unit_graph
 
 -- | 'doMkDependHS' needs this to be enabled.
-addDynWay :: HscEnv -> HscEnv
-addDynWay =
-  hscUpdateFlags \ d -> d { targetWays_ = addWay WayDyn (targetWays_ d) }
+metadataTempSession :: HscEnv -> HscEnv
+metadataTempSession =
+  hscUpdateFlags \ d -> d {ghcMode = MkDepend, targetWays_ = addWay WayDyn (targetWays_ d)}
+
+insertHomeUnit ::
+  UnitId ->
+  DynFlags ->
+  [UnitDatabase UnitId] ->
+  UnitState ->
+  HomeUnit ->
+  UnitEnv ->
+  UnitEnv
+insertHomeUnit unit dflags dbs unit_state home_unit unit_env =
+  (updateHug (unitEnv_insert unit hue) unit_env) {
+    ue_platform = targetPlatform dflags,
+    ue_namever = ghcNameVersion dflags
+  }
+  where
+    hue = HomeUnitEnv {
+      homeUnitEnv_units = unit_state,
+      homeUnitEnv_unit_dbs = Just dbs,
+      homeUnitEnv_dflags = dflags,
+      homeUnitEnv_hpt = emptyHomePackageTable,
+      homeUnitEnv_home_unit = Just home_unit
+    }
+
+initHomeUnit :: DynFlags -> Logger -> UnitId -> UnitEnv -> IO UnitEnv
+initHomeUnit dflags0 logger unit unit_env = do
+  (dbs, unit_state, home_unit, mconstants) <- initUnits logger dflags0 Nothing allUnitIds
+  dflags1 <- updatePlatformConstants dflags0 mconstants
+  pure (insertHomeUnit unit dflags1 dbs unit_state home_unit unit_env)
+  where
+    allUnitIds = unitEnv_keys (ue_home_unit_graph unit_env)
+
+-- | Add a new home unit to the current session using the provided 'DynFlags'.
+-- The flags have been constructed from Buck CLI args passed to the metadata step, which, crucially, contain the package
+-- DB arguments for dependencies.
+addHomeUnit :: DynFlags -> Ghc UnitId
+addHomeUnit dflags = do
+  modifySessionM \ hsc_env -> do
+    unit_env <- liftIO $ initHomeUnit dflags hsc_env.hsc_logger unit hsc_env.hsc_unit_env
+    pure hsc_env {hsc_unit_env = unit_env}
+  pure unit
+  where
+    unit = dflags.homeUnitId_
+
+-- | Initialize the home unit env for this target and restore the module graphs computed previously for other units.
+--
+-- This part is the most significant difference that the make worker has from GHC make mode, since it never happens
+-- natively that units are added incrementally.
+-- Therefore, this is a relatively delicate procedure that hasn't been fully optimized yet.
+--
+-- We especially want to take care that the command line flags aren't applied to the base session before we initialize
+-- the home unit in order to replicate what GHC does in @initMulti@.
+prepareMetadataSession :: Env -> DynFlags -> Ghc ()
+prepareMetadataSession env dflags = do
+  cache <- liftIO $ readMVar env.cache
+  modifySession (restoreEnv cache)
+  unit <- addHomeUnit dflags
+  setActiveUnit unit
+  storeNewUnit
+  where
+    setActiveUnit unit = modifySession (hscUpdateLoggerFlags . hscSetActiveUnitId unit)
+
+    storeNewUnit = withSession \ hsc_env -> liftIO $ modifyMVar_ env.cache (pure . insertUnitEnv hsc_env)
+
+-- | Run 'doMkDependHS' to write the metadata JSON file and exfiltrate the module graph.
+-- We need to use a temporary session because 'doMkDependHS' uses some custom settings that we don't want to leak,
+-- though it's not been thoroughly tested what precisely the impact is.
+writeMetadata :: [String] -> Ghc ModuleGraph
+writeMetadata srcs = do
+  initializeSessionPlugins
+  withTempSession metadataTempSession do
+    doMkDependHS srcs
 
 -- | Run downsweep and merge the resulting module graph into the cached graph.
 -- This is executed for the metadata step, which natively only calls 'doMkDependHS'.
@@ -38,28 +112,14 @@ addDynWay =
 -- This is usually not necessary (in fact, 'doMkDependHS' clears the target ways) but since we're keeping the module
 -- graph the target way will be reflected in the stored @ModSummary@ nodes.
 --
--- Before downsweep, we also create a fresh @Finder@ for some reason and restore the previous unit env so dependencies
--- are visible.
-computeMetadataInSession :: Ghc () -> Env -> [String] -> Ghc (Maybe Bool)
-computeMetadataInSession setup env srcs = do
-  initializeSessionPlugins
-  cache <- liftIO $ readMVar env.cache
-  modifySession (restoreEnv cache)
-  setup
-  withSession \ hsc_env ->
-    liftIO $ modifyMVar_ env.cache $ (pure . insertUnitEnv hsc_env)
-  module_graph <- withTempSession addDynWay do
-    modifySessionM \ hsc_env -> do
-      hsc_FC <- liftIO $ newFinderCache env.cache cache (Target "metadata")
-      pure hsc_env {hsc_FC}
-    modifySession $ hscUpdateFlags \ d -> d {ghcMode = MkDepend}
-    doMkDependHS srcs
-  liftIO $ updateModuleGraph env.cache module_graph
-  pure (Just True)
-
--- | Run 'computeMetadataInSession' without extra session initialization code, which is just used in tests.
+-- Before downsweep, we also create a fresh @Finder@ to prevent 'doMkDependHS' from polluting the cache with entries
+-- with different compilation ways and restore the previous unit env so dependencies are visible.
 computeMetadata :: Env -> IO Bool
 computeMetadata env = do
-  res <- fmap (fromMaybe False) $ runSession False env $ withGhcInSession env \ srcs ->
-    computeMetadataInSession (pure ()) env (fst <$> srcs)
+  res <- fmap isJust $ runSession True env $ withDynFlags env \ dflags srcs -> do
+    prepareMetadataSession env dflags
+    module_graph <- writeMetadata (fst <$> srcs)
+    liftIO do
+      updateModuleGraph env.cache module_graph
+    pure (Just ())
   res <$ logMemStats "after metadata" env.log
diff --git a/internal/src/Internal/Session.hs b/internal/src/Internal/Session.hs
index 014aeb00..e375e4d5 100644
--- a/internal/src/Internal/Session.hs
+++ b/internal/src/Internal/Session.hs
@@ -104,6 +104,22 @@ parseFlags argv = do
   (dflags, fileish_args, dynamicFlagWarnings) <- parseDynamicFlags logger2 dflags1 argv
   pure (dflags, setLogFlags logger2 (initLogFlags dflags), fileish_args, dynamicFlagWarnings)
 
+-- | Parse CLI args and initialize 'DynFlags'.
+-- Returns the subset of args that have not been recognized as options.
+initDynFlags ::
+  DynFlags ->
+  Logger ->
+  [Located String] ->
+  DriverMessages ->
+  Ghc (DynFlags, [(String, Maybe Phase)])
+initDynFlags dflags0 logger fileish_args dynamicFlagWarnings = do
+  liftIO $ printOrThrowDiagnostics logger (initPrintConfig dflags0) (initDiagOpts dflags0) flagWarnings'
+  let (dflags1, srcs, objs) = parseTargetFiles dflags0 (map unLoc fileish_args)
+  unless (null objs) $ throwGhcException (UsageError ("Targets contain object files: " ++ show objs))
+  pure (dflags1, srcs)
+  where
+    flagWarnings' = GhcDriverMessage <$> dynamicFlagWarnings
+
 -- | Parse CLI args and set up the GHC session.
 -- Returns the subset of args that have not been recognized as options.
 initGhc ::
@@ -113,29 +129,35 @@ initGhc ::
   DriverMessages ->
   Ghc [(String, Maybe Phase)]
 initGhc dflags0 logger fileish_args dynamicFlagWarnings = do
-  liftIO $ printOrThrowDiagnostics logger (initPrintConfig dflags0) (initDiagOpts dflags0) flagWarnings'
-  let (dflags1, srcs, objs) = parseTargetFiles dflags0 (map unLoc fileish_args)
-  unless (null objs) $ throwGhcException (UsageError ("Targets contain object files: " ++ show objs))
+  (dflags1, srcs) <- initDynFlags dflags0 logger fileish_args dynamicFlagWarnings
   setSessionDynFlags dflags1
   pure srcs
-  where
-    flagWarnings' = GhcDriverMessage <$> dynamicFlagWarnings
 
--- | Run a program with a fresh session constructed from command line args.
--- Passes the unprocessed args to the callback, which usually consist of the file or module names intended for
--- compilation.
+-- | Run a program with fresh 'DynFlags' constructed from command line args.
+-- Passes the flags and the unprocessed args to the callback, which usually consist of the file or module names intended
+-- for compilation.
 -- In a Buck compile step these should always be a single path, but in the metadata step they enumerate an entire unit.
-withGhcInSession :: Env -> ([(String, Maybe Phase)] -> Ghc a) -> [Located String] -> Ghc a
-withGhcInSession env prog argv = do
+withDynFlags :: Env -> (DynFlags -> [(String, Maybe Phase)] -> Ghc a) -> [Located String] -> Ghc a
+withDynFlags env prog argv = do
   let !log = env.log
   pushLogHookM (const (logToState log))
   cache <- liftIO $ readMVar env.cache
   (dflags0, logger, fileish_args, dynamicFlagWarnings) <- parseFlags (argv ++ map instrumentLocation (words cache.options.extraGhcOptions))
   result <- prettyPrintGhcErrors logger do
-    srcs <- initGhc dflags0 logger fileish_args dynamicFlagWarnings
-    prog srcs
+    (dflags, srcs) <- initDynFlags dflags0 logger fileish_args dynamicFlagWarnings
+    prog dflags srcs
   result <$ popLogHookM
 
+-- | Run a program with a fresh session constructed from command line args.
+-- Passes the unprocessed args to the callback, which usually consist of the file or module names intended for
+-- compilation.
+-- In a Buck compile step these should always be a single path, but in the metadata step they enumerate an entire unit.
+withGhcInSession :: Env -> ([(String, Maybe Phase)] -> Ghc a) -> [Located String] -> Ghc a
+withGhcInSession env prog =
+  withDynFlags env \ dflags srcs -> do
+    setSessionDynFlags dflags
+    prog srcs
+
 -- | Create a base session and store it in the cache.
 -- On subsequent calls, return the cached session, unless the cache is disabled or @reuse@ is true.
 -- This will at some point be replaced by more deliberate methods.

From e19633618b0f372cf200da2740c0bb3c3391b2a4 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Tue, 3 Jun 2025 19:09:15 +0200
Subject: [PATCH 03/27] Use newtype `TargetId` in `TestSetup`

Change-Id: I1ab88f19d45c80792e0d1113db3004362d4a95b9
---
 ghc-worker/test/CompileHptTest.hs | 2 --
 ghc-worker/test/TestSetup.hs      | 4 ++--
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/ghc-worker/test/CompileHptTest.hs b/ghc-worker/test/CompileHptTest.hs
index a50fac19..a9d3fc6c 100644
--- a/ghc-worker/test/CompileHptTest.hs
+++ b/ghc-worker/test/CompileHptTest.hs
@@ -1,6 +1,5 @@
 module CompileHptTest where
 
-import Control.Concurrent (readMVar, threadDelay)
 import Control.Monad (unless, when)
 import Data.Char (toUpper)
 import Data.Foldable (fold, for_, traverse_)
@@ -126,7 +125,6 @@ stepCompile Conf {cache, tmp, args0} Module {unit, src} = do
     dbg ""
     dbg (">>> compiling " ++ takeFileName target.get)
     modifySession $ hscUpdateFlags \ d -> d {ghcMode = CompManager}
-    cache' <- liftIO $ readMVar env.cache
     compileModuleWithDepsInHpt target
   when (isNothing result) do
       liftIO $ throwGhcExceptionIO (ProgramError "Compile failed")
diff --git a/ghc-worker/test/TestSetup.hs b/ghc-worker/test/TestSetup.hs
index 69a1cf31..a7223ce4 100644
--- a/ghc-worker/test/TestSetup.hs
+++ b/ghc-worker/test/TestSetup.hs
@@ -14,7 +14,7 @@ import System.Environment (getEnv)
 import System.FilePath ((<.>), (</>))
 import System.IO.Temp (withSystemTempDirectory)
 import System.Process.Typed (proc, runProcess_)
-import Types.Args (Args (..))
+import Types.Args (Args (..), TargetId (..))
 
 -- | Global configuration for a worker compilation test.
 data Conf =
@@ -128,7 +128,7 @@ baseArgs :: FilePath -> FilePath -> Args
 baseArgs topdir tmp =
   Args {
     topdir = Just topdir,
-    workerTargetId = Just "test",
+    workerTargetId = Just (TargetId "test"),
     env = mempty,
     binPath = [],
     tempDir = Nothing,

From 8b4222ecf6a7f65cf017a00cea703088c46b0d88 Mon Sep 17 00:00:00 2001
From: Sjoerd Visscher <sjoerd@w3future.com>
Date: Wed, 4 Jun 2025 17:34:11 +0200
Subject: [PATCH 04/27] Instrument: support multiple workers per session

Change-Id: I33d1f48a82cfd32e1d79189f50afbd6f68a907b0
---
 instrument/Main.hs               | 26 +++++++-------
 instrument/UI.hs                 | 20 +++++++----
 instrument/UI/ModuleSelector.hs  | 13 +++----
 instrument/UI/Session.hs         | 61 +++++++++++++++++++++-----------
 instrument/UI/SessionSelector.hs | 33 ++++++++++++-----
 5 files changed, 98 insertions(+), 55 deletions(-)

diff --git a/instrument/Main.hs b/instrument/Main.hs
index 2a74d8fd..1d376ef2 100644
--- a/instrument/Main.hs
+++ b/instrument/Main.hs
@@ -5,7 +5,7 @@ import BuckWorker (Instrument)
 import Control.Concurrent (forkIO, threadDelay)
 import Control.Exception (SomeException, catch)
 import Control.Monad (filterM, void, when, forever)
-import Data.List (dropWhileEnd, isSuffixOf)
+import Data.List (isInfixOf)
 import Data.Maybe (fromMaybe)
 import Data.Text qualified as Text
 import Data.Time (getCurrentTime)
@@ -20,7 +20,7 @@ import Proto.Instrument qualified as Instr
 import Proto.Instrument_Fields qualified as Fields
 import System.Directory (doesPathExist, getModificationTime, listDirectory)
 import System.Environment (lookupEnv)
-import System.FSNotify (Event(..), watchTree, withManager)
+import System.FSNotify (Event(..), EventIsDirectory(..), watchDir, withManager)
 import UI qualified
 import UI.SessionSelector qualified as SS
 import UI.Session qualified as Session
@@ -36,8 +36,12 @@ listen :: BChan UI.Event -> FilePath -> IO ()
 listen eventChan instrPath = do
   void $ forkIO $ go 5
  where
+  -- TODO: This is a hack, ids should be sent over grpc
+  (sessionId', workerId') = break (== '_') instrPath
+  sessionId = Session.Id $ Text.pack sessionId'
+  workerId = Session.WorkerId $ Text.pack workerId'
   go :: Int -> IO ()
-  go 0 = getCurrentTime >>= writeBChan eventChan . UI.SessionSelectorEvent . SS.EndSession instrPath
+  go 0 = writeBChan eventChan $ UI.SessionSelectorEvent $ SS.RemoveWorker sessionId workerId
   go n =
     catch @SomeException
       ( withConnection def (ServerUnix instrPath) $ \conn -> do
@@ -47,9 +51,9 @@ listen eventChan instrPath = do
                     mkOptions options
           serverStreaming conn (rpc @(Protobuf Instrument "notifyMe")) defMessage $ \recv -> do
             time <- getModificationTime instrPath
-            writeBChan eventChan $ UI.SessionSelectorEvent $ SS.StartSession instrPath time sendOptions
-            writeBChan eventChan UI.SendOptions
-            whileNext_ recv $ writeBChan eventChan . UI.SessionSelectorEvent . SS.SessionEvent instrPath . Session.InstrEvent
+            writeBChan eventChan $ UI.SessionSelectorEvent $ SS.AddWorker sessionId workerId time sendOptions
+            writeBChan eventChan (UI.SendOptions (Just workerId))
+            whileNext_ recv $ writeBChan eventChan . UI.SessionSelectorEvent . SS.SessionEvent sessionId . Session.InstrEvent workerId
       )
       (const $ threadDelay 100_000 >> go (n - 1))
   mkOptions :: Options -> Proto Instr.Options
@@ -74,16 +78,14 @@ main = do
   when workerPathExists do
     primaryDirs <- do
       dirs <- listDirectory workers.path
-      filterM (\dir -> doesPathExist (workers.path ++ dir ++ "/primary")) dirs
+      filterM (\dir -> doesPathExist (workers.path ++ dir ++ "/instrument")) dirs
     mapM_ (listen eventChan . (++ "/instrument") . (workers.path ++)) primaryDirs
 
   -- Detect new workers
   withManager $ \mgr -> do
-    void $ watchTree mgr workers.path (const True) $ \case
-      Added file _ _ | "/primary" `isSuffixOf` file -> do
-        listen eventChan $ dropWhileEnd (/= '/') file ++ "instrument"
-      Modified file _ _ | "/primary" `isSuffixOf` file -> do
-        listen eventChan $ dropWhileEnd (/= '/') file ++ "instrument"
+    void $ watchDir mgr workers.path (const True) $ \case
+      Added dir _ IsDirectory | not ("/log" `isInfixOf` dir) -> do
+        listen eventChan $ dir ++ "instrument"
       _ -> pure ()
 
     (_, vty) <- UI.customMainWithDefaultVty (Just eventChan) UI.app UI.initialState
diff --git a/instrument/UI.hs b/instrument/UI.hs
index 1fc4eb17..88c3b4e0 100644
--- a/instrument/UI.hs
+++ b/instrument/UI.hs
@@ -13,6 +13,7 @@ import Brick.Widgets.Center (center)
 import Brick.Widgets.Core (joinBorders, modifyDefAttr, str, vBox, withBorderStyle, (<+>))
 import Brick.Widgets.Edit (editFocusedAttr)
 import Brick.Widgets.List (handleListEvent, listSelectedAttr, listSelectedElement, listSelectedElementL, listSelectedFocusedAttr)
+import Control.Exception (handle)
 import Control.Monad.IO.Class (liftIO)
 import Data.Foldable (for_)
 import Data.Text qualified as Text
@@ -20,7 +21,7 @@ import Data.Time (UTCTime (..), fromGregorian)
 import Graphics.Vty qualified as V
 import Graphics.Vty.Attributes.Color
 import Internal.Cache (Options (..), defaultOptions)
-import Lens.Micro.Platform (Lens', lens, makeLenses, packed, preuse, use, zoom, (.=), _2)
+import Lens.Micro.Platform (Lens', lens, makeLenses, packed, use, zoom, (.=), _2)
 import UI.ActiveTasks qualified as ActiveTasks
 import UI.ModuleSelector qualified as ModuleSelector
 import UI.Session qualified as Session
@@ -29,7 +30,7 @@ import UI.Types (Name (..))
 import UI.Utils (popup)
 
 data Event
-  = SendOptions
+  = SendOptions (Maybe Session.WorkerId)
   | SetTime UTCTime
   | SessionSelectorEvent SessionSelector.Event
 
@@ -90,11 +91,16 @@ drawOptionsEditor form = popup 50 "Session Options" $ renderForm form
 
 handleEvent :: BrickEvent Name Event -> EventM Name State ()
 handleEvent (AppEvent (SetTime t)) = currentTime .= t
-handleEvent (AppEvent SendOptions) = do
+handleEvent (AppEvent (SendOptions mwid)) = do
   opts <- use options
-  sendOpts <- preuse (sessions . listSelectedElementL)
-  for_ sendOpts $ \(_, s) -> do
-    liftIO $ Session._sendOptions s (formState opts)
+  workers <- use (sessions . listSelectedElementL . _2 . Session.workers)
+  let workers' = case mwid of
+        Nothing -> workers
+        Just wid -> filter (\w -> w._workerId == wid) workers
+  for_ workers' $ \worker -> do
+    liftIO $
+      handle @IOError (\_ -> pure ()) $
+        Session._sendOptions worker (formState opts)
 handleEvent (AppEvent (SessionSelectorEvent evt)) =
   zoom sessions (SessionSelector.handleEvent evt)
 handleEvent (VtyEvent evt) = do
@@ -110,7 +116,7 @@ handleEvent (VtyEvent evt) = do
     OptionsEditor -> do
       let hide = do
             currentFocus .= ModuleSelector
-            handleEvent (AppEvent SendOptions)
+            handleEvent (AppEvent (SendOptions Nothing))
       case evt of
         V.EvKey V.KEsc [] -> hide
         V.EvKey V.KEnter [] -> hide
diff --git a/instrument/UI/ModuleSelector.hs b/instrument/UI/ModuleSelector.hs
index 16c4b825..68d72701 100644
--- a/instrument/UI/ModuleSelector.hs
+++ b/instrument/UI/ModuleSelector.hs
@@ -1,13 +1,13 @@
 module UI.ModuleSelector where
 
-import Brick.Types (Widget, EventM)
-import Brick.Widgets.Core (Padding (..), padRight, str, (<+>), vBox, strWrap)
-import Brick.Widgets.List (GenericList, list, renderList, listElementsL, listSelectedL)
-import Data.Fixed (Pico, Fixed (..))
+import Brick.Types (EventM, Widget)
+import Brick.Widgets.Core (Padding (..), padRight, str, strWrap, vBox, (<+>))
+import Brick.Widgets.List (GenericList, list, listElementsL, listSelectedL, renderList)
+import Data.Fixed (Fixed (..), Pico)
 import Data.Sequence qualified as Seq
-import Lens.Micro.Platform (use, modifying, (.=))
+import Lens.Micro.Platform (modifying, use, (.=))
 import UI.Types (Name (ModuleSelector))
-import UI.Utils (formatPico, popup, formatPs, upsertAscSeq)
+import UI.Utils (formatPico, formatPs, popup, upsertAscSeq)
 
 type State = GenericList Name Seq.Seq Module
 
@@ -35,6 +35,7 @@ drawModuleDetails Module{..} =
       ]
 
 addModule :: String -> String -> Maybe Pico -> EventM Name State ()
+addModule "" _ _ = pure () -- TODO: Filter out earlier
 addModule name content compileTime = do
   mods <- use listElementsL
   let (i, mods') = upsertAscSeq _modName (Module name content compileTime) mods
diff --git a/instrument/UI/Session.hs b/instrument/UI/Session.hs
index d34e403f..dad3d5b5 100644
--- a/instrument/UI/Session.hs
+++ b/instrument/UI/Session.hs
@@ -5,11 +5,12 @@ module UI.Session where
 import Brick.Types (EventM, Widget)
 import Brick.Widgets.Border (borderWithLabel, hBorder)
 import Brick.Widgets.Core (str, vBox, vLimitPercent)
+import Control.Monad.IO.Class (liftIO)
 import Data.Map qualified as Map
 import Data.Text qualified as Text
-import Data.Time (UTCTime, getCurrentTime, nominalDiffTimeToSeconds, diffUTCTime)
+import Data.Time (UTCTime, diffUTCTime, getCurrentTime, nominalDiffTimeToSeconds)
 import Internal.Cache (Options (..))
-import Lens.Micro.Platform (makeLenses, modifying, zoom)
+import Lens.Micro.Platform (each, filtered, makeLenses, modifying, use, zoom)
 import Network.GRPC.Common.Protobuf (Proto, (^.))
 import Proto.Instrument qualified as Instr
 import Proto.Instrument_Fields qualified as Instr
@@ -17,43 +18,56 @@ import UI.ActiveTasks qualified as ActiveTasks
 import UI.ModuleSelector qualified as ModuleSelector
 import UI.Types (Name)
 import UI.Utils (formatBytes, formatPs, stripEscSeqs)
-import Control.Monad.IO.Class (liftIO)
 
-type Id = FilePath
+newtype Id = Id { unId :: Text.Text }
+  deriving stock (Eq, Ord, Show)
+newtype WorkerId = WorkerId { unWorkerId :: Text.Text }
+  deriving stock (Eq, Ord, Show)
 
 data State = Session
   { _title :: String
+  , _workers :: [Worker]
   , _activeTasks :: ActiveTasks.State
   , _modules :: ModuleSelector.State
   , _sesStartTime :: UTCTime
   , _sesEndTime :: Maybe UTCTime
-  , _stats :: Stats
+  , _finishedWorkerStats :: Stats
+  }
+
+data Worker = Worker
+  { _workerId :: WorkerId
   , _sendOptions :: Options -> IO ()
+  , _stats :: Stats
   }
 
 data Stats = Stats
   { _memory :: Map.Map Text.Text Int -- in bytes
   , _gc_cpu_ns :: Int
   , _cpu_ns :: Int
-  , _activeTaskCount :: Int
   }
+instance Semigroup Stats where
+  Stats m1 gc1 cpu1 <> Stats m2 gc2 cpu2 =
+    Stats (Map.unionWith (+) m1 m2) (gc1 + gc2) (cpu1 + cpu2)
+instance Monoid Stats where
+  mempty = Stats mempty 0 0
 
 makeLenses ''State
+makeLenses ''Worker
 makeLenses ''Stats
 
 data Event
-  = InstrEvent (Proto Instr.Event)
+  = InstrEvent WorkerId (Proto Instr.Event)
 
-mkSession :: String -> UTCTime -> (Options -> IO ()) -> State
-mkSession _title _startTime _sendOptions =
+mkSession :: String -> UTCTime -> State
+mkSession _title _startTime =
   Session
     { _title
+    , _workers = []
     , _activeTasks = ActiveTasks.initialState
     , _modules = ModuleSelector.initialState
     , _sesStartTime = _startTime
     , _sesEndTime = Nothing
-    , _stats = Stats mempty 0 0 0
-    , _sendOptions
+    , _finishedWorkerStats = mempty
     }
 
 draw :: Name -> UTCTime -> State -> Widget Name
@@ -64,14 +78,16 @@ draw current now Session{..} =
       , hBorder
       , ModuleSelector.draw current _modules
       , hBorder
-      , drawStats _stats
+      , drawStats (length _workers) (foldMap _stats _workers <> _finishedWorkerStats)
       ]
 
-drawStats :: Stats -> Widget Name
-drawStats Stats{..} =
+drawStats :: Int -> Stats -> Widget Name
+drawStats workerCount Stats{..} =
   vBox
     [ str $
-        "Memory:"
+        " Worker count: "
+          ++ show workerCount
+          ++ " | Memory:"
           ++ concatMap
             (\(k, v) -> " " ++ Text.unpack k ++ "=" ++ formatBytes v)
             (Map.toList _memory)
@@ -80,18 +96,15 @@ drawStats Stats{..} =
           ++ formatPs (1000 * _cpu_ns)
           ++ " | GC Time: "
           ++ formatPs (1000 * _gc_cpu_ns)
-          ++ (if _activeTaskCount > 0 then " | Active jobs: " ++ show _activeTaskCount else "")
     ]
 
 handleEvent :: Event -> EventM Name State ()
-handleEvent (InstrEvent evt) =
+handleEvent (InstrEvent wid evt) =
   case evt ^. Instr.maybe'compileStart of
     Just cs -> do
       zoom activeTasks $ ActiveTasks.addTask (Text.unpack $ cs ^. Instr.target)
-      modifying (stats . activeTaskCount) succ
     _ -> case evt ^. Instr.maybe'compileEnd of
       Just ce -> do
-        modifying (stats . activeTaskCount) pred
         let content = stripEscSeqs (Text.unpack $ ce ^. Instr.stderr)
             target' = Text.unpack $ ce ^. Instr.target
             target = if target' == "" then takeWhile (/= ':') content else target'
@@ -105,10 +118,16 @@ handleEvent (InstrEvent evt) =
             zoom activeTasks $ ActiveTasks.taskFailure target content
       _ -> case evt ^. Instr.maybe'stats of
         Just msg -> do
-          modifying stats \st ->
+          modifying (workers . each . filtered (\w -> w._workerId == wid) . stats) \st ->
             st
               { _memory = fromIntegral <$> msg ^. Instr.memory
               , _gc_cpu_ns = fromIntegral $ msg ^. Instr.gcCpuNs
               , _cpu_ns = fromIntegral $ msg ^. Instr.cpuNs
               }
-        _ -> pure ()
\ No newline at end of file
+        _ -> pure ()
+
+removeWorker :: WorkerId -> EventM Name State ()
+removeWorker wid = do
+  st <- use (workers . each . filtered (\w -> w._workerId == wid) . stats)
+  modifying finishedWorkerStats (<> st{_memory = mempty})
+  modifying workers (filter (\w -> w._workerId /= wid))
diff --git a/instrument/UI/SessionSelector.hs b/instrument/UI/SessionSelector.hs
index 35c2a50f..860eb7af 100644
--- a/instrument/UI/SessionSelector.hs
+++ b/instrument/UI/SessionSelector.hs
@@ -3,11 +3,12 @@ module UI.SessionSelector where
 import Brick.Types (EventM, Widget)
 import Brick.Widgets.Core (str)
 import Brick.Widgets.List (GenericList, list, listElementsL, listSelectedL, renderList)
+import Control.Monad.IO.Class (liftIO)
 import Data.Sequence qualified as Seq
-import Data.Time (UTCTime)
+import Data.Time (UTCTime, getCurrentTime)
 import Data.Time.Format.ISO8601 (iso8601Show)
 import Internal.Cache (Options)
-import Lens.Micro.Platform (Traversal', each, filtered, modifying, zoom, (.=), _2)
+import Lens.Micro.Platform (Traversal', each, filtered, modifying, preuse, zoom, (.=), _2)
 import UI.Session qualified as Session
 import UI.Types (Name (SessionSelector))
 import UI.Utils (popup)
@@ -15,9 +16,11 @@ import UI.Utils (popup)
 type State = GenericList Name Seq.Seq (Session.Id, Session.State)
 
 data Event
-  = StartSession Session.Id UTCTime (Options -> IO ())
-  | EndSession Session.Id UTCTime
+  = StartSession Session.Id UTCTime
+  | EndSession Session.Id
   | SessionEvent Session.Id Session.Event
+  | AddWorker Session.Id Session.WorkerId UTCTime (Options -> IO ())
+  | RemoveWorker Session.Id Session.WorkerId
 
 initialState :: State
 initialState = list SessionSelector Seq.empty 1
@@ -32,7 +35,8 @@ draw ss =
         [ if isSel then "> " else "  "
         , _title
         , " - "
-        , maybe "Running..." (take 19 . iso8601Show) _sesEndTime
+        , show (length _workers)
+        , " workers"
         ]
 
 sessionLens :: Session.Id -> Traversal' State (Session.State)
@@ -40,15 +44,26 @@ sessionLens sid =
   listElementsL . each . filtered ((== sid) . fst) . _2
 
 handleEvent :: Event -> EventM Name State ()
-handleEvent (StartSession sid start sendOpts) = do
+handleEvent (AddWorker sid wid time sendOpts) = do
+  session <- preuse (sessionLens sid)
+  case session of
+    Nothing -> handleEvent (StartSession sid time)
+    _ -> pure ()
+  zoom (sessionLens sid) $ do
+    modifying Session.workers (Session.Worker wid sendOpts mempty :)
+    modifying Session.sesStartTime (min time)
+handleEvent (RemoveWorker sid wid) = do
+  zoom (sessionLens sid) $ Session.removeWorker wid
+handleEvent (StartSession sid start) = do
   modifying
-    (listElementsL)
+    listElementsL
     ( \m ->
         let i = Seq.length m + 1
             stitle = "Session " ++ show i ++ "  " ++ take 19 (iso8601Show start)
-         in Seq.insertAt 0 (sid, Session.mkSession stitle start sendOpts) m
+         in Seq.insertAt 0 (sid, Session.mkSession stitle start) m
     )
   listSelectedL .= Just 0
-handleEvent (EndSession sid end) = do
+handleEvent (EndSession sid) = do
+  end <- liftIO getCurrentTime
   modifying (sessionLens sid . Session.sesEndTime) (const $ Just end)
 handleEvent (SessionEvent sid evt) = zoom (sessionLens sid) (Session.handleEvent evt)
\ No newline at end of file

From e11e39e09428e3619ce067cd6b95bf04e99f0d15 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 28 May 2025 19:09:03 +0200
Subject: [PATCH 05/27] Remove some obsolete code and comments

Abandoning the unimplemented function `storeIface` that was intended to
be an experiment with caching.

Change-Id: I11187e3293738b065a72ec84724144f0943a97ee
---
 internal/src/Internal/Cache.hs | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/internal/src/Internal/Cache.hs b/internal/src/Internal/Cache.hs
index b7ad9452..c0aaa384 100644
--- a/internal/src/Internal/Cache.hs
+++ b/internal/src/Internal/Cache.hs
@@ -248,7 +248,6 @@ data CacheFeatures =
 newCacheFeatures :: CacheFeatures
 newCacheFeatures = CacheFeatures {enable = True, loader = True, names = True, finder = True, eps = True, hpt = False}
 
--- TODO the name cache could in principle be shared directly – try it out
 data Cache =
   Cache {
     features :: CacheFeatures,
@@ -338,7 +337,6 @@ basicSymbolsStats base update =
     new = sizeUFM (minusUFM update.get base.get)
   }
 
--- TODO complicated
 basicNamesStats :: OrigNameCache -> OrigNameCache -> NamesStats
 basicNamesStats _ _ =
   NamesStats {
@@ -425,7 +423,6 @@ restoreCache target initialLoaderState initialSymbolCache initialNames cache
   | otherwise
   = pure (initialNames, (initialSymbolCache, (initialLoaderState, cache)))
 
--- TODO filter all cached items to include only external Names if possible
 initCache ::
   LoaderState ->
   SymbolCache ->
@@ -728,9 +725,6 @@ prepareCache cacheVar target hsc_env0 cache0 = do
   let (hsc_env1, cache1) = fromMaybe (hsc_env0, cache0 {features = cache0.features {loader = False}}) result
   pure (cache1, (hsc_env1, cache1.features.enable))
 
-storeIface :: HscEnv -> ModIface -> IO ()
-storeIface _ _ =
-  pure ()
 
 storeHug :: HscEnv -> Cache -> IO Cache
 storeHug hsc_env cache = do
@@ -758,7 +752,7 @@ finalizeCache ::
   Maybe ModuleArtifacts ->
   Cache ->
   IO Cache
-finalizeCache logVar workerId hsc_env target artifacts cache0 = do
+finalizeCache logVar workerId hsc_env target _ cache0 = do
   cache1 <-
     if cache0.features.enable
     then do
@@ -776,8 +770,6 @@ finalizeCache logVar workerId hsc_env target artifacts cache0 = do
         then do
           storeHug hsc_env cache1
         else pure cache1
-      for_ artifacts \ ModuleArtifacts {iface} ->
-        storeIface hsc_env iface
       pure cache2
     else pure cache0
   report logVar workerId target cache1

From 50e383ec129a60cf5565316c16190b2ae7cbce6d Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 17:09:33 +0200
Subject: [PATCH 06/27] Handle Interp sharing in make mode by storing it in the
 cache

Replacing the prior ad hoc solution

Change-Id: I6f99b04c9decc3814c5dd28d037c5e0d166fc8a4
---
 internal/src/Internal/Cache.hs   | 42 +++++++++++++++++++++++---------
 internal/src/Internal/Session.hs |  4 ---
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/internal/src/Internal/Cache.hs b/internal/src/Internal/Cache.hs
index c0aaa384..824c70d7 100644
--- a/internal/src/Internal/Cache.hs
+++ b/internal/src/Internal/Cache.hs
@@ -20,7 +20,7 @@ import Data.Traversable (for)
 import GHC (Ghc, ModIface, ModuleName, mi_module, moduleName, moduleNameString, setSession)
 import GHC.Data.FastString (FastString)
 import GHC.Driver.Env (HscEnv (..))
-import GHC.Driver.Monad (withSession)
+import GHC.Driver.Monad (modifySessionM, withSession)
 import GHC.Linker.Types (Linkable, LinkerEnv (..), Loader (..), LoaderState (..))
 import GHC.Ptr (Ptr)
 import GHC.Runtime.Interpreter (Interp (..))
@@ -251,7 +251,7 @@ newCacheFeatures = CacheFeatures {enable = True, loader = True, names = True, fi
 data Cache =
   Cache {
     features :: CacheFeatures,
-    interp :: Maybe InterpCache,
+    interpCache :: Maybe InterpCache,
     names :: OrigNameCache,
     stats :: Map Target CacheStats,
     path :: BinPath,
@@ -259,6 +259,7 @@ data Cache =
     eps :: ExternalUnitCache,
     hug :: Maybe HomeUnitGraph,
     moduleGraph :: Maybe ModuleGraph,
+    interp :: Maybe Interp,
     baseSession :: Maybe HscEnv,
     options :: Options
   }
@@ -275,7 +276,7 @@ emptyCacheWith features = do
   eps <- initExternalUnitCache
   newMVar Cache {
     features,
-    interp = Nothing,
+    interpCache = Nothing,
     names = emptyModuleEnv,
     stats = mempty,
     path = BinPath {
@@ -286,6 +287,7 @@ emptyCacheWith features = do
     eps,
     hug = Nothing,
     moduleGraph = Nothing,
+    interp = Nothing,
     baseSession = Nothing,
     options = defaultOptions
   }
@@ -403,7 +405,7 @@ restoreCache ::
   Cache ->
   IO (OrigNameCache, (SymbolCache, (Maybe LoaderState, Cache)))
 restoreCache target initialLoaderState initialSymbolCache initialNames cache
-  | Just InterpCache {..} <- cache.interp
+  | Just InterpCache {..} <- cache.interpCache
   = do
     (restoredLs, loaderStats) <- case initialLoaderState of
       Just sessionLs ->
@@ -430,7 +432,7 @@ initCache ::
   Cache ->
   IO Cache
 initCache loaderState symbols names Cache {names = _, ..} =
-  pure Cache {interp = Just InterpCache {..}, ..}
+  pure Cache {interpCache = Just InterpCache {..}, ..}
 
 updateLinkerEnv :: LinkerEnv -> LinkerEnv -> (LinkerEnv, LinkerStats)
 updateLinkerEnv cached session =
@@ -483,7 +485,7 @@ updateCache ::
 updateCache target InterpCache {..} newLoaderState newSymbols newNames cache = do
   (updatedLs, stats) <- updateLoaderState loaderState newLoaderState
   pure $ pushStats False target stats symbolsStats namesStats cache {
-    interp = Just InterpCache {
+    interpCache = Just InterpCache {
       loaderState = updatedLs,
       symbols = symbols <> newSymbols
     },
@@ -763,7 +765,7 @@ finalizeCache logVar workerId hsc_env target _ cache0 = do
             readMVar loaderStateVar >>= traverse \ newLoaderState -> do
               newSymbols <- readMVar symbolCacheVar
               newNames <- readMVar nsNames
-              maybe initCache (updateCache target) cache0.interp newLoaderState (SymbolCache newSymbols) newNames cache0
+              maybe initCache (updateCache target) cache0.interpCache newLoaderState (SymbolCache newSymbols) newNames cache0
         else pure cache0
       cache2 <-
         if cache0.features.hpt
@@ -811,18 +813,34 @@ logMemStats step logVar = do
   logMem "Max mem in use" s.max_mem_in_use_bytes
   logMem "Max live bytes" s.max_live_bytes
 
--- | Restore the shared state used by @compileHpt@ from the cache, consisting of the module graph and the HPT.
+-- | Restore the shared state used by @compileHpt@ from the cache, consisting of the module graph, the HPT, and the
+-- loader state and symbol cache that's contained in 'Interp'.
 -- The module graph is only modified by @computeMetadata@, so it will not be written back to the cache after
 -- compilation.
+--
+-- Managing 'Interp' is a bit difficult: The field 'hsc_interp' isn't initialized with everything else in 'newHscEnv',
+-- but only after parsing the command line arguments in 'setTopSessionDynFlags', since it needs to know the Ways of the
+-- session if an external interpreter is used.
+-- Therefore we grab the 'Interp' from the session when the cached value is absent, which amounts to the first
+-- compilation session of the build.
+-- When the cached value is present, on the other hand, we instead restore it into the session, making all subsequent
+-- sessions share the first one's 'Interp'.
+-- Both fields of 'Interp' are 'MVar's, so the state is shared immediately and concurrently.
 loadCacheMake ::
   MVar Log ->
   HscEnv ->
   Cache ->
-  IO (Cache, (HscEnv, ()))
+  IO (Cache, HscEnv)
 loadCacheMake logVar hsc_env cache = do
   logMemStats "load cache" logVar
-  pure (cache, (restoreModuleGraph (restoreHug hsc_env), ()))
+  pure (restoreModuleGraph . restoreHug <$> ensureInterp)
   where
+    ensureInterp = maybe storeInterp restoreInterp cache.interp
+
+    storeInterp = (cache {interp = hsc_env.hsc_interp}, hsc_env)
+
+    restoreInterp interp = (cache, hsc_env {hsc_interp = Just interp})
+
     restoreModuleGraph =
       maybe id (\ mg e -> e {hsc_mod_graph = mg}) cache.moduleGraph
 
@@ -852,9 +870,9 @@ withCacheMake ::
   Ghc (Maybe (Maybe ModuleArtifacts, a)) ->
   Ghc (Maybe (Maybe ModuleArtifacts, a))
 withCacheMake logVar cacheVar prog = do
-  _ <- withSessionM restore
+  modifySessionM restore
   prog <* withSession store
   where
-    restore hsc_env = modifyMVar cacheVar (loadCacheMake logVar hsc_env)
+    restore hsc_env = liftIO (modifyMVar cacheVar (loadCacheMake logVar hsc_env))
 
     store hsc_env = liftIO (modifyMVar_ cacheVar (storeCacheMake logVar hsc_env))
diff --git a/internal/src/Internal/Session.hs b/internal/src/Internal/Session.hs
index e375e4d5..51413210 100644
--- a/internal/src/Internal/Session.hs
+++ b/internal/src/Internal/Session.hs
@@ -198,10 +198,6 @@ runSession reuse Env {log, args, cache} prog = do
     cleanup session =
       flip unGhc session do
         hsc_env <- getSession
-        liftIO $ modifyMVar_ cache \case
-          Cache {baseSession = Just cachedEnv@HscEnv {hsc_interp = Nothing}, ..} ->
-            pure Cache {baseSession = Just cachedEnv {hsc_interp = hsc_env.hsc_interp}, ..}
-          c -> pure c
         liftIO $ unless (gopt Opt_KeepTmpFiles (hsc_dflags hsc_env)) do
           let tmpfs = hsc_tmpfs hsc_env
               logger = hsc_logger hsc_env

From e9a5118ba05db734d6db817ce5051a5a672f44fb Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 20:03:42 +0200
Subject: [PATCH 07/27] Move make mode related state into a separate type

Change-Id: I47834b71bc00c42222f0eee95ca40438aa175ae4
---
 internal/src/Internal/Cache.hs    | 108 +++++++++++++++---------------
 internal/src/Internal/Metadata.hs |  18 +++--
 2 files changed, 64 insertions(+), 62 deletions(-)

diff --git a/internal/src/Internal/Cache.hs b/internal/src/Internal/Cache.hs
index 824c70d7..5adf16e2 100644
--- a/internal/src/Internal/Cache.hs
+++ b/internal/src/Internal/Cache.hs
@@ -7,7 +7,6 @@ import Control.Monad (join)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Data.Bifunctor (first)
 import Data.Coerce (coerce)
-import Data.Foldable (for_)
 import Data.IORef (readIORef)
 import Data.List (sortBy)
 import qualified Data.Map.Strict as Map
@@ -17,7 +16,7 @@ import Data.Ord (comparing)
 import qualified Data.Set as Set
 import Data.Set (Set, (\\))
 import Data.Traversable (for)
-import GHC (Ghc, ModIface, ModuleName, mi_module, moduleName, moduleNameString, setSession)
+import GHC (Ghc, ModIface, ModuleName, emptyMG, mi_module, moduleName, moduleNameString, setSession)
 import GHC.Data.FastString (FastString)
 import GHC.Driver.Env (HscEnv (..))
 import GHC.Driver.Monad (modifySessionM, withSession)
@@ -34,7 +33,7 @@ import GHC.Unit.Env (
   UnitEnv (..),
   unitEnv_insert,
   unitEnv_lookup,
-  unitEnv_singleton,
+  unitEnv_new,
   unitEnv_union,
   )
 import GHC.Unit.External (ExternalUnitCache (..), initExternalUnitCache)
@@ -234,6 +233,29 @@ finderEnv FinderState {cache = FinderCache {fcModuleCache}} =
 
 #endif
 
+-- | Data extracted from 'HscEnv' for the purpose of persisting it across sessions.
+--
+-- While many parts of the session are either contained in mutable variables or trivially reinitialized, some components
+-- must be handled explicitly: The module graph and home unit graph are pure fields that need to be shared, and the
+-- interpreter state for TH execution is only initialized when the flags are parsed.
+data MakeState =
+  MakeState {
+    -- | The module graph for a specific unit is computed in its metadata step, after which it's extracted and merged
+    -- into the existing graph.
+    moduleGraph :: ModuleGraph,
+
+    -- | The unit environment for a specific unit is inserted into the shared home unit graph at the beginning of the
+    -- metadata step, constructed from the dependency specifications provided by Buck.
+    -- After compilation of a module, its 'HomeUnitInfo' is inserted into the home package table contained in its unit's
+    -- unit environment.
+    hug :: HomeUnitGraph,
+
+    -- | While the interpreter state contains a mutable variable that would be shared across sessions, it isn't
+    -- initialized properly until the first module compilation's flags have been parsed, so we store it in the shared
+    -- state for consistency.
+    interp :: Maybe Interp
+  }
+
 data CacheFeatures =
   CacheFeatures {
     enable :: Bool,
@@ -257,11 +279,9 @@ data Cache =
     path :: BinPath,
     finder :: FinderState,
     eps :: ExternalUnitCache,
-    hug :: Maybe HomeUnitGraph,
-    moduleGraph :: Maybe ModuleGraph,
-    interp :: Maybe Interp,
     baseSession :: Maybe HscEnv,
-    options :: Options
+    options :: Options,
+    make :: MakeState
   }
 
 data Options =
@@ -285,11 +305,13 @@ emptyCacheWith features = do
     },
     finder,
     eps,
-    hug = Nothing,
-    moduleGraph = Nothing,
-    interp = Nothing,
     baseSession = Nothing,
-    options = defaultOptions
+    options = defaultOptions,
+    make = MakeState {
+      moduleGraph = emptyMG,
+      hug = unitEnv_new mempty,
+      interp = Nothing
+    }
   }
 
 emptyCache :: Bool -> IO (MVar Cache)
@@ -654,21 +676,6 @@ setTarget cacheVar cache target hsc_env = do
       hsc_FC <- newFinderCache cacheVar cache target
       pure e {hsc_FC}
 
--- | Restore cache parts related to make mode.
-restoreHptCache :: Cache -> HscEnv -> IO HscEnv
-restoreHptCache cache hsc_env = do
-  let
-    -- If the cache contains a module graph stored by @compileModuleWithDepsInHpt@, restore it
-    hsc_env1 = maybe id restoreHug cache.hug hsc_env
-
-    -- If the cache contains a module graph stored by @computeMetadata@, restore it
-    hsc_env2 = maybe id restoreModuleGraph cache.moduleGraph hsc_env1
-  pure hsc_env2
-  where
-    restoreModuleGraph mg e = e {hsc_mod_graph = mg}
-
-    restoreHug hug e = e {hsc_unit_env = e.hsc_unit_env {ue_home_unit_graph = hug}}
-
 -- | Restore cache parts related to oneshot mode.
 restoreOneshotCache :: Cache -> HscEnv -> IO HscEnv
 restoreOneshotCache cache hsc_env = do
@@ -689,7 +696,8 @@ updateModuleGraph :: MVar Cache -> ModuleGraph -> IO ()
 updateModuleGraph cacheVar new =
   modifyMVar_ cacheVar \ cache -> do
 #if defined(MWB)
-    pure cache {moduleGraph = Just (maybe new merge cache.moduleGraph)}
+    let !merged = merge cache.make.moduleGraph
+    pure cache {make = cache.make {moduleGraph = merged}}
   where
     merge old =
       mkModuleGraph (Map.elems (Map.unionWith mergeNodes oldMap newMap))
@@ -704,7 +712,7 @@ updateModuleGraph cacheVar new =
 
         newMap = Map.fromList $ [(mkNodeKey n, n) | n <- mgModSummaries' new]
 #else
-    pure cache {moduleGraph = Just (maybe id unionMG cache.moduleGraph new)}
+    pure cache {make = cache.make {moduleGraph = unionMG cache.make.moduleGraph new}}
 #endif
 
 prepareCache :: MVar Cache -> Target -> HscEnv -> Cache -> IO (Cache, (HscEnv, Bool))
@@ -712,7 +720,7 @@ prepareCache cacheVar target hsc_env0 cache0 = do
   result <-
     if cache0.features.enable
     then do
-      hsc_env1 <- restoreOneshotCache cache0 =<< restoreHptCache cache0 =<< setTarget cacheVar cache0 target hsc_env0
+      hsc_env1 <- restoreOneshotCache cache0 =<< setTarget cacheVar cache0 target hsc_env0
       if cache0.features.loader
       then do
         withHscState hsc_env1 \ nsNames loaderStateVar symbolCacheVar -> do
@@ -730,17 +738,16 @@ prepareCache cacheVar target hsc_env0 cache0 = do
 
 storeHug :: HscEnv -> Cache -> IO Cache
 storeHug hsc_env cache = do
-  pure cache {hug = Just merged}
-  where
-    merged = maybe id (unitEnv_union mergeHugs) cache.hug hsc_env.hsc_unit_env.ue_home_unit_graph
+  let !new = hsc_env.hsc_unit_env.ue_home_unit_graph
+      !hug = unitEnv_union mergeHugs cache.make.hug new
+  pure cache {make = cache.make {hug}}
 
 -- | Extract the unit env of the currently active unit and store it in the cache.
 -- This is used by the make mode worker after the metadata step has initialized the new unit.
 insertUnitEnv :: HscEnv -> Cache -> Cache
 insertUnitEnv hsc_env cache =
-  cache {hug = Just (maybe fresh update cache.hug)}
+  cache {make = cache.make {hug = update cache.make.hug}}
   where
-    fresh = unitEnv_singleton current ue
     ue = unitEnv_lookup current hsc_env.hsc_unit_env.ue_home_unit_graph
     current = hsc_env.hsc_unit_env.ue_current_unit
     update = unitEnv_insert current ue
@@ -758,21 +765,14 @@ finalizeCache logVar workerId hsc_env target _ cache0 = do
   cache1 <-
     if cache0.features.enable
     then do
-      cache1 <-
-        if cache0.features.loader
-        then do
-          fromMaybe cache0 . join <$> withHscState hsc_env \ nsNames loaderStateVar symbolCacheVar ->
-            readMVar loaderStateVar >>= traverse \ newLoaderState -> do
-              newSymbols <- readMVar symbolCacheVar
-              newNames <- readMVar nsNames
-              maybe initCache (updateCache target) cache0.interpCache newLoaderState (SymbolCache newSymbols) newNames cache0
-        else pure cache0
-      cache2 <-
-        if cache0.features.hpt
-        then do
-          storeHug hsc_env cache1
-        else pure cache1
-      pure cache2
+      if cache0.features.loader
+      then do
+        fromMaybe cache0 . join <$> withHscState hsc_env \ nsNames loaderStateVar symbolCacheVar ->
+          readMVar loaderStateVar >>= traverse \ newLoaderState -> do
+            newSymbols <- readMVar symbolCacheVar
+            newNames <- readMVar nsNames
+            maybe initCache (updateCache target) cache0.interpCache newLoaderState (SymbolCache newSymbols) newNames cache0
+      else pure cache0
     else pure cache0
   report logVar workerId target cache1
   pure cache1
@@ -835,17 +835,15 @@ loadCacheMake logVar hsc_env cache = do
   logMemStats "load cache" logVar
   pure (restoreModuleGraph . restoreHug <$> ensureInterp)
   where
-    ensureInterp = maybe storeInterp restoreInterp cache.interp
+    ensureInterp = maybe storeInterp restoreInterp cache.make.interp
 
-    storeInterp = (cache {interp = hsc_env.hsc_interp}, hsc_env)
+    storeInterp = (cache {make = cache.make {interp = hsc_env.hsc_interp}}, hsc_env)
 
     restoreInterp interp = (cache, hsc_env {hsc_interp = Just interp})
 
-    restoreModuleGraph =
-      maybe id (\ mg e -> e {hsc_mod_graph = mg}) cache.moduleGraph
+    restoreModuleGraph e = e {hsc_mod_graph = cache.make.moduleGraph}
 
-    restoreHug =
-      maybe id (\ hug e -> e {hsc_unit_env = e.hsc_unit_env {ue_home_unit_graph = hug}}) cache.hug
+    restoreHug e = e {hsc_unit_env = e.hsc_unit_env {ue_home_unit_graph = cache.make.hug}}
 
 -- | Store the changes made to the HUG by @compileHpt@ in the cache, which usually consists of adding a single
 -- 'HomeModInfo'.
diff --git a/internal/src/Internal/Metadata.hs b/internal/src/Internal/Metadata.hs
index 9a3d7a32..8f767bc9 100644
--- a/internal/src/Internal/Metadata.hs
+++ b/internal/src/Internal/Metadata.hs
@@ -12,21 +12,25 @@ import GHC.Runtime.Loader (initializeSessionPlugins)
 import GHC.Unit (HomeUnit, UnitDatabase, UnitId, UnitState, initUnits)
 import GHC.Unit.Env (HomeUnitEnv (..), UnitEnv (..), unitEnv_insert, unitEnv_keys, unitEnv_union, updateHug)
 import GHC.Unit.Home.ModInfo (emptyHomePackageTable)
-import Internal.Cache (Cache (..), insertUnitEnv, mergeHugs, updateModuleGraph, logMemStats)
+import Internal.Cache (Cache (..), MakeState (..), insertUnitEnv, logMemStats, mergeHugs, updateModuleGraph)
 import Internal.MakeFile (doMkDependHS)
 import Internal.Session (Env (..), runSession, withDynFlags)
 
 -- | Copy the cached unit env and module graph to the given session.
 restoreEnv :: Cache -> HscEnv -> HscEnv
-restoreEnv cache hsc_env = do
-  maybe id restoreMg cache.moduleGraph $ maybe hsc_env restore cache.hug
+restoreEnv cache hsc_env =
+  restoreMg cache.make.moduleGraph withHug
   where
-    restoreMg new e = e {hsc_mod_graph = new}
+    !withHug =
+      hsc_env {
+        hsc_unit_env = hsc_env.hsc_unit_env {
+          ue_home_unit_graph = unitEnv_union mergeHugs cache.make.hug current
+        }
+      }
 
-    restore hug =
-      hsc_env {hsc_unit_env = hsc_env.hsc_unit_env {ue_home_unit_graph = unitEnv_union mergeHugs hug current}}
+    restoreMg !new e = e {hsc_mod_graph = new}
 
-    current = hsc_env.hsc_unit_env.ue_home_unit_graph
+    !current = hsc_env.hsc_unit_env.ue_home_unit_graph
 
 -- | 'doMkDependHS' needs this to be enabled.
 metadataTempSession :: HscEnv -> HscEnv

From aec0ec111e34788be5baf93105df4a4614a178a0 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 20:23:26 +0200
Subject: [PATCH 08/27] Refactor `Metadata` to use the same function as
 `CompileHpt` to restore state

Change-Id: Icfd7a7737f4cfd1161c655aab5d1c857e9704d49
---
 internal/src/Internal/Cache.hs    | 32 ++++++++++++++++++++-----------
 internal/src/Internal/Metadata.hs | 22 +++------------------
 2 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/internal/src/Internal/Cache.hs b/internal/src/Internal/Cache.hs
index 5adf16e2..798ff213 100644
--- a/internal/src/Internal/Cache.hs
+++ b/internal/src/Internal/Cache.hs
@@ -813,6 +813,21 @@ logMemStats step logVar = do
   logMem "Max mem in use" s.max_mem_in_use_bytes
   logMem "Max live bytes" s.max_live_bytes
 
+-- | Restore the shared state used by both @computeMetadata@ and @compileHpt@ from the cache.
+-- See 'loadCacheMakeCompile' for details.
+loadCacheMake ::
+  MVar Log ->
+  HscEnv ->
+  Cache ->
+  IO HscEnv
+loadCacheMake logVar hsc_env cache = do
+  logMemStats "load cache" logVar
+  pure (restoreHug (restoreModuleGraph hsc_env))
+  where
+    restoreModuleGraph e = e {hsc_mod_graph = cache.make.moduleGraph}
+
+    restoreHug e = e {hsc_unit_env = e.hsc_unit_env {ue_home_unit_graph = cache.make.hug}}
+
 -- | Restore the shared state used by @compileHpt@ from the cache, consisting of the module graph, the HPT, and the
 -- loader state and symbol cache that's contained in 'Interp'.
 -- The module graph is only modified by @computeMetadata@, so it will not be written back to the cache after
@@ -826,24 +841,19 @@ logMemStats step logVar = do
 -- When the cached value is present, on the other hand, we instead restore it into the session, making all subsequent
 -- sessions share the first one's 'Interp'.
 -- Both fields of 'Interp' are 'MVar's, so the state is shared immediately and concurrently.
-loadCacheMake ::
+loadCacheMakeCompile ::
   MVar Log ->
   HscEnv ->
   Cache ->
   IO (Cache, HscEnv)
-loadCacheMake logVar hsc_env cache = do
-  logMemStats "load cache" logVar
-  pure (restoreModuleGraph . restoreHug <$> ensureInterp)
+loadCacheMakeCompile logVar hsc_env0 cache = do
+  ensureInterp <$> loadCacheMake logVar hsc_env0 cache
   where
     ensureInterp = maybe storeInterp restoreInterp cache.make.interp
 
-    storeInterp = (cache {make = cache.make {interp = hsc_env.hsc_interp}}, hsc_env)
+    storeInterp hsc_env = (cache {make = cache.make {interp = hsc_env.hsc_interp}}, hsc_env)
 
-    restoreInterp interp = (cache, hsc_env {hsc_interp = Just interp})
-
-    restoreModuleGraph e = e {hsc_mod_graph = cache.make.moduleGraph}
-
-    restoreHug e = e {hsc_unit_env = e.hsc_unit_env {ue_home_unit_graph = cache.make.hug}}
+    restoreInterp interp hsc_env = (cache, hsc_env {hsc_interp = Just interp})
 
 -- | Store the changes made to the HUG by @compileHpt@ in the cache, which usually consists of adding a single
 -- 'HomeModInfo'.
@@ -871,6 +881,6 @@ withCacheMake logVar cacheVar prog = do
   modifySessionM restore
   prog <* withSession store
   where
-    restore hsc_env = liftIO (modifyMVar cacheVar (loadCacheMake logVar hsc_env))
+    restore hsc_env = liftIO (modifyMVar cacheVar (loadCacheMakeCompile logVar hsc_env))
 
     store hsc_env = liftIO (modifyMVar_ cacheVar (storeCacheMake logVar hsc_env))
diff --git a/internal/src/Internal/Metadata.hs b/internal/src/Internal/Metadata.hs
index 8f767bc9..f9312c7a 100644
--- a/internal/src/Internal/Metadata.hs
+++ b/internal/src/Internal/Metadata.hs
@@ -10,28 +10,12 @@ import GHC.Driver.Session (updatePlatformConstants)
 import GHC.Platform.Ways (Way (WayDyn), addWay)
 import GHC.Runtime.Loader (initializeSessionPlugins)
 import GHC.Unit (HomeUnit, UnitDatabase, UnitId, UnitState, initUnits)
-import GHC.Unit.Env (HomeUnitEnv (..), UnitEnv (..), unitEnv_insert, unitEnv_keys, unitEnv_union, updateHug)
+import GHC.Unit.Env (HomeUnitEnv (..), UnitEnv (..), unitEnv_insert, unitEnv_keys, updateHug)
 import GHC.Unit.Home.ModInfo (emptyHomePackageTable)
-import Internal.Cache (Cache (..), MakeState (..), insertUnitEnv, logMemStats, mergeHugs, updateModuleGraph)
+import Internal.Cache (insertUnitEnv, loadCacheMake, logMemStats, updateModuleGraph)
 import Internal.MakeFile (doMkDependHS)
 import Internal.Session (Env (..), runSession, withDynFlags)
 
--- | Copy the cached unit env and module graph to the given session.
-restoreEnv :: Cache -> HscEnv -> HscEnv
-restoreEnv cache hsc_env =
-  restoreMg cache.make.moduleGraph withHug
-  where
-    !withHug =
-      hsc_env {
-        hsc_unit_env = hsc_env.hsc_unit_env {
-          ue_home_unit_graph = unitEnv_union mergeHugs cache.make.hug current
-        }
-      }
-
-    restoreMg !new e = e {hsc_mod_graph = new}
-
-    !current = hsc_env.hsc_unit_env.ue_home_unit_graph
-
 -- | 'doMkDependHS' needs this to be enabled.
 metadataTempSession :: HscEnv -> HscEnv
 metadataTempSession =
@@ -90,7 +74,7 @@ addHomeUnit dflags = do
 prepareMetadataSession :: Env -> DynFlags -> Ghc ()
 prepareMetadataSession env dflags = do
   cache <- liftIO $ readMVar env.cache
-  modifySession (restoreEnv cache)
+  modifySessionM \ hsc_env -> liftIO (loadCacheMake env.log hsc_env cache)
   unit <- addHomeUnit dflags
   setActiveUnit unit
   storeNewUnit

From 6411817160a237cba34e4fc79b102247511163ab Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 17:01:15 +0200
Subject: [PATCH 09/27] Remove legacy code for initializing the unit env in
 tests

Change-Id: I1d3f13e7d34e0b06a9e9f6968cdde6ac9dbbf5db
---
 ghc-worker/test/CompileHptTest.hs | 53 ++-----------------------------
 1 file changed, 3 insertions(+), 50 deletions(-)

diff --git a/ghc-worker/test/CompileHptTest.hs b/ghc-worker/test/CompileHptTest.hs
index a9d3fc6c..4a3e2f56 100644
--- a/ghc-worker/test/CompileHptTest.hs
+++ b/ghc-worker/test/CompileHptTest.hs
@@ -2,24 +2,20 @@ module CompileHptTest where
 
 import Control.Monad (unless, when)
 import Data.Char (toUpper)
-import Data.Foldable (fold, for_, traverse_)
+import Data.Foldable (for_, traverse_)
 import Data.List (intercalate)
 import qualified Data.List.NonEmpty as NonEmpty
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.Map.Strict as Map
 import Data.Map.Strict ((!?))
 import Data.Maybe (catMaybes, isNothing)
-import qualified Data.Set as Set
-import Data.Set (Set)
-import GHC (DynFlags (..), Ghc, GhcException (..), GhcMode (..), GhcMonad (..), mkGeneralLocated, setUnitDynFlags)
+import GHC (DynFlags (..), Ghc, GhcException (..), GhcMode (..), mkGeneralLocated)
 import GHC.Driver.Config.Diagnostic (initDiagOpts, initPrintConfig)
-import GHC.Driver.Env (HscEnv (..), hscSetActiveUnitId, hscUpdateFlags)
+import GHC.Driver.Env (HscEnv (..), hscUpdateFlags)
 import GHC.Driver.Errors (printOrThrowDiagnostics)
 import GHC.Driver.Errors.Types (GhcMessage (..))
 import GHC.Driver.Monad (modifySession)
 import GHC.Driver.Session (parseDynamicFlagsCmdLine)
-import GHC.Unit (UnitId, UnitState (..), stringToUnitId, unitIdString)
-import GHC.Unit.Env (HomeUnitEnv (..), HomeUnitGraph, UnitEnv (..), UnitEnvGraph (..), unitEnv_lookup_maybe)
 import GHC.Utils.Monad (MonadIO (..))
 import GHC.Utils.Outputable (ppr, showPprUnsafe, text, (<+>))
 import GHC.Utils.Panic (throwGhcExceptionIO)
@@ -41,49 +37,6 @@ unitFlags args HscEnv {hsc_logger, hsc_dflags = dflags0} = do
   liftIO $ printOrThrowDiagnostics hsc_logger (initPrintConfig dflags) (initDiagOpts dflags) (GhcDriverMessage <$> warns)
   pure dflags
 
--- | Turn each @-package@ argument in the given list into a @-package-id@ argument if the following package name refers
--- to a home unit, and create a set of all mentioned home unit names.
--- This is needed for GHC to recognize home unit dependencies.
-adaptHp :: HomeUnitGraph -> [String] -> (Set UnitId, [String])
-adaptHp (UnitEnvGraph ueg) =
-  spin mempty
-  where
-    spin seen = \case
-      "-package" : p : rest
-        | Map.member (stringToUnitId p) ueg
-        -> (["-package-id", p] ++) <$> spin (Set.insert (stringToUnitId p) seen) rest
-      arg : rest -> (arg :) <$> spin seen rest
-      [] -> (seen, [])
-
--- | Return the given unit's dependencies.
-homeUnitDeps :: HscEnv -> UnitId -> Maybe [UnitId]
-homeUnitDeps hsc_env target = do
-  HomeUnitEnv {homeUnitEnv_units = UnitState {homeUnitDepends}} <- unitEnv_lookup_maybe target hug
-  pure homeUnitDepends
-  where
-    hug = hsc_env.hsc_unit_env.ue_home_unit_graph
-
--- | Assemble @-package-id@ arguments for the current unit's dependencies, omitting those that are present in the
--- provided set, which are the current module's deps specified by Buck.
-homeUnitDepFlags :: HscEnv -> Set UnitId -> UnitId -> [String]
-homeUnitDepFlags hsc_env explicit target =
-  concat [["-package-id", unitIdString u] | u <- fold prev_deps, not (Set.member u explicit)]
-  where
-    prev_deps = homeUnitDeps hsc_env target
-
--- | Update the current unit's @DynFlags@ stored in the unit env, and reinitialize its unit state.
--- Since different modules in the same unit may have arbitrary subsets of the unit's package dependencies when Buck
--- compiles them, we take the union of existing and new dependencies.
-initUnit :: [String] -> Ghc ()
-initUnit specific = do
-  hsc_env0 <- getSession
-  let current = hsc_env0.hsc_unit_env.ue_current_unit
-      (explicit, withPackageId) = adaptHp hsc_env0.hsc_unit_env.ue_home_unit_graph specific
-      unitOptions = withPackageId ++ homeUnitDepFlags hsc_env0 explicit current
-  dflags <- unitFlags unitOptions hsc_env0
-  setUnitDynFlags current dflags
-  modifySession (hscSetActiveUnitId current)
-
 stepMetadata :: Conf -> Unit -> [Unit] -> IO ()
 stepMetadata Conf {cache, tmp, args0} unit deps = do
   log <- newLog True

From 32dff98d2cce668c4d49ffddb2120c174d515878 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 20:04:49 +0200
Subject: [PATCH 10/27] Remove obsolete function

Change-Id: Iaa2e40f934576be9375c9bf7add5d8da35c7de2e
---
 internal/src/Internal/Session.hs | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/internal/src/Internal/Session.hs b/internal/src/Internal/Session.hs
index 51413210..0b4e4ac6 100644
--- a/internal/src/Internal/Session.hs
+++ b/internal/src/Internal/Session.hs
@@ -336,11 +336,3 @@ withGhcMhu env f =
           a <- res
           pure (Nothing, a)
       pure (snd <$> result)
-
--- | Like @withGhcMhu@, specialized to @ModuleArtifacts@.
-withGhcMhuDefault ::
-  Env ->
-  ([String] -> Target -> Ghc (Maybe (Maybe ModuleArtifacts, a))) ->
-  IO (Maybe (Maybe ModuleArtifacts, a))
-withGhcMhuDefault env =
-  withGhcUsingCacheMhu (withCache env.log env.args.workerTargetId env.cache) env

From 9405d8186a1a15592fb3dac4ee3f2e6ca337ca76 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 28 May 2025 18:37:16 +0200
Subject: [PATCH 11/27] Disable finder cache statistics

I doubt that these are very helpful unless you're looking for something,
so let's exclude this potential source of leaks in the default settings.

Change-Id: I3ee75cd9e7c701a4c1f985b4e4d826c2125e4650
---
 ghc-worker/lib/GhcWorker/Run.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ghc-worker/lib/GhcWorker/Run.hs b/ghc-worker/lib/GhcWorker/Run.hs
index d9092216..381481c7 100644
--- a/ghc-worker/lib/GhcWorker/Run.hs
+++ b/ghc-worker/lib/GhcWorker/Run.hs
@@ -84,7 +84,7 @@ runWorker CliOptions {workerMode, serve} = do
           loader = False,
           enable = True,
           names = False,
-          finder = True,
+          finder = False,
           eps = False
         }
       WorkerOneshotMode -> emptyCache True

From d3fbe536db0462f3c2ce8c1768a3d206457784de Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 28 May 2025 18:32:45 +0200
Subject: [PATCH 12/27] Remove the command env from `Args`

Change-Id: I0735488a9c3fa58c11b52ece9c2c61091c7ab27e
---
 ghc-worker/test/TestSetup.hs | 1 -
 types/src/Types/Args.hs      | 2 --
 types/src/Types/BuckArgs.hs  | 1 -
 3 files changed, 4 deletions(-)

diff --git a/ghc-worker/test/TestSetup.hs b/ghc-worker/test/TestSetup.hs
index a7223ce4..ba6e7327 100644
--- a/ghc-worker/test/TestSetup.hs
+++ b/ghc-worker/test/TestSetup.hs
@@ -129,7 +129,6 @@ baseArgs topdir tmp =
   Args {
     topdir = Just topdir,
     workerTargetId = Just (TargetId "test"),
-    env = mempty,
     binPath = [],
     tempDir = Nothing,
     ghcPath = Nothing,
diff --git a/types/src/Types/Args.hs b/types/src/Types/Args.hs
index 9afa5f11..ce89e601 100644
--- a/types/src/Types/Args.hs
+++ b/types/src/Types/Args.hs
@@ -10,7 +10,6 @@ data Args =
   Args {
     topdir :: Maybe String,
     workerTargetId :: Maybe TargetId,
-    env :: Map String String,
     binPath :: [String],
     tempDir :: Maybe String,
     ghcPath :: Maybe String,
@@ -23,7 +22,6 @@ emptyArgs env =
   Args {
     topdir = Nothing,
     workerTargetId = Nothing,
-    env,
     binPath = [],
     tempDir = env !? "TMPDIR",
     ghcPath = Nothing,
diff --git a/types/src/Types/BuckArgs.hs b/types/src/Types/BuckArgs.hs
index f003697f..f8e40a3e 100644
--- a/types/src/Types/BuckArgs.hs
+++ b/types/src/Types/BuckArgs.hs
@@ -151,7 +151,6 @@ toGhcArgs args = do
   pure Args {
     topdir,
     workerTargetId = args.workerTargetId,
-    env = args.env,
     binPath = args.binPath,
     tempDir = args.tempDir,
     ghcPath = args.ghcPath,

From 51d1bee2a2880df68cbfbc238d122d9ed54846e3 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 28 May 2025 18:51:10 +0200
Subject: [PATCH 13/27] Fix some names

Change-Id: I6adb726cc89c3ef1294481910a0806f50c694ed3
---
 ghc-worker/lib/GhcWorker/Orchestration.hs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ghc-worker/lib/GhcWorker/Orchestration.hs b/ghc-worker/lib/GhcWorker/Orchestration.hs
index ac3b9bb7..b9573130 100644
--- a/ghc-worker/lib/GhcWorker/Orchestration.hs
+++ b/ghc-worker/lib/GhcWorker/Orchestration.hs
@@ -209,14 +209,14 @@ runOrProxyCentralGhc socketDir runServer = do
 
 -- | Start a gRPC server that either runs GHC (primary server) or a proxy that forwards requests to the primary.
 serveOrProxyCentralGhc :: CreateMethods -> ServerSocketPath -> IO ()
-serveOrProxyCentralGhc mode socket = do
+serveOrProxyCentralGhc methods socket = do
   runOrProxyCentralGhc socketDir run >>= \case
     Right (_, thread) -> onException (wait thread) (cancel thread)
     Left primary -> proxyServer primary socket
   where
     run primaryFile = do
       let primary = PrimarySocketPath socket.path
-      thread <- async (runCentralGhc mode primaryFile socket instr)
+      thread <- async (runCentralGhc methods primaryFile socket instr)
       waitPoll primary
       pure (primary, thread)
 

From 4a73ba4b6b37cd843614efae59115eb190bbe0ac Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 21:01:00 +0200
Subject: [PATCH 14/27] Enable temp file cleanup after compiling a module

Change-Id: Ib6e23b542463694685ed21259b94c07a2ad6c2bc
---
 internal/src/Internal/CompileHpt.hs | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/internal/src/Internal/CompileHpt.hs b/internal/src/Internal/CompileHpt.hs
index 962c50ba..e00d3a02 100644
--- a/internal/src/Internal/CompileHpt.hs
+++ b/internal/src/Internal/CompileHpt.hs
@@ -2,7 +2,6 @@
 
 module Internal.CompileHpt where
 
-import Control.Monad (when)
 import GHC (DynFlags (..), GeneralFlag (..), Ghc, GhcMonad (..), Logger, ModLocation (..), ModSummary (..), gopt)
 import GHC.Driver.Env (HscEnv (..), hscUpdateHUG)
 import GHC.Driver.Errors.Types (GhcMessage (..))
@@ -57,9 +56,7 @@ compileModuleWithDepsInHpt (Target src) = do
     summResult <- summariseFile hsc_env (ue_unsafeHomeUnit (hsc_unit_env hsc_env)) mempty src Nothing Nothing
     summary <- setHiLocation hsc_env <$> eitherMessages GhcDriverMessage summResult
     result <- compileOne hsc_env summary 1 100000 Nothing (HomeModLinkable Nothing Nothing)
-    -- This deletes assembly files too early
-    when False do
-      cleanCurrentModuleTempFilesMaybe (hsc_logger hsc_env) (hsc_tmpfs hsc_env) summary.ms_hspp_opts
+    cleanCurrentModuleTempFilesMaybe (hsc_logger hsc_env) (hsc_tmpfs hsc_env) summary.ms_hspp_opts
     pure result
   modifySession (addDepsToHscEnv [hmi])
   pure (Just ModuleArtifacts {iface, bytecode = homeMod_bytecode hm_linkable})

From 98a0baf8069420f8793eb771b3570af496d34405 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 20:54:16 +0200
Subject: [PATCH 15/27] Refactor make mode state into a separate module

Change-Id: Ib55e9ccc95cd5c0dcea794a45a4868c8b3c654ed
---
 internal/buck-worker-internal.cabal |   1 +
 internal/src/Internal/Cache.hs      | 173 +++-------------------------
 internal/src/Internal/Metadata.hs   |  12 +-
 internal/src/Internal/State/Make.hs | 156 +++++++++++++++++++++++++
 4 files changed, 182 insertions(+), 160 deletions(-)
 create mode 100644 internal/src/Internal/State/Make.hs

diff --git a/internal/buck-worker-internal.cabal b/internal/buck-worker-internal.cabal
index 90691653..6ae1ee63 100644
--- a/internal/buck-worker-internal.cabal
+++ b/internal/buck-worker-internal.cabal
@@ -32,6 +32,7 @@ library
                       Internal.MakeFile.JSON
                       Internal.Metadata
                       Internal.Session
+                      Internal.State.Make
 
     if (impl(ghc >= 9.9))
         build-depends: base ^>=4.20
diff --git a/internal/src/Internal/Cache.hs b/internal/src/Internal/Cache.hs
index 798ff213..e5ca2f05 100644
--- a/internal/src/Internal/Cache.hs
+++ b/internal/src/Internal/Cache.hs
@@ -27,15 +27,7 @@ import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats)
 import GHC.Types.Name.Cache (NameCache (..), OrigNameCache)
 import GHC.Types.Unique.DFM (plusUDFM)
 import GHC.Types.Unique.FM (UniqFM, minusUFM, nonDetEltsUFM, sizeUFM)
-import GHC.Unit.Env (
-  HomeUnitEnv (..),
-  HomeUnitGraph,
-  UnitEnv (..),
-  unitEnv_insert,
-  unitEnv_lookup,
-  unitEnv_new,
-  unitEnv_union,
-  )
+import GHC.Unit.Env (UnitEnv (..), unitEnv_new)
 import GHC.Unit.External (ExternalUnitCache (..), initExternalUnitCache)
 import GHC.Unit.Finder (InstalledFindResult (..))
 import GHC.Unit.Finder.Types (FinderCache (..))
@@ -44,6 +36,8 @@ import GHC.Unit.Module.Graph (ModuleGraph)
 import qualified GHC.Utils.Outputable as Outputable
 import GHC.Utils.Outputable (SDoc, comma, doublePrec, fsep, hang, nest, punctuate, text, vcat, ($$), (<+>))
 import Internal.Log (Log, logd)
+import qualified Internal.State.Make as Make
+import Internal.State.Make (MakeState (..))
 import System.Environment (lookupEnv)
 import Types.Args (TargetId (..))
 
@@ -63,16 +57,6 @@ import GHC.Unit.Finder (initFinderCache)
 
 #endif
 
-#if defined(MWB)
-
-import GHC.Unit.Module.Graph (ModuleGraphNode (..), mgModSummaries', mkModuleGraph, mkNodeKey)
-
-#else
-
-import GHC.Unit.Module.Graph (unionMG)
-
-#endif
-
 data ModuleArtifacts =
   ModuleArtifacts {
     iface :: ModIface,
@@ -233,29 +217,6 @@ finderEnv FinderState {cache = FinderCache {fcModuleCache}} =
 
 #endif
 
--- | Data extracted from 'HscEnv' for the purpose of persisting it across sessions.
---
--- While many parts of the session are either contained in mutable variables or trivially reinitialized, some components
--- must be handled explicitly: The module graph and home unit graph are pure fields that need to be shared, and the
--- interpreter state for TH execution is only initialized when the flags are parsed.
-data MakeState =
-  MakeState {
-    -- | The module graph for a specific unit is computed in its metadata step, after which it's extracted and merged
-    -- into the existing graph.
-    moduleGraph :: ModuleGraph,
-
-    -- | The unit environment for a specific unit is inserted into the shared home unit graph at the beginning of the
-    -- metadata step, constructed from the dependency specifications provided by Buck.
-    -- After compilation of a module, its 'HomeUnitInfo' is inserted into the home package table contained in its unit's
-    -- unit environment.
-    hug :: HomeUnitGraph,
-
-    -- | While the interpreter state contains a mutable variable that would be shared across sessions, it isn't
-    -- initialized properly until the first module compilation's flags have been parsed, so we store it in the shared
-    -- state for consistency.
-    interp :: Maybe Interp
-  }
-
 data CacheFeatures =
   CacheFeatures {
     enable :: Bool,
@@ -405,10 +366,15 @@ restoreLoaderState cached session =
       }
 
     (linker_env, linkerStats) = restoreLinkerEnv cached.linker_env session.linker_env
+-- | Update the 'MakeState' field in the 'Cache'.
+updateMakeState :: (MakeState -> MakeState) -> Cache -> Cache
+updateMakeState f cache = cache {make = f cache.make}
 
 modifyStats :: Target -> (CacheStats -> CacheStats) -> Cache -> Cache
 modifyStats target f cache =
   cache {stats = Map.alter (Just . f . fromMaybe emptyStats) target cache.stats}
+updateMakeStateVar :: MVar Cache -> (MakeState -> MakeState) -> IO ()
+updateMakeStateVar var f = modifyMVar_ var (pure . updateMakeState f)
 
 pushStats :: Bool -> Target -> Maybe LoaderStats -> SymbolsStats -> NamesStats -> Cache -> Cache
 pushStats restoring target (Just new) symbols names =
@@ -655,13 +621,6 @@ withHscState HscEnv {hsc_interp, hsc_NC = NameCache {nsNames}} use =
     use nsNames loader_state symbolCacheVar
 #endif
 
-mergeHugs ::
-  HomeUnitEnv ->
-  HomeUnitEnv ->
-  HomeUnitEnv
-mergeHugs old new =
-  new {homeUnitEnv_hpt = plusUDFM old.homeUnitEnv_hpt new.homeUnitEnv_hpt}
-
 -- | Restore cache parts that depend on the 'Target'.
 setTarget :: MVar Cache -> Cache -> Target -> HscEnv -> IO HscEnv
 setTarget cacheVar cache target hsc_env = do
@@ -690,30 +649,11 @@ restoreOneshotCache cache hsc_env = do
   where
     restoreCachedEps e = e {hsc_unit_env = e.hsc_unit_env {ue_eps = cache.eps}}
 
--- | Merge the given module graph into the cached graph, or initialize it it doesn't exist yet.
+-- | Merge the given module graph into the cached graph.
 -- This is used by the make mode worker after the metadata step has computed the module graph.
 updateModuleGraph :: MVar Cache -> ModuleGraph -> IO ()
 updateModuleGraph cacheVar new =
-  modifyMVar_ cacheVar \ cache -> do
-#if defined(MWB)
-    let !merged = merge cache.make.moduleGraph
-    pure cache {make = cache.make {moduleGraph = merged}}
-  where
-    merge old =
-      mkModuleGraph (Map.elems (Map.unionWith mergeNodes oldMap newMap))
-      where
-        mergeNodes = \cases
-          (ModuleNode oldDeps _) (ModuleNode newDeps summ) -> ModuleNode (mergeDeps oldDeps newDeps) summ
-          _ newNode -> newNode
-
-        mergeDeps oldDeps newDeps = Set.toList (Set.fromList oldDeps <> Set.fromList newDeps)
-
-        oldMap = Map.fromList $ [(mkNodeKey n, n) | n <- mgModSummaries' old]
-
-        newMap = Map.fromList $ [(mkNodeKey n, n) | n <- mgModSummaries' new]
-#else
-    pure cache {make = cache.make {moduleGraph = unionMG cache.make.moduleGraph new}}
-#endif
+  updateMakeStateVar cacheVar (Make.storeModuleGraph new)
 
 prepareCache :: MVar Cache -> Target -> HscEnv -> Cache -> IO (Cache, (HscEnv, Bool))
 prepareCache cacheVar target hsc_env0 cache0 = do
@@ -735,23 +675,6 @@ prepareCache cacheVar target hsc_env0 cache0 = do
   let (hsc_env1, cache1) = fromMaybe (hsc_env0, cache0 {features = cache0.features {loader = False}}) result
   pure (cache1, (hsc_env1, cache1.features.enable))
 
-
-storeHug :: HscEnv -> Cache -> IO Cache
-storeHug hsc_env cache = do
-  let !new = hsc_env.hsc_unit_env.ue_home_unit_graph
-      !hug = unitEnv_union mergeHugs cache.make.hug new
-  pure cache {make = cache.make {hug}}
-
--- | Extract the unit env of the currently active unit and store it in the cache.
--- This is used by the make mode worker after the metadata step has initialized the new unit.
-insertUnitEnv :: HscEnv -> Cache -> Cache
-insertUnitEnv hsc_env cache =
-  cache {make = cache.make {hug = update cache.make.hug}}
-  where
-    ue = unitEnv_lookup current hsc_env.hsc_unit_env.ue_home_unit_graph
-    current = hsc_env.hsc_unit_env.ue_current_unit
-    update = unitEnv_insert current ue
-
 finalizeCache ::
   MVar Log ->
   -- | A description of the current worker process.
@@ -802,70 +725,6 @@ withCache logVar workerId cacheVar target prog = do
       withSession \ hsc_env ->
         liftIO (modifyMVar_ cacheVar (finalizeCache logVar workerId hsc_env target art))
 
-------------------------------------------------------------------------------------------------------------------------
-
-logMemStats :: String -> MVar Log -> IO ()
-logMemStats step logVar = do
-  s <- liftIO getRTSStats
-  let logMem desc value = logd logVar (text (desc ++ ":") <+> doublePrec 2 (fromIntegral value / 1_000_000) <+> text "MB")
-  logd logVar (text ("-------------- " ++ step))
-  logMem "Mem in use" s.gc.gcdetails_mem_in_use_bytes
-  logMem "Max mem in use" s.max_mem_in_use_bytes
-  logMem "Max live bytes" s.max_live_bytes
-
--- | Restore the shared state used by both @computeMetadata@ and @compileHpt@ from the cache.
--- See 'loadCacheMakeCompile' for details.
-loadCacheMake ::
-  MVar Log ->
-  HscEnv ->
-  Cache ->
-  IO HscEnv
-loadCacheMake logVar hsc_env cache = do
-  logMemStats "load cache" logVar
-  pure (restoreHug (restoreModuleGraph hsc_env))
-  where
-    restoreModuleGraph e = e {hsc_mod_graph = cache.make.moduleGraph}
-
-    restoreHug e = e {hsc_unit_env = e.hsc_unit_env {ue_home_unit_graph = cache.make.hug}}
-
--- | Restore the shared state used by @compileHpt@ from the cache, consisting of the module graph, the HPT, and the
--- loader state and symbol cache that's contained in 'Interp'.
--- The module graph is only modified by @computeMetadata@, so it will not be written back to the cache after
--- compilation.
---
--- Managing 'Interp' is a bit difficult: The field 'hsc_interp' isn't initialized with everything else in 'newHscEnv',
--- but only after parsing the command line arguments in 'setTopSessionDynFlags', since it needs to know the Ways of the
--- session if an external interpreter is used.
--- Therefore we grab the 'Interp' from the session when the cached value is absent, which amounts to the first
--- compilation session of the build.
--- When the cached value is present, on the other hand, we instead restore it into the session, making all subsequent
--- sessions share the first one's 'Interp'.
--- Both fields of 'Interp' are 'MVar's, so the state is shared immediately and concurrently.
-loadCacheMakeCompile ::
-  MVar Log ->
-  HscEnv ->
-  Cache ->
-  IO (Cache, HscEnv)
-loadCacheMakeCompile logVar hsc_env0 cache = do
-  ensureInterp <$> loadCacheMake logVar hsc_env0 cache
-  where
-    ensureInterp = maybe storeInterp restoreInterp cache.make.interp
-
-    storeInterp hsc_env = (cache {make = cache.make {interp = hsc_env.hsc_interp}}, hsc_env)
-
-    restoreInterp interp hsc_env = (cache, hsc_env {hsc_interp = Just interp})
-
--- | Store the changes made to the HUG by @compileHpt@ in the cache, which usually consists of adding a single
--- 'HomeModInfo'.
-storeCacheMake ::
-  MVar Log ->
-  HscEnv ->
-  Cache ->
-  IO Cache
-storeCacheMake logVar hsc_env cache = do
-  logMemStats "store cache" logVar
-  storeHug hsc_env cache
-
 -- | This reduced version of 'withCache' is tailored specifically to make mode, only restoring the HUG and module graph
 -- from the cache, since those are the only two components modified by the worker that aren't already shared by the base
 -- session.
@@ -881,6 +740,12 @@ withCacheMake logVar cacheVar prog = do
   modifySessionM restore
   prog <* withSession store
   where
-    restore hsc_env = liftIO (modifyMVar cacheVar (loadCacheMakeCompile logVar hsc_env))
-
-    store hsc_env = liftIO (modifyMVar_ cacheVar (storeCacheMake logVar hsc_env))
+    restore hsc_env =
+      liftIO $ modifyMVar cacheVar \ cache -> do
+        (make, hsc_env1) <- Make.loadStateCompile logVar hsc_env cache.make
+        pure (cache {make}, hsc_env1)
+
+    store hsc_env =
+      liftIO $ modifyMVar_ cacheVar \ cache -> do
+        make <- Make.storeState logVar hsc_env cache.make
+        pure cache {make}
diff --git a/internal/src/Internal/Metadata.hs b/internal/src/Internal/Metadata.hs
index f9312c7a..e88653d8 100644
--- a/internal/src/Internal/Metadata.hs
+++ b/internal/src/Internal/Metadata.hs
@@ -1,6 +1,6 @@
 module Internal.Metadata where
 
-import Control.Concurrent (modifyMVar_, readMVar)
+import Control.Concurrent (readMVar)
 import Control.Monad.IO.Class (liftIO)
 import Data.Maybe (isJust)
 import GHC (DynFlags (..), Ghc, GhcMode (..), Logger, ModuleGraph)
@@ -12,9 +12,10 @@ import GHC.Runtime.Loader (initializeSessionPlugins)
 import GHC.Unit (HomeUnit, UnitDatabase, UnitId, UnitState, initUnits)
 import GHC.Unit.Env (HomeUnitEnv (..), UnitEnv (..), unitEnv_insert, unitEnv_keys, updateHug)
 import GHC.Unit.Home.ModInfo (emptyHomePackageTable)
-import Internal.Cache (insertUnitEnv, loadCacheMake, logMemStats, updateModuleGraph)
+import Internal.Cache (Cache (..), updateMakeStateVar)
 import Internal.MakeFile (doMkDependHS)
 import Internal.Session (Env (..), runSession, withDynFlags)
+import Internal.State.Make (insertUnitEnv, loadState, logMemStats, storeModuleGraph)
 
 -- | 'doMkDependHS' needs this to be enabled.
 metadataTempSession :: HscEnv -> HscEnv
@@ -74,14 +75,14 @@ addHomeUnit dflags = do
 prepareMetadataSession :: Env -> DynFlags -> Ghc ()
 prepareMetadataSession env dflags = do
   cache <- liftIO $ readMVar env.cache
-  modifySessionM \ hsc_env -> liftIO (loadCacheMake env.log hsc_env cache)
+  modifySessionM \ hsc_env -> liftIO (loadState env.log hsc_env cache.make)
   unit <- addHomeUnit dflags
   setActiveUnit unit
   storeNewUnit
   where
     setActiveUnit unit = modifySession (hscUpdateLoggerFlags . hscSetActiveUnitId unit)
 
-    storeNewUnit = withSession \ hsc_env -> liftIO $ modifyMVar_ env.cache (pure . insertUnitEnv hsc_env)
+    storeNewUnit = withSession \ hsc_env -> liftIO $ updateMakeStateVar env.cache (insertUnitEnv hsc_env)
 
 -- | Run 'doMkDependHS' to write the metadata JSON file and exfiltrate the module graph.
 -- We need to use a temporary session because 'doMkDependHS' uses some custom settings that we don't want to leak,
@@ -107,7 +108,6 @@ computeMetadata env = do
   res <- fmap isJust $ runSession True env $ withDynFlags env \ dflags srcs -> do
     prepareMetadataSession env dflags
     module_graph <- writeMetadata (fst <$> srcs)
-    liftIO do
-      updateModuleGraph env.cache module_graph
+    liftIO $ updateMakeStateVar env.cache (storeModuleGraph module_graph)
     pure (Just ())
   res <$ logMemStats "after metadata" env.log
diff --git a/internal/src/Internal/State/Make.hs b/internal/src/Internal/State/Make.hs
new file mode 100644
index 00000000..f95ef3d1
--- /dev/null
+++ b/internal/src/Internal/State/Make.hs
@@ -0,0 +1,156 @@
+{-# LANGUAGE CPP #-}
+
+module Internal.State.Make where
+
+import Control.Concurrent.MVar (MVar)
+import Control.Monad.IO.Class (liftIO)
+import GHC.Driver.Env (HscEnv (..))
+import GHC.Runtime.Interpreter (Interp (..))
+import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats)
+import GHC.Types.Unique.DFM (plusUDFM)
+import GHC.Unit.Env (HomeUnitEnv (..), HomeUnitGraph, UnitEnv (..), unitEnv_insert, unitEnv_lookup, unitEnv_union)
+import GHC.Unit.Module.Graph (ModuleGraph)
+import GHC.Utils.Outputable (doublePrec, text, (<+>))
+import Internal.Log (Log, logd)
+
+#if defined(MWB)
+
+import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
+import GHC.Unit.Module.Graph (ModuleGraphNode (..), mgModSummaries', mkModuleGraph, mkNodeKey)
+
+#else
+
+import GHC.Unit.Module.Graph (unionMG)
+
+#endif
+
+-- | Data extracted from 'HscEnv' for the purpose of persisting it across sessions.
+--
+-- While many parts of the session are either contained in mutable variables or trivially reinitialized, some components
+-- must be handled explicitly: The module graph and home unit graph are pure fields that need to be shared, and the
+-- interpreter state for TH execution is only initialized when the flags are parsed.
+data MakeState =
+  MakeState {
+    -- | The module graph for a specific unit is computed in its metadata step, after which it's extracted and merged
+    -- into the existing graph.
+    moduleGraph :: ModuleGraph,
+
+    -- | The unit environment for a specific unit is inserted into the shared home unit graph at the beginning of the
+    -- metadata step, constructed from the dependency specifications provided by Buck.
+    -- After compilation of a module, its 'HomeUnitInfo' is inserted into the home package table contained in its unit's
+    -- unit environment.
+    hug :: HomeUnitGraph,
+
+    -- | While the interpreter state contains a mutable variable that would be shared across sessions, it isn't
+    -- initialized properly until the first module compilation's flags have been parsed, so we store it in the shared
+    -- state for consistency.
+    interp :: Maybe Interp
+  }
+
+logMemStats :: String -> MVar Log -> IO ()
+logMemStats step logVar = do
+  s <- liftIO getRTSStats
+  let logMem desc value = logd logVar (text (desc ++ ":") <+> doublePrec 2 (fromIntegral value / 1_000_000) <+> text "MB")
+  logd logVar (text ("-------------- " ++ step))
+  logMem "Mem in use" s.gc.gcdetails_mem_in_use_bytes
+  logMem "Max mem in use" s.max_mem_in_use_bytes
+  logMem "Max live bytes" s.max_live_bytes
+
+-- | Restore the shared state used by both @computeMetadata@ and @compileHpt@ from the cache.
+-- See 'loadCacheMakeCompile' for details.
+loadState ::
+  MVar Log ->
+  HscEnv ->
+  MakeState ->
+  IO HscEnv
+loadState logVar hsc_env state = do
+  logMemStats "load state" logVar
+  pure (restoreHug (restoreModuleGraph hsc_env))
+  where
+    restoreModuleGraph e = e {hsc_mod_graph = state.moduleGraph}
+
+    restoreHug e = e {hsc_unit_env = e.hsc_unit_env {ue_home_unit_graph = state.hug}}
+
+-- | Restore the shared state used by @compileHpt@ from the state, consisting of the module graph, the HPT, and the
+-- loader state and symbol cache that's contained in 'Interp'.
+-- The module graph is only modified by @computeMetadata@, so it will not be written back to the state after
+-- compilation.
+--
+-- Managing 'Interp' is a bit difficult: The field 'hsc_interp' isn't initialized with everything else in 'newHscEnv',
+-- but only after parsing the command line arguments in 'setTopSessionDynFlags', since it needs to know the Ways of the
+-- session if an external interpreter is used.
+-- Therefore we grab the 'Interp' from the session when the cached value is absent, which amounts to the first
+-- compilation session of the build.
+-- When the cached value is present, on the other hand, we instead restore it into the session, making all subsequent
+-- sessions share the first one's 'Interp'.
+-- Both fields of 'Interp' are 'MVar's, so the state is shared immediately and concurrently.
+loadStateCompile ::
+  MVar Log ->
+  HscEnv ->
+  MakeState ->
+  IO (MakeState, HscEnv)
+loadStateCompile logVar hsc_env0 state = do
+  ensureInterp <$> loadState logVar hsc_env0 state
+  where
+    ensureInterp = maybe storeInterp restoreInterp state.interp
+
+    storeInterp hsc_env = (state {interp = hsc_env.hsc_interp}, hsc_env)
+
+    restoreInterp interp hsc_env = (state, hsc_env {hsc_interp = Just interp})
+
+-- | Merge the given module graph into the cached graph.
+-- In more recent versions of GHC, the function for merging graphs is not exposed anymore.
+-- There was also some issue with node duplication, which is why this function is so convoluted.
+storeModuleGraph :: ModuleGraph -> MakeState -> MakeState
+storeModuleGraph new state =
+#if defined(MWB)
+  state {moduleGraph = merged}
+  where
+    !merged = merge state.moduleGraph
+
+    merge old =
+      mkModuleGraph (Map.elems (Map.unionWith mergeNodes oldMap newMap))
+      where
+        mergeNodes = \cases
+          (ModuleNode oldDeps _) (ModuleNode newDeps summ) -> ModuleNode (mergeDeps oldDeps newDeps) summ
+          _ newNode -> newNode
+
+        mergeDeps oldDeps newDeps = Set.toList (Set.fromList oldDeps <> Set.fromList newDeps)
+
+        oldMap = Map.fromList $ [(mkNodeKey n, n) | n <- mgModSummaries' old]
+
+        newMap = Map.fromList $ [(mkNodeKey n, n) | n <- mgModSummaries' new]
+#else
+    state {moduleGraph = unionMG state.moduleGraph new}
+#endif
+
+-- | Extract the unit env of the currently active unit and store it in the cache.
+-- This is used by the make mode worker after the metadata step has initialized the new unit.
+insertUnitEnv :: HscEnv -> MakeState -> MakeState
+insertUnitEnv hsc_env state =
+  state {hug = update state.hug}
+  where
+    ue = unitEnv_lookup current hsc_env.hsc_unit_env.ue_home_unit_graph
+    current = hsc_env.hsc_unit_env.ue_current_unit
+    update = unitEnv_insert current ue
+
+mergeHugs ::
+  HomeUnitEnv ->
+  HomeUnitEnv ->
+  HomeUnitEnv
+mergeHugs old new =
+  new {homeUnitEnv_hpt = plusUDFM old.homeUnitEnv_hpt new.homeUnitEnv_hpt}
+
+-- | Store the changes made to the HUG by @compileHpt@ in the state, which usually consists of adding a single
+-- 'HomeModInfo'.
+storeState ::
+  MVar Log ->
+  HscEnv ->
+  MakeState ->
+  IO MakeState
+storeState logVar hsc_env state = do
+  logMemStats "store make state" logVar
+  let !new = hsc_env.hsc_unit_env.ue_home_unit_graph
+      !hug = unitEnv_union mergeHugs state.hug new
+  pure state {hug}

From 2b90b00d2b3a50594dd1956472ad5784ab5827a4 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 21:54:37 +0200
Subject: [PATCH 16/27] Refactor statistics collection logic into a separate
 module

Change-Id: I93e2175ab6ef8ee2f8d4dea927d726f505a922bf
---
 ghc-worker/lib/GhcWorker/GhcHandler.hs      |   3 +-
 ghc-worker/lib/GhcWorker/Instrumentation.hs |   7 +-
 ghc-worker/test/CompileHptTest.hs           |   5 +-
 internal/buck-worker-internal.cabal         |   1 +
 internal/src/Internal/Cache.hs              | 249 +++-----------------
 internal/src/Internal/Compile.hs            |   3 +-
 internal/src/Internal/CompileHpt.hs         |   3 +-
 internal/src/Internal/Metadata.hs           |   3 +-
 internal/src/Internal/Session.hs            |   2 +-
 internal/src/Internal/State/Make.hs         |  15 +-
 internal/src/Internal/State/Stats.hs        | 220 +++++++++++++++++
 types/buck-worker-types.cabal               |   2 +
 types/src/Types/State.hs                    |  18 ++
 13 files changed, 287 insertions(+), 244 deletions(-)
 create mode 100644 internal/src/Internal/State/Stats.hs
 create mode 100644 types/src/Types/State.hs

diff --git a/ghc-worker/lib/GhcWorker/GhcHandler.hs b/ghc-worker/lib/GhcWorker/GhcHandler.hs
index f5c02eb3..3ab30016 100644
--- a/ghc-worker/lib/GhcWorker/GhcHandler.hs
+++ b/ghc-worker/lib/GhcWorker/GhcHandler.hs
@@ -15,7 +15,7 @@ import GHC.Driver.Monad (modifySession)
 import GhcWorker.CompileResult (CompileResult (..), writeCloseOutput, writeResult)
 import GhcWorker.Instrumentation (Hooks (..), InstrumentedHandler (..))
 import Internal.AbiHash (AbiHash (..), showAbiHash)
-import Internal.Cache (Cache (..), ModuleArtifacts (..), Target (..))
+import Internal.Cache (Cache (..), ModuleArtifacts (..))
 import Internal.Compile (compileModuleWithDepsInEps)
 import Internal.CompileHpt (compileModuleWithDepsInHpt)
 import Internal.Log (LogName (..), dbg, logFlush, newLog)
@@ -27,6 +27,7 @@ import System.Posix.Process (exitImmediately)
 import Types.BuckArgs (BuckArgs, Mode (..), parseBuckArgs, toGhcArgs)
 import qualified Types.BuckArgs
 import Types.GhcHandler (WorkerMode (..))
+import Types.State (Target (Target))
 
 data LockState = LockStart | LockFreeze Int | LockThaw Int | LockEnd
   deriving stock (Eq, Show)
diff --git a/ghc-worker/lib/GhcWorker/Instrumentation.hs b/ghc-worker/lib/GhcWorker/Instrumentation.hs
index 82c4d7a9..935edc8f 100644
--- a/ghc-worker/lib/GhcWorker/Instrumentation.hs
+++ b/ghc-worker/lib/GhcWorker/Instrumentation.hs
@@ -8,12 +8,13 @@ import Data.Foldable (traverse_)
 import Data.Int (Int32)
 import Data.Text qualified as Text
 import GhcWorker.Grpc (mkStats)
-import Internal.Cache (Cache, Target (..))
+import Internal.Cache (Cache)
 import Internal.Log (dbg)
 import Network.GRPC.Common.Protobuf (Proto, defMessage, (&), (.~))
 import Prelude hiding (log)
 import qualified Proto.Instrument as Instr
 import Proto.Instrument_Fields qualified as Instr
+import Types.State (Target (..))
 
 -- | Rudimentary dummy state for instrumentation, counting concurrently compiling sessions.
 data WorkerStatus =
@@ -111,12 +112,12 @@ withInstrumentation instrChan status cacheVar handler =
         writeChan instrChan $
           defMessage &
             Instr.compileStart .~
-              messageCompileStart target.get
+              messageCompileStart target.path
 
     -- Note: This is WIP.
     compileFinish =
       traverse_ \ (target, output, exitCode) -> do
-        let tgt = maybe "" (.get) target
+        let tgt = maybe "" (.path) target
         writeChan instrChan $
           defMessage &
             Instr.compileEnd .~
diff --git a/ghc-worker/test/CompileHptTest.hs b/ghc-worker/test/CompileHptTest.hs
index 4a3e2f56..00d695bd 100644
--- a/ghc-worker/test/CompileHptTest.hs
+++ b/ghc-worker/test/CompileHptTest.hs
@@ -19,16 +19,17 @@ import GHC.Driver.Session (parseDynamicFlagsCmdLine)
 import GHC.Utils.Monad (MonadIO (..))
 import GHC.Utils.Outputable (ppr, showPprUnsafe, text, (<+>))
 import GHC.Utils.Panic (throwGhcExceptionIO)
-import Internal.Cache (Target (..), logMemStats)
 import Internal.CompileHpt (compileModuleWithDepsInHpt)
 import Internal.Log (dbg, dbgp, dbgs, newLog)
 import Internal.Metadata (computeMetadata)
 import Internal.Session (Env (..), withGhcMhu)
+import Internal.State.Stats (logMemStats)
 import Prelude hiding (log)
 import System.Directory (createDirectoryIfMissing, listDirectory, removeDirectoryRecursive)
 import System.FilePath (dropExtension, takeBaseName, takeExtension, takeFileName, (</>))
 import TestSetup (Conf (..), Module (..), ModuleSpec (..), Unit (..), UnitSpec (..), withProject)
 import Types.Args (Args (..))
+import Types.State (Target (..))
 
 -- | Parse command line flags, used to create unit-specific @DynFlags@.
 unitFlags :: [String] -> HscEnv -> Ghc DynFlags
@@ -76,7 +77,7 @@ stepCompile Conf {cache, tmp, args0} Module {unit, src} = do
   liftIO $ createDirectoryIfMissing False sessionTmpDir
   result <- liftIO $ withGhcMhu env \ _ target -> do
     dbg ""
-    dbg (">>> compiling " ++ takeFileName target.get)
+    dbg (">>> compiling " ++ takeFileName target.path)
     modifySession $ hscUpdateFlags \ d -> d {ghcMode = CompManager}
     compileModuleWithDepsInHpt target
   when (isNothing result) do
diff --git a/internal/buck-worker-internal.cabal b/internal/buck-worker-internal.cabal
index 6ae1ee63..67075898 100644
--- a/internal/buck-worker-internal.cabal
+++ b/internal/buck-worker-internal.cabal
@@ -33,6 +33,7 @@ library
                       Internal.Metadata
                       Internal.Session
                       Internal.State.Make
+                      Internal.State.Stats
 
     if (impl(ghc >= 9.9))
         build-depends: base ^>=4.20
diff --git a/internal/src/Internal/Cache.hs b/internal/src/Internal/Cache.hs
index e5ca2f05..38fc08e2 100644
--- a/internal/src/Internal/Cache.hs
+++ b/internal/src/Internal/Cache.hs
@@ -8,38 +8,43 @@ import Control.Monad.IO.Class (MonadIO, liftIO)
 import Data.Bifunctor (first)
 import Data.Coerce (coerce)
 import Data.IORef (readIORef)
-import Data.List (sortBy)
 import qualified Data.Map.Strict as Map
-import Data.Map.Strict (Map, (!?))
+import Data.Map.Strict (Map)
 import Data.Maybe (fromMaybe)
-import Data.Ord (comparing)
-import qualified Data.Set as Set
-import Data.Set (Set, (\\))
+import Data.Set (Set)
 import Data.Traversable (for)
-import GHC (Ghc, ModIface, ModuleName, emptyMG, mi_module, moduleName, moduleNameString, setSession)
-import GHC.Data.FastString (FastString)
+import GHC (Ghc, ModIface, emptyMG, mi_module, moduleName, moduleNameString, setSession)
 import GHC.Driver.Env (HscEnv (..))
 import GHC.Driver.Monad (modifySessionM, withSession)
 import GHC.Linker.Types (Linkable, LinkerEnv (..), Loader (..), LoaderState (..))
-import GHC.Ptr (Ptr)
 import GHC.Runtime.Interpreter (Interp (..))
-import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats)
 import GHC.Types.Name.Cache (NameCache (..), OrigNameCache)
 import GHC.Types.Unique.DFM (plusUDFM)
-import GHC.Types.Unique.FM (UniqFM, minusUFM, nonDetEltsUFM, sizeUFM)
 import GHC.Unit.Env (UnitEnv (..), unitEnv_new)
 import GHC.Unit.External (ExternalUnitCache (..), initExternalUnitCache)
 import GHC.Unit.Finder (InstalledFindResult (..))
 import GHC.Unit.Finder.Types (FinderCache (..))
-import GHC.Unit.Module.Env (InstalledModuleEnv, emptyModuleEnv, moduleEnvKeys, plusModuleEnv)
+import GHC.Unit.Module.Env (InstalledModuleEnv, emptyModuleEnv, plusModuleEnv)
 import GHC.Unit.Module.Graph (ModuleGraph)
-import qualified GHC.Utils.Outputable as Outputable
-import GHC.Utils.Outputable (SDoc, comma, doublePrec, fsep, hang, nest, punctuate, text, vcat, ($$), (<+>))
-import Internal.Log (Log, logd)
+import Internal.Log (Log)
 import qualified Internal.State.Make as Make
 import Internal.State.Make (MakeState (..))
+import qualified Internal.State.Stats as Stats
+import Internal.State.Stats (
+  CacheStats (..),
+  LinkerStats (..),
+  LoaderStats (..),
+  StatsUpdate (..),
+  SymbolsStats (..),
+  basicLinkerStats,
+  basicLoaderStats,
+  basicSymbolsStats,
+  emptyLinkerStats,
+  emptyStats,
+  )
 import System.Environment (lookupEnv)
 import Types.Args (TargetId (..))
+import Types.State (SymbolCache (..), SymbolMap, Target)
 
 #if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
 
@@ -50,6 +55,7 @@ import GHC.IORef (atomicModifyIORef')
 import GHC.Unit (InstalledModule, emptyInstalledModuleEnv, extendInstalledModuleEnv, lookupInstalledModuleEnv)
 import GHC.Unit.Finder (FinderCache (..), InstalledFindResult (..))
 import GHC.Utils.Panic (panic)
+import Internal.State.Stats (FinderStats (..))
 
 #else
 
@@ -67,111 +73,12 @@ instance Show ModuleArtifacts where
   show ModuleArtifacts {iface} =
     "ModuleArtifacts { iface = " ++ moduleNameString (moduleName (mi_module iface)) ++ " }"
 
-type SymbolMap = UniqFM FastString (Ptr ())
-
-newtype SymbolCache =
-  SymbolCache { get :: SymbolMap }
-  deriving newtype (Semigroup, Monoid)
-
-data LinkerStats =
-  LinkerStats {
-    newClosures :: Int,
-    newItables :: Int
-  }
-  deriving stock (Eq, Show)
-
-emptyLinkerStats :: LinkerStats
-emptyLinkerStats =
-  LinkerStats {
-    newClosures = 0,
-    newItables = 0
-  }
-
-data LoaderStats =
-  LoaderStats {
-    newBcos :: [String],
-    sameBcos :: Int,
-    linker :: LinkerStats
-  }
-  deriving stock (Eq, Show)
-
-emptyLoaderStats :: LoaderStats
-emptyLoaderStats =
-  LoaderStats {
-    newBcos = mempty,
-    sameBcos = 0,
-    linker = emptyLinkerStats
-  }
-
-data SymbolsStats =
-  SymbolsStats {
-    new :: Int
-  }
-  deriving stock (Eq, Show)
-
-data NamesStats =
-  NamesStats {
-    new :: Int
-  }
-  deriving stock (Eq, Show)
-
-data StatsUpdate =
-  StatsUpdate {
-    loaderStats :: LoaderStats,
-    symbols :: SymbolsStats,
-    names :: NamesStats
-  }
-  deriving stock (Eq, Show)
-
-emptyStatsUpdate :: StatsUpdate
-emptyStatsUpdate =
-  StatsUpdate {
-    loaderStats = emptyLoaderStats,
-    symbols = SymbolsStats {new = 0},
-    names = NamesStats {new = 0}
-  }
-
-data FinderStats =
-  FinderStats {
-    hits :: Map ModuleName Int,
-    misses :: Map ModuleName Int
-  }
-  deriving stock (Eq, Show)
-
-emptyFinderStats :: FinderStats
-emptyFinderStats =
-  FinderStats {
-    hits = mempty,
-    misses = mempty
-  }
-
-data CacheStats =
-  CacheStats {
-    restore :: StatsUpdate,
-    update :: StatsUpdate,
-    finder :: FinderStats
-  }
-  deriving stock (Eq, Show)
-
-emptyStats :: CacheStats
-emptyStats =
-  CacheStats {
-    restore = emptyStatsUpdate,
-    update = emptyStatsUpdate,
-    finder = emptyFinderStats
-  }
-
 data InterpCache =
   InterpCache {
     loaderState :: LoaderState,
     symbols :: SymbolCache
   }
 
-newtype Target =
-  Target { get :: String }
-  deriving stock (Eq, Show)
-  deriving newtype (Ord)
-
 data BinPath =
   BinPath {
     initial :: Maybe String,
@@ -285,50 +192,6 @@ defaultOptions =
     extraGhcOptions = ""
   }
 
-basicLinkerStats :: LinkerEnv -> LinkerEnv -> LinkerStats
-basicLinkerStats base update =
-  LinkerStats {
-    newClosures = Set.size (updateClosures \\ baseClosures),
-    newItables = Set.size (updateItables \\ baseItables)
-  }
-  where
-    updateClosures = names update.closure_env
-    baseClosures = names base.closure_env
-    updateItables = names update.itbl_env
-    baseItables = names base.itbl_env
-
-    names = Set.fromList . fmap fst . nonDetEltsUFM
-
-basicLoaderStats ::
-  LoaderState ->
-  LoaderState ->
-  LinkerStats ->
-  LoaderStats
-basicLoaderStats base update linker =
-  LoaderStats {
-    newBcos = modStr <$> Set.toList (updateBcos \\ baseBcos),
-    sameBcos = Set.size bcoSame,
-    linker
-  }
-  where
-    modStr = moduleNameString . moduleName
-    bcoSame = Set.intersection updateBcos baseBcos
-    updateBcos = Set.fromList (moduleEnvKeys update.bcos_loaded)
-    baseBcos = Set.fromList (moduleEnvKeys base.bcos_loaded)
-
-basicSymbolsStats :: SymbolCache -> SymbolCache -> SymbolsStats
-basicSymbolsStats base update =
-  SymbolsStats {
-    new = sizeUFM (minusUFM update.get base.get)
-  }
-
-basicNamesStats :: OrigNameCache -> OrigNameCache -> NamesStats
-basicNamesStats _ _ =
-  NamesStats {
-    new = 0
-  }
-  where
-
 restoreLinkerEnv :: LinkerEnv -> LinkerEnv -> (LinkerEnv, LinkerStats)
 restoreLinkerEnv cached session =
   (merged, basicLinkerStats session cached)
@@ -376,13 +239,13 @@ modifyStats target f cache =
 updateMakeStateVar :: MVar Cache -> (MakeState -> MakeState) -> IO ()
 updateMakeStateVar var f = modifyMVar_ var (pure . updateMakeState f)
 
-pushStats :: Bool -> Target -> Maybe LoaderStats -> SymbolsStats -> NamesStats -> Cache -> Cache
-pushStats restoring target (Just new) symbols names =
+pushStats :: Bool -> Target -> Maybe LoaderStats -> SymbolsStats -> Cache -> Cache
+pushStats restoring target (Just new) symbols =
   modifyStats target add
   where
-    add old | restoring = old {restore = old.restore {loaderStats = new, symbols, names}}
-            | otherwise = old {update = old.update {loaderStats = new, symbols, names}}
-pushStats _ _ _ _ _ =
+    add old | restoring = old {restore = old.restore {loaderStats = new, symbols}}
+            | otherwise = old {update = old.update {loaderStats = new, symbols}}
+pushStats _ _ _ _ =
   id
 
 restoreCache ::
@@ -403,8 +266,7 @@ restoreCache target initialLoaderState initialSymbolCache initialNames cache
     let
       newSymbols = initialSymbolCache <> symbols
       symbolsStats = basicSymbolsStats initialSymbolCache symbols
-      namesStats = basicNamesStats initialNames cache.names
-      newCache = pushStats True target loaderStats symbolsStats namesStats cache
+      newCache = pushStats True target loaderStats symbolsStats cache
       -- this overwrites entire modules, since OrigNameCache is a three-level map.
       -- eventually we'll want to merge properly.
       names = plusModuleEnv initialNames cache.names
@@ -472,7 +334,7 @@ updateCache ::
   IO Cache
 updateCache target InterpCache {..} newLoaderState newSymbols newNames cache = do
   (updatedLs, stats) <- updateLoaderState loaderState newLoaderState
-  pure $ pushStats False target stats symbolsStats namesStats cache {
+  pure $ pushStats False target stats symbolsStats cache {
     interpCache = Just InterpCache {
       loaderState = updatedLs,
       symbols = symbols <> newSymbols
@@ -484,55 +346,6 @@ updateCache target InterpCache {..} newLoaderState newSymbols newNames cache = d
   }
   where
     symbolsStats = basicSymbolsStats symbols newSymbols
-    namesStats = basicNamesStats cache.names newNames
-
-moduleColumns :: Show a => Map ModuleName a -> SDoc
-moduleColumns m =
-  vcat [text n Outputable.<> text ":" $$ nest offset (text (show h)) | (n, h) <- kvs]
-  where
-    offset = length (fst (last kvs)) + 2
-    kvs = sortBy (comparing (length . fst)) (first moduleNameString <$> Map.toList m)
-
--- | Assemble log messages about cache statistics.
-statsMessages :: CacheStats -> SDoc
-statsMessages CacheStats {restore, update, finder} =
-  hang (text "Restore:") 2 restoreStats $$
-  hang (text "Update:") 2 updateStats $$
-  hang (text "Finder:") 2 finderStats
-  where
-      restoreStats =
-        text (show (length restore.loaderStats.newBcos)) <+> text "BCOs" $$
-        text (show restore.loaderStats.linker.newClosures) <+> text "closures" $$
-        text (show restore.symbols.new) <+> text "symbols" $$
-        text (show restore.loaderStats.sameBcos) <+> text "BCOs already in cache"
-
-      newBcos = text <$> update.loaderStats.newBcos
-
-      updateStats =
-        (if null newBcos then text "No new BCOs" else text "New BCOs:" <+> fsep (punctuate comma newBcos)) $$
-        text (show update.loaderStats.linker.newClosures) <+> text "new closures" $$
-        text (show update.symbols.new) <+> text "new symbols" $$
-        text (show update.loaderStats.sameBcos) <+> text "BCOs already in cache"
-
-      finderStats =
-        hang (text "Hits:") 2 (moduleColumns finder.hits) $$
-        hang (text "Misses:") 2 (moduleColumns finder.misses)
-
--- | Assemble report messages, consisting of:
---
--- - Cache statistics, if the feature is enabled
--- - Current RTS memory usage
-reportMessages :: Target -> Cache -> Double -> SDoc
-reportMessages target Cache {stats, features} memory =
-  statsPart $$
-  memoryPart
-  where
-    statsPart =
-      if features.enable
-      then maybe (text "Cache unused for this module.") statsMessages (stats !? target)
-      else text "Cache disabled."
-
-    memoryPart = text "Memory:" <+> doublePrec 2 memory <+> text "MB"
 
 -- | Log a report for a completed compilation, using 'reportMessages' to assemble the content.
 report ::
@@ -544,13 +357,7 @@ report ::
   Cache ->
   m ()
 report logVar workerId target cache = do
-  s <- liftIO getRTSStats
-  let memory = fromIntegral (s.gc.gcdetails_mem_in_use_bytes) / 1000000
-  logd logVar (hang header 2 (reportMessages target cache memory))
-  where
-    header = text target.get Outputable.<> maybe (text "") workerDesc workerId Outputable.<> text ":"
-
-    workerDesc wid = text (" (" ++ wid.string ++ ")")
+  Stats.report logVar workerId target (if cache.features.enable then Just cache.stats else Nothing)
 
 #if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
 
diff --git a/internal/src/Internal/Compile.hs b/internal/src/Internal/Compile.hs
index df25373e..d43bc260 100644
--- a/internal/src/Internal/Compile.hs
+++ b/internal/src/Internal/Compile.hs
@@ -28,8 +28,9 @@ import GHC.Types.SourceFile (HscSource)
 import GHC.Unit.Home.ModInfo (HomeModLinkable (..))
 import GHC.Utils.Monad (MonadIO (..), unlessM)
 import GHC.Utils.Panic (panic, throwGhcExceptionIO)
-import Internal.Cache (ModuleArtifacts (..), Target (..))
+import Internal.Cache (ModuleArtifacts (..))
 import System.Directory (doesFileExist)
+import Types.State (Target (Target))
 
 type P m = TPipelineClass TPhase m
 
diff --git a/internal/src/Internal/CompileHpt.hs b/internal/src/Internal/CompileHpt.hs
index e00d3a02..5ce58aad 100644
--- a/internal/src/Internal/CompileHpt.hs
+++ b/internal/src/Internal/CompileHpt.hs
@@ -13,8 +13,9 @@ import GHC.Unit.Env (addHomeModInfoToHug, ue_unsafeHomeUnit)
 import GHC.Unit.Home.ModInfo (HomeModInfo (..), HomeModLinkable (..))
 import GHC.Utils.Monad (MonadIO (..))
 import GHC.Utils.TmpFs (TmpFs, cleanCurrentModuleTempFiles, keepCurrentModuleTempFiles)
-import Internal.Cache (ModuleArtifacts (..), Target (..))
+import Internal.Cache (ModuleArtifacts (..))
 import Internal.Error (eitherMessages)
+import Types.State (Target (Target))
 
 -- | Insert a compilation result into the current unit's home package table, as it is done by upsweep.
 addDepsToHscEnv :: [HomeModInfo] -> HscEnv -> HscEnv
diff --git a/internal/src/Internal/Metadata.hs b/internal/src/Internal/Metadata.hs
index e88653d8..8738042a 100644
--- a/internal/src/Internal/Metadata.hs
+++ b/internal/src/Internal/Metadata.hs
@@ -15,7 +15,8 @@ import GHC.Unit.Home.ModInfo (emptyHomePackageTable)
 import Internal.Cache (Cache (..), updateMakeStateVar)
 import Internal.MakeFile (doMkDependHS)
 import Internal.Session (Env (..), runSession, withDynFlags)
-import Internal.State.Make (insertUnitEnv, loadState, logMemStats, storeModuleGraph)
+import Internal.State.Make (insertUnitEnv, loadState, storeModuleGraph)
+import Internal.State.Stats (logMemStats)
 
 -- | 'doMkDependHS' needs this to be enabled.
 metadataTempSession :: HscEnv -> HscEnv
diff --git a/internal/src/Internal/Session.hs b/internal/src/Internal/Session.hs
index 0b4e4ac6..4c013c7a 100644
--- a/internal/src/Internal/Session.hs
+++ b/internal/src/Internal/Session.hs
@@ -45,7 +45,6 @@ import Internal.Cache (
   CacheFeatures (..),
   ModuleArtifacts,
   Options (..),
-  Target (..),
   withCache,
   withCacheMake,
   )
@@ -54,6 +53,7 @@ import Internal.Log (Log (..), logToState)
 import Prelude hiding (log)
 import System.Environment (setEnv)
 import Types.Args (Args (..))
+import Types.State (Target (Target))
 
 -- | Worker state.
 data Env =
diff --git a/internal/src/Internal/State/Make.hs b/internal/src/Internal/State/Make.hs
index f95ef3d1..b4d8572d 100644
--- a/internal/src/Internal/State/Make.hs
+++ b/internal/src/Internal/State/Make.hs
@@ -3,15 +3,13 @@
 module Internal.State.Make where
 
 import Control.Concurrent.MVar (MVar)
-import Control.Monad.IO.Class (liftIO)
 import GHC.Driver.Env (HscEnv (..))
 import GHC.Runtime.Interpreter (Interp (..))
-import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats)
 import GHC.Types.Unique.DFM (plusUDFM)
 import GHC.Unit.Env (HomeUnitEnv (..), HomeUnitGraph, UnitEnv (..), unitEnv_insert, unitEnv_lookup, unitEnv_union)
 import GHC.Unit.Module.Graph (ModuleGraph)
-import GHC.Utils.Outputable (doublePrec, text, (<+>))
-import Internal.Log (Log, logd)
+import Internal.Log (Log)
+import Internal.State.Stats (logMemStats)
 
 #if defined(MWB)
 
@@ -48,15 +46,6 @@ data MakeState =
     interp :: Maybe Interp
   }
 
-logMemStats :: String -> MVar Log -> IO ()
-logMemStats step logVar = do
-  s <- liftIO getRTSStats
-  let logMem desc value = logd logVar (text (desc ++ ":") <+> doublePrec 2 (fromIntegral value / 1_000_000) <+> text "MB")
-  logd logVar (text ("-------------- " ++ step))
-  logMem "Mem in use" s.gc.gcdetails_mem_in_use_bytes
-  logMem "Max mem in use" s.max_mem_in_use_bytes
-  logMem "Max live bytes" s.max_live_bytes
-
 -- | Restore the shared state used by both @computeMetadata@ and @compileHpt@ from the cache.
 -- See 'loadCacheMakeCompile' for details.
 loadState ::
diff --git a/internal/src/Internal/State/Stats.hs b/internal/src/Internal/State/Stats.hs
new file mode 100644
index 00000000..8f37084f
--- /dev/null
+++ b/internal/src/Internal/State/Stats.hs
@@ -0,0 +1,220 @@
+module Internal.State.Stats where
+
+import Control.Concurrent.MVar (MVar)
+import Control.Monad.IO.Class (MonadIO, liftIO)
+import Data.Bifunctor (first)
+import Data.List (sortBy)
+import qualified Data.Map.Strict as Map
+import Data.Map.Strict (Map, (!?))
+import Data.Ord (comparing)
+import qualified Data.Set as Set
+import Data.Set ((\\))
+import GHC (ModuleName, moduleName, moduleNameString)
+import GHC.Linker.Types (LinkerEnv (..), LoaderState (..))
+import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats)
+import GHC.Types.Unique.FM (minusUFM, nonDetEltsUFM, sizeUFM)
+import GHC.Unit.Module.Env (moduleEnvKeys)
+import qualified GHC.Utils.Outputable as Outputable
+import GHC.Utils.Outputable (SDoc, comma, doublePrec, fsep, hang, nest, punctuate, text, vcat, ($$), (<+>))
+import Internal.Log (Log, logd)
+import Types.Args (TargetId (..))
+import Types.State (SymbolCache (..), Target (..))
+
+data LinkerStats =
+  LinkerStats {
+    newClosures :: Int,
+    newItables :: Int
+  }
+  deriving stock (Eq, Show)
+
+emptyLinkerStats :: LinkerStats
+emptyLinkerStats =
+  LinkerStats {
+    newClosures = 0,
+    newItables = 0
+  }
+
+data LoaderStats =
+  LoaderStats {
+    newBcos :: [String],
+    sameBcos :: Int,
+    linker :: LinkerStats
+  }
+  deriving stock (Eq, Show)
+
+emptyLoaderStats :: LoaderStats
+emptyLoaderStats =
+  LoaderStats {
+    newBcos = mempty,
+    sameBcos = 0,
+    linker = emptyLinkerStats
+  }
+
+data SymbolsStats =
+  SymbolsStats {
+    new :: Int
+  }
+  deriving stock (Eq, Show)
+
+data StatsUpdate =
+  StatsUpdate {
+    loaderStats :: LoaderStats,
+    symbols :: SymbolsStats
+  }
+  deriving stock (Eq, Show)
+
+emptyStatsUpdate :: StatsUpdate
+emptyStatsUpdate =
+  StatsUpdate {
+    loaderStats = emptyLoaderStats,
+    symbols = SymbolsStats {new = 0}
+  }
+
+data FinderStats =
+  FinderStats {
+    hits :: Map ModuleName Int,
+    misses :: Map ModuleName Int
+  }
+  deriving stock (Eq, Show)
+
+emptyFinderStats :: FinderStats
+emptyFinderStats =
+  FinderStats {
+    hits = mempty,
+    misses = mempty
+  }
+
+data CacheStats =
+  CacheStats {
+    restore :: StatsUpdate,
+    update :: StatsUpdate,
+    finder :: FinderStats
+  }
+  deriving stock (Eq, Show)
+
+emptyStats :: CacheStats
+emptyStats =
+  CacheStats {
+    restore = emptyStatsUpdate,
+    update = emptyStatsUpdate,
+    finder = emptyFinderStats
+  }
+
+----------------------------------------------------------------------------------------------------
+
+basicLinkerStats :: LinkerEnv -> LinkerEnv -> LinkerStats
+basicLinkerStats base update =
+  LinkerStats {
+    newClosures = Set.size (updateClosures \\ baseClosures),
+    newItables = Set.size (updateItables \\ baseItables)
+  }
+  where
+    updateClosures = names update.closure_env
+    baseClosures = names base.closure_env
+    updateItables = names update.itbl_env
+    baseItables = names base.itbl_env
+
+    names = Set.fromList . fmap fst . nonDetEltsUFM
+
+basicLoaderStats ::
+  LoaderState ->
+  LoaderState ->
+  LinkerStats ->
+  LoaderStats
+basicLoaderStats base update linker =
+  LoaderStats {
+    newBcos = modStr <$> Set.toList (updateBcos \\ baseBcos),
+    sameBcos = Set.size bcoSame,
+    linker
+  }
+  where
+    modStr = moduleNameString . moduleName
+    bcoSame = Set.intersection updateBcos baseBcos
+    updateBcos = Set.fromList (moduleEnvKeys update.bcos_loaded)
+    baseBcos = Set.fromList (moduleEnvKeys base.bcos_loaded)
+
+basicSymbolsStats :: SymbolCache -> SymbolCache -> SymbolsStats
+basicSymbolsStats base update =
+  SymbolsStats {
+    new = sizeUFM (minusUFM update.symbols base.symbols)
+  }
+
+----------------------------------------------------------------------------------------------------
+
+moduleColumns :: Show a => Map ModuleName a -> SDoc
+moduleColumns m =
+  vcat [text n Outputable.<> text ":" $$ nest offset (text (show h)) | (n, h) <- kvs]
+  where
+    offset = length (fst (last kvs)) + 2
+    kvs = sortBy (comparing (length . fst)) (first moduleNameString <$> Map.toList m)
+
+-- | Assemble log messages about cache statistics.
+statsMessages :: CacheStats -> SDoc
+statsMessages CacheStats {restore, update, finder} =
+  hang (text "Restore:") 2 restoreStats $$
+  hang (text "Update:") 2 updateStats $$
+  hang (text "Finder:") 2 finderStats
+  where
+      restoreStats =
+        text (show (length restore.loaderStats.newBcos)) <+> text "BCOs" $$
+        text (show restore.loaderStats.linker.newClosures) <+> text "closures" $$
+        text (show restore.symbols.new) <+> text "symbols" $$
+        text (show restore.loaderStats.sameBcos) <+> text "BCOs already in cache"
+
+      newBcos = text <$> update.loaderStats.newBcos
+
+      updateStats =
+        (if null newBcos then text "No new BCOs" else text "New BCOs:" <+> fsep (punctuate comma newBcos)) $$
+        text (show update.loaderStats.linker.newClosures) <+> text "new closures" $$
+        text (show update.symbols.new) <+> text "new symbols" $$
+        text (show update.loaderStats.sameBcos) <+> text "BCOs already in cache"
+
+      finderStats =
+        hang (text "Hits:") 2 (moduleColumns finder.hits) $$
+        hang (text "Misses:") 2 (moduleColumns finder.misses)
+
+-- | Assemble report messages, consisting of:
+--
+-- - Cache statistics, if the feature is enabled
+-- - Current RTS memory usage
+reportMessages ::
+  Target ->
+  Maybe (Map Target CacheStats) ->
+  Double ->
+  SDoc
+reportMessages target mb_stats memory =
+  statsPart $$
+  memoryPart
+  where
+    statsPart = case mb_stats of
+      Just stats -> maybe (text "Cache unused for this module.") statsMessages (stats !? target)
+      Nothing -> text "Cache disabled."
+
+    memoryPart = text "Memory:" <+> doublePrec 2 memory <+> text "MB"
+
+-- | Log a report for a completed compilation, using 'reportMessages' to assemble the content.
+report ::
+  MonadIO m =>
+  MVar Log ->
+  -- | A description of the current worker process.
+  Maybe TargetId ->
+  Target ->
+  Maybe (Map Target CacheStats) ->
+  m ()
+report logVar workerId target stats = do
+  s <- liftIO getRTSStats
+  let memory = fromIntegral (s.gc.gcdetails_mem_in_use_bytes) / 1000000
+  logd logVar (hang header 2 (reportMessages target stats memory))
+  where
+    header = text target.path Outputable.<> maybe (text "") workerDesc workerId Outputable.<> text ":"
+
+    workerDesc wid = text (" (" ++ wid.string ++ ")")
+
+logMemStats :: String -> MVar Log -> IO ()
+logMemStats step logVar = do
+  s <- liftIO getRTSStats
+  let logMem desc value = logd logVar (text (desc ++ ":") <+> doublePrec 2 (fromIntegral value / 1_000_000) <+> text "MB")
+  logd logVar (text ("-------------- " ++ step))
+  logMem "Mem in use" s.gc.gcdetails_mem_in_use_bytes
+  logMem "Max mem in use" s.max_mem_in_use_bytes
+  logMem "Max live bytes" s.max_live_bytes
diff --git a/types/buck-worker-types.cabal b/types/buck-worker-types.cabal
index 951fb46b..d1dd0193 100644
--- a/types/buck-worker-types.cabal
+++ b/types/buck-worker-types.cabal
@@ -29,8 +29,10 @@ library
     Types.GhcHandler
     Types.Grpc
     Types.Orchestration
+    Types.State
   build-depends:
     base ==4.*,
     containers,
     filepath,
+    ghc,
     split
diff --git a/types/src/Types/State.hs b/types/src/Types/State.hs
new file mode 100644
index 00000000..4da5dc8e
--- /dev/null
+++ b/types/src/Types/State.hs
@@ -0,0 +1,18 @@
+module Types.State where
+
+import GHC.Data.FastString (FastString)
+import GHC.Ptr (Ptr)
+import GHC.Types.Unique.FM (UniqFM)
+
+-- | The path to the source file the worker is currently compiling.
+-- used primarily to index maps in the state and for logging.
+newtype Target =
+  Target { path :: FilePath }
+  deriving stock (Eq, Show)
+  deriving newtype (Ord)
+
+type SymbolMap = UniqFM FastString (Ptr ())
+
+newtype SymbolCache =
+  SymbolCache { symbols :: SymbolMap }
+  deriving newtype (Semigroup, Monoid)

From 3148e77c119121266430db5f668597403840e396 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Wed, 4 Jun 2025 22:37:33 +0200
Subject: [PATCH 17/27] Refactor oneshot mode state into a separate module

Change-Id: Icc47c9f67788210555476179279ac5a5fc3e0d73
---
 ghc-worker/lib/GhcWorker/Run.hs        |  11 +-
 ghc-worker/test/TestSetup.hs           |   6 +-
 internal/buck-worker-internal.cabal    |   1 +
 internal/src/Internal/Cache.hs         | 444 ++----------------------
 internal/src/Internal/Session.hs       |  19 +-
 internal/src/Internal/State/Oneshot.hs | 451 +++++++++++++++++++++++++
 6 files changed, 498 insertions(+), 434 deletions(-)
 create mode 100644 internal/src/Internal/State/Oneshot.hs

diff --git a/ghc-worker/lib/GhcWorker/Run.hs b/ghc-worker/lib/GhcWorker/Run.hs
index 381481c7..f80eecf0 100644
--- a/ghc-worker/lib/GhcWorker/Run.hs
+++ b/ghc-worker/lib/GhcWorker/Run.hs
@@ -9,11 +9,9 @@ import Control.Exception (throwIO)
 import GhcWorker.GhcHandler (LockState (..), ghcHandler)
 import GhcWorker.Grpc (instrumentMethods)
 import GhcWorker.Instrumentation (WorkerStatus (..), toGrpcHandler)
-import GhcWorker.Orchestration (
-  CreateMethods (..),
-  runCentralGhcSpawned,
-  )
-import Internal.Cache (Cache (..), CacheFeatures (..), emptyCache, emptyCacheWith)
+import GhcWorker.Orchestration (CreateMethods (..), runCentralGhcSpawned)
+import Internal.Cache (Cache (..), emptyCache, emptyCacheWith)
+import Internal.State.Oneshot (OneshotCacheFeatures (..))
 import Network.GRPC.Common.Protobuf (Proto)
 import Network.GRPC.Server.Protobuf (ProtobufMethodsOf)
 import Network.GRPC.Server.StreamType (Methods)
@@ -79,8 +77,7 @@ runWorker CliOptions {workerMode, serve} = do
   cache <-
     case workerMode of
       WorkerMakeMode ->
-        emptyCacheWith CacheFeatures {
-          hpt = True,
+        emptyCacheWith OneshotCacheFeatures {
           loader = False,
           enable = True,
           names = False,
diff --git a/ghc-worker/test/TestSetup.hs b/ghc-worker/test/TestSetup.hs
index ba6e7327..93715f9a 100644
--- a/ghc-worker/test/TestSetup.hs
+++ b/ghc-worker/test/TestSetup.hs
@@ -6,8 +6,9 @@ import Data.Functor ((<&>))
 import Data.List.NonEmpty (NonEmpty (..))
 import Data.Traversable (for)
 import GHC.Unit (UnitId, stringToUnitId, unitIdString)
-import Internal.Cache (Cache (..), CacheFeatures (..), emptyCacheWith)
+import Internal.Cache (Cache (..), emptyCacheWith)
 import Internal.Log (dbg)
+import Internal.State.Oneshot (OneshotCacheFeatures (..))
 import Prelude hiding (log)
 import System.Directory (createDirectoryIfMissing, listDirectory, withCurrentDirectory)
 import System.Environment (getEnv)
@@ -202,8 +203,7 @@ withProject mkTargets use =
     withCurrentDirectory tmp do
       for_ @[] ["src", "tmp", "out"] \ dir ->
         createDirectoryIfMissing False (tmp </> dir)
-      cache <- emptyCacheWith CacheFeatures {
-        hpt = True,
+      cache <- emptyCacheWith OneshotCacheFeatures {
         loader = False,
         enable = True,
         names = False,
diff --git a/internal/buck-worker-internal.cabal b/internal/buck-worker-internal.cabal
index 67075898..5684453d 100644
--- a/internal/buck-worker-internal.cabal
+++ b/internal/buck-worker-internal.cabal
@@ -33,6 +33,7 @@ library
                       Internal.Metadata
                       Internal.Session
                       Internal.State.Make
+                      Internal.State.Oneshot
                       Internal.State.Stats
 
     if (impl(ghc >= 9.9))
diff --git a/internal/src/Internal/Cache.hs b/internal/src/Internal/Cache.hs
index 38fc08e2..8ab6b768 100644
--- a/internal/src/Internal/Cache.hs
+++ b/internal/src/Internal/Cache.hs
@@ -2,66 +2,24 @@
 
 module Internal.Cache where
 
-import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar, readMVar)
-import Control.Monad (join)
+import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar)
 import Control.Monad.IO.Class (MonadIO, liftIO)
-import Data.Bifunctor (first)
-import Data.Coerce (coerce)
-import Data.IORef (readIORef)
-import qualified Data.Map.Strict as Map
-import Data.Map.Strict (Map)
-import Data.Maybe (fromMaybe)
 import Data.Set (Set)
-import Data.Traversable (for)
 import GHC (Ghc, ModIface, emptyMG, mi_module, moduleName, moduleNameString, setSession)
 import GHC.Driver.Env (HscEnv (..))
 import GHC.Driver.Monad (modifySessionM, withSession)
-import GHC.Linker.Types (Linkable, LinkerEnv (..), Loader (..), LoaderState (..))
-import GHC.Runtime.Interpreter (Interp (..))
-import GHC.Types.Name.Cache (NameCache (..), OrigNameCache)
-import GHC.Types.Unique.DFM (plusUDFM)
-import GHC.Unit.Env (UnitEnv (..), unitEnv_new)
-import GHC.Unit.External (ExternalUnitCache (..), initExternalUnitCache)
-import GHC.Unit.Finder (InstalledFindResult (..))
-import GHC.Unit.Finder.Types (FinderCache (..))
-import GHC.Unit.Module.Env (InstalledModuleEnv, emptyModuleEnv, plusModuleEnv)
+import GHC.Linker.Types (Linkable)
+import GHC.Unit.Env (unitEnv_new)
 import GHC.Unit.Module.Graph (ModuleGraph)
 import Internal.Log (Log)
 import qualified Internal.State.Make as Make
 import Internal.State.Make (MakeState (..))
+import qualified Internal.State.Oneshot as Oneshot
+import Internal.State.Oneshot (OneshotCacheFeatures (..), OneshotState, newOneshotCacheFeatures, newOneshotStateWith)
 import qualified Internal.State.Stats as Stats
-import Internal.State.Stats (
-  CacheStats (..),
-  LinkerStats (..),
-  LoaderStats (..),
-  StatsUpdate (..),
-  SymbolsStats (..),
-  basicLinkerStats,
-  basicLoaderStats,
-  basicSymbolsStats,
-  emptyLinkerStats,
-  emptyStats,
-  )
 import System.Environment (lookupEnv)
 import Types.Args (TargetId (..))
-import Types.State (SymbolCache (..), SymbolMap, Target)
-
-#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
-
-import Data.IORef (IORef, newIORef)
-import qualified Data.Map.Lazy as LazyMap
-import GHC.Fingerprint (Fingerprint, getFileHash)
-import GHC.IORef (atomicModifyIORef')
-import GHC.Unit (InstalledModule, emptyInstalledModuleEnv, extendInstalledModuleEnv, lookupInstalledModuleEnv)
-import GHC.Unit.Finder (FinderCache (..), InstalledFindResult (..))
-import GHC.Utils.Panic (panic)
-import Internal.State.Stats (FinderStats (..))
-
-#else
-
-import GHC.Unit.Finder (initFinderCache)
-
-#endif
+import Types.State (Target)
 
 data ModuleArtifacts =
   ModuleArtifacts {
@@ -73,12 +31,6 @@ instance Show ModuleArtifacts where
   show ModuleArtifacts {iface} =
     "ModuleArtifacts { iface = " ++ moduleNameString (moduleName (mi_module iface)) ++ " }"
 
-data InterpCache =
-  InterpCache {
-    loaderState :: LoaderState,
-    symbols :: SymbolCache
-  }
-
 data BinPath =
   BinPath {
     initial :: Maybe String,
@@ -86,70 +38,13 @@ data BinPath =
   }
   deriving stock (Eq, Show)
 
-#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
-
-data FinderState =
-  FinderState {
-     modules :: IORef (InstalledModuleEnv InstalledFindResult),
-     files :: IORef (Map String Fingerprint)
-  }
-
-emptyFinderState :: MonadIO m => m FinderState
-emptyFinderState =
-  liftIO do
-    modules <- newIORef emptyInstalledModuleEnv
-    files <- newIORef LazyMap.empty
-    pure FinderState {modules, files}
-
-finderEnv :: FinderState -> IO (InstalledModuleEnv InstalledFindResult)
-finderEnv FinderState {modules} =
-  readIORef modules
-
-#else
-
-data FinderState =
-  FinderState {
-    cache :: FinderCache
-  }
-
-emptyFinderState :: MonadIO m => m FinderState
-emptyFinderState =
-  liftIO do
-    cache <- initFinderCache
-    pure FinderState {cache}
-
-finderEnv :: FinderState -> IO (InstalledModuleEnv InstalledFindResult)
-finderEnv FinderState {cache = FinderCache {fcModuleCache}} =
-  readIORef fcModuleCache
-
-#endif
-
-data CacheFeatures =
-  CacheFeatures {
-    enable :: Bool,
-    loader :: Bool,
-    names :: Bool,
-    finder :: Bool,
-    eps :: Bool,
-    hpt :: Bool
-  }
-  deriving stock (Eq, Show)
-
-newCacheFeatures :: CacheFeatures
-newCacheFeatures = CacheFeatures {enable = True, loader = True, names = True, finder = True, eps = True, hpt = False}
-
 data Cache =
   Cache {
-    features :: CacheFeatures,
-    interpCache :: Maybe InterpCache,
-    names :: OrigNameCache,
-    stats :: Map Target CacheStats,
     path :: BinPath,
-    finder :: FinderState,
-    eps :: ExternalUnitCache,
     baseSession :: Maybe HscEnv,
     options :: Options,
-    make :: MakeState
+    make :: MakeState,
+    oneshot :: OneshotState
   }
 
 data Options =
@@ -157,195 +52,47 @@ data Options =
     extraGhcOptions :: String
   }
 
-emptyCacheWith :: CacheFeatures -> IO (MVar Cache)
+emptyCacheWith :: OneshotCacheFeatures -> IO (MVar Cache)
 emptyCacheWith features = do
   initialPath <- lookupEnv "PATH"
-  finder <- emptyFinderState
-  eps <- initExternalUnitCache
+  oneshot <- newOneshotStateWith features
   newMVar Cache {
-    features,
-    interpCache = Nothing,
-    names = emptyModuleEnv,
-    stats = mempty,
     path = BinPath {
       initial = initialPath,
       extra = mempty
     },
-    finder,
-    eps,
     baseSession = Nothing,
     options = defaultOptions,
     make = MakeState {
       moduleGraph = emptyMG,
       hug = unitEnv_new mempty,
       interp = Nothing
-    }
+    },
+    oneshot
   }
 
 emptyCache :: Bool -> IO (MVar Cache)
 emptyCache enable = do
-  emptyCacheWith newCacheFeatures {enable}
-
-defaultOptions :: Options
-defaultOptions =
-  Options {
-    extraGhcOptions = ""
-  }
-
-restoreLinkerEnv :: LinkerEnv -> LinkerEnv -> (LinkerEnv, LinkerStats)
-restoreLinkerEnv cached session =
-  (merged, basicLinkerStats session cached)
-  where
-    merged =
-      LinkerEnv {
-        -- UniqFM, <> right-biased
-        closure_env =
-          cached.closure_env
-          <>
-          session.closure_env,
-        -- UniqFM, <> right-biased
-        itbl_env = cached.itbl_env <> session.itbl_env,
-        -- UniqFM, <> right-biased
-        addr_env = cached.addr_env <> session.addr_env
-      }
-
-restoreLoaderState ::
-  LoaderState ->
-  LoaderState ->
-  IO (LoaderState, Maybe LoaderStats)
-restoreLoaderState cached session =
-  pure (merged, Just (basicLoaderStats session cached linkerStats))
-  where
-    merged =
-      LoaderState {
-        linker_env,
-        -- ModuleEnv, left-biased
-        bcos_loaded = plusModuleEnv session.bcos_loaded cached.bcos_loaded,
-        -- ModuleEnv, left-biased
-        objs_loaded = plusModuleEnv session.objs_loaded cached.objs_loaded,
-        -- UniqDFM, depends on the elements in the maps
-        pkgs_loaded = plusUDFM cached.pkgs_loaded session.pkgs_loaded,
-        temp_sos = session.temp_sos
-      }
+  emptyCacheWith newOneshotCacheFeatures {enable}
 
-    (linker_env, linkerStats) = restoreLinkerEnv cached.linker_env session.linker_env
 -- | Update the 'MakeState' field in the 'Cache'.
 updateMakeState :: (MakeState -> MakeState) -> Cache -> Cache
 updateMakeState f cache = cache {make = f cache.make}
 
-modifyStats :: Target -> (CacheStats -> CacheStats) -> Cache -> Cache
-modifyStats target f cache =
-  cache {stats = Map.alter (Just . f . fromMaybe emptyStats) target cache.stats}
 updateMakeStateVar :: MVar Cache -> (MakeState -> MakeState) -> IO ()
 updateMakeStateVar var f = modifyMVar_ var (pure . updateMakeState f)
 
-pushStats :: Bool -> Target -> Maybe LoaderStats -> SymbolsStats -> Cache -> Cache
-pushStats restoring target (Just new) symbols =
-  modifyStats target add
-  where
-    add old | restoring = old {restore = old.restore {loaderStats = new, symbols}}
-            | otherwise = old {update = old.update {loaderStats = new, symbols}}
-pushStats _ _ _ _ =
-  id
-
-restoreCache ::
-  Target ->
-  Maybe LoaderState ->
-  SymbolCache ->
-  OrigNameCache ->
-  Cache ->
-  IO (OrigNameCache, (SymbolCache, (Maybe LoaderState, Cache)))
-restoreCache target initialLoaderState initialSymbolCache initialNames cache
-  | Just InterpCache {..} <- cache.interpCache
-  = do
-    (restoredLs, loaderStats) <- case initialLoaderState of
-      Just sessionLs ->
-        restoreLoaderState loaderState sessionLs
-      Nothing ->
-        pure (loaderState, Nothing)
-    let
-      newSymbols = initialSymbolCache <> symbols
-      symbolsStats = basicSymbolsStats initialSymbolCache symbols
-      newCache = pushStats True target loaderStats symbolsStats cache
-      -- this overwrites entire modules, since OrigNameCache is a three-level map.
-      -- eventually we'll want to merge properly.
-      names = plusModuleEnv initialNames cache.names
-    pure (names, (newSymbols, (Just restoredLs, newCache)))
+updateOneshotState :: (OneshotState -> OneshotState) -> Cache -> Cache
+updateOneshotState f cache = cache {oneshot = f cache.oneshot}
 
-  | otherwise
-  = pure (initialNames, (initialSymbolCache, (initialLoaderState, cache)))
+updateOneshotStateVar :: MVar Cache -> (OneshotState -> OneshotState) -> IO ()
+updateOneshotStateVar var f = modifyMVar_ var (pure . updateOneshotState f)
 
-initCache ::
-  LoaderState ->
-  SymbolCache ->
-  OrigNameCache ->
-  Cache ->
-  IO Cache
-initCache loaderState symbols names Cache {names = _, ..} =
-  pure Cache {interpCache = Just InterpCache {..}, ..}
-
-updateLinkerEnv :: LinkerEnv -> LinkerEnv -> (LinkerEnv, LinkerStats)
-updateLinkerEnv cached session =
-  (merged, basicLinkerStats cached session)
-  where
-    merged =
-      LinkerEnv {
-        -- UniqFM, <> right-biased
-        closure_env =
-          cached.closure_env
-          <>
-          session.closure_env,
-        -- UniqFM, <> right-biased
-        itbl_env = cached.itbl_env <> session.itbl_env,
-        -- UniqFM, <> right-biased
-        addr_env = cached.addr_env <> session.addr_env
-      }
-
-updateLoaderState ::
-  LoaderState ->
-  LoaderState ->
-  IO (LoaderState, Maybe LoaderStats)
-updateLoaderState cached session = do
-  pure (merged, Just stats {linker = linkerStats})
-  where
-    merged =
-      LoaderState {
-        linker_env,
-        -- ModuleEnv, left-biased
-        bcos_loaded = plusModuleEnv session.bcos_loaded cached.bcos_loaded,
-        -- ModuleEnv, left-biased
-        objs_loaded = plusModuleEnv session.objs_loaded cached.objs_loaded,
-        -- UniqDFM, depends on the elements in the maps
-        pkgs_loaded = plusUDFM cached.pkgs_loaded session.pkgs_loaded,
-        temp_sos = session.temp_sos
-      }
-
-    (linker_env, linkerStats) = updateLinkerEnv cached.linker_env session.linker_env
-
-    stats = basicLoaderStats cached session emptyLinkerStats
-
-updateCache ::
-  Target ->
-  InterpCache ->
-  LoaderState ->
-  SymbolCache ->
-  OrigNameCache ->
-  Cache ->
-  IO Cache
-updateCache target InterpCache {..} newLoaderState newSymbols newNames cache = do
-  (updatedLs, stats) <- updateLoaderState loaderState newLoaderState
-  pure $ pushStats False target stats symbolsStats cache {
-    interpCache = Just InterpCache {
-      loaderState = updatedLs,
-      symbols = symbols <> newSymbols
-    },
-    -- for now: when a module is compiled, its names are definitely complete, so when a downstream module uses it as a
-    -- dep, we don't want to overwrite the previous entry.
-    -- but when we recompile parts of the tree this is different, so wel'll want to merge properly.
-    names = plusModuleEnv newNames cache.names
+defaultOptions :: Options
+defaultOptions =
+  Options {
+    extraGhcOptions = ""
   }
-  where
-    symbolsStats = basicSymbolsStats symbols newSymbols
 
 -- | Log a report for a completed compilation, using 'reportMessages' to assemble the content.
 report ::
@@ -357,104 +104,7 @@ report ::
   Cache ->
   m ()
 report logVar workerId target cache = do
-  Stats.report logVar workerId target (if cache.features.enable then Just cache.stats else Nothing)
-
-#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
-
--- | This replacement of the Finder implementation has the sole purpose of recording some cache stats, for now.
--- While its mutable state is allocated separately and shared across sessions, this doesn't really make a difference at
--- the moment since we're also initializing each compilation session with a shared @HscEnv@.
--- Ultimately this might be used to exert some more control over what modules GHC is allowed to access by using Buck's
--- deps, or some additional optimization.
-newFinderCache :: MVar Cache -> Cache -> Target -> IO FinderCache
-newFinderCache cacheVar Cache {finder = FinderState {modules, files}} target = do
-  let flushFinderCaches :: UnitEnv -> IO ()
-      flushFinderCaches _ = panic "GHC attempted to flush finder caches, which shouldn't happen in worker mode"
-
-      addToFinderCache :: InstalledModule -> InstalledFindResult -> IO ()
-      addToFinderCache key val =
-        atomicModifyIORef' modules $ \c ->
-          case (lookupInstalledModuleEnv c key, val) of
-            (Just InstalledFound{}, InstalledNotFound{}) -> (c, ())
-            _ -> (extendInstalledModuleEnv c key val, ())
-
-      lookupFinderCache :: InstalledModule -> IO (Maybe InstalledFindResult)
-      lookupFinderCache key = do
-        c <- readIORef modules
-        let result = lookupInstalledModuleEnv c key
-        case result of
-          Just _ -> cacheHit key
-          Nothing -> cacheMiss key
-        pure $! result
-
-      lookupFileCache :: FilePath -> IO Fingerprint
-      lookupFileCache key = do
-         fc <- readIORef files
-         case LazyMap.lookup key fc of
-           Nothing -> do
-             hash <- getFileHash key
-             atomicModifyIORef' files $ \c -> (LazyMap.insert key hash c, ())
-             return hash
-           Just fp -> return fp
-  return FinderCache {..}
-  where
-    cacheHit m =
-      updateStats \ FinderStats {hits, ..} -> FinderStats {hits = incStat m hits, ..}
-
-    cacheMiss m =
-      updateStats \ FinderStats {misses, ..} -> FinderStats {misses = incStat m misses, ..}
-
-    incStat m = Map.alter (Just . succ . fromMaybe 0) (moduleName m)
-
-    updateStats f =
-      modifyMVar_ cacheVar $ pure . modifyStats target \ CacheStats {..} -> CacheStats {finder = f finder, ..}
-
-#else
-
-newFinderCache :: MVar Cache -> Cache -> Target -> IO FinderCache
-newFinderCache _ Cache {finder = FinderState {cache}} _ = pure cache
-
-#endif
-
-withHscState :: HscEnv -> (MVar OrigNameCache -> MVar (Maybe LoaderState) -> MVar SymbolMap -> IO a) -> IO (Maybe a)
-withHscState HscEnv {hsc_interp, hsc_NC = NameCache {nsNames}} use =
-#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) || defined(MWB)
-  for hsc_interp \ Interp {interpLoader = Loader {loader_state}, interpLookupSymbolCache} ->
-    liftIO $ use nsNames loader_state interpLookupSymbolCache
-#else
-  for hsc_interp \ Interp {interpLoader = Loader {loader_state}} ->
-    liftIO do
-    symbolCacheVar <- newMVar mempty
-    use nsNames loader_state symbolCacheVar
-#endif
-
--- | Restore cache parts that depend on the 'Target'.
-setTarget :: MVar Cache -> Cache -> Target -> HscEnv -> IO HscEnv
-setTarget cacheVar cache target hsc_env = do
-  -- The Finder cache is already shared by the base session, but with this, we can additionally inject the target file
-  -- for stats collection.
-  -- Only has an effect if the patch that abstracts the 'FinderCache' interface is in GHC, so >= 9.12.
-  if cache.features.finder
-  then restoreFinderCache hsc_env
-  else pure hsc_env
-  where
-    restoreFinderCache e = do
-      hsc_FC <- newFinderCache cacheVar cache target
-      pure e {hsc_FC}
-
--- | Restore cache parts related to oneshot mode.
-restoreOneshotCache :: Cache -> HscEnv -> IO HscEnv
-restoreOneshotCache cache hsc_env = do
-  -- If the feature is enabled, restore the EPS.
-  -- This is only relevant for the oneshot worker, but even there the base session should already be sharing the EPS
-  -- across modules.
-  -- Might be removed soon.
-  pure
-    if cache.features.eps
-    then restoreCachedEps hsc_env
-    else hsc_env
-  where
-    restoreCachedEps e = e {hsc_unit_env = e.hsc_unit_env {ue_eps = cache.eps}}
+  Stats.report logVar workerId target (if cache.oneshot.features.enable then Just cache.oneshot.stats else Nothing)
 
 -- | Merge the given module graph into the cached graph.
 -- This is used by the make mode worker after the metadata step has computed the module graph.
@@ -462,26 +112,6 @@ updateModuleGraph :: MVar Cache -> ModuleGraph -> IO ()
 updateModuleGraph cacheVar new =
   updateMakeStateVar cacheVar (Make.storeModuleGraph new)
 
-prepareCache :: MVar Cache -> Target -> HscEnv -> Cache -> IO (Cache, (HscEnv, Bool))
-prepareCache cacheVar target hsc_env0 cache0 = do
-  result <-
-    if cache0.features.enable
-    then do
-      hsc_env1 <- restoreOneshotCache cache0 =<< setTarget cacheVar cache0 target hsc_env0
-      if cache0.features.loader
-      then do
-        withHscState hsc_env1 \ nsNames loaderStateVar symbolCacheVar -> do
-          cache1 <- modifyMVar loaderStateVar \ initialLoaderState ->
-            modifyMVar symbolCacheVar \ initialSymbolCache ->
-              (first coerce) <$>
-              modifyMVar nsNames \ names ->
-                restoreCache target initialLoaderState (SymbolCache initialSymbolCache) names cache0
-          pure (hsc_env1, cache1)
-      else pure (Just (hsc_env1, cache0))
-    else pure Nothing
-  let (hsc_env1, cache1) = fromMaybe (hsc_env0, cache0 {features = cache0.features {loader = False}}) result
-  pure (cache1, (hsc_env1, cache1.features.enable))
-
 finalizeCache ::
   MVar Log ->
   -- | A description of the current worker process.
@@ -492,18 +122,8 @@ finalizeCache ::
   Cache ->
   IO Cache
 finalizeCache logVar workerId hsc_env target _ cache0 = do
-  cache1 <-
-    if cache0.features.enable
-    then do
-      if cache0.features.loader
-      then do
-        fromMaybe cache0 . join <$> withHscState hsc_env \ nsNames loaderStateVar symbolCacheVar ->
-          readMVar loaderStateVar >>= traverse \ newLoaderState -> do
-            newSymbols <- readMVar symbolCacheVar
-            newNames <- readMVar nsNames
-            maybe initCache (updateCache target) cache0.interpCache newLoaderState (SymbolCache newSymbols) newNames cache0
-      else pure cache0
-    else pure cache0
+  oneshot <- Oneshot.storeState hsc_env target cache0.oneshot
+  let cache1 = cache0 {oneshot}
   report logVar workerId target cache1
   pure cache1
 
@@ -514,7 +134,7 @@ withSessionM use =
     setSession new_env
     pure a
 
-withCache ::
+withCacheOneshot ::
   MVar Log ->
   -- | A description of the current worker process.
   Maybe TargetId ->
@@ -522,8 +142,10 @@ withCache ::
   Target ->
   Ghc (Maybe (Maybe ModuleArtifacts, a)) ->
   Ghc (Maybe (Maybe ModuleArtifacts, a))
-withCache logVar workerId cacheVar target prog = do
-  _ <- withSessionM \ hsc_env -> modifyMVar cacheVar (prepareCache cacheVar target hsc_env)
+withCacheOneshot logVar workerId cacheVar target prog = do
+  _ <- withSessionM \ hsc_env -> modifyMVar cacheVar \ cache -> do
+    (oneshot, result) <- Oneshot.loadState (updateOneshotStateVar cacheVar) target hsc_env cache.oneshot
+    pure (cache {oneshot}, result)
   result <- prog
   finalize (fst =<< result)
   pure result
@@ -532,11 +154,11 @@ withCache logVar workerId cacheVar target prog = do
       withSession \ hsc_env ->
         liftIO (modifyMVar_ cacheVar (finalizeCache logVar workerId hsc_env target art))
 
--- | This reduced version of 'withCache' is tailored specifically to make mode, only restoring the HUG and module graph
--- from the cache, since those are the only two components modified by the worker that aren't already shared by the base
--- session.
+-- | This reduced version of 'withCache' is tailored specifically to make mode, only restoring the HUG, module graph and
+-- interpreter state from the cache, since those are the only two components modified by the worker that aren't already
+-- shared by the base session.
 --
--- The mechanisms in in 'withCache' are partially legacy experiments whose purpose was to explore which data can be
+-- The mechanisms in 'withCache' are partially legacy experiments whose purpose was to explore which data can be
 -- shared manually in oneshot mode, so this variant will be improved more deliberately.
 withCacheMake ::
   MVar Log ->
diff --git a/internal/src/Internal/Session.hs b/internal/src/Internal/Session.hs
index 4c013c7a..970b7921 100644
--- a/internal/src/Internal/Session.hs
+++ b/internal/src/Internal/Session.hs
@@ -38,18 +38,11 @@ import GHC.Runtime.Loader (initializeSessionPlugins)
 import GHC.Types.SrcLoc (Located, mkGeneralLocated, unLoc)
 import GHC.Utils.Logger (Logger, getLogger, setLogFlags)
 import GHC.Utils.Panic (GhcException (UsageError), panic, throwGhcException)
-import GHC.Utils.TmpFs (TempDir (..), initTmpFs, cleanTempFiles, cleanTempDirs)
-import Internal.Cache (
-  BinPath (..),
-  Cache (..),
-  CacheFeatures (..),
-  ModuleArtifacts,
-  Options (..),
-  withCache,
-  withCacheMake,
-  )
+import GHC.Utils.TmpFs (TempDir (..), cleanTempDirs, cleanTempFiles, initTmpFs)
+import Internal.Cache (BinPath (..), Cache (..), ModuleArtifacts, Options (..), withCacheMake, withCacheOneshot)
 import Internal.Error (handleExceptions)
 import Internal.Log (Log (..), logToState)
+import Internal.State.Oneshot (OneshotCacheFeatures (..), OneshotState (..))
 import Prelude hiding (log)
 import System.Environment (setEnv)
 import Types.Args (Args (..))
@@ -167,7 +160,7 @@ withGhcInSession env prog =
 ensureSession :: Bool -> MVar Cache -> Args -> IO HscEnv
 ensureSession reuse cacheVar args =
   modifyMVar cacheVar \ cache -> do
-    if cache.features.enable && reuse
+    if cache.oneshot.features.enable && reuse
     then do
       newEnv <- maybe (initHscEnv args.topdir) prepReused cache.baseSession
       pure (cache {baseSession = Just newEnv}, newEnv)
@@ -229,7 +222,7 @@ withGhc env =
   withGhcUsingCache cacheHandler env
   where
     cacheHandler target prog = do
-      result <- withCache env.log env.args.workerTargetId env.cache target do
+      result <- withCacheOneshot env.log env.args.workerTargetId env.cache target do
         res <- prog
         pure do
           a <- res
@@ -240,7 +233,7 @@ withGhc env =
 -- Return the interface and bytecode.
 withGhcDefault :: Env -> (Target -> Ghc (Maybe (Maybe ModuleArtifacts, a))) -> IO (Maybe (Maybe ModuleArtifacts, a))
 withGhcDefault env =
-  withGhcUsingCache (withCache env.log env.args.workerTargetId env.cache) env
+  withGhcUsingCache (withCacheOneshot env.log env.args.workerTargetId env.cache) env
 
 -- | Command line args that have to be stored in the current home unit env.
 -- These are specified as a single program argument with their option argument, without whitespace in between.
diff --git a/internal/src/Internal/State/Oneshot.hs b/internal/src/Internal/State/Oneshot.hs
new file mode 100644
index 00000000..2b6bf15b
--- /dev/null
+++ b/internal/src/Internal/State/Oneshot.hs
@@ -0,0 +1,451 @@
+{-# LANGUAGE CPP, NoFieldSelectors #-}
+
+module Internal.State.Oneshot where
+
+import Control.Concurrent.MVar (MVar, modifyMVar, readMVar)
+import Control.Monad (join)
+import Control.Monad.IO.Class (MonadIO, liftIO)
+import Data.Bifunctor (first)
+import Data.Coerce (coerce)
+import Data.IORef (readIORef)
+import qualified Data.Map.Strict as Map
+import Data.Map.Strict (Map)
+import Data.Maybe (fromMaybe)
+import Data.Traversable (for)
+import GHC.Driver.Env (HscEnv (..))
+import GHC.Linker.Types (LinkerEnv (..), Loader (..), LoaderState (..))
+import GHC.Runtime.Interpreter (Interp (..))
+import GHC.Types.Name.Cache (NameCache (..), OrigNameCache)
+import GHC.Types.Unique.DFM (plusUDFM)
+import GHC.Unit.Env (UnitEnv (..))
+import GHC.Unit.External (ExternalUnitCache (..), initExternalUnitCache)
+import GHC.Unit.Finder (InstalledFindResult (..))
+import GHC.Unit.Finder.Types (FinderCache (..))
+import GHC.Unit.Module.Env (InstalledModuleEnv, emptyModuleEnv, plusModuleEnv)
+import Internal.State.Stats (
+  CacheStats (..),
+  LinkerStats (..),
+  LoaderStats (..),
+  StatsUpdate (..),
+  SymbolsStats (..),
+  basicLinkerStats,
+  basicLoaderStats,
+  basicSymbolsStats,
+  emptyLinkerStats,
+  emptyStats,
+  )
+import Types.State (SymbolCache (..), SymbolMap, Target)
+
+#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) || defined(MWB)
+
+import Control.Exception (evaluate)
+import Data.IORef (IORef, newIORef)
+import qualified Data.Map.Lazy as LazyMap
+import GHC.Fingerprint (Fingerprint, getFileHash)
+import GHC.IORef (atomicModifyIORef')
+import GHC.Unit (
+  InstalledModule,
+  emptyInstalledModuleEnv,
+  extendInstalledModuleEnv,
+  lookupInstalledModuleEnv,
+  moduleName,
+  )
+import GHC.Utils.Panic (panic)
+import Internal.Debug (forceInstalledFindResult)
+import Internal.State.Stats (FinderStats (..))
+
+#else
+
+import Control.Concurrent.MVar (newMVar)
+import GHC.Unit.Finder (initFinderCache)
+
+#endif
+
+#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) || defined(MWB)
+
+data FinderState =
+  FinderState {
+     modules :: IORef (InstalledModuleEnv InstalledFindResult),
+     files :: IORef (Map String Fingerprint)
+  }
+
+emptyFinderState :: MonadIO m => m FinderState
+emptyFinderState =
+  liftIO do
+    modules <- newIORef emptyInstalledModuleEnv
+    files <- newIORef LazyMap.empty
+    pure FinderState {modules, files}
+
+finderEnv :: FinderState -> IO (InstalledModuleEnv InstalledFindResult)
+finderEnv FinderState {modules} =
+  readIORef modules
+
+#else
+
+data FinderState =
+  FinderState {
+    cache :: FinderCache
+  }
+
+emptyFinderState :: MonadIO m => m FinderState
+emptyFinderState =
+  liftIO do
+    cache <- initFinderCache
+    pure FinderState {cache}
+
+finderEnv :: FinderState -> IO (InstalledModuleEnv InstalledFindResult)
+finderEnv FinderState {cache = FinderCache {fcModuleCache}} =
+  readIORef fcModuleCache
+
+#endif
+
+data OneshotCacheFeatures =
+  OneshotCacheFeatures {
+    enable :: Bool,
+    loader :: Bool,
+    names :: Bool,
+    finder :: Bool,
+    eps :: Bool
+  }
+  deriving stock (Eq, Show)
+
+newOneshotCacheFeatures :: OneshotCacheFeatures
+newOneshotCacheFeatures = OneshotCacheFeatures {enable = True, loader = True, names = True, finder = True, eps = True}
+
+data InterpCache =
+  InterpCache {
+    loaderState :: LoaderState,
+    symbols :: SymbolCache
+  }
+
+data OneshotState =
+  OneshotState {
+    features :: OneshotCacheFeatures,
+    interpCache :: Maybe InterpCache,
+    names :: OrigNameCache,
+    finder :: FinderState,
+    eps :: ExternalUnitCache,
+    stats :: Map Target CacheStats
+  }
+
+newOneshotStateWith :: OneshotCacheFeatures -> IO OneshotState
+newOneshotStateWith features = do
+  finder <- emptyFinderState
+  eps <- initExternalUnitCache
+  pure OneshotState {
+    features,
+    interpCache = Nothing,
+    names = emptyModuleEnv,
+    stats = mempty,
+    finder,
+    eps
+  }
+
+newOneshotState :: Bool -> IO OneshotState
+newOneshotState enable = do
+  newOneshotStateWith newOneshotCacheFeatures {enable}
+
+restoreLinkerEnv :: LinkerEnv -> LinkerEnv -> (LinkerEnv, LinkerStats)
+restoreLinkerEnv cached session =
+  (merged, basicLinkerStats session cached)
+  where
+    merged =
+      LinkerEnv {
+        -- UniqFM, <> right-biased
+        closure_env =
+          cached.closure_env
+          <>
+          session.closure_env,
+        -- UniqFM, <> right-biased
+        itbl_env = cached.itbl_env <> session.itbl_env,
+        -- UniqFM, <> right-biased
+        addr_env = cached.addr_env <> session.addr_env
+      }
+
+restoreLoaderState ::
+  LoaderState ->
+  LoaderState ->
+  IO (LoaderState, Maybe LoaderStats)
+restoreLoaderState cached session =
+  pure (merged, Just (basicLoaderStats session cached linkerStats))
+  where
+    merged =
+      LoaderState {
+        linker_env,
+        -- ModuleEnv, left-biased
+        bcos_loaded = plusModuleEnv session.bcos_loaded cached.bcos_loaded,
+        -- ModuleEnv, left-biased
+        objs_loaded = plusModuleEnv session.objs_loaded cached.objs_loaded,
+        -- UniqDFM, depends on the elements in the maps
+        pkgs_loaded = plusUDFM cached.pkgs_loaded session.pkgs_loaded,
+        temp_sos = session.temp_sos
+      }
+
+    (linker_env, linkerStats) = restoreLinkerEnv cached.linker_env session.linker_env
+
+modifyStats :: Target -> (CacheStats -> CacheStats) -> OneshotState -> OneshotState
+modifyStats target f cache =
+  cache {stats = Map.alter (Just . f . fromMaybe emptyStats) target cache.stats}
+
+pushStats :: Bool -> Target -> Maybe LoaderStats -> SymbolsStats -> OneshotState -> OneshotState
+pushStats restoring target (Just new) symbols =
+  modifyStats target add
+  where
+    add old | restoring = old {restore = old.restore {loaderStats = new, symbols}}
+            | otherwise = old {update = old.update {loaderStats = new, symbols}}
+pushStats _ _ _ _ =
+  id
+
+loadCache ::
+  Target ->
+  Maybe LoaderState ->
+  SymbolCache ->
+  OrigNameCache ->
+  OneshotState ->
+  IO (OrigNameCache, (SymbolCache, (Maybe LoaderState, OneshotState)))
+loadCache target initialLoaderState initialSymbolCache initialNames cache
+  | Just InterpCache {..} <- cache.interpCache
+  = do
+    (restoredLs, loaderStats) <- case initialLoaderState of
+      Just sessionLs ->
+        restoreLoaderState loaderState sessionLs
+      Nothing ->
+        pure (loaderState, Nothing)
+    let
+      newSymbols = initialSymbolCache <> symbols
+      symbolsStats = basicSymbolsStats initialSymbolCache symbols
+      newCache = pushStats True target loaderStats symbolsStats cache
+      -- this overwrites entire modules, since OrigNameCache is a three-level map.
+      -- eventually we'll want to merge properly.
+      names = plusModuleEnv initialNames cache.names
+    pure (names, (newSymbols, (Just restoredLs, newCache)))
+
+  | otherwise
+  = pure (initialNames, (initialSymbolCache, (initialLoaderState, cache)))
+
+initState ::
+  LoaderState ->
+  SymbolCache ->
+  OrigNameCache ->
+  OneshotState ->
+  IO OneshotState
+initState loaderState symbols names OneshotState {names = _, ..} =
+  pure OneshotState {interpCache = Just InterpCache {..}, ..}
+
+updateLinkerEnv :: LinkerEnv -> LinkerEnv -> (LinkerEnv, LinkerStats)
+updateLinkerEnv cached session =
+  (merged, basicLinkerStats cached session)
+  where
+    merged =
+      LinkerEnv {
+        -- UniqFM, <> right-biased
+        closure_env =
+          cached.closure_env
+          <>
+          session.closure_env,
+        -- UniqFM, <> right-biased
+        itbl_env = cached.itbl_env <> session.itbl_env,
+        -- UniqFM, <> right-biased
+        addr_env = cached.addr_env <> session.addr_env
+      }
+
+updateLoaderState ::
+  LoaderState ->
+  LoaderState ->
+  IO (LoaderState, Maybe LoaderStats)
+updateLoaderState cached session = do
+  pure (merged, Just stats {linker = linkerStats})
+  where
+    merged =
+      LoaderState {
+        linker_env,
+        -- ModuleEnv, left-biased
+        bcos_loaded = plusModuleEnv session.bcos_loaded cached.bcos_loaded,
+        -- ModuleEnv, left-biased
+        objs_loaded = plusModuleEnv session.objs_loaded cached.objs_loaded,
+        -- UniqDFM, depends on the elements in the maps
+        pkgs_loaded = plusUDFM cached.pkgs_loaded session.pkgs_loaded,
+        temp_sos = session.temp_sos
+      }
+
+    (linker_env, linkerStats) = updateLinkerEnv cached.linker_env session.linker_env
+
+    stats = basicLoaderStats cached session emptyLinkerStats
+
+updateState ::
+  Target ->
+  InterpCache ->
+  LoaderState ->
+  SymbolCache ->
+  OrigNameCache ->
+  OneshotState ->
+  IO OneshotState
+updateState target InterpCache {..} newLoaderState newSymbols newNames cache = do
+  (updatedLs, stats) <- updateLoaderState loaderState newLoaderState
+  pure $ pushStats False target stats symbolsStats cache {
+    interpCache = Just InterpCache {
+      loaderState = updatedLs,
+      symbols = symbols <> newSymbols
+    },
+    -- for now: when a module is compiled, its names are definitely complete, so when a downstream module uses it as a
+    -- dep, we don't want to overwrite the previous entry.
+    -- but when we recompile parts of the tree this is different, so wel'll want to merge properly.
+    names = plusModuleEnv newNames cache.names
+  }
+  where
+    symbolsStats = basicSymbolsStats symbols newSymbols
+
+#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) || defined(MWB)
+
+-- | This replacement of the Finder implementation has the sole purpose of recording some cache stats, for now.
+-- While its mutable state is allocated separately and shared across sessions, this doesn't really make a difference at
+-- the moment since we're also initializing each compilation session with a shared @HscEnv@.
+-- Ultimately this might be used to exert some more control over what modules GHC is allowed to access by using Buck's
+-- deps, or some additional optimization.
+newFinderCache ::
+  ((OneshotState -> OneshotState) -> IO ()) ->
+  OneshotState ->
+  Target ->
+  IO FinderCache
+newFinderCache updateOneshot OneshotState {finder = FinderState {modules, files}} target = do
+  let flushFinderCaches :: UnitEnv -> IO ()
+      flushFinderCaches _ = panic "GHC attempted to flush finder caches, which shouldn't happen in worker mode"
+
+      addToFinderCache :: InstalledModule -> InstalledFindResult -> IO ()
+      addToFinderCache key val = do
+        !() <- evaluate (forceInstalledFindResult val)
+        atomicModifyIORef' modules $ \c ->
+          case (lookupInstalledModuleEnv c key, val) of
+            (Just InstalledFound{}, InstalledNotFound{}) -> (c, ())
+            _ -> (extendInstalledModuleEnv c key val, ())
+
+      lookupFinderCache :: InstalledModule -> IO (Maybe InstalledFindResult)
+      lookupFinderCache key = do
+        c <- readIORef modules
+        let result = lookupInstalledModuleEnv c key
+        case result of
+          Just _ -> cacheHit key
+          Nothing -> cacheMiss key
+        pure $! result
+
+      lookupFileCache :: FilePath -> IO Fingerprint
+      lookupFileCache key = do
+         fc <- readIORef files
+         case LazyMap.lookup key fc of
+           Nothing -> do
+             hash <- getFileHash key
+             atomicModifyIORef' files $ \c -> (LazyMap.insert key hash c, ())
+             return hash
+           Just fp -> return fp
+  return FinderCache {..}
+  where
+    cacheHit m =
+      updateStats \ FinderStats {hits, ..} -> FinderStats {hits = incStat m hits, ..}
+
+    cacheMiss m =
+      updateStats \ FinderStats {misses, ..} -> FinderStats {misses = incStat m misses, ..}
+
+    incStat m = Map.alter (Just . succ . fromMaybe 0) (moduleName m)
+
+    updateStats f =
+      updateOneshot $ modifyStats target \ CacheStats {..} -> CacheStats {finder = f finder, ..}
+
+#else
+
+newFinderCache ::
+  a ->
+  OneshotState ->
+  Target ->
+  IO FinderCache
+newFinderCache _ OneshotState {finder = FinderState {cache}} _ = pure cache
+
+#endif
+
+withHscState :: HscEnv -> (MVar OrigNameCache -> MVar (Maybe LoaderState) -> MVar SymbolMap -> IO a) -> IO (Maybe a)
+withHscState HscEnv {hsc_interp, hsc_NC = NameCache {nsNames}} use =
+#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) || defined(MWB)
+  for hsc_interp \ Interp {interpLoader = Loader {loader_state}, interpLookupSymbolCache} ->
+    liftIO $ use nsNames loader_state interpLookupSymbolCache
+#else
+  for hsc_interp \ Interp {interpLoader = Loader {loader_state}} ->
+    liftIO do
+    symbolCacheVar <- newMVar mempty
+    use nsNames loader_state symbolCacheVar
+#endif
+
+-- | Restore cache parts that depend on the 'Target'.
+setTarget ::
+  ((OneshotState -> OneshotState) -> IO ()) ->
+  OneshotState ->
+  Target ->
+  HscEnv ->
+  IO HscEnv
+setTarget update cache target hsc_env = do
+  -- The Finder cache is already shared by the base session, but with this, we can additionally inject the target file
+  -- for stats collection.
+  -- Only has an effect if the patch that abstracts the 'FinderCache' interface is in GHC, so >= 9.12.
+  if cache.features.finder
+  then restoreFinderCache hsc_env
+  else pure hsc_env
+  where
+    restoreFinderCache e = do
+      hsc_FC <- newFinderCache update cache target
+      pure e {hsc_FC}
+
+-- | Restore cache parts related to oneshot mode.
+restoreOneshotCache :: OneshotState -> HscEnv -> IO HscEnv
+restoreOneshotCache cache hsc_env = do
+  -- If the feature is enabled, restore the EPS.
+  -- This is only relevant for the oneshot worker, but even there the base session should already be sharing the EPS
+  -- across modules.
+  -- Might be removed soon.
+  pure
+    if cache.features.eps
+    then restoreCachedEps hsc_env
+    else hsc_env
+  where
+    restoreCachedEps e = e {hsc_unit_env = e.hsc_unit_env {ue_eps = cache.eps}}
+
+loadState ::
+  ((OneshotState -> OneshotState) -> IO ()) ->
+  Target ->
+  HscEnv ->
+  OneshotState ->
+  IO (OneshotState, (HscEnv, Bool))
+loadState update target hsc_env0 state0 = do
+  result <-
+    if state0.features.enable
+    then do
+      hsc_env1 <- restoreOneshotCache state0 =<< setTarget update state0 target hsc_env0
+      if state0.features.loader
+      then do
+        withHscState hsc_env1 \ nsNames loaderStateVar symbolCacheVar -> do
+          cache1 <- modifyMVar loaderStateVar \ initialLoaderState ->
+            modifyMVar symbolCacheVar \ initialSymbolCache ->
+              (first coerce) <$>
+              modifyMVar nsNames \ names ->
+                loadCache target initialLoaderState (SymbolCache initialSymbolCache) names state0
+          pure (hsc_env1, cache1)
+      else pure (Just (hsc_env1, state0))
+    else pure Nothing
+  let (hsc_env1, cache1) = fromMaybe (hsc_env0, state0 {features = state0.features {loader = False}}) result
+  pure (cache1, (hsc_env1, cache1.features.enable))
+
+storeState ::
+  HscEnv ->
+  Target ->
+  OneshotState ->
+  IO OneshotState
+storeState hsc_env target state = do
+  if state.features.enable
+  then do
+    if state.features.loader
+    then do
+      fromMaybe state . join <$> withHscState hsc_env \ nsNames loaderStateVar symbolCacheVar ->
+        readMVar loaderStateVar >>= traverse \ newLoaderState -> do
+          newSymbols <- readMVar symbolCacheVar
+          newNames <- readMVar nsNames
+          let store = maybe initState (updateState target) state.interpCache
+          store newLoaderState (SymbolCache newSymbols) newNames state
+    else pure state
+  else pure state

From 5732e449f3f1aa241460d00674419344aa1a7009 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Thu, 5 Jun 2025 15:45:58 +0200
Subject: [PATCH 18/27] rename `Cache` to `WorkerState`

Change-Id: I6befaff50f63b66ead6d5b0db755ad0eac96c893
---
 ghc-worker/lib/GhcWorker/CompileResult.hs    |  2 +-
 ghc-worker/lib/GhcWorker/GhcHandler.hs       |  8 +--
 ghc-worker/lib/GhcWorker/Grpc.hs             | 28 ++++----
 ghc-worker/lib/GhcWorker/Instrumentation.hs  | 16 ++---
 ghc-worker/lib/GhcWorker/Run.hs              | 16 ++---
 ghc-worker/test/CompileHptTest.hs            |  8 +--
 ghc-worker/test/TestSetup.hs                 | 10 +--
 instrument/Main.hs                           |  2 +-
 instrument/UI.hs                             |  4 +-
 instrument/UI/Session.hs                     |  2 +-
 instrument/UI/SessionSelector.hs             |  4 +-
 internal/buck-worker-internal.cabal          |  2 +-
 internal/src/Internal/Compile.hs             |  2 +-
 internal/src/Internal/CompileHpt.hs          |  2 +-
 internal/src/Internal/Metadata.hs            | 10 +--
 internal/src/Internal/Session.hs             | 40 +++++------
 internal/src/Internal/{Cache.hs => State.hs} | 75 ++++++++++----------
 17 files changed, 115 insertions(+), 116 deletions(-)
 rename internal/src/Internal/{Cache.hs => State.hs} (70%)

diff --git a/ghc-worker/lib/GhcWorker/CompileResult.hs b/ghc-worker/lib/GhcWorker/CompileResult.hs
index 03c84a97..cdf54cd8 100644
--- a/ghc-worker/lib/GhcWorker/CompileResult.hs
+++ b/ghc-worker/lib/GhcWorker/CompileResult.hs
@@ -3,7 +3,7 @@ module GhcWorker.CompileResult where
 import Data.Foldable (for_)
 import Data.Int (Int32)
 import Internal.AbiHash (AbiHash (..))
-import Internal.Cache (ModuleArtifacts)
+import Internal.State (ModuleArtifacts)
 import Types.BuckArgs (BuckArgs (..))
 
 -- | Right now the 'Maybe' just corresponds to the presence of the CLI argument @--abi-out@ – errors occuring while
diff --git a/ghc-worker/lib/GhcWorker/GhcHandler.hs b/ghc-worker/lib/GhcWorker/GhcHandler.hs
index 3ab30016..0c97b5b8 100644
--- a/ghc-worker/lib/GhcWorker/GhcHandler.hs
+++ b/ghc-worker/lib/GhcWorker/GhcHandler.hs
@@ -15,7 +15,7 @@ import GHC.Driver.Monad (modifySession)
 import GhcWorker.CompileResult (CompileResult (..), writeCloseOutput, writeResult)
 import GhcWorker.Instrumentation (Hooks (..), InstrumentedHandler (..))
 import Internal.AbiHash (AbiHash (..), showAbiHash)
-import Internal.Cache (Cache (..), ModuleArtifacts (..))
+import Internal.State (WorkerState (..), ModuleArtifacts (..))
 import Internal.Compile (compileModuleWithDepsInEps)
 import Internal.CompileHpt (compileModuleWithDepsInHpt)
 import Internal.Log (LogName (..), dbg, logFlush, newLog)
@@ -126,15 +126,15 @@ dispatch lock workerMode hooks env args =
 ghcHandler ::
   -- | first req lock hack
   TVar LockState ->
-  MVar Cache ->
+  MVar WorkerState ->
   WorkerMode ->
   InstrumentedHandler
-ghcHandler lock cache workerMode =
+ghcHandler lock state workerMode =
   InstrumentedHandler \ hooks -> GrpcHandler \ commandEnv argv -> do
     buckArgs <- either (throwIO . userError) pure (parseBuckArgs commandEnv argv)
     args <- toGhcArgs buckArgs
     log <- newLog True
-    let env = Env {log, cache, args}
+    let env = Env {log, state, args}
     onException
       do
         (result, target) <- dispatch lock workerMode hooks env buckArgs
diff --git a/ghc-worker/lib/GhcWorker/Grpc.hs b/ghc-worker/lib/GhcWorker/Grpc.hs
index 5629955a..ea56a320 100644
--- a/ghc-worker/lib/GhcWorker/Grpc.hs
+++ b/ghc-worker/lib/GhcWorker/Grpc.hs
@@ -7,7 +7,7 @@ import Control.Monad (forever)
 import Data.Map.Strict qualified as Map
 import Data.Text qualified as Text
 import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats)
-import Internal.Cache (Cache (..), Options (..))
+import Internal.State (WorkerState (..), Options (..))
 import Network.GRPC.Common (NextElem (..))
 import Network.GRPC.Common.Protobuf (Proto, defMessage, (&), (.~))
 import Network.GRPC.Server.Protobuf (ProtobufMethodsOf)
@@ -22,7 +22,7 @@ import Proto.Instrument (Instrument)
 import Proto.Instrument_Fields qualified as Instr
 
 -- | Fetch statistics about the current state of the RTS for instrumentation.
-mkStats :: Cache -> IO (Proto Instr.Stats)
+mkStats :: WorkerState -> IO (Proto Instr.Stats)
 mkStats _ = do
   s <- getRTSStats
   pure $
@@ -36,14 +36,14 @@ mkStats _ = do
 -- | Implementation of a streaming grapesy handler that sends instrumentation statistics pulled from the provided
 -- channel to the client.
 notifyMe ::
-  MVar Cache ->
+  MVar WorkerState ->
   Chan (Proto Instr.Event) ->
   (NextElem (Proto Instr.Event) -> IO ()) ->
   IO ()
-notifyMe cacheVar chan callback = do
-  cache <- readMVar cacheVar
+notifyMe stateVar chan callback = do
+  state <- readMVar stateVar
   myChan <- dupChan chan
-  stats <- mkStats cache
+  stats <- mkStats state
   callback $ NextElem $
     defMessage
       & Instr.stats .~ stats
@@ -53,12 +53,12 @@ notifyMe cacheVar chan callback = do
 
 -- | Set the options for the server.
 setOptions ::
-  MVar Cache ->
+  MVar WorkerState ->
   Proto  Instr.Options ->
   IO (Proto Instr.Empty)
-setOptions cacheVar opts = do
-  modifyMVar_ cacheVar $ \cache ->
-    pure cache {
+setOptions stateVar opts = do
+  modifyMVar_ stateVar $ \state ->
+    pure state {
       options = Options {
         extraGhcOptions = Text.unpack opts.extraGhcOptions
       }
@@ -68,9 +68,9 @@ setOptions cacheVar opts = do
 -- | A grapesy server that streams instrumentation data from the provided channel.
 instrumentMethods ::
   Chan (Proto Instr.Event) ->
-  MVar Cache ->
+  MVar WorkerState ->
   Methods IO (ProtobufMethodsOf Instrument)
-instrumentMethods chan cacheVar =
+instrumentMethods chan stateVar =
   simpleMethods
-    (mkServerStreaming (const (notifyMe cacheVar chan)))
-    (mkNonStreaming (setOptions cacheVar))
+    (mkServerStreaming (const (notifyMe stateVar chan)))
+    (mkNonStreaming (setOptions stateVar))
diff --git a/ghc-worker/lib/GhcWorker/Instrumentation.hs b/ghc-worker/lib/GhcWorker/Instrumentation.hs
index 935edc8f..93d257b6 100644
--- a/ghc-worker/lib/GhcWorker/Instrumentation.hs
+++ b/ghc-worker/lib/GhcWorker/Instrumentation.hs
@@ -8,7 +8,7 @@ import Data.Foldable (traverse_)
 import Data.Int (Int32)
 import Data.Text qualified as Text
 import GhcWorker.Grpc (mkStats)
-import Internal.Cache (Cache)
+import Internal.State (WorkerState)
 import Internal.Log (dbg)
 import Network.GRPC.Common.Protobuf (Proto, defMessage, (&), (.~))
 import Prelude hiding (log)
@@ -90,15 +90,15 @@ messageCompileEnd target exitCode err =
 withInstrumentation ::
   Chan (Proto Instr.Event) ->
   MVar WorkerStatus ->
-  MVar Cache ->
+  MVar WorkerState ->
   InstrumentedHandler ->
   GrpcHandler
-withInstrumentation instrChan status cacheVar handler =
+withInstrumentation instrChan status stateVar handler =
   GrpcHandler \ commandEnv argv -> do
-    cache <- readMVar cacheVar
+    state <- readMVar stateVar
     bracket_ (startJob status) (finishJob status) do
       result <- (handler.create hooks).run commandEnv argv
-      stats <- mkStats cache
+      stats <- mkStats state
       writeChan instrChan (defMessage & Instr.stats .~ stats)
       pure result
   where
@@ -128,9 +128,9 @@ withInstrumentation instrChan status cacheVar handler =
 toGrpcHandler ::
   InstrumentedHandler ->
   MVar WorkerStatus ->
-  MVar Cache ->
+  MVar WorkerState ->
   Maybe (Chan (Proto Instr.Event)) ->
   GrpcHandler
-toGrpcHandler createHandler status cacheVar = \case
+toGrpcHandler createHandler status stateVar = \case
   Nothing -> createHandler.create hooksNoop
-  Just instrChan -> withInstrumentation instrChan status cacheVar createHandler
+  Just instrChan -> withInstrumentation instrChan status stateVar createHandler
diff --git a/ghc-worker/lib/GhcWorker/Run.hs b/ghc-worker/lib/GhcWorker/Run.hs
index f80eecf0..ebd5b9a9 100644
--- a/ghc-worker/lib/GhcWorker/Run.hs
+++ b/ghc-worker/lib/GhcWorker/Run.hs
@@ -10,7 +10,7 @@ import GhcWorker.GhcHandler (LockState (..), ghcHandler)
 import GhcWorker.Grpc (instrumentMethods)
 import GhcWorker.Instrumentation (WorkerStatus (..), toGrpcHandler)
 import GhcWorker.Orchestration (CreateMethods (..), runCentralGhcSpawned)
-import Internal.Cache (Cache (..), emptyCache, emptyCacheWith)
+import Internal.State (WorkerState (..), newState, newStateWith)
 import Internal.State.Oneshot (OneshotCacheFeatures (..))
 import Network.GRPC.Common.Protobuf (Proto)
 import Network.GRPC.Server.Protobuf (ProtobufMethodsOf)
@@ -55,15 +55,15 @@ parseOptions =
 -- events to a client.
 --
 -- Returns the channel so that a GHC server can use it to send events.
-createInstrumentMethods :: MVar Cache -> IO (Chan (Proto Instr.Event), Methods IO (ProtobufMethodsOf Instrument))
-createInstrumentMethods cacheVar = do
+createInstrumentMethods :: MVar WorkerState -> IO (Chan (Proto Instr.Event), Methods IO (ProtobufMethodsOf Instrument))
+createInstrumentMethods stateVar = do
   instrChan <- newChan
-  pure (instrChan, instrumentMethods instrChan cacheVar)
+  pure (instrChan, instrumentMethods instrChan stateVar)
 
 -- | Construct a gRPC server handler for the main part of the persistent worker.
 createGhcMethods ::
   TVar LockState ->
-  MVar Cache ->
+  MVar WorkerState ->
   WorkerMode ->
   MVar WorkerStatus ->
   Maybe (Chan (Proto Instr.Event)) ->
@@ -77,15 +77,15 @@ runWorker CliOptions {workerMode, serve} = do
   cache <-
     case workerMode of
       WorkerMakeMode ->
-        emptyCacheWith OneshotCacheFeatures {
+        newStateWith OneshotCacheFeatures {
           loader = False,
           enable = True,
           names = False,
           finder = False,
           eps = False
         }
-      WorkerOneshotMode -> emptyCache True
-  lock <- newTVarIO LockStart 
+      WorkerOneshotMode -> newState True
+  lock <- newTVarIO LockStart
   status <- newMVar WorkerStatus {active = 0}
   let
     methods = CreateMethods {
diff --git a/ghc-worker/test/CompileHptTest.hs b/ghc-worker/test/CompileHptTest.hs
index 00d695bd..7ba5bbaf 100644
--- a/ghc-worker/test/CompileHptTest.hs
+++ b/ghc-worker/test/CompileHptTest.hs
@@ -39,12 +39,12 @@ unitFlags args HscEnv {hsc_logger, hsc_dflags = dflags0} = do
   pure dflags
 
 stepMetadata :: Conf -> Unit -> [Unit] -> IO ()
-stepMetadata Conf {cache, tmp, args0} unit deps = do
+stepMetadata Conf {state, tmp, args0} unit deps = do
   log <- newLog True
   createDirectoryIfMissing False sessionTmpDir
   names <- listDirectory unit.dir
   let srcs = [unit.dir </> name | name <- names, takeExtension name == ".hs"]
-      env = Env {log, cache, args = args srcs}
+      env = Env {log, state, args = args srcs}
   dbgp (text ">>> metadata for" <+> ppr unit.uid)
   success <- computeMetadata env
   unless success do
@@ -71,9 +71,9 @@ stepMetadata Conf {cache, tmp, args0} unit deps = do
     sessionTmpDir = tmp </> "tmp" </> unit.name
 
 stepCompile :: Conf -> Module -> IO ()
-stepCompile Conf {cache, tmp, args0} Module {unit, src} = do
+stepCompile Conf {state, tmp, args0} Module {unit, src} = do
   log <- newLog True
-  let env = Env {log, cache, args}
+  let env = Env {log, state, args}
   liftIO $ createDirectoryIfMissing False sessionTmpDir
   result <- liftIO $ withGhcMhu env \ _ target -> do
     dbg ""
diff --git a/ghc-worker/test/TestSetup.hs b/ghc-worker/test/TestSetup.hs
index 93715f9a..1987debb 100644
--- a/ghc-worker/test/TestSetup.hs
+++ b/ghc-worker/test/TestSetup.hs
@@ -6,8 +6,8 @@ import Data.Functor ((<&>))
 import Data.List.NonEmpty (NonEmpty (..))
 import Data.Traversable (for)
 import GHC.Unit (UnitId, stringToUnitId, unitIdString)
-import Internal.Cache (Cache (..), emptyCacheWith)
 import Internal.Log (dbg)
+import Internal.State (WorkerState (..), newStateWith)
 import Internal.State.Oneshot (OneshotCacheFeatures (..))
 import Prelude hiding (log)
 import System.Directory (createDirectoryIfMissing, listDirectory, withCurrentDirectory)
@@ -23,8 +23,8 @@ data Conf =
     -- | Root directory of the test in @/tmp@.
     tmp :: FilePath,
 
-    -- | The worker cache.
-    cache :: MVar Cache,
+    -- | The worker state.
+    state :: MVar WorkerState,
 
     -- | The base cli args used for all modules.
     args0 :: Args,
@@ -203,7 +203,7 @@ withProject mkTargets use =
     withCurrentDirectory tmp do
       for_ @[] ["src", "tmp", "out"] \ dir ->
         createDirectoryIfMissing False (tmp </> dir)
-      cache <- emptyCacheWith OneshotCacheFeatures {
+      state <- newStateWith OneshotCacheFeatures {
         loader = False,
         enable = True,
         names = False,
@@ -215,7 +215,7 @@ withProject mkTargets use =
         [d] -> "lib" </> d </> "lib"
         ds -> error ("weird GHC lib dir contains /= 1 entries: " ++ show ds)
       let topdir = ghcDir </> libPath
-          conf = Conf {tmp, cache, args0 = baseArgs topdir tmp, ..}
+          conf = Conf {tmp, state, args0 = baseArgs topdir tmp, ..}
       targets <- mkTargets conf
       units <- for targets \ unit -> do
         let dir = tmp </> "src" </> unit.name
diff --git a/instrument/Main.hs b/instrument/Main.hs
index 1d376ef2..69dcc93b 100644
--- a/instrument/Main.hs
+++ b/instrument/Main.hs
@@ -10,7 +10,7 @@ import Data.Maybe (fromMaybe)
 import Data.Text qualified as Text
 import Data.Time (getCurrentTime)
 import Graphics.Vty (Vty (shutdown))
-import Internal.Cache (Options (..))
+import Internal.State (Options (..))
 import Network.GRPC.Client (Server (ServerUnix), rpc, withConnection)
 import Network.GRPC.Client.StreamType.IO (nonStreaming, serverStreaming)
 import Network.GRPC.Common (def)
diff --git a/instrument/UI.hs b/instrument/UI.hs
index 88c3b4e0..f401e03f 100644
--- a/instrument/UI.hs
+++ b/instrument/UI.hs
@@ -20,7 +20,7 @@ import Data.Text qualified as Text
 import Data.Time (UTCTime (..), fromGregorian)
 import Graphics.Vty qualified as V
 import Graphics.Vty.Attributes.Color
-import Internal.Cache (Options (..), defaultOptions)
+import Internal.State (Options (..), defaultOptions)
 import Lens.Micro.Platform (Lens', lens, makeLenses, packed, use, zoom, (.=), _2)
 import UI.ActiveTasks qualified as ActiveTasks
 import UI.ModuleSelector qualified as ModuleSelector
@@ -172,4 +172,4 @@ app =
             , (listSelectedFocusedAttr, brightWhite `on` blue)
             ]
     , appChooseCursor = showFirstCursor
-    }
\ No newline at end of file
+    }
diff --git a/instrument/UI/Session.hs b/instrument/UI/Session.hs
index dad3d5b5..723b9e62 100644
--- a/instrument/UI/Session.hs
+++ b/instrument/UI/Session.hs
@@ -9,7 +9,7 @@ import Control.Monad.IO.Class (liftIO)
 import Data.Map qualified as Map
 import Data.Text qualified as Text
 import Data.Time (UTCTime, diffUTCTime, getCurrentTime, nominalDiffTimeToSeconds)
-import Internal.Cache (Options (..))
+import Internal.State (Options (..))
 import Lens.Micro.Platform (each, filtered, makeLenses, modifying, use, zoom)
 import Network.GRPC.Common.Protobuf (Proto, (^.))
 import Proto.Instrument qualified as Instr
diff --git a/instrument/UI/SessionSelector.hs b/instrument/UI/SessionSelector.hs
index 860eb7af..4f37b16d 100644
--- a/instrument/UI/SessionSelector.hs
+++ b/instrument/UI/SessionSelector.hs
@@ -7,7 +7,7 @@ import Control.Monad.IO.Class (liftIO)
 import Data.Sequence qualified as Seq
 import Data.Time (UTCTime, getCurrentTime)
 import Data.Time.Format.ISO8601 (iso8601Show)
-import Internal.Cache (Options)
+import Internal.State (Options)
 import Lens.Micro.Platform (Traversal', each, filtered, modifying, preuse, zoom, (.=), _2)
 import UI.Session qualified as Session
 import UI.Types (Name (SessionSelector))
@@ -66,4 +66,4 @@ handleEvent (StartSession sid start) = do
 handleEvent (EndSession sid) = do
   end <- liftIO getCurrentTime
   modifying (sessionLens sid . Session.sesEndTime) (const $ Just end)
-handleEvent (SessionEvent sid evt) = zoom (sessionLens sid) (Session.handleEvent evt)
\ No newline at end of file
+handleEvent (SessionEvent sid evt) = zoom (sessionLens sid) (Session.handleEvent evt)
diff --git a/internal/buck-worker-internal.cabal b/internal/buck-worker-internal.cabal
index 5684453d..88585abd 100644
--- a/internal/buck-worker-internal.cabal
+++ b/internal/buck-worker-internal.cabal
@@ -22,7 +22,6 @@ library
     ghc-options:      -Wall
     hs-source-dirs:   src
     exposed-modules:  Internal.AbiHash
-                      Internal.Cache
                       Internal.Compile
                       Internal.CompileHpt
                       Internal.Debug
@@ -32,6 +31,7 @@ library
                       Internal.MakeFile.JSON
                       Internal.Metadata
                       Internal.Session
+                      Internal.State
                       Internal.State.Make
                       Internal.State.Oneshot
                       Internal.State.Stats
diff --git a/internal/src/Internal/Compile.hs b/internal/src/Internal/Compile.hs
index d43bc260..1419ecde 100644
--- a/internal/src/Internal/Compile.hs
+++ b/internal/src/Internal/Compile.hs
@@ -28,7 +28,7 @@ import GHC.Types.SourceFile (HscSource)
 import GHC.Unit.Home.ModInfo (HomeModLinkable (..))
 import GHC.Utils.Monad (MonadIO (..), unlessM)
 import GHC.Utils.Panic (panic, throwGhcExceptionIO)
-import Internal.Cache (ModuleArtifacts (..))
+import Internal.State (ModuleArtifacts (..))
 import System.Directory (doesFileExist)
 import Types.State (Target (Target))
 
diff --git a/internal/src/Internal/CompileHpt.hs b/internal/src/Internal/CompileHpt.hs
index 5ce58aad..5b89b357 100644
--- a/internal/src/Internal/CompileHpt.hs
+++ b/internal/src/Internal/CompileHpt.hs
@@ -13,7 +13,7 @@ import GHC.Unit.Env (addHomeModInfoToHug, ue_unsafeHomeUnit)
 import GHC.Unit.Home.ModInfo (HomeModInfo (..), HomeModLinkable (..))
 import GHC.Utils.Monad (MonadIO (..))
 import GHC.Utils.TmpFs (TmpFs, cleanCurrentModuleTempFiles, keepCurrentModuleTempFiles)
-import Internal.Cache (ModuleArtifacts (..))
+import Internal.State (ModuleArtifacts (..))
 import Internal.Error (eitherMessages)
 import Types.State (Target (Target))
 
diff --git a/internal/src/Internal/Metadata.hs b/internal/src/Internal/Metadata.hs
index 8738042a..612c0900 100644
--- a/internal/src/Internal/Metadata.hs
+++ b/internal/src/Internal/Metadata.hs
@@ -12,7 +12,7 @@ import GHC.Runtime.Loader (initializeSessionPlugins)
 import GHC.Unit (HomeUnit, UnitDatabase, UnitId, UnitState, initUnits)
 import GHC.Unit.Env (HomeUnitEnv (..), UnitEnv (..), unitEnv_insert, unitEnv_keys, updateHug)
 import GHC.Unit.Home.ModInfo (emptyHomePackageTable)
-import Internal.Cache (Cache (..), updateMakeStateVar)
+import Internal.State (WorkerState (..), updateMakeStateVar)
 import Internal.MakeFile (doMkDependHS)
 import Internal.Session (Env (..), runSession, withDynFlags)
 import Internal.State.Make (insertUnitEnv, loadState, storeModuleGraph)
@@ -75,15 +75,15 @@ addHomeUnit dflags = do
 -- the home unit in order to replicate what GHC does in @initMulti@.
 prepareMetadataSession :: Env -> DynFlags -> Ghc ()
 prepareMetadataSession env dflags = do
-  cache <- liftIO $ readMVar env.cache
-  modifySessionM \ hsc_env -> liftIO (loadState env.log hsc_env cache.make)
+  state <- liftIO $ readMVar env.state
+  modifySessionM \ hsc_env -> liftIO (loadState env.log hsc_env state.make)
   unit <- addHomeUnit dflags
   setActiveUnit unit
   storeNewUnit
   where
     setActiveUnit unit = modifySession (hscUpdateLoggerFlags . hscSetActiveUnitId unit)
 
-    storeNewUnit = withSession \ hsc_env -> liftIO $ updateMakeStateVar env.cache (insertUnitEnv hsc_env)
+    storeNewUnit = withSession \ hsc_env -> liftIO $ updateMakeStateVar env.state (insertUnitEnv hsc_env)
 
 -- | Run 'doMkDependHS' to write the metadata JSON file and exfiltrate the module graph.
 -- We need to use a temporary session because 'doMkDependHS' uses some custom settings that we don't want to leak,
@@ -109,6 +109,6 @@ computeMetadata env = do
   res <- fmap isJust $ runSession True env $ withDynFlags env \ dflags srcs -> do
     prepareMetadataSession env dflags
     module_graph <- writeMetadata (fst <$> srcs)
-    liftIO $ updateMakeStateVar env.cache (storeModuleGraph module_graph)
+    liftIO $ updateMakeStateVar env.state (storeModuleGraph module_graph)
     pure (Just ())
   res <$ logMemStats "after metadata" env.log
diff --git a/internal/src/Internal/Session.hs b/internal/src/Internal/Session.hs
index 970b7921..d224430c 100644
--- a/internal/src/Internal/Session.hs
+++ b/internal/src/Internal/Session.hs
@@ -39,7 +39,7 @@ import GHC.Types.SrcLoc (Located, mkGeneralLocated, unLoc)
 import GHC.Utils.Logger (Logger, getLogger, setLogFlags)
 import GHC.Utils.Panic (GhcException (UsageError), panic, throwGhcException)
 import GHC.Utils.TmpFs (TempDir (..), cleanTempDirs, cleanTempFiles, initTmpFs)
-import Internal.Cache (BinPath (..), Cache (..), ModuleArtifacts, Options (..), withCacheMake, withCacheOneshot)
+import Internal.State (BinPath (..), WorkerState (..), ModuleArtifacts, Options (..), withCacheMake, withCacheOneshot)
 import Internal.Error (handleExceptions)
 import Internal.Log (Log (..), logToState)
 import Internal.State.Oneshot (OneshotCacheFeatures (..), OneshotState (..))
@@ -48,14 +48,14 @@ import System.Environment (setEnv)
 import Types.Args (Args (..))
 import Types.State (Target (Target))
 
--- | Worker state.
+-- | Data used by a single worker request session, consisting of a logger, shared state, and request arguments.
 data Env =
   Env {
     -- | Logger used to receive messages from GHC and relay them to Buck.
     log :: MVar Log,
 
-    -- | Parts of @HscEnv@ we share between sessions.
-    cache :: MVar Cache,
+    -- | The entirety of the persistent state of a worker thats's shared across sessions.
+    state :: MVar WorkerState,
 
     -- | Preprocessed command line args from Buck.
     args :: Args
@@ -64,7 +64,7 @@ data Env =
 -- | Add all the directories passed by Buck in @--bin-path@ options to the global @$PATH@.
 -- Although Buck intends these to be module specific, all subsequent compile jobs will see all previous jobs' entries,
 -- since we only have one process environment.
-setupPath :: Args -> Cache -> IO Cache
+setupPath :: Args -> WorkerState -> IO WorkerState
 setupPath args old = do
   setEnv "PATH" (intercalate ":" (toList path.extra ++ maybeToList path.initial))
   pure new
@@ -134,8 +134,8 @@ withDynFlags :: Env -> (DynFlags -> [(String, Maybe Phase)] -> Ghc a) -> [Locate
 withDynFlags env prog argv = do
   let !log = env.log
   pushLogHookM (const (logToState log))
-  cache <- liftIO $ readMVar env.cache
-  (dflags0, logger, fileish_args, dynamicFlagWarnings) <- parseFlags (argv ++ map instrumentLocation (words cache.options.extraGhcOptions))
+  state <- liftIO $ readMVar env.state
+  (dflags0, logger, fileish_args, dynamicFlagWarnings) <- parseFlags (argv ++ map instrumentLocation (words state.options.extraGhcOptions))
   result <- prettyPrintGhcErrors logger do
     (dflags, srcs) <- initDynFlags dflags0 logger fileish_args dynamicFlagWarnings
     prog dflags srcs
@@ -157,16 +157,16 @@ withGhcInSession env prog =
 --
 -- When reusing the base session, create a new @TmpFs@ to avoid keeping old entries around after Buck deletes the
 -- directories.
-ensureSession :: Bool -> MVar Cache -> Args -> IO HscEnv
-ensureSession reuse cacheVar args =
-  modifyMVar cacheVar \ cache -> do
-    if cache.oneshot.features.enable && reuse
+ensureSession :: Bool -> MVar WorkerState -> Args -> IO HscEnv
+ensureSession reuse stateVar args =
+  modifyMVar stateVar \ state -> do
+    if state.oneshot.features.enable && reuse
     then do
-      newEnv <- maybe (initHscEnv args.topdir) prepReused cache.baseSession
-      pure (cache {baseSession = Just newEnv}, newEnv)
+      newEnv <- maybe (initHscEnv args.topdir) prepReused state.baseSession
+      pure (state {baseSession = Just newEnv}, newEnv)
     else do
       newEnv <- initHscEnv args.topdir
-      pure (cache, newEnv)
+      pure (state, newEnv)
   where
     prepReused hsc_env = do
       hsc_tmpfs <- initTmpFs
@@ -177,9 +177,9 @@ ensureSession reuse cacheVar args =
 --
 -- Delete all temporary files on completion.
 runSession :: Bool -> Env -> ([Located String] -> Ghc (Maybe a)) -> IO (Maybe a)
-runSession reuse Env {log, args, cache} prog = do
-  modifyMVar_ cache (setupPath args)
-  hsc_env <- ensureSession reuse cache args
+runSession reuse Env {log, args, state} prog = do
+  modifyMVar_ state (setupPath args)
+  hsc_env <- ensureSession reuse state args
   session <- Session <$> newIORef hsc_env
   finally (run session) (cleanup session)
   where
@@ -222,7 +222,7 @@ withGhc env =
   withGhcUsingCache cacheHandler env
   where
     cacheHandler target prog = do
-      result <- withCacheOneshot env.log env.args.workerTargetId env.cache target do
+      result <- withCacheOneshot env.log env.args.workerTargetId env.state target do
         res <- prog
         pure do
           a <- res
@@ -233,7 +233,7 @@ withGhc env =
 -- Return the interface and bytecode.
 withGhcDefault :: Env -> (Target -> Ghc (Maybe (Maybe ModuleArtifacts, a))) -> IO (Maybe (Maybe ModuleArtifacts, a))
 withGhcDefault env =
-  withGhcUsingCache (withCacheOneshot env.log env.args.workerTargetId env.cache) env
+  withGhcUsingCache (withCacheOneshot env.log env.args.workerTargetId env.state) env
 
 -- | Command line args that have to be stored in the current home unit env.
 -- These are specified as a single program argument with their option argument, without whitespace in between.
@@ -323,7 +323,7 @@ withGhcMhu env f =
   withGhcUsingCacheMhu cacheHandler env f
   where
     cacheHandler _ prog = do
-      result <- withCacheMake env.log env.cache do
+      result <- withCacheMake env.log env.state do
         res <- prog
         pure do
           a <- res
diff --git a/internal/src/Internal/Cache.hs b/internal/src/Internal/State.hs
similarity index 70%
rename from internal/src/Internal/Cache.hs
rename to internal/src/Internal/State.hs
index 8ab6b768..33457e4b 100644
--- a/internal/src/Internal/Cache.hs
+++ b/internal/src/Internal/State.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE CPP, NoFieldSelectors #-}
 
-module Internal.Cache where
+module Internal.State where
 
 import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar)
 import Control.Monad.IO.Class (MonadIO, liftIO)
@@ -38,8 +38,8 @@ data BinPath =
   }
   deriving stock (Eq, Show)
 
-data Cache =
-  Cache {
+data WorkerState =
+  WorkerState {
     path :: BinPath,
     baseSession :: Maybe HscEnv,
     options :: Options,
@@ -52,11 +52,11 @@ data Options =
     extraGhcOptions :: String
   }
 
-emptyCacheWith :: OneshotCacheFeatures -> IO (MVar Cache)
-emptyCacheWith features = do
+newStateWith :: OneshotCacheFeatures -> IO (MVar WorkerState)
+newStateWith features = do
   initialPath <- lookupEnv "PATH"
   oneshot <- newOneshotStateWith features
-  newMVar Cache {
+  newMVar WorkerState {
     path = BinPath {
       initial = initialPath,
       extra = mempty
@@ -71,21 +71,20 @@ emptyCacheWith features = do
     oneshot
   }
 
-emptyCache :: Bool -> IO (MVar Cache)
-emptyCache enable = do
-  emptyCacheWith newOneshotCacheFeatures {enable}
+newState :: Bool -> IO (MVar WorkerState)
+newState enable = newStateWith newOneshotCacheFeatures {enable}
 
--- | Update the 'MakeState' field in the 'Cache'.
-updateMakeState :: (MakeState -> MakeState) -> Cache -> Cache
-updateMakeState f cache = cache {make = f cache.make}
+-- | Update the 'MakeState' field in the 'WorkerState'.
+updateMakeState :: (MakeState -> MakeState) -> WorkerState -> WorkerState
+updateMakeState f state = state {make = f state.make}
 
-updateMakeStateVar :: MVar Cache -> (MakeState -> MakeState) -> IO ()
+updateMakeStateVar :: MVar WorkerState -> (MakeState -> MakeState) -> IO ()
 updateMakeStateVar var f = modifyMVar_ var (pure . updateMakeState f)
 
-updateOneshotState :: (OneshotState -> OneshotState) -> Cache -> Cache
-updateOneshotState f cache = cache {oneshot = f cache.oneshot}
+updateOneshotState :: (OneshotState -> OneshotState) -> WorkerState -> WorkerState
+updateOneshotState f state = state {oneshot = f state.oneshot}
 
-updateOneshotStateVar :: MVar Cache -> (OneshotState -> OneshotState) -> IO ()
+updateOneshotStateVar :: MVar WorkerState -> (OneshotState -> OneshotState) -> IO ()
 updateOneshotStateVar var f = modifyMVar_ var (pure . updateOneshotState f)
 
 defaultOptions :: Options
@@ -101,16 +100,16 @@ report ::
   -- | A description of the current worker process.
   Maybe TargetId ->
   Target ->
-  Cache ->
+  WorkerState ->
   m ()
-report logVar workerId target cache = do
-  Stats.report logVar workerId target (if cache.oneshot.features.enable then Just cache.oneshot.stats else Nothing)
+report logVar workerId target state = do
+  Stats.report logVar workerId target (if state.oneshot.features.enable then Just state.oneshot.stats else Nothing)
 
 -- | Merge the given module graph into the cached graph.
 -- This is used by the make mode worker after the metadata step has computed the module graph.
-updateModuleGraph :: MVar Cache -> ModuleGraph -> IO ()
-updateModuleGraph cacheVar new =
-  updateMakeStateVar cacheVar (Make.storeModuleGraph new)
+updateModuleGraph :: MVar WorkerState -> ModuleGraph -> IO ()
+updateModuleGraph stateVar new =
+  updateMakeStateVar stateVar (Make.storeModuleGraph new)
 
 finalizeCache ::
   MVar Log ->
@@ -119,8 +118,8 @@ finalizeCache ::
   HscEnv ->
   Target ->
   Maybe ModuleArtifacts ->
-  Cache ->
-  IO Cache
+  WorkerState ->
+  IO WorkerState
 finalizeCache logVar workerId hsc_env target _ cache0 = do
   oneshot <- Oneshot.storeState hsc_env target cache0.oneshot
   let cache1 = cache0 {oneshot}
@@ -138,21 +137,21 @@ withCacheOneshot ::
   MVar Log ->
   -- | A description of the current worker process.
   Maybe TargetId ->
-  MVar Cache ->
+  MVar WorkerState ->
   Target ->
   Ghc (Maybe (Maybe ModuleArtifacts, a)) ->
   Ghc (Maybe (Maybe ModuleArtifacts, a))
-withCacheOneshot logVar workerId cacheVar target prog = do
-  _ <- withSessionM \ hsc_env -> modifyMVar cacheVar \ cache -> do
-    (oneshot, result) <- Oneshot.loadState (updateOneshotStateVar cacheVar) target hsc_env cache.oneshot
-    pure (cache {oneshot}, result)
+withCacheOneshot logVar workerId stateVar target prog = do
+  _ <- withSessionM \ hsc_env -> modifyMVar stateVar \ state -> do
+    (oneshot, result) <- Oneshot.loadState (updateOneshotStateVar stateVar) target hsc_env state.oneshot
+    pure (state {oneshot}, result)
   result <- prog
   finalize (fst =<< result)
   pure result
   where
     finalize art =
       withSession \ hsc_env ->
-        liftIO (modifyMVar_ cacheVar (finalizeCache logVar workerId hsc_env target art))
+        liftIO (modifyMVar_ stateVar (finalizeCache logVar workerId hsc_env target art))
 
 -- | This reduced version of 'withCache' is tailored specifically to make mode, only restoring the HUG, module graph and
 -- interpreter state from the cache, since those are the only two components modified by the worker that aren't already
@@ -162,19 +161,19 @@ withCacheOneshot logVar workerId cacheVar target prog = do
 -- shared manually in oneshot mode, so this variant will be improved more deliberately.
 withCacheMake ::
   MVar Log ->
-  MVar Cache ->
+  MVar WorkerState ->
   Ghc (Maybe (Maybe ModuleArtifacts, a)) ->
   Ghc (Maybe (Maybe ModuleArtifacts, a))
-withCacheMake logVar cacheVar prog = do
+withCacheMake logVar stateVar prog = do
   modifySessionM restore
   prog <* withSession store
   where
     restore hsc_env =
-      liftIO $ modifyMVar cacheVar \ cache -> do
-        (make, hsc_env1) <- Make.loadStateCompile logVar hsc_env cache.make
-        pure (cache {make}, hsc_env1)
+      liftIO $ modifyMVar stateVar \ state -> do
+        (make, hsc_env1) <- Make.loadStateCompile logVar hsc_env state.make
+        pure (state {make}, hsc_env1)
 
     store hsc_env =
-      liftIO $ modifyMVar_ cacheVar \ cache -> do
-        make <- Make.storeState logVar hsc_env cache.make
-        pure cache {make}
+      liftIO $ modifyMVar_ stateVar \ state -> do
+        make <- Make.storeState logVar hsc_env state.make
+        pure state {make}

From 73387113e5f1182be628c179d4ee7d06fac4ad20 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Thu, 15 May 2025 13:48:33 +0200
Subject: [PATCH 19/27] Add CLI option to enable instrumentation

Change-Id: Ie0f58102f5084cc6562f392bab1a7cba3d71555d
---
 ghc-worker/lib/GhcWorker/Orchestration.hs | 33 ++++++++++++++---------
 ghc-worker/lib/GhcWorker/Run.hs           | 16 ++++++-----
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/ghc-worker/lib/GhcWorker/Orchestration.hs b/ghc-worker/lib/GhcWorker/Orchestration.hs
index b9573130..6352aefa 100644
--- a/ghc-worker/lib/GhcWorker/Orchestration.hs
+++ b/ghc-worker/lib/GhcWorker/Orchestration.hs
@@ -42,10 +42,14 @@ import Types.Orchestration (
 -- The 'Instrument' component is intended to be optional.
 data CreateMethods where
   CreateMethods :: {
-    createInstrumentation :: IO (instr, Methods IO (ProtobufMethodsOf Instrument)),
-    createGhc :: Maybe instr -> IO (Methods IO (ProtobufMethodsOf Worker))
+    createInstrumentation :: IO (instrumentSocket, Methods IO (ProtobufMethodsOf Instrument)),
+    createGhc :: Maybe instrumentSocket -> IO (Methods IO (ProtobufMethodsOf Worker))
   } -> CreateMethods
 
+newtype FeatureInstrument =
+  FeatureInstrument { flag :: Bool }
+  deriving stock (Eq, Show)
+
 -- | Start a gRPC server that dispatches requests to GHC handlers.
 runLocalGhc ::
   CreateMethods ->
@@ -54,10 +58,10 @@ runLocalGhc ::
   IO ()
 runLocalGhc CreateMethods {..} socket minstr = do
   dbg ("Starting ghc server on " ++ socket.path)
-  instrResource <- for minstr \instr -> do
-    dbg ("Instrumentation info available on " ++ instr.path)
+  instrResource <- for minstr \instrumentSocket -> do
+    dbg ("Instrumentation info available on " ++ instrumentSocket.path)
     (resource, methods) <- createInstrumentation
-    _instrThread <- async $ runServerWithHandlers def (grpcServerConfig instr.path) (fromMethods methods)
+    _instrThread <- async $ runServerWithHandlers def (grpcServerConfig instrumentSocket.path) (fromMethods methods)
     pure resource
   methods <- createGhc instrResource
   runServerWithHandlers def (grpcServerConfig socket.path) (fromMethods methods)
@@ -69,8 +73,8 @@ runCentralGhc ::
   ServerSocketPath ->
   Maybe InstrumentSocketPath ->
   IO ()
-runCentralGhc mode discovery socket instr =
-  finally (runLocalGhc mode socket instr) do
+runCentralGhc mode discovery socket instrumentSocket =
+  finally (runLocalGhc mode socket instrumentSocket) do
     dbg ("Shutting down ghc server on " ++ socket.path)
     removeFile discovery.path
 
@@ -163,11 +167,14 @@ waitForCentralGhc proc socket = do
     dbg "Spawned process for the GHC server exited after starting up."
 
 -- | Run a GHC server synchronously.
-runCentralGhcSpawned :: CreateMethods -> ServerSocketPath -> IO ()
-runCentralGhcSpawned methods socket =
-  runCentralGhc methods primaryFile socket instr
+runCentralGhcSpawned :: CreateMethods -> FeatureInstrument -> ServerSocketPath -> IO ()
+runCentralGhcSpawned methods featureInstrument socket =
+  runCentralGhc methods primaryFile socket instrumentSocket
   where
-    instr = Just (instrumentSocketIn dir)
+    instrumentSocket =
+      if featureInstrument.flag
+      then Just (instrumentSocketIn dir)
+      else Nothing
 
     primaryFile = primarySocketDiscoveryIn dir
 
@@ -216,10 +223,10 @@ serveOrProxyCentralGhc methods socket = do
   where
     run primaryFile = do
       let primary = PrimarySocketPath socket.path
-      thread <- async (runCentralGhc methods primaryFile socket instr)
+      thread <- async (runCentralGhc methods primaryFile socket instrumentSocket)
       waitPoll primary
       pure (primary, thread)
 
-    instr = Just (instrumentSocketIn socketDir)
+    instrumentSocket = Just (instrumentSocketIn socketDir)
 
     socketDir = SocketDirectory (init (dropWhileEnd ('-' /=) (takeDirectory socket.path)))
diff --git a/ghc-worker/lib/GhcWorker/Run.hs b/ghc-worker/lib/GhcWorker/Run.hs
index ebd5b9a9..2134464a 100644
--- a/ghc-worker/lib/GhcWorker/Run.hs
+++ b/ghc-worker/lib/GhcWorker/Run.hs
@@ -9,7 +9,7 @@ import Control.Exception (throwIO)
 import GhcWorker.GhcHandler (LockState (..), ghcHandler)
 import GhcWorker.Grpc (instrumentMethods)
 import GhcWorker.Instrumentation (WorkerStatus (..), toGrpcHandler)
-import GhcWorker.Orchestration (CreateMethods (..), runCentralGhcSpawned)
+import GhcWorker.Orchestration (CreateMethods (..), FeatureInstrument (..), runCentralGhcSpawned)
 import Internal.State (WorkerState (..), newState, newStateWith)
 import Internal.State.Oneshot (OneshotCacheFeatures (..))
 import Network.GRPC.Common.Protobuf (Proto)
@@ -29,8 +29,10 @@ data CliOptions =
     -- | The worker implementation: Make mode or oneshot mode.
     workerMode :: WorkerMode,
 
-    -- | listening on the given path.
-    serve :: ServerSocketPath
+    -- | If this is given, the app should start a GHC server synchronously, listening on the given path.
+    serve :: ServerSocketPath,
+
+    instrument :: FeatureInstrument
   }
   deriving stock (Eq, Show)
 
@@ -38,7 +40,8 @@ defaultCliOptions :: CliOptions
 defaultCliOptions =
   CliOptions {
     workerMode = WorkerOneshotMode,
-    serve = ServerSocketPath "" "" ""
+    serve = ServerSocketPath "" "" "",
+    instrument = FeatureInstrument False
   }
 
 parseOptions :: [String] -> IO CliOptions
@@ -49,6 +52,7 @@ parseOptions =
       [] -> pure z
       "--make" : rest -> spin z {workerMode = WorkerMakeMode} rest
       "--serve" : socket : rest -> spin z {serve = serverSocketFromPath socket} rest
+      "--instrument" : rest -> spin z {instrument = FeatureInstrument True} rest
       arg -> throwIO (userError ("Invalid worker CLI args: " ++ unwords arg))
 
 -- | Allocate a communication channel for instrumentation events and construct a gRPC server handler that streams said
@@ -73,7 +77,7 @@ createGhcMethods lock cache workerMode status instrChan =
 
 -- | Main function for running the default persistent worker using the provided server socket path and CLI options.
 runWorker :: CliOptions -> IO ()
-runWorker CliOptions {workerMode, serve} = do
+runWorker CliOptions {workerMode, serve, instrument} = do
   cache <-
     case workerMode of
       WorkerMakeMode ->
@@ -92,4 +96,4 @@ runWorker CliOptions {workerMode, serve} = do
       createInstrumentation = createInstrumentMethods cache,
       createGhc = createGhcMethods lock cache workerMode status
     }
-  runCentralGhcSpawned methods serve
+  runCentralGhcSpawned methods instrument serve

From d944b509a8acab3f258cda4ce30b012ff5130eea Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Fri, 6 Jun 2025 13:45:33 +0200
Subject: [PATCH 20/27] Use BUCK_BUILD_ID for the primary socket path if it
 exists

Change-Id: Ie036864b000f1d47ad60fa9c73495fa81d69431d
---
 buck-proxy/lib/BuckProxy/Orchestration.hs | 11 +++++++----
 types/src/Types/Grpc.hs                   |  2 +-
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/buck-proxy/lib/BuckProxy/Orchestration.hs b/buck-proxy/lib/BuckProxy/Orchestration.hs
index cac23743..d96a24ea 100644
--- a/buck-proxy/lib/BuckProxy/Orchestration.hs
+++ b/buck-proxy/lib/BuckProxy/Orchestration.hs
@@ -13,9 +13,9 @@ import Control.Concurrent (threadDelay)
 import Control.Concurrent.MVar (MVar, modifyMVar)
 import Control.Exception (throwIO, try)
 import Control.Monad (void, when)
-import Data.Map.Strict (Map)
+import Data.Map.Strict (Map, (!?))
 import Data.Map.Strict qualified as Map
-import Data.Maybe (isJust)
+import Data.Maybe (fromMaybe, isJust)
 import Data.Text qualified as Text
 import Data.Text.Encoding (decodeUtf8Lenient)
 import Network.GRPC.Client (Connection, Server (..), recvNextOutput, sendFinalInput, withConnection, withRPC)
@@ -37,7 +37,7 @@ import System.Process (ProcessHandle, getProcessExitCode, spawnProcess)
 import Types.Args (TargetId)
 import Types.BuckArgs (BuckArgs (workerTargetId), parseBuckArgs)
 import Types.GhcHandler (WorkerMode (..))
-import Types.Grpc (RequestArgs (..))
+import Types.Grpc (CommandEnv (..), RequestArgs (..))
 import Types.Orchestration (
   PrimarySocketPath (..),
   ServerSocketPath (..),
@@ -83,6 +83,9 @@ proxyHandler ::
 proxyHandler workerMap exe wmode basePath req = do
   let cmdEnv = commandEnv req.env
       argv = Text.unpack . decodeUtf8Lenient <$> req.argv
+      -- Get the build ID for the primary socket path from the command environment, and fall back to the value extracted
+      -- from the gRPC socket path if the key is absent from the env.
+      socketId = fromMaybe basePath (cmdEnv.values !? "BUCK_BUILD_ID")
   buckArgs <- either (throwIO . userError) pure (parseBuckArgs cmdEnv (RequestArgs argv))
   case buckArgs.workerTargetId of
     Nothing -> throwIO (userError "No --worker-target-id passed")
@@ -91,7 +94,7 @@ proxyHandler workerMap exe wmode basePath req = do
         modifyMVar workerMap \wmap -> do
           case Map.lookup targetId wmap of
             Nothing -> do
-              let workerSocketDir = projectSocketDirectory basePath targetId
+              let workerSocketDir = projectSocketDirectory socketId targetId
               void $ try @IOError (createDirectoryIfMissing True workerSocketDir.path)
               resource <- spawnGhcWorker exe wmode workerSocketDir
               dbg $ "No primary socket for " ++ show targetId ++ ", so created it on " ++ resource.primarySocket.path
diff --git a/types/src/Types/Grpc.hs b/types/src/Types/Grpc.hs
index f59d3ec6..c0cce4dc 100644
--- a/types/src/Types/Grpc.hs
+++ b/types/src/Types/Grpc.hs
@@ -4,7 +4,7 @@ import Data.Map (Map)
 
 -- | The environment variables sent by Buck.
 newtype CommandEnv =
-  CommandEnv (Map String String)
+  CommandEnv { values :: Map String String }
   deriving stock (Eq, Show)
 
 -- | The command line arguments sent by Buck.

From 201e6f121d3d7de18d9aa72c0b547658643b239c Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Thu, 12 Jun 2025 14:55:20 +0200
Subject: [PATCH 21/27] Remove accidental debug code introduced in 6f97ca36

Change-Id: Iea8bae23a92ad6143982d55a33629b262da8beaf
---
 internal/src/Internal/State/Oneshot.hs | 2 --
 1 file changed, 2 deletions(-)

diff --git a/internal/src/Internal/State/Oneshot.hs b/internal/src/Internal/State/Oneshot.hs
index 2b6bf15b..4552bd13 100644
--- a/internal/src/Internal/State/Oneshot.hs
+++ b/internal/src/Internal/State/Oneshot.hs
@@ -51,7 +51,6 @@ import GHC.Unit (
   moduleName,
   )
 import GHC.Utils.Panic (panic)
-import Internal.Debug (forceInstalledFindResult)
 import Internal.State.Stats (FinderStats (..))
 
 #else
@@ -313,7 +312,6 @@ newFinderCache updateOneshot OneshotState {finder = FinderState {modules, files}
 
       addToFinderCache :: InstalledModule -> InstalledFindResult -> IO ()
       addToFinderCache key val = do
-        !() <- evaluate (forceInstalledFindResult val)
         atomicModifyIORef' modules $ \c ->
           case (lookupInstalledModuleEnv c key, val) of
             (Just InstalledFound{}, InstalledNotFound{}) -> (c, ())

From f2a689cc49c88ffa29905470bc5f693d1cb5465b Mon Sep 17 00:00:00 2001
From: Sjoerd Visscher <sjoerd@w3future.com>
Date: Thu, 12 Jun 2025 17:10:58 +0200
Subject: [PATCH 22/27] Grpc: We're not going to use grpc metadata, so declare
 NoMetadata for all methods in one go.

Change-Id: I23a413748c53830cf00f7d24ef7d15a21bf156aa
---
 proto/src/BuckWorker.hs | 20 ++++++--------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/proto/src/BuckWorker.hs b/proto/src/BuckWorker.hs
index 41469661..83d9fc2b 100644
--- a/proto/src/BuckWorker.hs
+++ b/proto/src/BuckWorker.hs
@@ -12,18 +12,10 @@ import Network.GRPC.Common.Protobuf (Protobuf)
 import Proto.Worker
 import Proto.Instrument
 
-type instance RequestMetadata (Protobuf Worker "exec") = NoMetadata
-type instance ResponseInitialMetadata (Protobuf Worker "exec") = NoMetadata
-type instance ResponseTrailingMetadata (Protobuf Worker "exec") = NoMetadata
+type instance RequestMetadata (Protobuf Worker _) = NoMetadata
+type instance ResponseInitialMetadata (Protobuf Worker _) = NoMetadata
+type instance ResponseTrailingMetadata (Protobuf Worker _) = NoMetadata
 
-type instance RequestMetadata (Protobuf Worker "execute") = NoMetadata
-type instance ResponseInitialMetadata (Protobuf Worker "execute") = NoMetadata
-type instance ResponseTrailingMetadata (Protobuf Worker "execute") = NoMetadata
-
-type instance RequestMetadata (Protobuf Instrument "notifyMe") = NoMetadata
-type instance ResponseInitialMetadata (Protobuf Instrument "notifyMe") = NoMetadata
-type instance ResponseTrailingMetadata (Protobuf Instrument "notifyMe") = NoMetadata
-
-type instance RequestMetadata (Protobuf Instrument "setOptions") = NoMetadata
-type instance ResponseInitialMetadata (Protobuf Instrument "setOptions") = NoMetadata
-type instance ResponseTrailingMetadata (Protobuf Instrument "setOptions") = NoMetadata
+type instance RequestMetadata (Protobuf Instrument _) = NoMetadata
+type instance ResponseInitialMetadata (Protobuf Instrument _) = NoMetadata
+type instance ResponseTrailingMetadata (Protobuf Instrument _) = NoMetadata

From 63a55ea87a7f99ba7e52efc4e331b593dcdbe965 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Tue, 17 Jun 2025 18:15:01 +0200
Subject: [PATCH 23/27] Remove accidental inclusion of -DMWB in the Finder
 interface CPP

Change-Id: I45411bec41097b683976743d5c3c871eb991f982
---
 internal/src/Internal/State/Oneshot.hs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/internal/src/Internal/State/Oneshot.hs b/internal/src/Internal/State/Oneshot.hs
index 4552bd13..925c4598 100644
--- a/internal/src/Internal/State/Oneshot.hs
+++ b/internal/src/Internal/State/Oneshot.hs
@@ -36,7 +36,7 @@ import Internal.State.Stats (
   )
 import Types.State (SymbolCache (..), SymbolMap, Target)
 
-#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) || defined(MWB)
+#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
 
 import Control.Exception (evaluate)
 import Data.IORef (IORef, newIORef)
@@ -60,7 +60,7 @@ import GHC.Unit.Finder (initFinderCache)
 
 #endif
 
-#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) || defined(MWB)
+#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
 
 data FinderState =
   FinderState {
@@ -294,7 +294,7 @@ updateState target InterpCache {..} newLoaderState newSymbols newNames cache = d
   where
     symbolsStats = basicSymbolsStats symbols newSymbols
 
-#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) || defined(MWB)
+#if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
 
 -- | This replacement of the Finder implementation has the sole purpose of recording some cache stats, for now.
 -- While its mutable state is allocated separately and shared across sessions, this doesn't really make a difference at

From 92d1cc9f43d11926e6df15db8d72715e3fe1266e Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Thu, 12 Jun 2025 19:47:10 +0200
Subject: [PATCH 24/27] Replace the `debug` flag in `Log` with a simple log
 level

Change-Id: I6c68c6fa12f35331637341589b18971a03d08043
---
 ghc-worker/lib/GhcWorker/GhcHandler.hs |  2 +-
 ghc-worker/test/CompileHptTest.hs      |  6 ++--
 internal/src/Internal/Error.hs         |  4 +--
 internal/src/Internal/Log.hs           | 44 +++++++++++++++-----------
 4 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/ghc-worker/lib/GhcWorker/GhcHandler.hs b/ghc-worker/lib/GhcWorker/GhcHandler.hs
index 0c97b5b8..1ff5c5e2 100644
--- a/ghc-worker/lib/GhcWorker/GhcHandler.hs
+++ b/ghc-worker/lib/GhcWorker/GhcHandler.hs
@@ -133,7 +133,7 @@ ghcHandler lock state workerMode =
   InstrumentedHandler \ hooks -> GrpcHandler \ commandEnv argv -> do
     buckArgs <- either (throwIO . userError) pure (parseBuckArgs commandEnv argv)
     args <- toGhcArgs buckArgs
-    log <- newLog True
+    log <- newLog
     let env = Env {log, state, args}
     onException
       do
diff --git a/ghc-worker/test/CompileHptTest.hs b/ghc-worker/test/CompileHptTest.hs
index 7ba5bbaf..94cac562 100644
--- a/ghc-worker/test/CompileHptTest.hs
+++ b/ghc-worker/test/CompileHptTest.hs
@@ -40,7 +40,7 @@ unitFlags args HscEnv {hsc_logger, hsc_dflags = dflags0} = do
 
 stepMetadata :: Conf -> Unit -> [Unit] -> IO ()
 stepMetadata Conf {state, tmp, args0} unit deps = do
-  log <- newLog True
+  log <- newLog
   createDirectoryIfMissing False sessionTmpDir
   names <- listDirectory unit.dir
   let srcs = [unit.dir </> name | name <- names, takeExtension name == ".hs"]
@@ -72,7 +72,7 @@ stepMetadata Conf {state, tmp, args0} unit deps = do
 
 stepCompile :: Conf -> Module -> IO ()
 stepCompile Conf {state, tmp, args0} Module {unit, src} = do
-  log <- newLog True
+  log <- newLog
   let env = Env {log, state, args}
   liftIO $ createDirectoryIfMissing False sessionTmpDir
   result <- liftIO $ withGhcMhu env \ _ target -> do
@@ -312,7 +312,7 @@ runStep conf = \case
 
 testWorker :: (Conf -> NonEmpty UnitSpec) -> IO ()
 testWorker mkSpecs = do
-  log <- newLog True
+  log <- newLog
   logMemStats "initial" log
   withProject (pure . mkSpecs) \ conf units -> do
     let steps = testSteps units
diff --git a/internal/src/Internal/Error.hs b/internal/src/Internal/Error.hs
index e03078b3..bb58bc86 100644
--- a/internal/src/Internal/Error.hs
+++ b/internal/src/Internal/Error.hs
@@ -8,7 +8,7 @@ import qualified Control.Monad.Catch as MC
 import Control.Monad.IO.Class (liftIO, MonadIO)
 import GHC (Ghc, GhcException (..), printException)
 import GHC.Types.SourceError (SourceError, throwErrors)
-import Internal.Log (Log, logOther)
+import Internal.Log (Log, LogLevel (..), logOther)
 import System.Environment (getProgName)
 import System.Exit (ExitCode)
 import GHC.Driver.Errors.Types (GhcMessage)
@@ -48,7 +48,7 @@ handleExceptions logVar errResult =
       | otherwise
       = fm (show (Panic (show exception)))
 
-    fm = logOther logVar
+    fm = logOther logVar LogInfo
 
 eitherMessages ::
   MonadIO m =>
diff --git a/internal/src/Internal/Log.hs b/internal/src/Internal/Log.hs
index 2e7e5133..b9c696ad 100644
--- a/internal/src/Internal/Log.hs
+++ b/internal/src/Internal/Log.hs
@@ -1,7 +1,7 @@
 module Internal.Log where
 
 import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar)
-import Control.Monad (unless, when)
+import Control.Monad (unless)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Data.ByteString (ByteString)
 import Data.Foldable (traverse_)
@@ -10,7 +10,7 @@ import Data.Text.Encoding (encodeUtf8)
 import GHC (Ghc, Severity (SevIgnore), noSrcSpan)
 import GHC.Driver.Config.Diagnostic (initDiagOpts)
 import GHC.Driver.DynFlags (getDynFlags)
-import GHC.Driver.Errors.Types (DriverMessage (..), GhcMessage(GhcDriverMessage))
+import GHC.Driver.Errors.Types (DriverMessage (..), GhcMessage (GhcDriverMessage))
 import GHC.Driver.Monad qualified as GHC (logDiagnostics)
 import GHC.Types.Error (
   DiagnosticReason (WarningWithoutFlag),
@@ -47,17 +47,24 @@ newtype LogName =
   LogName { get :: String }
   deriving stock (Eq, Show)
 
+-- | Simple log level that decides whether non-diagnostic messages will be sent to Buck in addition to basic file
+-- logging.
+data LogLevel =
+  LogDebug
+  |
+  LogInfo
+  deriving stock (Eq, Show)
+
 data Log =
   Log {
     diagnostics :: [String],
-    other :: [String],
-    debug :: Bool
+    other :: [(String, LogLevel)]
   }
   deriving stock (Eq, Show)
 
-newLog :: MonadIO m => Bool -> m (MVar Log)
-newLog debug =
-  liftIO $ newMVar Log {diagnostics = [], other = [], debug}
+newLog :: MonadIO m => m (MVar Log)
+newLog =
+  liftIO $ newMVar Log {diagnostics = [], other = []}
 
 logDiagnostics ::
   MonadIO m =>
@@ -65,19 +72,18 @@ logDiagnostics ::
   String ->
   m ()
 logDiagnostics logVar msg =
-  liftIO $ modifyMVar_ logVar \ Log {diagnostics, ..} -> do
-    when debug (dbg msg)
+  liftIO $ modifyMVar_ logVar \ Log {diagnostics, ..} ->
     pure Log {diagnostics = msg : diagnostics, ..}
 
 logOther ::
   MonadIO m =>
   MVar Log ->
+  LogLevel ->
   String ->
   m ()
-logOther logVar msg =
-  liftIO $ modifyMVar_ logVar \ Log {other, ..} -> do
-    when debug (dbg msg)
-    pure Log {other = msg : other, ..}
+logOther logVar level msg =
+  liftIO $ modifyMVar_ logVar \ Log {other, ..} ->
+    pure Log {other = (msg, level) : other, ..}
 
 logDir :: FilePath
 logDir =
@@ -89,14 +95,14 @@ logDir =
 --
 -- If the session fails before the target could be determined, this is 'Nothing', so we choose @unknown@ for the file
 -- name.
-writeLogFile :: [String] -> LogName -> IO ()
+writeLogFile :: [(String, LogLevel)] -> LogName -> IO ()
 writeLogFile logLines (LogName logName) =
   either warn pure =<< tryIOError do
     createDirectoryIfMissing True (takeDirectory path)
     exists <- doesPathExist path
     unless exists do
       writeFile path ""
-    appendFile path (unlines logLines)
+    appendFile path (unlines (fst <$> logLines))
   where
     path = logDir </> addExtension logName "log"
 
@@ -106,9 +112,9 @@ writeLogFile logLines (LogName logName) =
 logFlush :: Maybe LogName -> MVar Log -> IO [String]
 logFlush logName var = do
   modifyMVar var \ Log {..} -> do
-    let logLines = reverse (other ++ diagnostics)
+    let logLines = reverse (other ++ [(msg, LogInfo) | msg <- diagnostics])
     traverse_ (writeLogFile logLines) logName
-    pure (Log {diagnostics = [], other = [], debug}, logLines)
+    pure (Log {diagnostics = [], other = [], ..}, [msg | (msg, level) <- logLines, LogInfo == level])
 
 logFlushBytes :: MVar Log -> IO ByteString
 logFlushBytes var = do
@@ -137,7 +143,7 @@ logToState logVar logflags msg_class srcSpan msg = case msg_class of
 
     diagnostic = logDiagnostics logVar . render
 
-    other = logOther logVar . render
+    other = logOther logVar LogInfo . render
 
     render d = renderWithContext (log_default_user_context logflags) d
 
@@ -157,7 +163,7 @@ logp ::
   a ->
   m ()
 logp logVar =
-  logOther logVar . showPprUnsafe
+  logOther logVar LogInfo . showPprUnsafe
 
 logd ::
   MonadIO m =>

From 2444b686627f9987af3858f52fcfcfa1f46c0a2c Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Thu, 12 Jun 2025 19:57:28 +0200
Subject: [PATCH 25/27] Use the Buck trace ID to prefix log directories

Change-Id: Ie63ca41f8d68492e24c25007e68e436307e05eb3
---
 ghc-worker/lib/GhcWorker/GhcHandler.hs | 18 +++++----
 ghc-worker/lib/GhcWorker/Run.hs        | 14 ++++---
 ghc-worker/test/CompileHptTest.hs      |  6 +--
 internal/src/Internal/Log.hs           | 51 +++++++++++++++++---------
 internal/src/Internal/State/Oneshot.hs |  1 -
 5 files changed, 56 insertions(+), 34 deletions(-)

diff --git a/ghc-worker/lib/GhcWorker/GhcHandler.hs b/ghc-worker/lib/GhcWorker/GhcHandler.hs
index 1ff5c5e2..b61c0891 100644
--- a/ghc-worker/lib/GhcWorker/GhcHandler.hs
+++ b/ghc-worker/lib/GhcWorker/GhcHandler.hs
@@ -18,7 +18,7 @@ import Internal.AbiHash (AbiHash (..), showAbiHash)
 import Internal.State (WorkerState (..), ModuleArtifacts (..))
 import Internal.Compile (compileModuleWithDepsInEps)
 import Internal.CompileHpt (compileModuleWithDepsInHpt)
-import Internal.Log (LogName (..), dbg, logFlush, newLog)
+import Internal.Log (TraceId, dbg, logFlush, newLog, setLogTarget)
 import Internal.Metadata (computeMetadata)
 import Internal.Session (Env (..), withGhc, withGhcMhu)
 import Prelude hiding (log)
@@ -91,10 +91,12 @@ dispatch lock workerMode hooks env args =
           pure (code, result)
       pure (code, snd <$> result)
     Just ModeMetadata -> do
+      let target = Target "metadata"
+      liftIO $ setLogTarget env.log target
       code <- computeMetadata env <&> \case
         True -> 0
         False -> 1
-      pure (code, Just (Target "metadata"))
+      pure (code, Just target)
     Just ModeClose -> do
       dbg "in dispatch. Mode Close"
       _ <- writeCloseOutput args
@@ -112,7 +114,8 @@ dispatch lock workerMode hooks env args =
         withGhcMhu env \ _ ->
           withTarget (compileAndReadAbiHash CompManager compileModuleWithDepsInHpt hooks args)
 
-    withTarget f target =
+    withTarget f target = do
+      liftIO $ setLogTarget env.log target
       f target <&> fmap \ r -> (r, target)
 
 -- | Default implementation of an 'InstrumentedHandler' using our custom persistent worker GHC mode, either using HPT or
@@ -128,20 +131,19 @@ ghcHandler ::
   TVar LockState ->
   MVar WorkerState ->
   WorkerMode ->
+  Maybe TraceId ->
   InstrumentedHandler
-ghcHandler lock state workerMode =
+ghcHandler lock state workerMode traceId =
   InstrumentedHandler \ hooks -> GrpcHandler \ commandEnv argv -> do
     buckArgs <- either (throwIO . userError) pure (parseBuckArgs commandEnv argv)
     args <- toGhcArgs buckArgs
-    log <- newLog
+    log <- newLog traceId
     let env = Env {log, state, args}
     onException
       do
         (result, target) <- dispatch lock workerMode hooks env buckArgs
-        output <- logFlush (logName <$> target) env.log
+        output <- logFlush env.log
         liftIO $ hooks.compileFinish (Just (target, output, result))
         pure (output, result)
       do
         liftIO $ hooks.compileFinish Nothing
-    where
-      logName (Target target) = LogName target
diff --git a/ghc-worker/lib/GhcWorker/Run.hs b/ghc-worker/lib/GhcWorker/Run.hs
index 2134464a..8e5c7aac 100644
--- a/ghc-worker/lib/GhcWorker/Run.hs
+++ b/ghc-worker/lib/GhcWorker/Run.hs
@@ -10,6 +10,7 @@ import GhcWorker.GhcHandler (LockState (..), ghcHandler)
 import GhcWorker.Grpc (instrumentMethods)
 import GhcWorker.Instrumentation (WorkerStatus (..), toGrpcHandler)
 import GhcWorker.Orchestration (CreateMethods (..), FeatureInstrument (..), runCentralGhcSpawned)
+import Internal.Log (TraceId (..))
 import Internal.State (WorkerState (..), newState, newStateWith)
 import Internal.State.Oneshot (OneshotCacheFeatures (..))
 import Network.GRPC.Common.Protobuf (Proto)
@@ -70,15 +71,16 @@ createGhcMethods ::
   MVar WorkerState ->
   WorkerMode ->
   MVar WorkerStatus ->
+  Maybe TraceId ->
   Maybe (Chan (Proto Instr.Event)) ->
   IO (Methods IO (ProtobufMethodsOf Worker))
-createGhcMethods lock cache workerMode status instrChan =
-  pure (fromGrpcHandler (toGrpcHandler (ghcHandler lock cache workerMode) status cache instrChan))
+createGhcMethods lock state workerMode status traceId instrChan =
+  pure (fromGrpcHandler (toGrpcHandler (ghcHandler lock state workerMode traceId) status state instrChan))
 
 -- | Main function for running the default persistent worker using the provided server socket path and CLI options.
 runWorker :: CliOptions -> IO ()
 runWorker CliOptions {workerMode, serve, instrument} = do
-  cache <-
+  state <-
     case workerMode of
       WorkerMakeMode ->
         newStateWith OneshotCacheFeatures {
@@ -93,7 +95,9 @@ runWorker CliOptions {workerMode, serve, instrument} = do
   status <- newMVar WorkerStatus {active = 0}
   let
     methods = CreateMethods {
-      createInstrumentation = createInstrumentMethods cache,
-      createGhc = createGhcMethods lock cache workerMode status
+      createInstrumentation = createInstrumentMethods state,
+      createGhc = createGhcMethods lock state workerMode status traceId
     }
   runCentralGhcSpawned methods instrument serve
+  where
+    traceId = if null serve.traceId then Nothing else Just (TraceId serve.traceId)
diff --git a/ghc-worker/test/CompileHptTest.hs b/ghc-worker/test/CompileHptTest.hs
index 94cac562..d9c0eea6 100644
--- a/ghc-worker/test/CompileHptTest.hs
+++ b/ghc-worker/test/CompileHptTest.hs
@@ -40,7 +40,7 @@ unitFlags args HscEnv {hsc_logger, hsc_dflags = dflags0} = do
 
 stepMetadata :: Conf -> Unit -> [Unit] -> IO ()
 stepMetadata Conf {state, tmp, args0} unit deps = do
-  log <- newLog
+  log <- newLog Nothing
   createDirectoryIfMissing False sessionTmpDir
   names <- listDirectory unit.dir
   let srcs = [unit.dir </> name | name <- names, takeExtension name == ".hs"]
@@ -72,7 +72,7 @@ stepMetadata Conf {state, tmp, args0} unit deps = do
 
 stepCompile :: Conf -> Module -> IO ()
 stepCompile Conf {state, tmp, args0} Module {unit, src} = do
-  log <- newLog
+  log <- newLog Nothing
   let env = Env {log, state, args}
   liftIO $ createDirectoryIfMissing False sessionTmpDir
   result <- liftIO $ withGhcMhu env \ _ target -> do
@@ -312,7 +312,7 @@ runStep conf = \case
 
 testWorker :: (Conf -> NonEmpty UnitSpec) -> IO ()
 testWorker mkSpecs = do
-  log <- newLog
+  log <- newLog Nothing
   logMemStats "initial" log
   withProject (pure . mkSpecs) \ conf units -> do
     let steps = testSteps units
diff --git a/internal/src/Internal/Log.hs b/internal/src/Internal/Log.hs
index b9c696ad..876074fd 100644
--- a/internal/src/Internal/Log.hs
+++ b/internal/src/Internal/Log.hs
@@ -4,7 +4,6 @@ import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar)
 import Control.Monad (unless)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Data.ByteString (ByteString)
-import Data.Foldable (traverse_)
 import Data.Text (pack)
 import Data.Text.Encoding (encodeUtf8)
 import GHC (Ghc, Severity (SevIgnore), noSrcSpan)
@@ -37,15 +36,12 @@ import GHC.Utils.Outputable (
   ($$),
   ($+$),
   )
+import Prelude hiding (log)
 import System.Directory (createDirectoryIfMissing, doesPathExist)
 import System.FilePath (addExtension, takeDirectory, (</>))
 import System.IO (hPutStrLn, stderr)
 import System.IO.Error (tryIOError)
-
--- | Name of the current session's target for the log file path.
-newtype LogName =
-  LogName { get :: String }
-  deriving stock (Eq, Show)
+import Types.State (Target (..))
 
 -- | Simple log level that decides whether non-diagnostic messages will be sent to Buck in addition to basic file
 -- logging.
@@ -55,16 +51,32 @@ data LogLevel =
   LogInfo
   deriving stock (Eq, Show)
 
+-- | The ID of the current build, usually generated by Buck.
+newtype TraceId =
+  TraceId { string :: String }
+  deriving stock (Eq, Show)
+
 data Log =
   Log {
     diagnostics :: [String],
-    other :: [(String, LogLevel)]
+    other :: [(String, LogLevel)],
+    traceId :: Maybe TraceId,
+    target :: Maybe Target
   }
   deriving stock (Eq, Show)
 
-newLog :: MonadIO m => m (MVar Log)
-newLog =
-  liftIO $ newMVar Log {diagnostics = [], other = []}
+newLog ::
+  MonadIO m =>
+  Maybe TraceId ->
+  m (MVar Log)
+newLog traceId =
+  liftIO $ newMVar Log {diagnostics = [], other = [], traceId, target = Nothing}
+
+-- | After the current request's target has been determined, the log state can be updated to generate more specific log
+-- file paths.
+setLogTarget :: MVar Log -> Target -> IO ()
+setLogTarget logVar target =
+  modifyMVar_ logVar \ log -> pure log {target = Just target}
 
 logDiagnostics ::
   MonadIO m =>
@@ -95,8 +107,8 @@ logDir =
 --
 -- If the session fails before the target could be determined, this is 'Nothing', so we choose @unknown@ for the file
 -- name.
-writeLogFile :: [(String, LogLevel)] -> LogName -> IO ()
-writeLogFile logLines (LogName logName) =
+writeLogFile :: Maybe TraceId -> Maybe Target -> [(String, LogLevel)] -> IO ()
+writeLogFile traceId target logLines =
   either warn pure =<< tryIOError do
     createDirectoryIfMissing True (takeDirectory path)
     exists <- doesPathExist path
@@ -104,21 +116,26 @@ writeLogFile logLines (LogName logName) =
       writeFile path ""
     appendFile path (unlines (fst <$> logLines))
   where
-    path = logDir </> addExtension logName "log"
+    path = targetIdDir </> addExtension logName "log"
+
+    targetIdDir | Just (TraceId wtId) <- traceId = logDir </> wtId
+                | otherwise = logDir
 
     warn err = dbg ("Failed to write log file for " ++ logName ++ ": " ++ show err)
 
+    logName = maybe "global" (.path) target
+
 -- | Write the current session's log to a file, clear the fields in the 'MVar' and return the log lines.
-logFlush :: Maybe LogName -> MVar Log -> IO [String]
-logFlush logName var = do
+logFlush :: MVar Log -> IO [String]
+logFlush var = do
   modifyMVar var \ Log {..} -> do
     let logLines = reverse (other ++ [(msg, LogInfo) | msg <- diagnostics])
-    traverse_ (writeLogFile logLines) logName
+    writeLogFile traceId target logLines
     pure (Log {diagnostics = [], other = [], ..}, [msg | (msg, level) <- logLines, LogInfo == level])
 
 logFlushBytes :: MVar Log -> IO ByteString
 logFlushBytes var = do
-  lns <- logFlush Nothing var
+  lns <- logFlush var
   pure (encodeUtf8 (pack (unlines lns)))
 
 logToState :: MVar Log -> LogAction
diff --git a/internal/src/Internal/State/Oneshot.hs b/internal/src/Internal/State/Oneshot.hs
index 925c4598..062ca5d3 100644
--- a/internal/src/Internal/State/Oneshot.hs
+++ b/internal/src/Internal/State/Oneshot.hs
@@ -38,7 +38,6 @@ import Types.State (SymbolCache (..), SymbolMap, Target)
 
 #if MIN_VERSION_GLASGOW_HASKELL(9,11,0,0)
 
-import Control.Exception (evaluate)
 import Data.IORef (IORef, newIORef)
 import qualified Data.Map.Lazy as LazyMap
 import GHC.Fingerprint (Fingerprint, getFileHash)

From fb785c3da02d73a352c054e7657b268418036628 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Thu, 12 Jun 2025 19:58:43 +0200
Subject: [PATCH 26/27] Dump the module graph and HUG to the log file on
 failure

Change-Id: I114d356a836c961e3b812476b83e2a63ee072dfd
---
 ghc-worker/lib/GhcWorker/GhcHandler.hs | 37 ++++++++++++++++++--------
 internal/src/Internal/Debug.hs         | 12 ++++++---
 internal/src/Internal/Log.hs           | 24 +++++++++++++++++
 internal/src/Internal/State.hs         | 26 ++++++++++++++++--
 4 files changed, 82 insertions(+), 17 deletions(-)

diff --git a/ghc-worker/lib/GhcWorker/GhcHandler.hs b/ghc-worker/lib/GhcWorker/GhcHandler.hs
index b61c0891..16f9f27b 100644
--- a/ghc-worker/lib/GhcWorker/GhcHandler.hs
+++ b/ghc-worker/lib/GhcWorker/GhcHandler.hs
@@ -3,8 +3,8 @@ module GhcWorker.GhcHandler where
 import Common.Grpc (GrpcHandler (..))
 import Control.Concurrent (MVar, forkIO, threadDelay)
 import Control.Concurrent.STM (TVar, atomically, modifyTVar', readTVar, retry, writeTVar)
-import Control.Exception (throwIO)
-import Control.Monad.Catch (onException)
+import Control.Exception (throwIO, try)
+import Control.Monad (when)
 import Control.Monad.IO.Class (liftIO)
 import Data.Functor ((<&>))
 import Data.Int (Int32)
@@ -15,12 +15,12 @@ import GHC.Driver.Monad (modifySession)
 import GhcWorker.CompileResult (CompileResult (..), writeCloseOutput, writeResult)
 import GhcWorker.Instrumentation (Hooks (..), InstrumentedHandler (..))
 import Internal.AbiHash (AbiHash (..), showAbiHash)
-import Internal.State (WorkerState (..), ModuleArtifacts (..))
 import Internal.Compile (compileModuleWithDepsInEps)
 import Internal.CompileHpt (compileModuleWithDepsInHpt)
 import Internal.Log (TraceId, dbg, logFlush, newLog, setLogTarget)
 import Internal.Metadata (computeMetadata)
 import Internal.Session (Env (..), withGhc, withGhcMhu)
+import Internal.State (ModuleArtifacts (..), WorkerState (..), dumpState)
 import Prelude hiding (log)
 import System.Exit (ExitCode (ExitSuccess))
 import System.Posix.Process (exitImmediately)
@@ -118,6 +118,27 @@ dispatch lock workerMode hooks env args =
       liftIO $ setLogTarget env.log target
       f target <&> fmap \ r -> (r, target)
 
+processResult ::
+  Hooks ->
+  Env ->
+  Either IOError (Int32, Maybe Target) ->
+  IO ([String], Int32)
+processResult hooks env result = do
+  when (exitCode /= 0) do
+    dumpState env.log env.state exception
+  output <- logFlush env.log
+  hooks.compileFinish (hookPayload output)
+  pure (output, exitCode)
+  where
+    hookPayload output =
+      if exitCode == 0
+      then Just (target, output, exitCode)
+      else Nothing
+
+    ((exitCode, target), exception) = case result of
+      Right out -> (out, Nothing)
+      Left err -> ((1, Nothing), Just ("Exception: " ++ show err))
+
 -- | Default implementation of an 'InstrumentedHandler' using our custom persistent worker GHC mode, either using HPT or
 -- EPS for local dependency lookup.
 --
@@ -139,11 +160,5 @@ ghcHandler lock state workerMode traceId =
     args <- toGhcArgs buckArgs
     log <- newLog traceId
     let env = Env {log, state, args}
-    onException
-      do
-        (result, target) <- dispatch lock workerMode hooks env buckArgs
-        output <- logFlush env.log
-        liftIO $ hooks.compileFinish (Just (target, output, result))
-        pure (output, result)
-      do
-        liftIO $ hooks.compileFinish Nothing
+    result <- try $ dispatch lock workerMode hooks env buckArgs
+    processResult hooks env result
diff --git a/internal/src/Internal/Debug.hs b/internal/src/Internal/Debug.hs
index f929a7e9..858253e3 100644
--- a/internal/src/Internal/Debug.hs
+++ b/internal/src/Internal/Debug.hs
@@ -12,7 +12,6 @@ import GHC.Unit.External (ExternalPackageState (..), eucEPS)
 import GHC.Unit.Home.ModInfo (HomeModInfo (..), HomePackageTable, hm_iface)
 import GHC.Unit.Module.Graph (ModuleGraph)
 import GHC.Utils.Outputable (Outputable, SDoc, comma, hang, hcat, ppr, punctuate, text, vcat, (<+>))
-import System.FilePath (takeDirectory, takeFileName)
 
 #if !MIN_VERSION_GLASGOW_HASKELL(9,11,0,0) && !defined(MWB)
 
@@ -21,7 +20,8 @@ import GHC.Unit.Module.Graph (mgTransDeps)
 
 #else
 
-import GHC.Unit.Module.Graph (mgModSummaries')
+import GHC (ms_mod_name)
+import GHC.Unit.Module.Graph (ModuleGraphNode (..), mgModSummaries')
 
 #endif
 
@@ -52,7 +52,11 @@ showModGraph g =
 
 showModGraph :: ModuleGraph -> SDoc
 showModGraph g =
-  ppr (mgModSummaries' g)
+  vcat (showOne <$> mgModSummaries' g)
+  where
+    showOne = \case
+      ModuleNode deps ms -> hang (ppr (ms_mod_name ms) <+> "->") 2 (vcat (ppr <$> deps))
+      _ -> ""
 
 #endif
 
@@ -85,7 +89,7 @@ showHpt hpt =
 
 showDbPath :: UnitDatabase UnitId -> SDoc
 showDbPath UnitDatabase {unitDatabasePath} =
-  text (takeFileName (takeDirectory unitDatabasePath))
+  text unitDatabasePath
 
 showHomeUnitEnvShort :: HomeUnitEnv -> SDoc
 showHomeUnitEnvShort HomeUnitEnv {..} =
diff --git a/internal/src/Internal/Log.hs b/internal/src/Internal/Log.hs
index 876074fd..c5def793 100644
--- a/internal/src/Internal/Log.hs
+++ b/internal/src/Internal/Log.hs
@@ -189,6 +189,30 @@ logd ::
   m ()
 logd = logp
 
+logDebug ::
+  MonadIO m =>
+  MVar Log ->
+  String ->
+  m ()
+logDebug logVar =
+  logOther logVar LogDebug
+
+logDebugP ::
+  Outputable a =>
+  MonadIO m =>
+  MVar Log ->
+  a ->
+  m ()
+logDebugP logVar =
+  logOther logVar LogDebug . showPprUnsafe
+
+logDebugD ::
+  MonadIO m =>
+  MVar Log ->
+  SDoc ->
+  m ()
+logDebugD =
+  logDebugP
 
 ghcLogd :: SDoc -> Ghc ()
 ghcLogd doc = do
diff --git a/internal/src/Internal/State.hs b/internal/src/Internal/State.hs
index 33457e4b..c7be0a5a 100644
--- a/internal/src/Internal/State.hs
+++ b/internal/src/Internal/State.hs
@@ -2,8 +2,9 @@
 
 module Internal.State where
 
-import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar)
+import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar, withMVar)
 import Control.Monad.IO.Class (MonadIO, liftIO)
+import Data.Foldable (traverse_)
 import Data.Set (Set)
 import GHC (Ghc, ModIface, emptyMG, mi_module, moduleName, moduleNameString, setSession)
 import GHC.Driver.Env (HscEnv (..))
@@ -11,7 +12,8 @@ import GHC.Driver.Monad (modifySessionM, withSession)
 import GHC.Linker.Types (Linkable)
 import GHC.Unit.Env (unitEnv_new)
 import GHC.Unit.Module.Graph (ModuleGraph)
-import Internal.Log (Log)
+import Internal.Debug (showHugShort, showModGraph)
+import Internal.Log (Log, logDebug, logDebugD)
 import qualified Internal.State.Make as Make
 import Internal.State.Make (MakeState (..))
 import qualified Internal.State.Oneshot as Oneshot
@@ -177,3 +179,23 @@ withCacheMake logVar stateVar prog = do
       liftIO $ modifyMVar_ stateVar \ state -> do
         make <- Make.storeState logVar hsc_env state.make
         pure state {make}
+
+dumpState ::
+  MVar Log ->
+  MVar WorkerState ->
+  Maybe String ->
+  IO ()
+dumpState logVar state exception =
+  withMVar state \ WorkerState {make = MakeState {moduleGraph, hug}} -> do
+    write "-----------------"
+    write "Request failed!"
+    traverse_ write exception
+    write "-----------------"
+    write "Module graph:"
+    writeD (showModGraph moduleGraph)
+    write "-----------------"
+    write "Home unit graph:"
+    writeD (showHugShort hug)
+  where
+    write = logDebug logVar
+    writeD = logDebugD logVar

From 3628e2cfe733928b610844950d78017d0223ecb6 Mon Sep 17 00:00:00 2001
From: Torsten Schmits <git@tryp.io>
Date: Fri, 13 Jun 2025 15:18:46 +0200
Subject: [PATCH 27/27] Use the unit ID as the target for the metadata step

Change-Id: I73d07735252288464a4a5354377f9f9e33018984
---
 ghc-worker/lib/GhcWorker/GhcHandler.hs | 15 +++++++--------
 ghc-worker/test/CompileHptTest.hs      |  2 +-
 internal/src/Internal/Metadata.hs      | 22 ++++++++++++++--------
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/ghc-worker/lib/GhcWorker/GhcHandler.hs b/ghc-worker/lib/GhcWorker/GhcHandler.hs
index 16f9f27b..420a475e 100644
--- a/ghc-worker/lib/GhcWorker/GhcHandler.hs
+++ b/ghc-worker/lib/GhcWorker/GhcHandler.hs
@@ -6,6 +6,7 @@ import Control.Concurrent.STM (TVar, atomically, modifyTVar', readTVar, retry, w
 import Control.Exception (throwIO, try)
 import Control.Monad (when)
 import Control.Monad.IO.Class (liftIO)
+import Data.Coerce (coerce)
 import Data.Functor ((<&>))
 import Data.Int (Int32)
 import GHC (DynFlags (..), Ghc, getSession)
@@ -17,7 +18,7 @@ import GhcWorker.Instrumentation (Hooks (..), InstrumentedHandler (..))
 import Internal.AbiHash (AbiHash (..), showAbiHash)
 import Internal.Compile (compileModuleWithDepsInEps)
 import Internal.CompileHpt (compileModuleWithDepsInHpt)
-import Internal.Log (TraceId, dbg, logFlush, newLog, setLogTarget)
+import Internal.Log (TraceId, dbg, logDebug, logFlush, newLog, setLogTarget)
 import Internal.Metadata (computeMetadata)
 import Internal.Session (Env (..), withGhc, withGhcMhu)
 import Internal.State (ModuleArtifacts (..), WorkerState (..), dumpState)
@@ -27,7 +28,8 @@ import System.Posix.Process (exitImmediately)
 import Types.BuckArgs (BuckArgs, Mode (..), parseBuckArgs, toGhcArgs)
 import qualified Types.BuckArgs
 import Types.GhcHandler (WorkerMode (..))
-import Types.State (Target (Target))
+import Types.Grpc (RequestArgs (..))
+import Types.State (Target)
 
 data LockState = LockStart | LockFreeze Int | LockThaw Int | LockEnd
   deriving stock (Eq, Show)
@@ -91,12 +93,8 @@ dispatch lock workerMode hooks env args =
           pure (code, result)
       pure (code, snd <$> result)
     Just ModeMetadata -> do
-      let target = Target "metadata"
-      liftIO $ setLogTarget env.log target
-      code <- computeMetadata env <&> \case
-        True -> 0
-        False -> 1
-      pure (code, Just target)
+      (success, target) <- computeMetadata env
+      pure (if success then 0 else 1, target)
     Just ModeClose -> do
       dbg "in dispatch. Mode Close"
       _ <- writeCloseOutput args
@@ -159,6 +157,7 @@ ghcHandler lock state workerMode traceId =
     buckArgs <- either (throwIO . userError) pure (parseBuckArgs commandEnv argv)
     args <- toGhcArgs buckArgs
     log <- newLog traceId
+    logDebug log (unlines (coerce argv))
     let env = Env {log, state, args}
     result <- try $ dispatch lock workerMode hooks env buckArgs
     processResult hooks env result
diff --git a/ghc-worker/test/CompileHptTest.hs b/ghc-worker/test/CompileHptTest.hs
index d9c0eea6..389e364a 100644
--- a/ghc-worker/test/CompileHptTest.hs
+++ b/ghc-worker/test/CompileHptTest.hs
@@ -46,7 +46,7 @@ stepMetadata Conf {state, tmp, args0} unit deps = do
   let srcs = [unit.dir </> name | name <- names, takeExtension name == ".hs"]
       env = Env {log, state, args = args srcs}
   dbgp (text ">>> metadata for" <+> ppr unit.uid)
-  success <- computeMetadata env
+  (success, _) <- computeMetadata env
   unless success do
     liftIO $ throwGhcExceptionIO (ProgramError "Metadata failed")
   where
diff --git a/internal/src/Internal/Metadata.hs b/internal/src/Internal/Metadata.hs
index 612c0900..c5ce163d 100644
--- a/internal/src/Internal/Metadata.hs
+++ b/internal/src/Internal/Metadata.hs
@@ -9,14 +9,16 @@ import GHC.Driver.Monad (modifySession, modifySessionM, withSession, withTempSes
 import GHC.Driver.Session (updatePlatformConstants)
 import GHC.Platform.Ways (Way (WayDyn), addWay)
 import GHC.Runtime.Loader (initializeSessionPlugins)
-import GHC.Unit (HomeUnit, UnitDatabase, UnitId, UnitState, initUnits)
+import GHC.Unit (HomeUnit, UnitDatabase, UnitId, UnitState, initUnits, unitIdString)
 import GHC.Unit.Env (HomeUnitEnv (..), UnitEnv (..), unitEnv_insert, unitEnv_keys, updateHug)
 import GHC.Unit.Home.ModInfo (emptyHomePackageTable)
-import Internal.State (WorkerState (..), updateMakeStateVar)
+import Internal.Log (setLogTarget)
 import Internal.MakeFile (doMkDependHS)
 import Internal.Session (Env (..), runSession, withDynFlags)
+import Internal.State (WorkerState (..), updateMakeStateVar)
 import Internal.State.Make (insertUnitEnv, loadState, storeModuleGraph)
 import Internal.State.Stats (logMemStats)
+import Types.State (Target (..))
 
 -- | 'doMkDependHS' needs this to be enabled.
 metadataTempSession :: HscEnv -> HscEnv
@@ -73,13 +75,14 @@ addHomeUnit dflags = do
 --
 -- We especially want to take care that the command line flags aren't applied to the base session before we initialize
 -- the home unit in order to replicate what GHC does in @initMulti@.
-prepareMetadataSession :: Env -> DynFlags -> Ghc ()
+prepareMetadataSession :: Env -> DynFlags -> Ghc UnitId
 prepareMetadataSession env dflags = do
   state <- liftIO $ readMVar env.state
   modifySessionM \ hsc_env -> liftIO (loadState env.log hsc_env state.make)
   unit <- addHomeUnit dflags
   setActiveUnit unit
   storeNewUnit
+  pure unit
   where
     setActiveUnit unit = modifySession (hscUpdateLoggerFlags . hscSetActiveUnitId unit)
 
@@ -104,11 +107,14 @@ writeMetadata srcs = do
 --
 -- Before downsweep, we also create a fresh @Finder@ to prevent 'doMkDependHS' from polluting the cache with entries
 -- with different compilation ways and restore the previous unit env so dependencies are visible.
-computeMetadata :: Env -> IO Bool
+computeMetadata :: Env -> IO (Bool, Maybe Target)
 computeMetadata env = do
-  res <- fmap isJust $ runSession True env $ withDynFlags env \ dflags srcs -> do
-    prepareMetadataSession env dflags
+  res <- runSession True env $ withDynFlags env \ dflags srcs -> do
+    unit <- prepareMetadataSession env dflags
+    let target = Target (unitIdString unit)
+    liftIO $ setLogTarget env.log target
     module_graph <- writeMetadata (fst <$> srcs)
     liftIO $ updateMakeStateVar env.state (storeModuleGraph module_graph)
-    pure (Just ())
-  res <$ logMemStats "after metadata" env.log
+    pure (Just target)
+  logMemStats "after metadata" env.log
+  pure (isJust res, res)