ModuleLoader fetch SourcePackage in main transaction (#27094)

Split behavior in ModuleLoader along better lines:

1. `get_module` looks up ModuleMetadata and SourcePackage and paths to prefetch (previously source package and paths to prefetch would happen later)
2. `get_module_with_metadata` just looks up the source code in storage. It does not need to read from the database at all, because everything from the database was already fetched.

Advantages:
1. We don't have to construct new transactions inside the cache
2. Since we don't need new transactions, we can make TransactionIngredients no longer implement Clone, we don't need a Database or system identity, we don't need to check the timestamp, and we don't need to hackily set entries in the ReadSet.
3. by passing in the set of paths to prefetch, we no longer need the hack to convert TabletId -> ComponentId.

Disadvantages:
1. in a few places application-level callers were calling `get_module_with_metadata` directly, and I changed these places to call `get_module`, which will return the same result but will do an extra fetch of SourcePackage and ModuleMetadata. Luckily these fetches should hit the system index cache, so I don't anticipate performance issues.
GitOrigin-RevId: 42e18a1ca8e2cffb597e89b86e506e1d5bfbe0ea

Author: ldanilek

crates/application/src/lib.rs

@@ -529,8 +529,7 @@ impl<RT: Runtime> Application<RT> {
         scheduled_jobs_pause_client: PauseClient,
         app_auth: Arc<ApplicationAuth>,
     ) -> anyhow::Result<Self> {
-        let module_cache =
-            ModuleCache::new(runtime.clone(), database.clone(), modules_storage.clone()).await;
+        let module_cache = ModuleCache::new(runtime.clone(), modules_storage.clone()).await;
         let module_loader = Arc::new(module_cache.clone());
 
         let system_env_vars = btreemap! {
@@ -818,10 +817,9 @@ impl<RT: Runtime> Application<RT> {
         let Some(source_index) = source_mapped.source_index else {
             return Ok(None);
         };
-        let full_source = self
-            .module_cache
-            .get_module_with_metadata(&mut tx, metadata)
-            .await?;
+        let Some(full_source) = self.module_cache.get_module(&mut tx, path).await? else {
+            return Ok(None);
+        };
         let Some(source_map_str) = &full_source.source_map else {
             return Ok(None);
         };
@@ -1646,13 +1644,14 @@ impl<RT: Runtime> Application<RT> {
             component: ComponentId::Root,
             module_path: AUTH_CONFIG_FILE_NAME.parse()?,
         };
-        let auth_config_metadata = ModuleModel::new(tx).get_metadata(path).await?;
+        let auth_config_metadata = ModuleModel::new(tx).get_metadata(path.clone()).await?;
         if let Some(auth_config_metadata) = auth_config_metadata {
             let environment = auth_config_metadata.environment;
             let auth_config_source = runner
                 .module_cache
-                .get_module_with_metadata(tx, auth_config_metadata)
-                .await?;
+                .get_module(tx, path)
+                .await?
+                .context("Module has metadata but no source")?;
             let auth_config_module = ModuleConfig {
                 path: AUTH_CONFIG_FILE_NAME.parse()?,
                 source: auth_config_source.source.clone(),

crates/application/src/module_cache/mod.rs

@@ -1,4 +1,7 @@
-use std::sync::Arc;
+use std::{
+    collections::BTreeMap,
+    sync::Arc,
+};
 
 use async_lru::async_lru::AsyncLru;
 use async_trait::async_trait;
@@ -10,38 +13,34 @@ use common::{
     },
     runtime::Runtime,
 };
-use database::{
-    Database,
-    Transaction,
-};
 use futures::FutureExt;
 use isolate::environment::helpers::module_loader::get_module_and_prefetch;
-use keybroker::Identity;
 use model::{
     config::module_loader::ModuleLoader,
     modules::{
         module_versions::FullModuleSource,
         types::ModuleMetadata,
-        ModuleModel,
     },
-    source_packages::types::SourcePackageId,
+    source_packages::types::{
+        SourcePackage,
+        SourcePackageId,
+    },
 };
 use storage::Storage;
+use sync_types::CanonicalizedModulePath;
 use value::ResolvedDocumentId;
 
 mod metrics;
 
 #[derive(Clone)]
 pub struct ModuleCache<RT: Runtime> {
-    database: Database<RT>,
-
     modules_storage: Arc<dyn Storage>,
 
     cache: AsyncLru<RT, (ResolvedDocumentId, SourcePackageId), FullModuleSource>,
 }
 
 impl<RT: Runtime> ModuleCache<RT> {
-    pub async fn new(rt: RT, database: Database<RT>, modules_storage: Arc<dyn Storage>) -> Self {
+    pub async fn new(rt: RT, modules_storage: Arc<dyn Storage>) -> Self {
         let cache = AsyncLru::new(
             rt.clone(),
             *MODULE_CACHE_MAX_SIZE_BYTES,
@@ -50,7 +49,6 @@ impl<RT: Runtime> ModuleCache<RT> {
         );
 
         Self {
-            database,
             modules_storage,
             cache,
         }
@@ -62,29 +60,27 @@ impl<RT: Runtime> ModuleLoader<RT> for ModuleCache<RT> {
     #[minitrace::trace]
     async fn get_module_with_metadata(
         &self,
-        tx: &mut Transaction<RT>,
         module_metadata: ParsedDocument<ModuleMetadata>,
+        source_package: ParsedDocument<SourcePackage>,
+        paths_to_prefetch: BTreeMap<ResolvedDocumentId, CanonicalizedModulePath>,
     ) -> anyhow::Result<Arc<FullModuleSource>> {
         let timer = metrics::module_cache_get_module_timer();
 
         let key = (module_metadata.id(), module_metadata.source_package_id);
-        let mut cache_tx = self.database.begin(Identity::system()).await?;
         let modules_storage = self.modules_storage.clone();
         let result = self
             .cache
             .get_and_prepopulate(
                 key,
-                async move {
-                    get_module_and_prefetch(&mut cache_tx, modules_storage, module_metadata).await
-                }
+                get_module_and_prefetch(
+                    modules_storage,
+                    module_metadata,
+                    source_package,
+                    paths_to_prefetch,
+                )
                 .boxed(),
             )
             .await?;
-        // Record read dependency on the module version so the transactions
-        // read same is the same regardless if we hit the cache or not.
-        // This is not technically needed since the module version is immutable,
-        // but better safe and consistent that sorry.
-        ModuleModel::new(tx).record_module_version_read_dependency(key.0)?;
 
         timer.finish();
         Ok(result)

crates/database/src/transaction.rs

@@ -971,18 +971,6 @@ impl<RT: Runtime> Transaction<RT> {
             .await
     }
 
-    /// Used when a system table is served from cache - to manually add a read
-    /// dependency on that system table.
-    pub fn record_system_table_cache_hit(
-        &mut self,
-        index_name: TabletIndexName,
-        fields: IndexedFields,
-        interval: Interval,
-    ) {
-        self.reads
-            .record_indexed_derived(index_name, fields, interval)
-    }
-
     #[cfg(any(test, feature = "testing"))]
     pub fn set_index_size_hard_limit(&mut self, size: usize) {
         self.index_size_override = Some(size);

crates/function_runner/src/in_memory_indexes.rs

@@ -87,9 +87,8 @@ use crate::{
     FunctionWrites,
 };
 
-/// Convenience struct for data required to create a [Transaction]. This struct
-/// is cloneable, while [Transaction] is not.
-#[derive(Clone)]
+/// Struct from which you can create a [Transaction].
+/// TODO: delete this in favor of using Transaction directly.
 pub struct TransactionIngredients<RT: Runtime> {
     pub ts: RepeatableTimestamp,
     pub identity: Identity,

crates/function_runner/src/module_cache.rs

@@ -1,4 +1,7 @@
-use std::sync::Arc;
+use std::{
+    collections::BTreeMap,
+    sync::Arc,
+};
 
 use async_lru::async_lru::AsyncLru;
 use async_trait::async_trait;
@@ -10,23 +13,23 @@ use common::{
     },
     runtime::Runtime,
 };
-use database::Transaction;
 use futures::FutureExt;
 use isolate::environment::helpers::module_loader::get_module_and_prefetch;
 use model::{
     config::module_loader::ModuleLoader,
     modules::{
         module_versions::FullModuleSource,
         types::ModuleMetadata,
-        ModuleModel,
     },
-    source_packages::types::SourcePackageId,
+    source_packages::types::{
+        SourcePackage,
+        SourcePackageId,
+    },
 };
 use storage::Storage;
+use sync_types::CanonicalizedModulePath;
 use value::ResolvedDocumentId;
 
-use crate::in_memory_indexes::TransactionIngredients;
-
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
 pub(crate) struct ModuleCacheKey {
     instance_name: String,
@@ -51,7 +54,6 @@ impl<RT: Runtime> ModuleCache<RT> {
 pub(crate) struct FunctionRunnerModuleLoader<RT: Runtime> {
     pub cache: ModuleCache<RT>,
     pub instance_name: String,
-    pub transaction_ingredients: TransactionIngredients<RT>,
     pub modules_storage: Arc<dyn Storage>,
 }
 
@@ -60,30 +62,30 @@ impl<RT: Runtime> ModuleLoader<RT> for FunctionRunnerModuleLoader<RT> {
     #[minitrace::trace]
     async fn get_module_with_metadata(
         &self,
-        tx: &mut Transaction<RT>,
         module_metadata: ParsedDocument<ModuleMetadata>,
+        source_package: ParsedDocument<SourcePackage>,
+        paths_to_prefetch: BTreeMap<ResolvedDocumentId, CanonicalizedModulePath>,
     ) -> anyhow::Result<Arc<FullModuleSource>> {
-        // The transaction we're getting modules for should be from the same ts as when
-        // this module loader was created.
-        assert_eq!(tx.begin_timestamp(), self.transaction_ingredients.ts);
-
         let instance_name = self.instance_name.clone();
         let key = ModuleCacheKey {
             instance_name: self.instance_name.clone(),
             module_id: module_metadata.id(),
             source_package_id: module_metadata.source_package_id,
         };
-        let mut transaction = self.transaction_ingredients.clone().try_into()?;
         let modules_storage = self.modules_storage.clone();
         let result = self
             .cache
             .0
             .get_and_prepopulate(
                 key.clone(),
                 async move {
-                    let modules =
-                        get_module_and_prefetch(&mut transaction, modules_storage, module_metadata)
-                            .await;
+                    let modules = get_module_and_prefetch(
+                        modules_storage,
+                        module_metadata,
+                        source_package,
+                        paths_to_prefetch,
+                    )
+                    .await;
                     modules
                         .into_iter()
                         .map(move |((module_id, source_package_id), source)| {
@@ -101,11 +103,6 @@ impl<RT: Runtime> ModuleLoader<RT> for FunctionRunnerModuleLoader<RT> {
                 .boxed(),
             )
             .await?;
-        // Record read dependency on the module version so the transactions
-        // read same is the same regardless if we hit the cache or not.
-        // This is not technically needed since the module version is immutable,
-        // but better safe and consistent that sorry.
-        ModuleModel::new(tx).record_module_version_read_dependency(key.module_id)?;
 
         Ok(result)
     }

crates/function_runner/src/server.rs

@@ -325,7 +325,7 @@ impl<RT: Runtime, S: StorageForInstance<RT>> FunctionRunnerCore<RT, S> {
                 retention_validator,
             )
             .await?;
-        let transaction = transaction_ingredients.clone().try_into()?;
+        let transaction = transaction_ingredients.try_into()?;
         Ok(transaction)
     }
 
@@ -387,7 +387,7 @@ impl<RT: Runtime, S: StorageForInstance<RT>> FunctionRunnerCore<RT, S> {
                 retention_validator,
             )
             .await?;
-        let mut transaction = transaction_ingredients.clone().try_into()?;
+        let mut transaction = transaction_ingredients.try_into()?;
         let storage = self
             .storage
             .storage_for_instance(&mut transaction, StorageUseCase::Files)
@@ -406,7 +406,6 @@ impl<RT: Runtime, S: StorageForInstance<RT>> FunctionRunnerCore<RT, S> {
             module_loader: Arc::new(FunctionRunnerModuleLoader {
                 instance_name: instance_name.clone(),
                 cache: self.module_cache.clone(),
-                transaction_ingredients,
                 modules_storage,
             }),
         };

crates/isolate/src/environment/action/phase.rs

@@ -4,9 +4,9 @@ use std::{
     sync::Arc,
 };
 
+use anyhow::Context;
 use common::{
     components::{
-        ComponentId,
         ComponentPath,
         Reference,
         Resource,
@@ -38,6 +38,7 @@ use model::{
         types::ModuleMetadata,
         ModuleModel,
     },
+    source_packages::SourcePackageModel,
     udf_config::UdfConfigModel,
 };
 use parking_lot::Mutex;
@@ -154,40 +155,45 @@ impl<RT: Runtime> ActionPhase<RT> {
             .map(|c| ChaCha12Rng::from_seed(c.import_phase_rng_seed));
         let import_time_unix_timestamp = udf_config.as_ref().map(|c| c.import_phase_unix_timestamp);
 
-        let module_metadata = with_release_permit(timeout, permit_slot, async {
-            let result = if !*COMPONENTS_ENABLED {
+        let (module_metadata, source_package) = with_release_permit(timeout, permit_slot, async {
+            let module_metadata = ModuleModel::new(&mut tx)
+                .get_all_metadata(component_id)
+                .await?;
+            let source_package = SourcePackageModel::new(&mut tx, component_id.into())
+                .get_latest()
+                .await?;
+            if !*COMPONENTS_ENABLED {
                 anyhow::ensure!(self.component.is_root());
-                ModuleModel::new(&mut tx)
-                    .get_all_metadata(ComponentId::Root)
-                    .await?
             } else {
-                let module_metadata = ModuleModel::new(&mut tx)
-                    .get_all_metadata(component_id)
-                    .await?;
-
                 let loaded_resources = ComponentsModel::new(&mut tx)
                     .preload_resources(component_id)
                     .await?;
                 {
                     let mut resources = resources.lock();
                     *resources = loaded_resources;
                 }
-
-                module_metadata
-            };
-            Ok(result)
+            }
+            Ok((module_metadata, source_package))
         })
         .await?;
 
         let modules = with_release_permit(timeout, permit_slot, async {
             let mut modules = BTreeMap::new();
+            let paths_to_prefetch: BTreeMap<_, _> = module_metadata
+                .iter()
+                .map(|metadata| (metadata.id(), metadata.path.clone()))
+                .collect();
             for metadata in module_metadata {
                 if metadata.path.is_system() {
                     continue;
                 }
                 let path = metadata.path.clone();
                 let module = module_loader
-                    .get_module_with_metadata(&mut tx, metadata.clone())
+                    .get_module_with_metadata(
+                        metadata.clone(),
+                        source_package.clone().context("source package not found")?,
+                        paths_to_prefetch.clone(),
+                    )
                     .await?;
                 modules.insert(path, (metadata.into_value(), module));
             }