parallel storage.getUrl (#24821)

allow `ctx.storage.getUrl` in queries and mutations to run in parallel with itself, as a batch.

note this removes a couple of the sanity checks, because i don't think they're necessary and with batching they're harder: (1) removes the assertion that there's at most one stored file with a given storage id, (2) removes the assertion that `entry.storage_id == storage_id` when `entry` was fetched with an equality index filter on `storage_id`

i wanted to make it also run in parallel with `db.get` and `db.query` but that turns out to be tricky because those need to use `DeveloperQuery` while file storage urls needs to use `ResolvedQuery`.

GitOrigin-RevId: 41e686efe0b36c60d9e91a99a5cfc8bc4d3f56d3

Author: ldanilek

Cargo.lock

@@ -17,6 +17,7 @@ errors = { path = "../errors" }
 futures = { workspace = true }
 headers = { workspace = true }
 keybroker = { path = "../keybroker" }
+maplit = { workspace = true }
 metrics = { path = "../metrics" }
 mime = { workspace = true }
 model = { path = "../model" }

crates/file_storage/Cargo.toml

@@ -1,4 +1,5 @@
 use std::{
+    collections::BTreeMap,
     ops::Bound,
     sync::Arc,
 };
@@ -32,12 +33,14 @@ use keybroker::{
     Identity,
     KeyBroker,
 };
+use maplit::btreemap;
 use mime::Mime;
 use model::file_storage::{
     types::{
         FileStorageEntry,
         StorageUuid,
     },
+    BatchKey,
     FileStorageId,
     FileStorageModel,
 };
@@ -92,15 +95,30 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
         tx: &mut Transaction<RT>,
         storage_id: FileStorageId,
     ) -> anyhow::Result<Option<String>> {
+        self.get_url_batch(tx, btreemap! { 0 => storage_id })
+            .await
+            .remove(&0)
+            .context("batch_key missing")?
+    }
+
+    pub async fn get_url_batch(
+        &self,
+        tx: &mut Transaction<RT>,
+        storage_ids: BTreeMap<BatchKey, FileStorageId>,
+    ) -> BTreeMap<BatchKey, anyhow::Result<Option<String>>> {
         let origin = &self.convex_origin;
-        let file: Option<FileStorageEntry> = self.get_file_entry(tx, storage_id.clone()).await?;
-        Ok(match file {
-            Some(entry) => {
-                let storage_id = entry.storage_id;
-                Some(format!("{origin}/api/storage/{storage_id}"))
-            },
-            None => None,
-        })
+        let files = self.get_file_entry_batch(tx, storage_ids).await;
+        files
+            .into_iter()
+            .map(|(batch_key, result)| {
+                (
+                    batch_key,
+                    result.map(|file| {
+                        file.map(|entry| format!("{origin}/api/storage/{}", entry.storage_id))
+                    }),
+                )
+            })
+            .collect()
     }
 
     pub async fn delete(
@@ -123,7 +141,23 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
         tx: &mut Transaction<RT>,
         storage_id: FileStorageId,
     ) -> anyhow::Result<Option<FileStorageEntry>> {
-        FileStorageModel::new(tx).get_file(storage_id).await
+        self.get_file_entry_batch(tx, btreemap! { 0 => storage_id })
+            .await
+            .remove(&0)
+            .context("batch_key missing")?
+    }
+
+    pub async fn get_file_entry_batch(
+        &self,
+        tx: &mut Transaction<RT>,
+        storage_ids: BTreeMap<BatchKey, FileStorageId>,
+    ) -> BTreeMap<BatchKey, anyhow::Result<Option<FileStorageEntry>>> {
+        FileStorageModel::new(tx)
+            .get_file_batch(storage_ids)
+            .await
+            .into_iter()
+            .map(|(batch_key, result)| (batch_key, result.map(|r| r.map(|r| r.into_value()))))
+            .collect()
     }
 
     pub async fn get_file_stream(

crates/file_storage/src/core.rs

@@ -48,6 +48,7 @@ use keybroker::KeyBroker;
 use model::{
     file_storage::{
         types::FileStorageEntry,
+        BatchKey,
         FileStorageId,
     },
     scheduled_jobs::VirtualSchedulerModel,
@@ -122,6 +123,7 @@ pub fn system_table_guard(name: &TableName, expect_system_table: bool) -> anyhow
 #[derive(Debug)]
 pub enum AsyncSyscallBatch {
     Reads(Vec<AsyncRead>),
+    StorageGetUrls(Vec<JsonValue>),
     Unbatched { name: String, args: JsonValue },
 }
 
@@ -136,6 +138,7 @@ impl AsyncSyscallBatch {
         match &*name {
             "1.0/get" => Self::Reads(vec![AsyncRead::Get(args)]),
             "1.0/queryStreamNext" => Self::Reads(vec![AsyncRead::QueryStreamNext(args)]),
+            "1.0/storageGetUrl" => Self::StorageGetUrls(vec![args]),
             _ => Self::Unbatched { name, args },
         }
     }
@@ -148,6 +151,8 @@ impl AsyncSyscallBatch {
             (Self::Reads(_), "1.0/get") => true,
             (Self::Reads(_), "1.0/queryStreamNext") => true,
             (Self::Reads(_), _) => false,
+            (Self::StorageGetUrls(_), "1.0/storageGetUrl") => true,
+            (Self::StorageGetUrls(_), _) => false,
             (Self::Unbatched { .. }, _) => false,
         }
     }
@@ -158,6 +163,9 @@ impl AsyncSyscallBatch {
             (Self::Reads(batch_args), "1.0/queryStreamNext") => {
                 batch_args.push(AsyncRead::QueryStreamNext(args))
             },
+            (Self::StorageGetUrls(batch_args), "1.0/storageGetUrl") => {
+                batch_args.push(args);
+            },
             _ => anyhow::bail!("cannot push {name} onto {self:?}"),
         }
         Ok(())
@@ -167,13 +175,15 @@ impl AsyncSyscallBatch {
         match self {
             // 1.0/get is grouped in with 1.0/queryStreamNext.
             Self::Reads(_) => "1.0/queryStreamNext",
+            Self::StorageGetUrls(_) => "1.0/storageGetUrl",
             Self::Unbatched { name, .. } => name,
         }
     }
 
     pub fn len(&self) -> usize {
         match self {
             Self::Reads(args) => args.len(),
+            Self::StorageGetUrls(args) => args.len(),
             Self::Unbatched { .. } => 1,
         }
     }
@@ -254,10 +264,10 @@ pub trait AsyncSyscallProvider<RT: Runtime> {
     ) -> anyhow::Result<(UdfPath, ConvexArray)>;
 
     fn file_storage_generate_upload_url(&self) -> anyhow::Result<String>;
-    async fn file_storage_get_url(
+    async fn file_storage_get_url_batch(
         &mut self,
-        storage_id: FileStorageId,
-    ) -> anyhow::Result<Option<String>>;
+        storage_ids: BTreeMap<BatchKey, FileStorageId>,
+    ) -> BTreeMap<BatchKey, anyhow::Result<Option<String>>>;
     async fn file_storage_delete(&mut self, storage_id: FileStorageId) -> anyhow::Result<()>;
     async fn file_storage_get_entry(
         &mut self,
@@ -350,13 +360,20 @@ impl<RT: Runtime> AsyncSyscallProvider<RT> for DatabaseUdfEnvironment<RT> {
         Ok(post_url)
     }
 
-    async fn file_storage_get_url(
+    async fn file_storage_get_url_batch(
         &mut self,
-        storage_id: FileStorageId,
-    ) -> anyhow::Result<Option<String>> {
-        self.file_storage
-            .get_url(self.phase.tx()?, storage_id)
-            .await
+        storage_ids: BTreeMap<BatchKey, FileStorageId>,
+    ) -> BTreeMap<BatchKey, anyhow::Result<Option<String>>> {
+        let tx = match self.phase.tx() {
+            Ok(tx) => tx,
+            Err(e) => {
+                return storage_ids
+                    .into_keys()
+                    .map(|batch_key| (batch_key, Err(e.clone().into())))
+                    .collect();
+            },
+        };
+        self.file_storage.get_url_batch(tx, storage_ids).await
     }
 
     async fn file_storage_delete(&mut self, storage_id: FileStorageId) -> anyhow::Result<()> {
@@ -389,7 +406,7 @@ impl<RT: Runtime, P: AsyncSyscallProvider<RT>> DatabaseSyscallsV1<RT, P> {
     pub async fn run_async_syscall_batch(
         provider: &mut P,
         batch: AsyncSyscallBatch,
-    ) -> anyhow::Result<Vec<anyhow::Result<String>>> {
+    ) -> Vec<anyhow::Result<String>> {
         let start = provider.rt().monotonic_now();
         let batch_name = batch.name().to_string();
         let timer = async_syscall_timer(&batch_name);
@@ -398,6 +415,9 @@ impl<RT: Runtime, P: AsyncSyscallProvider<RT>> DatabaseSyscallsV1<RT, P> {
         // errors.
         let results = match batch {
             AsyncSyscallBatch::Reads(batch_args) => Self::query_batch(provider, batch_args).await,
+            AsyncSyscallBatch::StorageGetUrls(batch_args) => {
+                Self::storage_get_url_batch(provider, batch_args).await
+            },
             AsyncSyscallBatch::Unbatched { name, args } => {
                 let result = match &name[..] {
                     // Database
@@ -415,7 +435,6 @@ impl<RT: Runtime, P: AsyncSyscallProvider<RT>> DatabaseSyscallsV1<RT, P> {
                     "1.0/storageGenerateUploadUrl" => {
                         Self::storage_generate_upload_url(provider, args).await
                     },
-                    "1.0/storageGetUrl" => Self::storage_get_url(provider, args).await,
                     // Scheduling
                     "1.0/schedule" => Self::schedule(provider, args).await,
                     "1.0/cancel_job" => Self::cancel_job(provider, args).await,
@@ -444,10 +463,10 @@ impl<RT: Runtime, P: AsyncSyscallProvider<RT>> DatabaseSyscallsV1<RT, P> {
             results.iter().all(|result| result.is_ok()),
         );
         timer.finish();
-        Ok(results
+        results
             .into_iter()
             .map(|result| anyhow::Ok(serde_json::to_string(&result?)?))
-            .collect())
+            .collect()
     }
 
     #[convex_macro::instrument_future]
@@ -493,18 +512,40 @@ impl<RT: Runtime, P: AsyncSyscallProvider<RT>> DatabaseSyscallsV1<RT, P> {
     }
 
     #[convex_macro::instrument_future]
-    async fn storage_get_url(provider: &mut P, args: JsonValue) -> anyhow::Result<JsonValue> {
+    async fn storage_get_url_batch(
+        provider: &mut P,
+        batch_args: Vec<JsonValue>,
+    ) -> Vec<anyhow::Result<JsonValue>> {
         #[derive(Deserialize)]
         #[serde(rename_all = "camelCase")]
         struct GetUrlArgs {
             storage_id: String,
         }
-        let storage_id: FileStorageId = with_argument_error("storage.getUrl", || {
-            let GetUrlArgs { storage_id } = serde_json::from_value(args)?;
-            storage_id.parse().context(ArgName("storageId"))
-        })?;
-        let url = provider.file_storage_get_url(storage_id).await?;
-        Ok(url.into())
+        let batch_size = batch_args.len();
+        let mut results = BTreeMap::new();
+        let mut storage_ids = BTreeMap::new();
+        for (idx, args) in batch_args.into_iter().enumerate() {
+            let storage_id_result = with_argument_error("storage.getUrl", || {
+                let GetUrlArgs { storage_id } = serde_json::from_value(args)?;
+                storage_id.parse().context(ArgName("storageId"))
+            });
+            match storage_id_result {
+                Ok(storage_id) => {
+                    storage_ids.insert(idx, storage_id);
+                },
+                Err(e) => {
+                    assert!(results.insert(idx, Err(e)).is_none());
+                },
+            }
+        }
+        let urls = provider.file_storage_get_url_batch(storage_ids).await;
+        for (batch_key, url) in urls {
+            assert!(results
+                .insert(batch_key, url.map(JsonValue::from))
+                .is_none());
+        }
+        assert_eq!(results.len(), batch_size);
+        results.into_values().collect()
     }
 
     #[convex_macro::instrument_future]

crates/isolate/src/environment/udf/async_syscall.rs

@@ -718,7 +718,9 @@ impl<RT: Runtime> DatabaseUdfEnvironment<RT> {
                     results = with_release_permit(
                         &mut state.timeout,
                         &mut state.permit,
-                        DatabaseSyscallsV1::run_async_syscall_batch(&mut state.environment, batch),
+                        DatabaseSyscallsV1::run_async_syscall_batch(
+                            &mut state.environment, batch,
+                        ).map(Ok),
                     ).fuse() => results?,
                 };
                 (resolvers, results)

crates/isolate/src/environment/udf/mod.rs

@@ -48,6 +48,7 @@ use model::{
     },
     file_storage::{
         types::FileStorageEntry,
+        BatchKey,
         FileStorageId,
     },
     udf_config::UdfConfigModel,
@@ -581,7 +582,7 @@ async fn run_request<RT: Runtime>(
 
                 if let Some(batch) = syscall_batch.take() {
                     let results =
-                        DatabaseSyscallsV1::run_async_syscall_batch(&mut provider, batch).await?;
+                        DatabaseSyscallsV1::run_async_syscall_batch(&mut provider, batch).await;
                     assert_eq!(results.len(), batch_promise_ids.len());
 
                     for (promise_id, result) in batch_promise_ids.drain(..).zip(results) {
@@ -597,7 +598,7 @@ async fn run_request<RT: Runtime>(
             }
             if let Some(batch) = syscall_batch {
                 let results =
-                    DatabaseSyscallsV1::run_async_syscall_batch(&mut provider, batch).await?;
+                    DatabaseSyscallsV1::run_async_syscall_batch(&mut provider, batch).await;
                 assert_eq!(results.len(), batch_promise_ids.len());
 
                 for (promise_id, result) in batch_promise_ids.into_iter().zip(results) {
@@ -882,10 +883,10 @@ impl<'a, RT: Runtime> AsyncSyscallProvider<RT> for Isolate2SyscallProvider<'a, R
         todo!()
     }
 
-    async fn file_storage_get_url(
+    async fn file_storage_get_url_batch(
         &mut self,
-        _storage_id: FileStorageId,
-    ) -> anyhow::Result<Option<String>> {
+        _storage_ids: BTreeMap<BatchKey, FileStorageId>,
+    ) -> BTreeMap<BatchKey, anyhow::Result<Option<String>>> {
         todo!()
     }
 

crates/isolate/src/isolate2/runner.rs

@@ -28,6 +28,7 @@ mod search;
 mod shapes;
 mod size_errors;
 mod source_maps;
+mod storage;
 mod system_udfs;
 mod unicode;
 mod user_error;