use SourcePackageId instead of ModuleVersion as cache key (#26934)

ldanilek · Convex, Inc. · commit 5e2c4a8116f2 · 2024-06-11T18:55:32.000Z
in my bid to get rid of module versions, we can now start using source package id as the key for cache lookups.
this is more reliable because source packages don't disappear when the module gets updated.
and once we do this we can allow REPL lookups to go through the cache. without the package id change, we can't because REPL module version might conflict with a real module version, but the package ids won't conflict.

GitOrigin-RevId: a4429407817c8fe1e2eda998c3f0ac06bbb19551
diff --git a/crates/application/src/module_cache/mod.rs b/crates/application/src/module_cache/mod.rs
@@ -20,20 +20,14 @@ use keybroker::Identity;
 use model::{
     config::module_loader::ModuleLoader,
     modules::{
-        module_versions::{
-            FullModuleSource,
-            ModuleVersion,
-        },
+        module_versions::FullModuleSource,
         types::ModuleMetadata,
         ModuleModel,
-        MODULE_VERSIONS_TABLE,
     },
+    source_packages::types::SourcePackageId,
 };
 use storage::Storage;
-use value::{
-    ResolvedDocumentId,
-    TableNamespace,
-};
+use value::ResolvedDocumentId;
 
 mod metrics;
 
@@ -43,7 +37,7 @@ pub struct ModuleCache<RT: Runtime> {
 
     modules_storage: Arc<dyn Storage>,
 
-    cache: AsyncLru<RT, (ResolvedDocumentId, ModuleVersion), FullModuleSource>,
+    cache: AsyncLru<RT, (ResolvedDocumentId, SourcePackageId), FullModuleSource>,
 }
 
 impl<RT: Runtime> ModuleCache<RT> {
@@ -73,20 +67,7 @@ impl<RT: Runtime> ModuleLoader<RT> for ModuleCache<RT> {
     ) -> anyhow::Result<Arc<FullModuleSource>> {
         let timer = metrics::module_cache_get_module_timer();
 
-        // If this transaction wrote to module_versions (true for REPLs), we cannot use
-        // the cache, load the module directly.
-        let module_versions_table_id = tx
-            .table_mapping()
-            .namespace(TableNamespace::by_component_definition_TODO())
-            .id(&MODULE_VERSIONS_TABLE)?;
-        if tx.writes().has_written_to(&module_versions_table_id) {
-            let source = ModuleModel::new(tx)
-                .get_source_from_db(module_metadata.id(), module_metadata.latest_version)
-                .await?;
-            return Ok(Arc::new(source));
-        }
-
-        let key = (module_metadata.id(), module_metadata.latest_version);
+        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
diff --git a/crates/database/src/writes.rs b/crates/database/src/writes.rs
@@ -66,8 +66,6 @@ pub struct Writes {
     user_tx_size: TransactionWriteSize,
     // Size of writes to system tables
     system_tx_size: TransactionWriteSize,
-    // When we write to module versions we cannot use the module cache.
-    written_tables: BTreeSet<TabletIdAndTableNumber>,
 }
 
 #[derive(Clone, Debug, Default, PartialEq)]
@@ -87,7 +85,6 @@ impl Writes {
             generated_ids: BTreeSet::new(),
             user_tx_size: TransactionWriteSize::default(),
             system_tx_size: TransactionWriteSize::default(),
-            written_tables: BTreeSet::new(),
         }
     }
 
@@ -96,10 +93,6 @@ impl Writes {
         self.updates.is_empty()
     }
 
-    pub fn has_written_to(&self, table_id: &TabletIdAndTableNumber) -> bool {
-        self.written_tables.contains(table_id)
-    }
-
     pub fn update(
         &mut self,
         bootstrap_tables: BootstrapTableIds,
@@ -121,8 +114,6 @@ impl Writes {
             .map(|d| d.value().size())
             .unwrap_or(0);
 
-        self.written_tables.insert(*document_id.table());
-
         let tx_size = if is_system_document {
             &mut self.system_tx_size
         } else {
diff --git a/crates/function_runner/src/module_cache.rs b/crates/function_runner/src/module_cache.rs
@@ -16,14 +16,11 @@ use isolate::environment::helpers::module_loader::get_module_and_prefetch;
 use model::{
     config::module_loader::ModuleLoader,
     modules::{
-        module_versions::{
-            FullModuleSource,
-            ModuleVersion,
-        },
+        module_versions::FullModuleSource,
         types::ModuleMetadata,
         ModuleModel,
-        MODULE_VERSIONS_TABLE,
     },
+    source_packages::types::SourcePackageId,
 };
 use storage::Storage;
 use value::ResolvedDocumentId;
@@ -34,7 +31,7 @@ use crate::in_memory_indexes::TransactionIngredients;
 pub(crate) struct ModuleCacheKey {
     instance_name: String,
     module_id: ResolvedDocumentId,
-    module_version: ModuleVersion,
+    source_package_id: SourcePackageId,
 }
 
 #[derive(Clone)]
@@ -70,27 +67,11 @@ impl<RT: Runtime> ModuleLoader<RT> for FunctionRunnerModuleLoader<RT> {
         // this module loader was created.
         assert_eq!(tx.begin_timestamp(), self.transaction_ingredients.ts);
 
-        let namespace = tx
-            .table_mapping()
-            .tablet_namespace(module_metadata.id().table().tablet_id)?;
-        // If this transaction wrote to module_versions (true for REPLs), we cannot use
-        // the cache, load the module directly.
-        let module_versions_table_id = tx
-            .table_mapping()
-            .namespace(namespace)
-            .id(&MODULE_VERSIONS_TABLE)?;
-        if tx.writes().has_written_to(&module_versions_table_id) {
-            let source = ModuleModel::new(tx)
-                .get_source_from_db(module_metadata.id(), module_metadata.latest_version)
-                .await?;
-            return Ok(Arc::new(source));
-        }
-
         let instance_name = self.instance_name.clone();
         let key = ModuleCacheKey {
             instance_name: self.instance_name.clone(),
             module_id: module_metadata.id(),
-            module_version: module_metadata.latest_version,
+            source_package_id: module_metadata.source_package_id,
         };
         let mut transaction = self.transaction_ingredients.clone().try_into()?;
         let modules_storage = self.modules_storage.clone();
@@ -105,12 +86,12 @@ impl<RT: Runtime> ModuleLoader<RT> for FunctionRunnerModuleLoader<RT> {
                             .await;
                     modules
                         .into_iter()
-                        .map(move |((module_id, module_version), source)| {
+                        .map(move |((module_id, source_package_id), source)| {
                             (
                                 ModuleCacheKey {
                                     instance_name: instance_name.clone(),
                                     module_id,
-                                    module_version,
+                                    source_package_id,
                                 },
                                 source,
                             )
diff --git a/crates/isolate/src/environment/helpers/module_loader.rs b/crates/isolate/src/environment/helpers/module_loader.rs
@@ -14,10 +14,7 @@ use database::Transaction;
 use deno_core::ModuleSpecifier;
 use model::{
     modules::{
-        module_versions::{
-            FullModuleSource,
-            ModuleVersion,
-        },
+        module_versions::FullModuleSource,
         types::ModuleMetadata,
         ModuleModel,
     },
@@ -45,7 +42,7 @@ pub async fn get_module_and_prefetch<RT: Runtime>(
     tx: &mut Transaction<RT>,
     modules_storage: Arc<dyn Storage>,
     module_metadata: ParsedDocument<ModuleMetadata>,
-) -> HashMap<(ResolvedDocumentId, ModuleVersion), anyhow::Result<FullModuleSource>> {
+) -> HashMap<(ResolvedDocumentId, SourcePackageId), anyhow::Result<FullModuleSource>> {
     let all_source_result = if *READ_MODULES_FROM_SOURCE_PACKAGE {
         let _timer = module_load_timer("package");
         download_module_source_from_package(
@@ -63,7 +60,7 @@ pub async fn get_module_and_prefetch<RT: Runtime>(
             .map(|source| {
                 let mut result = HashMap::new();
                 result.insert(
-                    (module_metadata.id(), module_metadata.latest_version),
+                    (module_metadata.id(), module_metadata.source_package_id),
                     source,
                 );
                 result
@@ -73,7 +70,7 @@ pub async fn get_module_and_prefetch<RT: Runtime>(
         Err(e) => {
             let mut result = HashMap::new();
             result.insert(
-                (module_metadata.id(), module_metadata.latest_version),
+                (module_metadata.id(), module_metadata.source_package_id),
                 Err(e),
             );
             result
@@ -91,7 +88,7 @@ async fn download_module_source_from_package<RT: Runtime>(
     modules_storage: Arc<dyn Storage>,
     modules_tablet: TabletId,
     source_package_id: SourcePackageId,
-) -> anyhow::Result<HashMap<(ResolvedDocumentId, ModuleVersion), FullModuleSource>> {
+) -> anyhow::Result<HashMap<(ResolvedDocumentId, SourcePackageId), FullModuleSource>> {
     let mut result = HashMap::new();
     let source_package = SourcePackageModel::new(tx).get(source_package_id).await?;
     let mut package = download_package(
@@ -119,7 +116,7 @@ async fn download_module_source_from_package<RT: Runtime>(
             },
             Some(source) => {
                 result.insert(
-                    (module_metadata.id(), module_metadata.latest_version),
+                    (module_metadata.id(), module_metadata.source_package_id),
                     FullModuleSource {
                         source: source.source,
                         source_map: source.source_map,
diff --git a/crates/model/src/modules/mod.rs b/crates/model/src/modules/mod.rs
@@ -311,7 +311,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
     pub async fn get_version(
         &mut self,
         module_id: ResolvedDocumentId,
-        version: ModuleVersion,
+        version: Option<ModuleVersion>,
     ) -> anyhow::Result<Option<ParsedDocument<ModuleVersionMetadata>>> {
         let timer = get_module_version_timer();
         let module_id_value: ConvexValue = module_id.into();
@@ -335,7 +335,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
             .map(TryFrom::try_from)
             .transpose()?;
         if let Some(module_version) = &module_version {
-            anyhow::ensure!(module_version.version == Some(version));
+            anyhow::ensure!(module_version.version == version);
         }
         timer.finish();
         Ok(module_version)
@@ -345,13 +345,13 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
     pub async fn get_source_from_db(
         &mut self,
         module_id: ResolvedDocumentId,
-        version: ModuleVersion,
+        version: Option<ModuleVersion>,
     ) -> anyhow::Result<FullModuleSource> {
         let module_version = self
             .get_version(module_id, version)
             .await?
             .context(format!(
-                "Dangling module version reference: {module_id}@{version}"
+                "Dangling module version reference: {module_id}@{version:?}"
             ))?
             .into_value();
         Ok(FullModuleSource {
@@ -426,7 +426,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
         analyze_result: Option<AnalyzedModule>,
         environment: ModuleEnvironment,
         sha256: Sha256Digest,
-    ) -> anyhow::Result<(ResolvedDocumentId, ModuleVersion)> {
+    ) -> anyhow::Result<(ResolvedDocumentId, Option<ModuleVersion>)> {
         let (module_id, version) = match self.module_metadata(path.clone()).await? {
             Some(module_metadata) => {
                 let previous_version = module_metadata.latest_version;
@@ -437,7 +437,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
                     .await?
                     .map(|version| version.id());
 
-                let latest_version = previous_version + 1;
+                let latest_version = previous_version.map(|v| v + 1);
                 let new_metadata = ModuleMetadata {
                     path: path.module_path,
                     latest_version,
@@ -459,7 +459,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
                 (module_metadata.id(), latest_version)
             },
             None => {
-                let version = 0;
+                let version = Some(0);
                 let new_metadata = ModuleMetadata {
                     path: path.module_path,
                     latest_version: version,
@@ -481,7 +481,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
     async fn put_module_source_into_db(
         &mut self,
         module_id: ResolvedDocumentId,
-        version: ModuleVersion,
+        version: Option<ModuleVersion>,
         source: ModuleSource,
         source_map: Option<SourceMap>,
         component: ComponentId,
@@ -490,7 +490,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
             module_id: module_id.into(),
             source,
             source_map,
-            version: Some(version),
+            version,
         }.try_into()
         .map_err(|e: anyhow::Error| e.map_error_metadata(|em| {
             if em.short_msg == VALUE_TOO_LARGE_SHORT_MSG {
diff --git a/crates/model/src/modules/types.rs b/crates/model/src/modules/types.rs
@@ -27,7 +27,7 @@ pub struct ModuleMetadata {
     /// Path stored as a "path" field.
     pub path: CanonicalizedModulePath,
     /// What is the latest version of the module?
-    pub latest_version: ModuleVersion,
+    pub latest_version: Option<ModuleVersion>,
 
     pub source_package_id: SourcePackageId,
     pub environment: ModuleEnvironment,
@@ -40,7 +40,7 @@ pub struct ModuleMetadata {
 #[serde(rename_all = "camelCase")]
 pub struct SerializedModuleMetadata {
     pub path: String,
-    pub latest_version: ModuleVersion,
+    pub latest_version: Option<ModuleVersion>,
     pub deleted: Option<bool>,
     pub source_package_id: String,
     pub environment: String,
diff --git a/crates/model/src/source_packages/types.rs b/crates/model/src/source_packages/types.rs
@@ -128,7 +128,7 @@ impl TryFrom<PackageSize> for ConvexObject {
 }
 
 #[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
-#[derive(Debug, Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
+#[derive(Debug, Clone, PartialEq, Eq, Copy, PartialOrd, Ord, Hash)]
 pub struct SourcePackageId(DeveloperDocumentId);
 
 impl HeapSize for SourcePackageId {

Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,7 @@ impl TryFrom<PackageSize> for ConvexObject {`
`128`	`128`	`}`
`129`	`129`
`130`	`130`	`#[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]`
`131`		`-#[derive(Debug, Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]`
	`131`	`+#[derive(Debug, Clone, PartialEq, Eq, Copy, PartialOrd, Ord, Hash)]`
`132`	`132`	`pub struct SourcePackageId(DeveloperDocumentId);`
`133`	`133`
`134`	`134`	`impl HeapSize for SourcePackageId {`