Merge pull request #257 from influxdata/crepererum/target-triplet

feat: allow setting a target triplet

Author: crepererum

Cargo.lock

@@ -42,6 +42,7 @@ sqlparser = { version = "0.58.0", default-features = false, features = [
   "visitor"
 ] }
 tar = { version = "0.4.44", default-features = false }
+target-lexicon = { version = "0.13", default-features = false }
 tempfile = { version = "3.23.0", default-features = false }
 tokio = { version = "1.48.0", default-features = false }
 uuid = { version = "1.19.0", default-features = false, features = ["v4"] }

Cargo.toml

@@ -32,8 +32,13 @@ datafusion-udf-wasm-bundle = { workspace = true, features = [
 ] }
 insta.workspace = true
 regex.workspace = true
+target-lexicon.workspace = true
 tokio = { workspace = true, features = ["fs", "macros"] }
 wiremock = "0.6.5"
 
+[features]
+# enable all architectures in cranelift/wasmtime
+all-arch = ["wasmtime/all-arch"]
+
 [lints]
 workspace = true

host/Cargo.toml

@@ -21,17 +21,37 @@ use crate::{
 };
 
 /// Create WASM engine.
-fn create_engine() -> DataFusionResult<Engine> {
-    Engine::new(
-        wasmtime::Config::new()
-            .async_support(true)
-            .epoch_interruption(true)
-            .memory_init_cow(true)
-            // Disable backtraces for now since debug info parsing doesn't seem to work and hence error
-            // messages are nondeterministic.
-            .wasm_backtrace(false),
-    )
-    .context("create WASM engine", None)
+fn create_engine(flags: &CompilationFlags) -> DataFusionResult<Engine> {
+    // TODO: Once https://github.com/bytecodealliance/wasmtime/pull/12089 is released, make this an `Option` and treat
+    //       `None` as `Config::without_compiler`.
+    let CompilationFlags { target } = flags;
+
+    let mut config = wasmtime::Config::new();
+    config.async_support(true);
+    config.epoch_interruption(true);
+    config.memory_init_cow(true);
+    // Disable backtraces for now since debug info parsing doesn't seem to work and hence error
+    // messages are nondeterministic.
+    config.wasm_backtrace(false);
+
+    if let Some(target) = &target {
+        config
+            .target(target)
+            .with_context(|_| format!("cannot set target: {target}"), None)?;
+    }
+
+    Engine::new(&config).context("create WASM engine", None)
+}
+
+/// Code compilation flags.
+///
+/// This is used when [compiling a component](WasmComponentPrecompiled::compile).
+#[derive(Debug, Default, Clone)]
+pub struct CompilationFlags {
+    /// Target (triplet).
+    ///
+    /// Set to [`None`] to use the host configuration. Note that this may lead to unportable compiled code.
+    pub target: Option<String>,
 }
 
 /// Pre-compiled WASM component.
@@ -51,11 +71,14 @@ impl WasmComponentPrecompiled {
     ///
     ///
     /// [binary format]: https://webassembly.github.io/spec/core/binary/index.html
-    pub async fn new(wasm_binary: Arc<[u8]>) -> DataFusionResult<Self> {
-        tokio::task::spawn_blocking(move || {
-            // Create temporary engine that we need for compilation.
-            let engine = create_engine()?;
+    pub async fn compile(
+        wasm_binary: Arc<[u8]>,
+        flags: &CompilationFlags,
+    ) -> DataFusionResult<Self> {
+        // Create temporary engine that we need for compilation.
+        let engine = create_engine(flags)?;
 
+        tokio::task::spawn_blocking(move || {
             let compiled_component = engine
                 .precompile_component(&wasm_binary)
                 .context("pre-compile component", None)?;
@@ -83,7 +106,7 @@ impl WasmComponentPrecompiled {
     ///
     /// # Exposure
     /// It is generally safe to leak/expose the pre-compiled data to the user that provided the WASM bytecode (see
-    /// [`new`](Self::new)). However, you must prevent the user from tampering the data, see "safety" section of
+    /// [`compile`](Self::compile)). However, you must prevent the user from tampering the data, see "safety" section of
     /// [`load`](Self::load).
     ///
     /// The exposed data is opaque and we make no guarantees about the internal structure of it.
@@ -139,7 +162,7 @@ impl WasmComponentPrecompiled {
         };
 
         // test hydration
-        let engine = create_engine()?;
+        let engine = create_engine(&CompilationFlags::default())?;
         this.hydrate(&engine)?;
 
         Ok(this)
@@ -187,7 +210,7 @@ impl WasmComponentInstance {
         io_rt: Handle,
         memory_pool: &Arc<dyn MemoryPool>,
     ) -> DataFusionResult<Self> {
-        let engine = create_engine()?;
+        let engine = create_engine(&CompilationFlags::default())?;
 
         // set up epoch timer
         let mut epoch_task = JoinSet::new();

host/src/component.rs

@@ -43,20 +43,43 @@ pub(crate) trait WasmToDataFusionResultExt {
     /// [`Ok`] payload.
     type T;
 
+    /// [`Err`] payload
+    type E;
+
     /// Add context to error.
     ///
     /// The context has:
     /// - `msg`: a human-readable context description
     /// - `stderr`: stderr output of the WASM payload if available
     fn context(self, msg: &str, stderr: Option<&[u8]>) -> Result<Self::T, DataFusionError>;
+
+    /// Add context to error.
+    ///
+    /// The context has:
+    /// - `msg`: a closure that generates a human-readable context description based on the error
+    /// - `stderr`: stderr output of the WASM payload if available
+    fn with_context<F>(self, msg: F, stderr: Option<&[u8]>) -> Result<Self::T, DataFusionError>
+    where
+        F: for<'a> FnOnce(&'a Self::E) -> String;
 }
 
 impl<T> WasmToDataFusionResultExt for Result<T, wasmtime::Error> {
     type T = T;
+    type E = wasmtime::Error;
 
     fn context(self, msg: &str, stderr: Option<&[u8]>) -> Result<Self::T, DataFusionError> {
         self.map_err(|err| WasmToDataFusionErrorExt::context(err, msg, stderr))
     }
+
+    fn with_context<F>(self, msg: F, stderr: Option<&[u8]>) -> Result<Self::T, DataFusionError>
+    where
+        F: for<'a> FnOnce(&'a Self::E) -> String,
+    {
+        self.map_err(|err| {
+            let msg = msg(&err);
+            WasmToDataFusionErrorExt::context(err, &msg, stderr)
+        })
+    }
 }
 
 /// Extension trait for [`Result`] containing a [`DataFusionError`].

host/src/error.rs

@@ -4,7 +4,7 @@
 //! [DataFusion]: https://datafusion.apache.org/
 
 pub use crate::{
-    component::WasmComponentPrecompiled,
+    component::{CompilationFlags, WasmComponentPrecompiled},
     conversion::limits::TrustedDataLimits,
     http::{
         AllowCertainHttpRequests, HttpMethod, HttpRequestMatcher, HttpRequestRejected,
@@ -22,6 +22,8 @@ use datafusion_udf_wasm_bundle as _;
 #[cfg(test)]
 use regex as _;
 #[cfg(test)]
+use target_lexicon as _;
+#[cfg(test)]
 use wiremock as _;
 
 mod bindings;

host/src/lib.rs

@@ -2,7 +2,9 @@ use std::sync::{Arc, LazyLock};
 
 use datafusion_common::DataFusionError;
 use datafusion_execution::memory_pool::GreedyMemoryPool;
-use datafusion_udf_wasm_host::{WasmComponentPrecompiled, WasmPermissions, WasmScalarUdf};
+use datafusion_udf_wasm_host::{
+    CompilationFlags, WasmComponentPrecompiled, WasmPermissions, WasmScalarUdf,
+};
 use tokio::{runtime::Runtime, sync::OnceCell};
 
 /// Static memory limit.
@@ -17,9 +19,12 @@ static COMPONENT: OnceCell<WasmComponentPrecompiled> = OnceCell::const_new();
 pub(crate) async fn component() -> &'static WasmComponentPrecompiled {
     COMPONENT
         .get_or_init(async || {
-            WasmComponentPrecompiled::new(datafusion_udf_wasm_bundle::BIN_EVIL.into())
-                .await
-                .unwrap()
+            WasmComponentPrecompiled::compile(
+                datafusion_udf_wasm_bundle::BIN_EVIL.into(),
+                &CompilationFlags::default(),
+            )
+            .await
+            .unwrap()
         })
         .await
 }

host/tests/integration_tests/evil/test_utils.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use datafusion_common::DataFusionError;
 use datafusion_execution::memory_pool::GreedyMemoryPool;
-use datafusion_udf_wasm_host::{WasmComponentPrecompiled, WasmScalarUdf};
+use datafusion_udf_wasm_host::{CompilationFlags, WasmComponentPrecompiled, WasmScalarUdf};
 use tokio::{runtime::Handle, sync::OnceCell};
 
 /// Memory limit in bytes.
@@ -17,9 +17,12 @@ static COMPONENT: OnceCell<WasmComponentPrecompiled> = OnceCell::const_new();
 pub(crate) async fn python_component() -> &'static WasmComponentPrecompiled {
     COMPONENT
         .get_or_init(async || {
-            WasmComponentPrecompiled::new(datafusion_udf_wasm_bundle::BIN_PYTHON.into())
-                .await
-                .unwrap()
+            WasmComponentPrecompiled::compile(
+                datafusion_udf_wasm_bundle::BIN_PYTHON.into(),
+                &CompilationFlags::default(),
+            )
+            .await
+            .unwrap()
         })
         .await
 }

host/tests/integration_tests/python/test_utils.rs

@@ -12,7 +12,8 @@ use datafusion_expr::{
     async_udf::AsyncScalarUDFImpl,
 };
 use datafusion_udf_wasm_host::{
-    StaticResourceLimits, WasmComponentPrecompiled, WasmPermissions, WasmScalarUdf,
+    CompilationFlags, StaticResourceLimits, WasmComponentPrecompiled, WasmPermissions,
+    WasmScalarUdf,
 };
 use tokio::{runtime::Handle, sync::OnceCell};
 
@@ -274,14 +275,95 @@ async fn test_limit_initial_n_memories() {
     );
 }
 
+#[tokio::test]
+async fn test_match_target() {
+    let component = WasmComponentPrecompiled::compile(
+        datafusion_udf_wasm_bundle::BIN_EXAMPLE.into(),
+        &CompilationFlags {
+            target: Some(target_lexicon::HOST.to_string()),
+        },
+    )
+    .await
+    .unwrap();
+
+    // instantiating works
+    WasmScalarUdf::new(
+        &component,
+        &Default::default(),
+        Handle::current(),
+        &(Arc::new(UnboundedMemoryPool::default()) as _),
+        "".to_owned(),
+    )
+    .await
+    .unwrap();
+
+    // and load->store also works
+    let data = component.store().to_vec();
+    // SAFETY: we just compiled that
+    let res = unsafe { WasmComponentPrecompiled::load(data) };
+    res.unwrap();
+}
+
+#[cfg(feature = "all-arch")]
+#[tokio::test]
+async fn test_mismatch_target() {
+    let component = WasmComponentPrecompiled::compile(
+        datafusion_udf_wasm_bundle::BIN_EXAMPLE.into(),
+        &CompilationFlags {
+            // It's unlikely that someone is gonna run the tests on a RISC-V 64bit host, but if they do, we need to
+            // make the test code smarter. It won't fail as expected.
+            target: Some("riscv64gc-unknown-linux-gnu".to_owned()),
+        },
+    )
+    .await
+    .unwrap();
+
+    // instantiating doesn't work
+    let err = WasmScalarUdf::new(
+        &component,
+        &Default::default(),
+        Handle::current(),
+        &(Arc::new(UnboundedMemoryPool::default()) as _),
+        "".to_owned(),
+    )
+    .await
+    .unwrap_err();
+
+    insta::assert_snapshot!(
+        err,
+        @r"
+    create WASM component
+    caused by
+    External error: Module was compiled for architecture 'riscv64gc'
+    "
+    );
+
+    // and load->store also fails
+    let data = component.store().to_vec();
+    // SAFETY: we just compiled that
+    let res = unsafe { WasmComponentPrecompiled::load(data) };
+
+    insta::assert_snapshot!(
+        res.unwrap_err(),
+        @r"
+    create WASM component
+    caused by
+    External error: Module was compiled for architecture 'riscv64gc'
+    "
+    );
+}
+
 async fn component() -> &'static WasmComponentPrecompiled {
     static COMPONENT: OnceCell<WasmComponentPrecompiled> = OnceCell::const_new();
 
     COMPONENT
         .get_or_init(async || {
-            WasmComponentPrecompiled::new(datafusion_udf_wasm_bundle::BIN_EXAMPLE.into())
-                .await
-                .unwrap()
+            WasmComponentPrecompiled::compile(
+                datafusion_udf_wasm_bundle::BIN_EXAMPLE.into(),
+                &CompilationFlags::default(),
+            )
+            .await
+            .unwrap()
         })
         .await
 }

host/tests/integration_tests/rust.rs

@@ -1,4 +1,4 @@
-use datafusion_udf_wasm_host::WasmComponentPrecompiled;
+use datafusion_udf_wasm_host::{CompilationFlags, WasmComponentPrecompiled};
 use tokio::sync::OnceCell;
 
 /// Static precompiled Python WASM component for tests
@@ -8,9 +8,12 @@ static COMPONENT: OnceCell<WasmComponentPrecompiled> = OnceCell::const_new();
 pub(crate) async fn python_component() -> &'static WasmComponentPrecompiled {
     COMPONENT
         .get_or_init(async || {
-            WasmComponentPrecompiled::new(datafusion_udf_wasm_bundle::BIN_PYTHON.into())
-                .await
-                .unwrap()
+            WasmComponentPrecompiled::compile(
+                datafusion_udf_wasm_bundle::BIN_PYTHON.into(),
+                &CompilationFlags::default(),
+            )
+            .await
+            .unwrap()
         })
         .await
 }