feat: Better transpiled module caching in SW

cdata · cdata · commit e9b69fb56d1c · 2024-05-23T20:41:22.000-07:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -9,6 +9,7 @@ members = [
 resolver = "2"
 
 [workspace.dependencies]
+anyhow = { version = "1" }
 async-trait = { version = "0.1" }
 axum = { version = "0.7" }
 blake3 = { version = "1.5" }
@@ -32,6 +33,7 @@ wasm-bindgen = { version = "0.2" }
 wasmtime-environ = { version = "20" }
 web-sys = { version = "0.3" }
 wit-bindgen = { version = "0.25" }
+wit-parser = { version = "0.208" }
 
 [profile.release]
 opt-level = 'z'
diff --git a/rust/usuba-compat/Cargo.toml b/rust/usuba-compat/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
+blake3 = { workspace = true }
 js-component-bindgen = { workspace = true }
 wit-bindgen = { workspace = true }
 wasmtime-environ = { workspace = true,  features = ["component-model", "compile"] }
diff --git a/rust/usuba-compat/src/lib.rs b/rust/usuba-compat/src/lib.rs
@@ -30,6 +30,10 @@ impl Guest for Polyfill {
             .map(|transpiled| transpiled.into())
             .map_err(|error| format!("{}", error))
     }
+
+    fn hash(bytes: Vec<u8>) -> String {
+        blake3::hash(&bytes).to_string()
+    }
 }
 
 export!(Polyfill);
diff --git a/rust/usuba-compat/wit/usuba-compat.wit b/rust/usuba-compat/wit/usuba-compat.wit
@@ -21,4 +21,5 @@ world usuba-compat {
   }
 
   export polyfill: func(component: list<u8>, options: polyfill-options) -> result<artifacts, string>;
+  export hash: func(bytes: list<u8>) -> string;
 }
diff --git a/rust/usuba/Cargo.toml b/rust/usuba/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+anyhow = { workspace = true }
 async-trait = { workspace = true }
 axum = { workspace = true, features = ["multipart"] }
 blake3 = { workspace = true }
@@ -23,5 +24,6 @@ tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 utoipa = { workspace = true, features = ["axum_extras"] }
 utoipa-swagger-ui = { workspace = true, features = ["axum"] }
+wit-parser = { workspace = true }
 # gcloud-sdk = { version = "0.24.6", features = ["google-cloud-aiplatform-v1"] }
 
diff --git a/rust/usuba/src/bake/bake.rs b/rust/usuba/src/bake/bake.rs
@@ -4,5 +4,10 @@ use bytes::Bytes;
 
 #[async_trait]
 pub trait Bake {
-    async fn bake(&self, wit: Bytes, source_code: Bytes) -> Result<Bytes, UsubaError>;
+    async fn bake(
+        &self,
+        world: &str,
+        wit: Vec<Bytes>,
+        source_code: Bytes,
+    ) -> Result<Bytes, UsubaError>;
 }
diff --git a/rust/usuba/src/bake/javascript.rs b/rust/usuba/src/bake/javascript.rs
@@ -1,20 +1,38 @@
 use std::io::Cursor;
+use std::path::{Path, PathBuf};
 use tracing::instrument;
 
+use crate::UsubaError;
+
 use super::Bake;
 use async_trait::async_trait;
 use bytes::Bytes;
 use tempfile::TempDir;
 
 use tokio::process::Command;
