Merge pull request #125 from unisoncomputing/cp/doc-eval-speedup

Speed up doc evaluation on Share

Author: ChrisPenner

share-api.cabal

@@ -45,6 +45,8 @@ library
       Share.Branch
       Share.ChatApps
       Share.Codebase
+      Share.Codebase.CodebaseRuntime
+      Share.Codebase.CodeCache
       Share.Codebase.Types
       Share.Contribution
       Share.Env

src/Share/Backend.hs

@@ -38,12 +38,15 @@ import Control.Lens hiding ((??))
 import Data.List (unzip4, zipWith4, zipWith5)
 import Data.List qualified as List
 import Data.Map qualified as Map
+import Data.Set qualified as Set
 import Share.Codebase qualified as Codebase
+import Share.Codebase.CodeCache qualified as CC
 import Share.Codebase.Types (CodebaseRuntime (CodebaseRuntime, cachedEvalResult))
 import Share.Postgres (QueryM)
 import Share.Postgres qualified as PG
 import Share.Postgres.Causal.Conversions (namespaceStatsPgToV2)
 import Share.Postgres.Causal.Queries qualified as CausalQ
+import Share.Postgres.Definitions.Queries qualified as DefnQ
 import Share.Prelude
 import Share.Utils.Lens (asListOfDeduped)
 import U.Codebase.Branch qualified as V2Branch
@@ -289,24 +292,35 @@ evalDocRef ::
   Codebase.CodebaseRuntime s IO ->
   V2.TermReference ->
   m (Doc.EvaluatedDoc Symbol)
-evalDocRef codebase (CodebaseRuntime {codeLookup, cachedEvalResult, unisonRuntime}) termRef = do
+evalDocRef codebase (CodebaseRuntime {codeLookup, codeCache, cachedEvalResult, unisonRuntime}) termRef = PG.transactionSpan "evalDocRef" mempty do
   let tm = Term.ref () termRef
-  -- TODO: batchify evalDoc.
-  Doc.evalDoc terms typeOf eval decls tm
+  case termRef of
+    Reference.Builtin _ -> Doc.evalDoc terms typeOf eval decls tm
+    Reference.DerivedId refId -> do
+      termId <- DefnQ.expectTermIdsByRefIdsOf id refId
+      (termDeps, typeDeps) <- DefnQ.termTransitiveDependencyRefs (Set.singleton termId)
+      -- Prime the cache with all the terms and types we know we'll need.
+      -- No need to store them manually, they'll be persistently cached automatically just by
+      -- fetching
+      _ <- CC.getTermsAndTypesByRefIdsOf codeCache traversed (Set.toList termDeps)
+      _ <- CC.getTypeDeclsByRefIdsOf codeCache traversed (Set.toList typeDeps)
+      -- TODO: batchify evalDoc within unison
+      Doc.evalDoc terms typeOf eval decls tm
   where
+    -- Loading one at a time is inefficient, so we prime the cache above.
     terms :: Reference -> m (Maybe (V1.Term Symbol ()))
-    terms termRef@(Reference.Builtin _) = pure (Just (Term.ref () termRef))
-    terms (Reference.DerivedId termRef) =
-      -- TODO: batchify properly
-      fmap (Term.unannotate . fst) <$> (Codebase.loadTermAndTypeByRefIdsOf codebase id termRef)
+    terms = CC.termsForRefsOf codeCache id
 
+    -- Loading one at a time is inefficient, so we prime the cache above.
     typeOf :: Referent.Referent -> m (Maybe (V1.Type Symbol ()))
-    typeOf termRef =
-      -- TODO: batchify properly
-      fmap void <$> Codebase.loadTypesOfReferentsOf codebase id (Cv.referent1to2 termRef)
+    typeOf = CC.typesOfReferentsOf codeCache id
+
+    -- Loading one at a time is inefficient, so we prime the cache above.
+    decls :: Reference -> m (Maybe (DD.Decl Symbol ()))
+    decls ref = fmap (DD.amap (const ())) <$> CC.getTypeDeclsByRefsOf codeCache id ref
 
     eval :: V1.Term Symbol a -> m (Maybe (V1.Term Symbol ()))
