pass modules_storage to module source fetch (#24963)

this PR does the plumbing to make the `modules_storage: Arc<dyn Storage>` available in the functions that fetch module source code. Module source code is (still) fetched only from `_module_versions`, but we eventually want it to be fetched from S3 instead.

this PR is piping the storage through the funrun and application structs, and also adding some Debug impls which i found useful while figuring this out.

in particular is the `StorageForInstanceImpl` naming and pattern okay?

GitOrigin-RevId: 3aaeab09e43c64ee76c3ea22bc2197a610363a93

Author: ldanilek

crates/application/src/lib.rs

@@ -464,7 +464,9 @@ impl<RT: Runtime> Application<RT> {
         snapshot_import_pause_client: PauseClient,
         scheduled_jobs_pause_client: PauseClient,
     ) -> anyhow::Result<Self> {
-        let module_cache = ModuleCacheWorker::start(runtime.clone(), database.clone()).await;
+        let module_cache =
+            ModuleCacheWorker::start(runtime.clone(), database.clone(), modules_storage.clone())
+                .await;
         let module_loader = Arc::new(module_cache.clone());
 
         let system_env_vars = btreemap! {

crates/application/src/module_cache/mod.rs

@@ -36,6 +36,7 @@ use model::modules::{
     MODULE_VERSIONS_TABLE,
 };
 use parking_lot::Mutex;
+use storage::Storage;
 use value::ResolvedDocumentId;
 
 mod metrics;
@@ -46,11 +47,16 @@ const MAX_BACKOFF: Duration = Duration::from_secs(30);
 pub struct ModuleCacheWorker<RT: Runtime> {
     rt: RT,
     database: Database<RT>,
+    modules_storage: Arc<dyn Storage>,
     cache: AsyncLru<RT, (ResolvedDocumentId, ModuleVersion), ModuleVersionMetadata>,
 }
 
 impl<RT: Runtime> ModuleCacheWorker<RT> {
-    pub async fn start(rt: RT, database: Database<RT>) -> ModuleCache<RT> {
+    pub async fn start(
+        rt: RT,
+        database: Database<RT>,
+        modules_storage: Arc<dyn Storage>,
+    ) -> ModuleCache<RT> {
         let cache = AsyncLru::new(
             rt.clone(),
             *MODULE_CACHE_MAX_SIZE_BYTES,
@@ -60,12 +66,14 @@ impl<RT: Runtime> ModuleCacheWorker<RT> {
         let worker = Self {
             rt: rt.clone(),
             database: database.clone(),
+            modules_storage: modules_storage.clone(),
             cache: cache.clone(),
         };
 
         let worker_handle = rt.spawn("module_cache_worker", worker.go());
         ModuleCache {
             database,
+            modules_storage,
             cache,
             worker: Arc::new(Mutex::new(worker_handle)),
         }
@@ -103,6 +111,7 @@ impl<RT: Runtime> ModuleCacheWorker<RT> {
             for key in referenced_versions {
                 let fetcher = ModuleVersionFetcher {
                     database: self.database.clone(),
+                    modules_storage: self.modules_storage.clone(),
                 };
                 self.cache
                     .get(key, fetcher.generate_value(key).boxed())
@@ -127,6 +136,9 @@ impl<RT: Runtime> ModuleCacheWorker<RT> {
 #[derive(Clone)]
 pub struct ModuleVersionFetcher<RT: Runtime> {
     database: Database<RT>,
+    // TODO(lee) read module source from storage.
+    #[allow(unused)]
+    modules_storage: Arc<dyn Storage>,
 }
 
 impl<RT: Runtime> ModuleVersionFetcher<RT> {
@@ -145,6 +157,8 @@ impl<RT: Runtime> ModuleVersionFetcher<RT> {
 pub struct ModuleCache<RT: Runtime> {
     database: Database<RT>,
 
+    modules_storage: Arc<dyn Storage>,
+
     cache: AsyncLru<RT, (ResolvedDocumentId, ModuleVersion), ModuleVersionMetadata>,
 
     worker: Arc<Mutex<RT::Handle>>,
@@ -160,6 +174,7 @@ impl<RT: Runtime> Clone for ModuleCache<RT> {
     fn clone(&self) -> Self {
         Self {
             database: self.database.clone(),
+            modules_storage: self.modules_storage.clone(),
             cache: self.cache.clone(),
             worker: self.worker.clone(),
         }
@@ -189,6 +204,7 @@ impl<RT: Runtime> ModuleLoader<RT> for ModuleCache<RT> {
         let key = (module_metadata.id(), module_metadata.latest_version);
         let fetcher = ModuleVersionFetcher {
             database: self.database.clone(),
+            modules_storage: self.modules_storage.clone(),
         };
         let result = self
             .cache

crates/application/src/test_helpers.rs

@@ -34,7 +34,10 @@ use file_storage::{
     FileStorage,
     TransactionalFileStorage,
 };
-use function_runner::server::InProcessFunctionRunner;
+use function_runner::server::{
+    InProcessFunctionRunner,
+    InstanceStorage,
+};
 use isolate::test_helpers::{
     TEST_SOURCE,
     TEST_SOURCE_ISOLATE_ONLY,
@@ -188,7 +191,10 @@ impl<RT: Runtime> ApplicationTestExt<RT> for Application<RT> {
                 convex_origin.clone(),
                 rt.clone(),
                 persistence.reader(),
-                files_storage.clone(),
+                InstanceStorage {
+                    files_storage: files_storage.clone(),
+                    modules_storage: modules_storage.clone(),
+                },
                 database.clone(),
                 fetch_client.clone(),
             )

crates/function_runner/src/module_cache.rs

@@ -25,6 +25,7 @@ use model::modules::{
     ModuleModel,
     MODULE_VERSIONS_TABLE,
 };
+use storage::Storage;
 use value::ResolvedDocumentId;
 
 use crate::in_memory_indexes::TransactionIngredients;
@@ -54,6 +55,7 @@ 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>,
 }
 
 #[async_trait]
@@ -89,7 +91,7 @@ impl<RT: Runtime> ModuleLoader<RT> for FunctionRunnerModuleLoader<RT> {
             .0
             .get(
                 key.clone(),
-                get_module(transaction, module_metadata).boxed(),
+                get_module(transaction, self.modules_storage.clone(), module_metadata).boxed(),
             )
             .await?;
         // Record read dependency on the module version so the transactions

crates/function_runner/src/server.rs

@@ -1,5 +1,6 @@
 use std::{
     collections::BTreeMap,
+    fmt::Debug,
     sync::{
         Arc,
         Weak,
@@ -91,7 +92,10 @@ use parking_lot::{
     Mutex,
     RwLock,
 };
-use storage::Storage;
+use storage::{
+    Storage,
+    StorageUseCase,
+};
 use sync_types::Timestamp;
 use usage_tracking::{
     FunctionUsageStats,
@@ -118,23 +122,35 @@ const MAX_ISOLATE_WORKERS: usize = 128;
 const ACTIVE_CONCURRENCY_PERMITS_LOG_FREQUENCY: Duration = Duration::from_secs(10);
 
 #[async_trait]
-pub trait StorageForInstance<RT: Runtime>: Clone + Send + Sync + 'static {
+pub trait StorageForInstance<RT: Runtime>: Debug + Clone + Send + Sync + 'static {
     /// Gets a storage impl for a instance. Agnostic to what kind of storage -
     /// local or s3, or how it was loaded (e.g. passed directly within backend,
     /// loaded from a transaction created in Funrun)
     async fn storage_for_instance(
         &self,
         transaction: &mut Transaction<RT>,
+        use_case: StorageUseCase,
     ) -> anyhow::Result<Arc<dyn Storage>>;
 }
 
+#[derive(Clone, Debug)]
+pub struct InstanceStorage {
+    pub files_storage: Arc<dyn Storage>,
+    pub modules_storage: Arc<dyn Storage>,
+}
+
 #[async_trait]
-impl<RT: Runtime> StorageForInstance<RT> for Arc<dyn Storage> {
+impl<RT: Runtime> StorageForInstance<RT> for InstanceStorage {
     async fn storage_for_instance(
         &self,
         _transaction: &mut Transaction<RT>,
+        use_case: StorageUseCase,
     ) -> anyhow::Result<Arc<dyn Storage>> {
-        Ok(self.clone())
+        match use_case {
+            StorageUseCase::Files => Ok(self.files_storage.clone()),
+            StorageUseCase::Modules => Ok(self.modules_storage.clone()),
+            _ => anyhow::bail!("function runner storage does not support {use_case}"),
+        }
     }
 }
 
@@ -357,8 +373,15 @@ impl<RT: Runtime, S: StorageForInstance<RT>> FunctionRunnerCore<RT, S> {
             )
             .await?;
         let mut transaction = transaction_ingredients.clone().try_into()?;
-        let storage = self.storage.storage_for_instance(&mut transaction).await?;
+        let storage = self
+            .storage
+            .storage_for_instance(&mut transaction, StorageUseCase::Files)
+            .await?;
         let file_storage = TransactionalFileStorage::new(self.rt.clone(), storage, convex_origin);
+        let modules_storage = self
+            .storage
+            .storage_for_instance(&mut transaction, StorageUseCase::Modules)
+            .await?;
 
         let key_broker = KeyBroker::new(&instance_name, instance_secret)?;
         let environment_data = EnvironmentData {
@@ -369,6 +392,7 @@ impl<RT: Runtime, S: StorageForInstance<RT>> FunctionRunnerCore<RT, S> {
                 instance_name: instance_name.clone(),
                 cache: self.module_cache.clone(),
                 transaction_ingredients,
+                modules_storage,
             }),
         };
 
@@ -434,7 +458,7 @@ impl<RT: Runtime, S: StorageForInstance<RT>> FunctionRunnerCore<RT, S> {
 }
 
 pub struct InProcessFunctionRunner<RT: Runtime> {
-    server: FunctionRunnerCore<RT, Arc<dyn Storage>>,
+    server: FunctionRunnerCore<RT, InstanceStorage>,
     persistence_reader: Arc<dyn PersistenceReader>,
 
     // Static information about the backend.
@@ -455,7 +479,7 @@ impl<RT: Runtime> InProcessFunctionRunner<RT> {
         convex_origin: ConvexOrigin,
         rt: RT,
         persistence_reader: Arc<dyn PersistenceReader>,
-        storage: Arc<dyn Storage>,
+        storage: InstanceStorage,
         database: Database<RT>,
         fetch_client: Arc<dyn FetchClient>,
     ) -> anyhow::Result<Self> {
@@ -567,16 +591,19 @@ mod tests {
     };
     use model::initialize_application_system_tables;
     use runtime::testing::TestRuntime;
-    use storage::{
-        LocalDirStorage,
-        Storage,
-    };
+    use storage::LocalDirStorage;
 
-    use crate::server::FunctionRunnerCore;
+    use crate::server::{
+        FunctionRunnerCore,
+        InstanceStorage,
+    };
     #[convex_macro::test_runtime]
     async fn test_scheduler_workers_limit_requests(rt: TestRuntime) -> anyhow::Result<()> {
         initialize_v8();
-        let storage = Arc::new(LocalDirStorage::new(rt.clone())?) as Arc<dyn Storage>;
+        let storage = InstanceStorage {
+            files_storage: Arc::new(LocalDirStorage::new(rt.clone())?),
+            modules_storage: Arc::new(LocalDirStorage::new(rt.clone())?),
+        };
         let function_runner_core = FunctionRunnerCore::_new(rt.clone(), storage, 100, 1).await?;
         let (mut pause1, pause_client1) = PauseController::new([PAUSE_REQUEST]);
         let DbFixtures { db, .. } = DbFixtures::new(&rt).await?;
@@ -591,7 +618,7 @@ mod tests {
         let request2 = bogus_udf_request(&db, client1, None, sender).await?;
         function_runner_core.send_request(request2)?;
         let response =
-            FunctionRunnerCore::<TestRuntime, Arc<dyn Storage>>::receive_response(rx2).await?;
+            FunctionRunnerCore::<TestRuntime, InstanceStorage>::receive_response(rx2).await?;
         let err = response.unwrap_err();
         assert!(err.is_rejected_before_execution());
         assert!(err.to_string().contains(NO_AVAILABLE_WORKERS));
@@ -603,7 +630,10 @@ mod tests {
         rt: TestRuntime,
     ) -> anyhow::Result<()> {
         initialize_v8();
-        let storage = Arc::new(LocalDirStorage::new(rt.clone())?) as Arc<dyn Storage>;
+        let storage = InstanceStorage {
+            files_storage: Arc::new(LocalDirStorage::new(rt.clone())?),
+            modules_storage: Arc::new(LocalDirStorage::new(rt.clone())?),
+        };
         let function_runner_core = FunctionRunnerCore::_new(rt.clone(), storage, 50, 2).await?;
         let (mut pause1, pause_client1) = PauseController::new([PAUSE_REQUEST]);
         let DbFixtures { db, .. } = DbFixtures::new(&rt).await?;
@@ -619,14 +649,17 @@ mod tests {
         let client2 = "client2";
         let request2 = bogus_udf_request(&db, client2, None, sender).await?;
         function_runner_core.send_request(request2)?;
-        FunctionRunnerCore::<TestRuntime, Arc<dyn Storage>>::receive_response(rx2).await??;
+        FunctionRunnerCore::<TestRuntime, InstanceStorage>::receive_response(rx2).await??;
         Ok(())
     }
 
     #[convex_macro::test_runtime]
     async fn test_scheduler_throttles_same_client(rt: TestRuntime) -> anyhow::Result<()> {
         initialize_v8();
-        let storage = Arc::new(LocalDirStorage::new(rt.clone())?) as Arc<dyn Storage>;
+        let storage = InstanceStorage {
+            files_storage: Arc::new(LocalDirStorage::new(rt.clone())?),
+            modules_storage: Arc::new(LocalDirStorage::new(rt.clone())?),
+        };
         let function_runner_core = FunctionRunnerCore::_new(rt.clone(), storage, 50, 2).await?;
         let (mut pause1, pause_client1) = PauseController::new([PAUSE_REQUEST]);
         let DbFixtures { db, .. } = DbFixtures::new(&rt).await?;
@@ -642,7 +675,7 @@ mod tests {
         let request2 = bogus_udf_request(&db, client, None, sender).await?;
         function_runner_core.send_request(request2)?;
         let response =
-            FunctionRunnerCore::<TestRuntime, Arc<dyn Storage>>::receive_response(rx2).await?;
+            FunctionRunnerCore::<TestRuntime, InstanceStorage>::receive_response(rx2).await?;
         let err = response.unwrap_err();
         assert!(err.is_rejected_before_execution());
         assert!(err.to_string().contains(NO_AVAILABLE_WORKERS));

crates/isolate/Cargo.toml

@@ -72,7 +72,7 @@ sha1 = { workspace = true }
 sha2 = { workspace = true }
 sourcemap = { workspace = true }
 spki = { workspace = true }
-storage = { path = "../storage", optional = true }
+storage = { path = "../storage" }
 sync_types = { package = "convex_sync_types", path = "../convex/sync_types" }
 thiserror = { workspace = true }
 tokio = { workspace = true }
@@ -128,7 +128,6 @@ testing = [
     "proptest-http",
     "runtime/testing",
     "search/testing",
-    "storage",
     "storage/testing",
     "usage_tracking/testing",
     "value/testing",

crates/isolate/src/environment/helpers/module_loader.rs

@@ -17,6 +17,7 @@ use model::modules::{
     types::ModuleMetadata,
     ModuleModel,
 };
+use storage::Storage;
 use sync_types::{
     CanonicalizedModulePath,
     CanonicalizedUdfPath,
@@ -114,6 +115,8 @@ impl<RT: Runtime> ModuleLoader<RT> for TransactionModuleLoader {
 
 pub async fn get_module<RT: Runtime>(
     mut tx: Transaction<RT>,
+    // TODO(lee) fetch from module storage
+    _modules_storage: Arc<dyn Storage>,
     module_metadata: ParsedDocument<ModuleMetadata>,
 ) -> anyhow::Result<ModuleVersionMetadata> {
     let _timer = module_load_timer();

crates/local_backend/src/lib.rs

@@ -45,7 +45,10 @@ use file_storage::{
     TransactionalFileStorage,
 };
 use function_runner::{
-    server::InProcessFunctionRunner,
+    server::{
+        InProcessFunctionRunner,
+        InstanceStorage,
+    },
     FunctionRunner,
 };
 use model::{
@@ -209,7 +212,10 @@ pub async fn make_app(
             config.convex_origin_url(),
             runtime.clone(),
             persistence.reader(),
-            files_storage.clone(),
+            InstanceStorage {
+                files_storage: files_storage.clone(),
+                modules_storage: modules_storage.clone(),
+            },
             database.clone(),
             fetch_client.clone(),
         )

crates/storage/src/lib.rs

@@ -521,7 +521,7 @@ impl StorageExt for Arc<dyn Storage> {
         let full_size = self
             .get_object_attributes(object_key)
             .await?
-            .context("object does not exist")?
+            .with_context(|| format!("object {object_key:?} does not exist in {self:?}"))?
             .size;
         let inner =
             StorageObjectReader::new_inner_reader_starting_at(self.clone(), object_key.clone(), 0);