Merge branch 'optimistic-presigned-urls'

Author: oscartbeaumont

apps/desktop/src-tauri/src/api.rs

@@ -49,7 +49,6 @@ pub async fn upload_multipart_presign_part(
     video_id: &str,
     upload_id: &str,
     part_number: u32,
-    md5_sum: &str,
 ) -> Result<String, AuthedApiError> {
     #[derive(Deserialize)]
     #[serde(rename_all = "camelCase")]
@@ -65,7 +64,6 @@ pub async fn upload_multipart_presign_part(
                     "videoId": video_id,
                     "uploadId": upload_id,
                     "partNumber": part_number,
-                    "md5Sum": md5_sum
                 }))
         })
         .await

apps/desktop/src-tauri/src/upload.rs

@@ -618,16 +618,49 @@ fn multipart_uploader(
     try_stream! {
         let mut stream = pin!(stream);
         let mut prev_part_number = None;
+
+        let mut optimistic_presigned_url_task: Option<tokio::task::JoinHandle<Result<String, AuthedApiError>>> = Some(
+            tokio::spawn({
+                let app = app.clone();
+                let video_id = video_id.clone();
+                let upload_id = upload_id.clone();
+
+                async move {
+                    api::upload_multipart_presign_part(&app, &video_id, &upload_id, 1)
+                        .await
+                }
+            })
+        );
+        let mut expected_part_number = 1u32;
+
         while let Some(item) = stream.next().await {
             let Chunk { total_size, part_number, chunk } = item.map_err(|err| format!("uploader/part/{:?}/fs: {err:?}", prev_part_number.map(|p| p + 1)))?;
             debug!("Uploading chunk {part_number} ({} bytes) for video {video_id:?}", chunk.len());
             prev_part_number = Some(part_number);
             let md5_sum = base64::encode(md5::compute(&chunk).0);
             let size = chunk.len();
 
-            let presigned_url =
-                api::upload_multipart_presign_part(&app, &video_id, &upload_id, part_number, &md5_sum)
-                    .await?;
+            let presigned_url = if expected_part_number == part_number {
+                // The optimistic presigned URL matches, wait for it
+                if let Some(task) = optimistic_presigned_url_task.take() {
+                    task.await
+                        .map_err(|e| format!("uploader/part/{part_number}/task_join: {e:?}"))?
+                        .map_err(|e| format!("uploader/part/{part_number}/presign: {e}"))?
+                } else {
+                    // Fallback if no task available
+                    api::upload_multipart_presign_part(&app, &video_id, &upload_id, part_number)
+                        .await?
+                }
+            } else {
+                // The optimistic presigned URL doesn't match, abort it and generate a new correct one
+                if let Some(task) = optimistic_presigned_url_task.take() {
+                    task.abort();
+                }
+                debug!("Throwing out optimistic presigned URL for part {expected_part_number} as part {part_number} was requested!");
+                expected_part_number = part_number;
+                api::upload_multipart_presign_part(&app, &video_id, &upload_id, part_number)
+                    .await?
+            };
 
             let url = Uri::from_str(&presigned_url).map_err(|err| format!("uploader/part/{part_number}/invalid_url: {err:?}"))?;
             let resp = retryable_client(url.host().unwrap_or("<unknown>").to_string())
@@ -655,6 +688,25 @@ fn multipart_uploader(
                 size,
                 total_size
             };
+
+            // We generate the presigned URL ahead of time for the next expected part.
+            // This means if the filesystem takes a while for the recording to reach previous total + CHUNK_SIZE, we aren't just doing nothing.
+            expected_part_number = part_number + 1;
+            optimistic_presigned_url_task = Some(tokio::spawn({
+                let app = app.clone();
+                let video_id = video_id.clone();
+                let upload_id = upload_id.clone();
+
+                async move {
+                    api::upload_multipart_presign_part(&app, &video_id, &upload_id, expected_part_number)
+                        .await
+                }
+            }));
+        }
+
+        // Clean up any remaining optimistic task
+        if let Some(task) = optimistic_presigned_url_task.take() {
+            task.abort();
         }
     }
 }

apps/web/app/api/upload/[...route]/multipart.ts

@@ -99,18 +99,19 @@ app.post(
 			.object({
 				uploadId: z.string(),
 				partNumber: z.number(),
-				md5Sum: z.string(),
 			})
 			.and(
 				z.union([
 					z.object({ videoId: z.string() }),
 					// deprecated
 					z.object({ fileKey: z.string() }),
+					// deprecated
+					// z.object({ md5Sum: z.string() }),
 				]),
 			),
 	),
 	async (c) => {
-		const { uploadId, partNumber, md5Sum, ...body } = c.req.valid("json");
+		const { uploadId, partNumber, ...body } = c.req.valid("json");
 		const user = c.get("user");
 
 		const fileKey = parseVideoIdOrFileKey(user.id, {
@@ -132,7 +133,6 @@ app.post(
 							fileKey,
 							uploadId,
 							partNumber,
-							{ ContentMD5: md5Sum },
 						);
 
 					return presignedUrl;