-    eval (Term.amap (const mempty) -> tm) = do
+    eval (Term.amap (const mempty) -> tm) = PG.transactionSpan "eval" mempty do
       -- We use an empty ppe for evalutation, it's only used for adding additional context to errors.
       let evalPPE = PPE.empty
       termRef <- fmap eitherToMaybe . PG.transactionUnsafeIO . liftIO $ Rt.evaluateTerm' codeLookup cachedEvalResult evalPPE unisonRuntime tm
@@ -320,12 +334,6 @@ evalDocRef codebase (CodebaseRuntime {codeLookup, cachedEvalResult, unisonRuntim
         _ -> pure ()
       pure $ termRef <&> Term.amap (const mempty) . snd
 
-    decls :: Reference -> m (Maybe (DD.Decl Symbol ()))
-    decls (Reference.DerivedId typeRef) =
-      -- TODO: batchify properly
-      fmap (DD.amap (const ())) <$> (Codebase.loadTypeDeclarationsByRefIdsOf codebase id typeRef)
-    decls _ = pure Nothing
-
 -- | Find all definitions and children reachable from the given 'V2Branch.Branch',
 lsBranch ::
   (QueryM m) =>
@@ -376,7 +384,7 @@ typeDeclHeader ::
   m (DisplayObject Syntax.SyntaxText Syntax.SyntaxText)
 typeDeclHeader codebase ppe r = case Reference.toId r of
   Just rid ->
-    Codebase.loadTypeDeclarationsByRefIdsOf codebase id rid <&> \case
+    Codebase.loadV1TypeDeclarationsByRefIdsOf codebase id rid <&> \case
       Nothing -> DisplayObject.MissingObject (Reference.toShortHash r)
       Just decl ->
         DisplayObject.UserObject $

src/Share/BackgroundJobs/Diffs/CausalDiffs.hs

@@ -9,6 +9,7 @@ import Share.BackgroundJobs.Errors (reportError)
 import Share.BackgroundJobs.Monad (Background)
 import Share.BackgroundJobs.Workers (newWorker)
 import Share.Codebase qualified as Codebase
+import Share.Codebase.CodebaseRuntime qualified as CR
 import Share.Env qualified as Env
 import Share.IDs qualified as IDs
 import Share.Metrics qualified as Metrics
@@ -98,8 +99,8 @@ maybeComputeAndStoreCausalDiff authZReceipt unisonRuntime (CausalDiffInfo {fromC
   let toCodebase = Codebase.codebaseEnv authZReceipt $ Codebase.codebaseLocationForUserCodebase toCodebaseOwner
   ContributionsQ.existsPrecomputedNamespaceDiff (fromCodebase, fromCausalId) (toCodebase, toCausalId) >>= \case
     True -> pure False
-    False -> Codebase.withCodebaseRuntime fromCodebase unisonRuntime \fromRuntime -> do
-      Codebase.withCodebaseRuntime toCodebase unisonRuntime \toRuntime -> do
+    False -> CR.withCodebaseRuntime fromCodebase unisonRuntime \fromRuntime -> do
+      CR.withCodebaseRuntime toCodebase unisonRuntime \toRuntime -> do
         _ <-
           Diffs.computeAndStoreCausalDiff
             authZReceipt

src/Share/BackgroundJobs/Search/DefinitionSync.hs

@@ -388,7 +388,7 @@ syncTypes codebase namesPerspective rootBranchHashId typesCursor = do
             Reference.Builtin _ -> pure (mempty, Arity 0)
             Reference.DerivedId refId -> do
               -- TODO: batchify this
-              decl <- lift (Codebase.loadTypeDeclarationsByRefIdsOf codebase id refId) `whenNothingM` throwError (NoDeclForType fqn ref)
+              decl <- lift (Codebase.loadV1TypeDeclarationsByRefIdsOf codebase id refId) `whenNothingM` throwError (NoDeclForType fqn ref)
               pure $ (tokensForDecl refId decl, Arity . fromIntegral . length . DD.bound $ DD.asDataDecl decl)
           let basicTokens = Set.fromList [NameToken fqn, HashToken $ Reference.toShortHash ref]
           typeSummary <- lift $ Summary.typeSummaryForReference codebase ref (Just fqn) Nothing

src/Share/Codebase.hs

@@ -9,24 +9,22 @@ module Share.Codebase
     codebaseOwner,
     CodebaseRuntime (..),
     codebaseEnv,
-    withCodebaseRuntime,
-    codebaseRuntimeTransaction,
     codebaseForProjectBranch,
     codebaseLocationForUserCodebase,
     codebaseLocationForProjectBranchCodebase,
     codebaseLocationForProjectRelease,
     CodebaseLocation (..),
 
     -- * Definitions
-    loadTermAndTypeByRefIdsOf,
+    loadV1TermAndTypeByRefIdsOf,
     expectTermsByRefIdsOf,
     loadTypesOfTermsOf,
     expectTypesOfTermsOf,
     expectTypesOfReferentsOf,
     expectTypesOfConstructorsOf,
     loadTypesOfConstructorsOf,
     loadTypesOfReferentsOf,
-    loadTypeDeclarationsByRefIdsOf,
+    loadV1TypeDeclarationsByRefIdsOf,
     expectTypeDeclarationsByRefIdsOf,
     loadDeclKindsOf,
     expectDeclKindsOf,
@@ -38,7 +36,6 @@ module Share.Codebase
     -- * Eval
     loadCachedEvalResult,
     saveCachedEvalResult,
-    codeLookupForUser,
 
     -- * Causals
     CausalQ.loadCausalNamespace,
@@ -62,7 +59,6 @@ module Share.Codebase
   )
 where
 
-import Control.Concurrent.STM (TVar, atomically, modifyTVar', newTVarIO, readTVarIO)
 import Control.Lens
 import Data.ByteString.Lazy.Char8 qualified as BL
 import Data.Map qualified as Map
@@ -105,8 +101,6 @@ import U.Codebase.Term qualified as TermV2
 import U.Codebase.Term qualified as V2.Term
 import Unison.Builtin qualified as Builtin
 import Unison.Builtin qualified as Builtins
-import Unison.Codebase.CodeLookup qualified as CL
-import Unison.Codebase.Runtime (Runtime)
 import Unison.Codebase.SqliteCodebase.Conversions qualified as Cv
 import Unison.ConstructorType qualified as CT
 import Unison.DataDeclaration qualified as DD
@@ -119,14 +113,11 @@ import Unison.Reference qualified as Reference
 import Unison.Reference qualified as V1
 import Unison.Referent qualified as V1
 import Unison.Referent qualified as V1Referent
-import Unison.Runtime.IOSource qualified as IOSource
 import Unison.ShortHash
 import Unison.ShortHash qualified as ShortHash
 import Unison.Symbol (Symbol)
-import Unison.Term qualified as Term
 import Unison.Term qualified as V1
 import Unison.Type qualified as V1
-import UnliftIO qualified
 
 data CodebaseError
   = MissingTypeForTerm Reference.Reference
@@ -170,30 +161,6 @@ codebaseEnv !_authZReceipt codebaseLoc = do
   let codebaseOwner = Codebase.codebaseOwnerUserId codebaseLoc
    in CodebaseEnv {codebaseOwner}
 
--- | Construct a Runtime linked to a specific codebase and transaction.
--- Don't use the runtime for one codebase with another codebase.
--- Don't use this runtime in any transaction other than the one where it's created.
-withCodebaseRuntime :: (Exception e) => CodebaseEnv -> Runtime Symbol -> (forall s. CodebaseRuntime s IO -> PG.Transaction e r) -> PG.Transaction e r
-withCodebaseRuntime codebaseEnv sandboxedRuntime f = do
-  rt <- PG.transactionUnsafeIO (codebaseRuntimeTransaction sandboxedRuntime codebaseEnv)
-  PG.asUnliftIOTransaction $ do
-    UnliftIO.withRunInIO \toIO -> do
-      toIO . PG.UnliftIOTransaction $ f $ hoistCodebaseRuntime (toIO . PG.UnliftIOTransaction) rt
-
--- | Ideally, we'd use this – a runtime with lookup actions in transaction, not IO. But that will require refactoring to
--- the runtime interface in ucm, so we can't use it for now. That's bad: we end up unsafely running separate
--- transactions for inner calls to 'codeLookup' / 'cachedEvalResult', which can lead to deadlock due to a starved
--- connection pool.
-codebaseRuntimeTransaction :: Runtime Symbol -> CodebaseEnv -> IO (CodebaseRuntime s (PG.Transaction e))
-codebaseRuntimeTransaction unisonRuntime codebase = do
-  cacheVar <- newTVarIO (CodeLookupCache mempty mempty)
-  pure
-    CodebaseRuntime
-      { codeLookup = codeLookupForUser cacheVar codebase,
-        cachedEvalResult = (fmap . fmap) Term.unannotate . loadCachedEvalResult codebase,
-        unisonRuntime
-      }
-
 -- | Wrap a response in caching.
 -- This combinator respects the cachability stored on the provided auth receipt.
 cachedCodebaseResponse ::
@@ -218,8 +185,8 @@ cachedCodebaseResponse authzReceipt codebaseOwner endpointName providedCachePara
     codebaseViewCacheKey = IDs.toText (codebaseOwnerUserId codebaseOwner)
 
 -- | Load terms and type.
-loadTermAndTypeByRefIdsOf :: (QueryM m) => CodebaseEnv -> Traversal s t TermReferenceId (Maybe (V1.Term Symbol Ann, V1.Type Symbol Ann)) -> s -> m t
-loadTermAndTypeByRefIdsOf codebase trav s = do
+loadV1TermAndTypeByRefIdsOf :: (QueryM m) => CodebaseEnv -> Traversal s t TermReferenceId (Maybe (V1.Term Symbol Ann, V1.Type Symbol Ann)) -> s -> m t
+loadV1TermAndTypeByRefIdsOf codebase trav s = do
   s
     & asListOf trav %%~ \refs -> do
       let hashes = refs <&> \(Reference.Id h _) -> h
@@ -257,7 +224,7 @@ convertTerms2to1Of trav s = do
 expectTermsByRefIdsOf :: (QueryM m) => CodebaseEnv -> Traversal s t TermReferenceId (V1.Term Symbol Ann, V1.Type Symbol Ann) -> s -> m t
 expectTermsByRefIdsOf codebase trav s = do
   s & asListOfDeduped trav \refIds -> do
-    termsAndTypes <- loadTermAndTypeByRefIdsOf codebase traversed refIds
+    termsAndTypes <- loadV1TermAndTypeByRefIdsOf codebase traversed refIds
     for (zip refIds termsAndTypes) \case
       (refId, Nothing) -> unrecoverableError (MissingTerm refId)
       (_, Just (term, typ)) -> pure (term, typ)
@@ -277,7 +244,7 @@ loadTypesOfTermsOf codebase trav s =
                         <&> \typ ->
                           (typ $> builtinAnnotation)
                  in Left builtinType
-      results <- loadTermAndTypeByRefIdsOf codebase (traversed . _Right) partitionedRefs
+      results <- loadV1TermAndTypeByRefIdsOf codebase (traversed . _Right) partitionedRefs
       pure $
         results <&> \case
           Left builtin -> builtin
@@ -329,8 +296,8 @@ loadDeclKindsOf trav s =
           & DefnQ.loadDeclKindsOf (traversed . _Right)
       pure (fmap (either id id) xs)
 
-loadTypeDeclarationsByRefIdsOf :: (QueryM m) => CodebaseEnv -> Traversal s t Reference.Id (Maybe (V1.Decl Symbol Ann)) -> s -> m t
-loadTypeDeclarationsByRefIdsOf codebase trav s =
+loadV1TypeDeclarationsByRefIdsOf :: (QueryM m) => CodebaseEnv -> Traversal s t Reference.Id (Maybe (V1.Decl Symbol Ann)) -> s -> m t
+loadV1TypeDeclarationsByRefIdsOf codebase trav s =
   s
     & asListOf trav %%~ \refIds -> do
       DefnQ.loadDeclsByRefIdsOf codebase traversed refIds
@@ -343,7 +310,7 @@ expectTypeDeclarationsByRefIdsOf :: (QueryM m) => CodebaseEnv -> Traversal s t R
 expectTypeDeclarationsByRefIdsOf codebase trav s = do
   s
     & asListOf trav %%~ \refIds -> do
-      results <- loadTypeDeclarationsByRefIdsOf codebase traversed refIds
+      results <- loadV1TypeDeclarationsByRefIdsOf codebase traversed refIds
       for (zip refIds results) \case
         (refId, Nothing) -> unrecoverableError (MissingDecl refId)
         (_, Just decl) -> pure decl
@@ -398,7 +365,7 @@ loadTypesOfConstructorsOf codebase trav s =
             refs <&> \case
               (Reference.Builtin _t, _conId) -> Left Nothing
               (Reference.DerivedId refId, conId) -> Right ((refId, conId), refId)
-      results <- loadTypeDeclarationsByRefIdsOf codebase (traversed . _Right . _2) partitionedRefs
+      results <- loadV1TypeDeclarationsByRefIdsOf codebase (traversed . _Right . _2) partitionedRefs
       for results \case
         (Right ((refId, _conId), Nothing)) -> (unrecoverableError (MissingDecl refId))
         (Right ((_refId, conId), (Just decl))) -> do
@@ -414,49 +381,6 @@ expectTypesOfConstructorsOf codebase trav s =
         ((ref, conId), Nothing) -> unrecoverableError (MissingTypeForConstructor ref conId)
         (_, Just r) -> pure r
 
-data CodeLookupCache = CodeLookupCache
-  { termCache :: Map Reference.Id (V1.Term Symbol Ann, V1.Type Symbol Ann),
-    typeCache :: Map Reference.Id (V1.Decl Symbol Ann)
-  }
-
--- | TODO: The code lookup should either be batchified or we should preload the cache with
--- everything we think we'll need.
-codeLookupForUser :: TVar CodeLookupCache -> CodebaseEnv -> CL.CodeLookup Symbol (PG.Transaction e) Ann
-codeLookupForUser cacheVar codebaseOwner = do
-  CL.CodeLookup (fmap (fmap fst) . getTermAndType) (fmap (fmap snd) . getTermAndType) getTypeDecl
-    <> Builtin.codeLookup
-    <> IOSource.codeLookupM
-  where
-    getTermAndType ::
-      Reference.Id ->
-      PG.Transaction e (Maybe (V1.Term Symbol Ann, V1.Type Symbol Ann))
-    getTermAndType r = do
-      CodeLookupCache {termCache} <- PG.transactionUnsafeIO (readTVarIO cacheVar)
-      case Map.lookup r termCache of
-        Just termAndType -> pure (Just termAndType)
-        Nothing -> do
-          maybeTermAndType <- loadTermAndTypeByRefIdsOf codebaseOwner id r
-          whenJust maybeTermAndType \termAndType -> do
-            PG.transactionUnsafeIO do
-              atomically do
-                modifyTVar' cacheVar \CodeLookupCache {termCache, ..} ->
-                  CodeLookupCache {termCache = Map.insert r termAndType termCache, ..}
-          pure maybeTermAndType
-
-    getTypeDecl :: Reference.Id -> PG.Transaction e (Maybe (V1.Decl Symbol Ann))
-    getTypeDecl r = do
-      CodeLookupCache {typeCache} <- PG.transactionUnsafeIO (readTVarIO cacheVar)
-      case Map.lookup r typeCache of
-        Just typ -> pure (Just typ)
-        Nothing -> do
-          maybeType <- loadTypeDeclarationsByRefIdsOf codebaseOwner id r
-          whenJust maybeType \typ ->
-            PG.transactionUnsafeIO do
-              atomically do
-                modifyTVar' cacheVar \CodeLookupCache {typeCache, ..} ->
-                  CodeLookupCache {typeCache = Map.insert r typ typeCache, ..}
-          pure maybeType
-
 -- | Look up the result of evaluating a term if we have it cached.
 --
 -- This is intentionally not in CodebaseM because it's used to build the CodebaseEnv.