preload branches into the branch cache when switching projects

parser-typechecker/src/Unison/Codebase.hs

@@ -59,6 +59,7 @@ module Unison.Codebase
     getShallowProjectRootByNames,
     expectProjectBranchRoot,
     getBranchAtProjectPath,
+    preloadProjectBranch,
 
     -- * Root branch
     SqliteCodebase.Operations.namesAtPath,
@@ -579,3 +580,11 @@ resolveProjectPathIds (PP.ProjectPath projectId projectBranchId path) = do
   proj <- Q.expectProject projectId
   projBranch <- Q.expectProjectBranch projectId projectBranchId
   pure $ PP.ProjectPath proj projBranch path
+
+-- | Starts loading the given project branch into cache in a background thread without blocking.
+preloadProjectBranch :: (MonadUnliftIO m) => Codebase m v a -> ProjectAndBranch Db.ProjectId Db.ProjectBranchId -> m ()
+preloadProjectBranch codebase (ProjectAndBranch projectId branchId) = do
+  ch <- runTransaction codebase $ do
+    causalHashId <- Q.expectProjectBranchHead projectId branchId
+    Q.expectCausalHash causalHashId
+  preloadProjectRoot codebase ch

parser-typechecker/src/Unison/Codebase/SqliteCodebase.hs

@@ -39,6 +39,7 @@ import Unison.Codebase.SqliteCodebase.Branch.Dependencies qualified as BD
 import Unison.Codebase.SqliteCodebase.Migrations qualified as Migrations
 import Unison.Codebase.SqliteCodebase.Operations qualified as CodebaseOps
 import Unison.Codebase.SqliteCodebase.Paths
+import Unison.Codebase.SqliteCodebase.ProjectRootCache qualified as ProjectRootCache
 import Unison.Codebase.SqliteCodebase.SyncEphemeral qualified as SyncEphemeral
 import Unison.Codebase.Type (LocalOrRemote (..))
 import Unison.Codebase.Type qualified as C
@@ -57,6 +58,8 @@ import Unison.Type (Type)
 import Unison.Util.Timing (time)
 import Unison.WatchKind qualified as UF
 import UnliftIO (UnliftIO (..), finally)
+import UnliftIO qualified as UnliftIO
+import UnliftIO.Concurrent qualified as UnliftIO
 import UnliftIO.Directory (createDirectoryIfMissing, doesDirectoryExist, doesFileExist)
 import UnliftIO.STM
 
@@ -162,6 +165,7 @@ sqliteCodebase ::
   m (Either Codebase1.OpenCodebaseError r)
 sqliteCodebase debugName root localOrRemote lockOption migrationStrategy action = handleLockOption do
   branchCache <- newBranchCache
+  projectRootCache <- ProjectRootCache.newProjectRootCache 5 {- Cache the last n project roots for quick switching. -}
   getDeclType <- CodebaseOps.makeCachedTransaction 2048 CodebaseOps.getDeclType
   -- The v1 codebase interface has operations to read and write individual definitions
   -- whereas the v2 codebase writes them as complete components.  These two fields buffer
@@ -242,6 +246,16 @@ sqliteCodebase debugName root localOrRemote lockOption migrationStrategy action
               withRunInIO \runInIO ->
                 runInIO (runTransaction (CodebaseOps.putBranch (Branch.transform (Sqlite.unsafeIO . runInIO) branch)))
 
+            preloadProjectRoot :: CausalHash -> m ()
+            preloadProjectRoot h = do
+              void . UnliftIO.forkIO $ void $ do
+                getBranchForHash h >>= \case
+                  Nothing -> pure ()
+                  Just b -> do
+                    ProjectRootCache.stashBranch projectRootCache b
+                    UnliftIO.evaluate b
+                    pure ()
+
             syncFromDirectory :: Codebase1.CodebasePath -> Branch m -> m ()
             syncFromDirectory srcRoot b =
               withConnection (debugName ++ ".sync.src") srcRoot \srcConn ->
@@ -307,7 +321,8 @@ sqliteCodebase debugName root localOrRemote lockOption migrationStrategy action
                   filterTermsByReferentIdHavingTypeImpl,
                   termReferentsByPrefix = referentsByPrefix,
                   withConnection = withConn,
-                  withConnectionIO = withConnection debugName root
+                  withConnectionIO = withConnection debugName root,
+                  preloadProjectRoot
                 }
         Right <$> action codebase
   where

parser-typechecker/src/Unison/Codebase/SqliteCodebase/ProjectRootCache.hs