+use tokio::task::JoinSet;
+
+use wit_parser::UnresolvedPackage;
+
+async fn write_file(path: PathBuf, bytes: Bytes) -> Result<(), UsubaError> {
+    let mut file = tokio::fs::File::create(&path).await?;
+    let mut cursor = Cursor::new(bytes.as_ref());
+    tokio::io::copy(&mut cursor, &mut file).await?;
+    Ok(())
+}
 
 #[derive(Debug)]
 pub struct JavaScriptBaker {}
 
 #[async_trait]
 impl Bake for JavaScriptBaker {
     #[instrument]
-    async fn bake(&self, wit: Bytes, source_code: Bytes) -> Result<Bytes, crate::UsubaError> {
+    async fn bake(
+        &self,
+        world: &str,
+        wit: Vec<Bytes>,
+        source_code: Bytes,
+    ) -> Result<Bytes, crate::UsubaError> {
         let workspace = TempDir::new()?;
         debug!(
             "Created temporary workspace in {}",
@@ -23,31 +41,32 @@ impl Bake for JavaScriptBaker {
 
         let wasm_path = workspace.path().join("module.wasm");
         let js_path = workspace.path().join("module.js");
-        let wit_path = workspace.path().join("module.wit");
 
-        let (mut wit_file, mut js_file) = tokio::try_join!(
-            tokio::fs::File::create(&wit_path),
-            tokio::fs::File::create(&js_path),
-        )?;
+        debug!(?workspace, "Created temporary workspace");
 
-        debug!(?wit_path, ?js_path, "Created temporary input files");
+        let mut writes = JoinSet::new();
 
-        let mut wit_cursor = Cursor::new(wit);
-        let mut js_cursor = Cursor::new(source_code);
+        wit.into_iter()
+            .enumerate()
+            .map(|(i, wit)| write_file(workspace.path().join(format!("module{}.wit", i)), wit))
+            .chain([write_file(js_path.clone(), source_code)])
+            .for_each(|fut| {
+                writes.spawn(fut);
+            });
 
-        tokio::try_join!(
-            tokio::io::copy(&mut wit_cursor, &mut wit_file),
-            tokio::io::copy(&mut js_cursor, &mut js_file),
-        )?;
+        while let Some(result) = writes.try_join_next() {
+            result??;
+            continue;
+        }
 
-        debug!(?wit_path, ?js_path, "Populated temporary input files");
+        debug!(?workspace, "Populated temporary input files");
 
         let mut command = Command::new("jco");
 
         command
             .arg("componentize")
             .arg("-w")
-            .arg(wit_path.display().to_string())
+            .arg(workspace.path())
             .arg("-o")
             .arg(wasm_path.display().to_string())
             .arg(js_path.display().to_string());
diff --git a/rust/usuba/src/bake/mod.rs b/rust/usuba/src/bake/mod.rs
@@ -13,9 +13,14 @@ pub enum Baker {
 
 #[async_trait]
 impl Bake for Baker {
-    async fn bake(&self, wit: Bytes, source_code: Bytes) -> Result<Bytes, crate::UsubaError> {
+    async fn bake(
+        &self,
+        world: &str,
+        wit: Vec<Bytes>,
+        source_code: Bytes,
+    ) -> Result<Bytes, crate::UsubaError> {
         match self {
-            Baker::JavaScript => (JavaScriptBaker {}).bake(wit, source_code).await,
+            Baker::JavaScript => (JavaScriptBaker {}).bake(world, wit, source_code).await,
         }
     }
 }
diff --git a/rust/usuba/src/error.rs b/rust/usuba/src/error.rs
@@ -5,6 +5,7 @@ use blake3::HexError;
 use redb::{CommitError, DatabaseError, StorageError, TableError, TransactionError};
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
+use tokio::task::JoinError;
 use tracing::subscriber::SetGlobalDefaultError;
 use utoipa::ToSchema;
 
@@ -16,6 +17,8 @@ pub enum UsubaError {
     BakeFailure(String),
     #[error("Invalid configuration: {0}")]
     InvalidConfiguration(String),
+    #[error("Invalid module: {0}")]
+    InvalidModule(String),
     #[error("Module not found")]
     ModuleNotFound,
     #[error("An internal error occurred")]
@@ -89,10 +92,23 @@ impl From<HexError> for UsubaError {
     }
 }
 
+impl From<JoinError> for UsubaError {
+    fn from(value: JoinError) -> Self {
+        UsubaError::Internal(format!("{}", value))
+    }
+}
+
+impl From<anyhow::Error> for UsubaError {
+    fn from(value: anyhow::Error) -> Self {
+        todo!()
+    }
+}
+
 impl IntoResponse for UsubaError {
     fn into_response(self) -> axum::response::Response {
         let status = match self {
             UsubaError::BadRequest => StatusCode::BAD_REQUEST,
+            UsubaError::InvalidModule(_) => StatusCode::BAD_REQUEST,
             UsubaError::BakeFailure(_) => StatusCode::INTERNAL_SERVER_ERROR,
             UsubaError::InvalidConfiguration(_) => StatusCode::BAD_REQUEST,
             UsubaError::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR,
diff --git a/rust/usuba/src/routes/module/build.rs b/rust/usuba/src/routes/module/build.rs
@@ -8,12 +8,16 @@ use axum::{
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use utoipa::ToSchema;
+use wit_parser::UnresolvedPackage;
 
 use crate::{Bake, Baker, HashStorage, UsubaError, UsubaState};
 
 #[derive(ToSchema)]
+/// A `multipart/form-data` payload that consists of module WIT + source code as
+/// well as additional (optional) library WIT files
 pub struct BuildModuleRequest {
-    pub files: Vec<Vec<u8>>,
+    pub module: Vec<Vec<u8>>,
+    pub library: Vec<Vec<u8>>,
 }
 
 #[derive(Serialize, Deserialize, ToSchema)]
@@ -41,36 +45,60 @@ pub async fn build_module(
     State(UsubaState { mut storage }): State<UsubaState>,
     mut form_data: Multipart,
 ) -> Result<BuildModuleResponse, UsubaError> {
-    let mut wit: Option<Bytes> = None;
+    let mut world_name: Option<String> = None;
+    let mut wit: Vec<Bytes> = Vec::new();
     let mut source_code: Option<Bytes> = None;
     let mut baker: Option<Baker> = None;
 
-    'collect_files: while let Some(field) = form_data.next_field().await? {
+    while let Some(field) = form_data.next_field().await? {
         if let Some(file_name) = field.file_name() {
             let file_name = PathBuf::from(file_name);
 
-            if let Some(extension) = file_name.extension() {
-                match extension.to_str() {
-                    Some("wit") => {
-                        wit = Some(field.bytes().await?);
-                    }
-                    Some("js") => {
-                        source_code = Some(field.bytes().await?);
-                        baker = Some(Baker::JavaScript);
-                    }
-                    _ => (),
-                };
-            }
-        }
+            match field.name() {
+                Some("module") => {
+                    if let Some(extension) = file_name.extension() {
+                        match extension.to_str() {
+                            Some("wit") => {
+                                let module_wit = field.bytes().await?;
+                                let wit_package = UnresolvedPackage::parse(
+                                    &PathBuf::from("module.wit"),
+                                    String::from_utf8_lossy(&module_wit).as_ref(),
+                                )?;
 
-        match (&wit, &source_code, &baker) {
-            (Some(_), Some(_), Some(_)) => break 'collect_files,
-            _ => (),
+                                wit.push(module_wit);
+
+                                world_name = Some(
+                                    wit_package
+                                        .worlds
+                                        .iter()
+                                        .nth(0)
+                                        .map(|(_, world)| world.name.clone())
+                                        .ok_or_else(|| {
+                                            UsubaError::InvalidModule(format!(
+                                                "Module WIT does not contain a world"
+                                            ))
+                                        })?,
+                                );
+                            }
+                            Some("js") => {
+                                source_code = Some(field.bytes().await?);
+                                baker = Some(Baker::JavaScript);
+                            }
+                            _ => (),
+                        };
+                    }
+                }
+                Some("library") => {
+                    wit.push(field.bytes().await?);
+                }
+                Some(name) => warn!("Unexpected multipart content: {name}"),
+                _ => warn!("Skipping unnamed multipart content"),
+            };
         }
     }
 
-    if let (Some(wit), Some(source_code), Some(baker)) = (wit, source_code, baker) {
-        let wasm = baker.bake(wit, source_code).await?;
+    if let (Some(world_name), Some(source_code), Some(baker)) = (world_name, source_code, baker) {
+        let wasm = baker.bake(&world_name, wit, source_code).await?;
         let hash = storage.write(wasm).await?;
 
         Ok(BuildModuleResponse {
diff --git a/typescript/packages/usuba-sw/openapi.json b/typescript/packages/usuba-sw/openapi.json
@@ -105,11 +105,20 @@
     "schemas": {
       "BuildModuleRequest": {
         "type": "object",
+        "description": "A `multipart/form-data` payload that consists of module WIT + source code as\nwell as additional (optional) library WIT files",
         "required": [
-          "files"
+          "module",
+          "library"
         ],
         "properties": {
-          "files": {
+          "library": {
+            "type": "array",
+            "items": {
+              "type": "string",
+              "format": "binary"
+            }
+          },
+          "module": {
             "type": "array",
             "items": {
               "type": "string",
diff --git a/typescript/packages/usuba-sw/package.json b/typescript/packages/usuba-sw/package.json
@@ -32,13 +32,19 @@
   },
   "wireit": {
     "update-openapi-spec": {
-      "command": "./scripts/update-openapi-spec.sh"
+      "command": "./scripts/update-openapi-spec.sh",
+      "output": [
+        "./openapi.json"
+      ]
     },
     "build:openapi-client": {
+      "command": "npx @hey-api/openapi-ts -i ./openapi.json -o ./src/openapi-client",
       "files": [
         "./openapi.json"
       ],
-      "command": "npx @hey-api/openapi-ts -i ./openapi.json -o ./src/openapi-client"
+      "output": [
+        "./src/openapi-client/**/*"
+      ]
     },
     "build:wasm": {
       "command": "../../../rust/usuba-compat/build-wasm-component.sh",
diff --git a/typescript/packages/usuba-sw/src/index.ts b/typescript/packages/usuba-sw/src/index.ts
diff --git a/typescript/packages/usuba-sw/tsconfig.tsbuildinfo b/typescript/packages/usuba-sw/tsconfig.tsbuildinfo

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,10 @@ impl Guest for Polyfill {`
`30`	`30`	`.map(\|transpiled\| transpiled.into())`
`31`	`31`	`.map_err(\|error\| format!("{}", error))`
`32`	`32`	`}`
	`33`	`+`
	`34`	`+ fn hash(bytes: Vec<u8>) -> String {`
	`35`	`+ blake3::hash(&bytes).to_string()`
	`36`	`+ }`
`33`	`37`	`}`
`34`	`38`
`35`	`39`	`export!(Polyfill);`
Original file line number	Diff line number	Diff line change
`@@ -21,4 +21,5 @@ world usuba-compat {`
`21`	`21`	`}`
`22`	`22`
`23`	`23`	`export polyfill: func(component: list<u8>, options: polyfill-options) -> result<artifacts, string>;`
	`24`	`+ export hash: func(bytes: list<u8>) -> string;`
`24`	`25`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,9 +13,14 @@ pub enum Baker {`
`13`	`13`
`14`	`14`	`#[async_trait]`
`15`	`15`	`impl Bake for Baker {`
`16`		`- async fn bake(&self, wit: Bytes, source_code: Bytes) -> Result<Bytes, crate::UsubaError> {`
	`16`	`+ async fn bake(`
	`17`	`+ &self,`
	`18`	`+ world: &str,`
	`19`	`+ wit: Vec<Bytes>,`
	`20`	`+ source_code: Bytes,`
	`21`	`+ ) -> Result<Bytes, crate::UsubaError> {`
`17`	`22`	`match self {`
`18`		`- Baker::JavaScript => (JavaScriptBaker {}).bake(wit, source_code).await,`
	`23`	`+ Baker::JavaScript => (JavaScriptBaker {}).bake(world, wit, source_code).await,`
`19`	`24`	`}`
`20`	`25`	`}`
`21`	`26`	`}`