@@ -0,0 +1,28 @@
+-- | Simple cache which just keeps the last n relevant project branches in memory.
+-- The Branch Cache handles all the lookups of the actual branch data by hash, this cache serves only to keep the last
+-- n accessed branches in memory so they don't get garbage collected. See the Branch Cache for more context.
+--
+-- This speeds up switching back and forth between project branches, and also serves to keep the current project branch
+-- in memory so it won't be cleaned up by the Branch Cache, since the Branch Cache only keeps
+-- a weak reference to the current branch and we no longer keep the actual branch in LoopState.
+module Unison.Codebase.SqliteCodebase.ProjectRootCache
+  ( newProjectRootCache,
+    stashBranch,
+  )
+where
+
+import Control.Concurrent.STM
+import Unison.Codebase.Branch
+import Unison.Prelude
+
+data ProjectRootCache m = ProjectRootCache {capacity :: Int, cached :: TVar [Branch m]}
+
+newProjectRootCache :: (MonadIO m) => Int -> m (ProjectRootCache n)
+newProjectRootCache capacity = do
+  var <- liftIO $ newTVarIO []
+  pure (ProjectRootCache capacity var)
+
+stashBranch :: (MonadIO n) => ProjectRootCache m -> Branch m -> n ()
+stashBranch ProjectRootCache {capacity, cached} branch = do
+  liftIO . atomically $ do
+    modifyTVar cached $ \branches -> take capacity (branch : filter (/= branch) branches)

parser-typechecker/src/Unison/Codebase/Type.hs

@@ -80,7 +80,12 @@ data Codebase m v a = Codebase
     -- | Acquire a new connection to the same underlying database file this codebase object connects to.
     withConnection :: forall x. (Sqlite.Connection -> m x) -> m x,
     -- | Acquire a new connection to the same underlying database file this codebase object connects to.
-    withConnectionIO :: forall x. (Sqlite.Connection -> IO x) -> IO x
+    withConnectionIO :: forall x. (Sqlite.Connection -> IO x) -> IO x,
+    -- | This optimization allows us to pre-fetch a branch from SQLite into the branch cache when we know we'll need it
+    -- soon, but not immediately. E.g. the user has switched a branch, but hasn't run any commands on it yet.
+    --
+    -- This combinator returns immediately, but warms the cache in the background with the desired branch.
+    preloadProjectRoot :: CausalHash -> m ()
   }
 
 -- | Whether a codebase is local or remote.

parser-typechecker/unison-parser-typechecker.cabal

@@ -81,6 +81,7 @@ library
       Unison.Codebase.SqliteCodebase.Migrations.MigrateSchema7To8
       Unison.Codebase.SqliteCodebase.Operations
       Unison.Codebase.SqliteCodebase.Paths
+      Unison.Codebase.SqliteCodebase.ProjectRootCache
       Unison.Codebase.SqliteCodebase.SyncEphemeral
       Unison.Codebase.TermEdit
       Unison.Codebase.TermEdit.Typing

unison-cli/src/Unison/Cli/Monad.hs

@@ -94,8 +94,6 @@ import Unison.Syntax.Parser qualified as Parser
 import Unison.Term (Term)
 import Unison.Type (Type)
 import Unison.UnisonFile qualified as UF
-import UnliftIO qualified
-import UnliftIO.Concurrent qualified as UnliftIO
 import Unsafe.Coerce (unsafeCoerce)
 
 -- | The main command-line app monad.
@@ -392,17 +390,13 @@ cd path = do
   #projectPathStack %= NonEmpty.cons newPP
 
 switchProject :: ProjectAndBranch ProjectId ProjectBranchId -> Cli ()
-switchProject (ProjectAndBranch projectId branchId) = do
+switchProject pab@(ProjectAndBranch projectId branchId) = do
   Env {codebase} <- ask
   let newPP = PP.ProjectPath projectId branchId Path.absoluteEmpty
   #projectPathStack %= NonEmpty.cons newPP
   runTransaction $ do Q.setMostRecentBranch projectId branchId
   setMostRecentProjectPath newPP
-  -- Prime the cache with the new project branch root so it's ready when a command needs it.
-  void . liftIO . UnliftIO.forkIO $ do
-    b <- Codebase.expectProjectBranchRoot codebase projectId branchId
-    -- Force the branch in the background thread to avoid delays later.
-    void $ UnliftIO.evaluate b
+  liftIO $ Codebase.preloadProjectBranch codebase pab
 
 -- | Pop the latest path off the stack, if it's not the only path in the stack.
 --

unison-cli/tests/Unison/Test/Cli/Monad.hs

@@ -36,8 +36,7 @@ dummyEnv = undefined
 dummyLoopState :: Cli.LoopState
 dummyLoopState =
   Cli.LoopState
-    { currentProjectRoot = undefined,
-      projectPathStack = undefined,
+    { projectPathStack = undefined,
       latestFile = Nothing,
       latestTypecheckedFile = Nothing,
       lastInput = Nothing,