feat: virtual file system

Author: crepererum

Cargo.lock

@@ -16,8 +16,9 @@ datafusion-expr.workspace = true
 datafusion-udf-wasm-arrow2bytes.workspace = true
 http.workspace = true
 hyper.workspace = true
+rand = { version = "0.9" }
+siphasher = { version = "1", default-features = false }
 tar.workspace = true
-tempfile.workspace = true
 tokio = { workspace = true, features = ["rt", "rt-multi-thread", "sync"] }
 wasmtime.workspace = true
 wasmtime-wasi.workspace = true

host/Cargo.toml

@@ -2,20 +2,17 @@
 //!
 //!
 //! [DataFusion]: https://datafusion.apache.org/
-use std::{any::Any, io::Cursor, ops::DerefMut, sync::Arc};
+use std::{any::Any, ops::DerefMut, sync::Arc};
 
 use arrow::datatypes::DataType;
 use datafusion_common::{DataFusionError, Result as DataFusionResult};
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
-use tempfile::TempDir;
 use tokio::sync::Mutex;
 use wasmtime::{
     Engine, Store,
     component::{Component, ResourceAny},
 };
-use wasmtime_wasi::{
-    DirPerms, FilePerms, ResourceTable, WasiCtx, WasiCtxView, WasiView, p2::pipe::MemoryOutputPipe,
-};
+use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxView, WasiView, p2::pipe::MemoryOutputPipe};
 use wasmtime_wasi_http::{
     HttpResult, WasiHttpCtx, WasiHttpView,
     bindings::http::types::ErrorCode as HttpErrorCode,
@@ -25,12 +22,12 @@ use wasmtime_wasi_http::{
 
 use crate::{
     bindings::exports::datafusion_udf_wasm::udf::types as wit_types,
-    error::DataFusionResultExt,
+    error::{DataFusionResultExt, WasmToDataFusionResultExt},
     http::{HttpRequestValidator, RejectAllHttpRequests},
     linker::link,
     tokio_helpers::async_in_sync_context,
+    vfs::{VfsCtxView, VfsState, VfsView},
 };
-use crate::{error::WasmToDataFusionResultExt, tokio_helpers::blocking_io};
 
 // unused-crate-dependencies false positives
 #[cfg(test)]
@@ -46,13 +43,14 @@ mod error;
 pub mod http;
 mod linker;
 mod tokio_helpers;
+mod vfs;
 
 /// State of the WASM payload.
 struct WasmStateImpl {
-    /// Temporary directory that holds the root filesystem.
+    /// Virtual filesystem for the WASM payload.
     ///
-    /// This filesystem is provided to the payload as read-only.
-    root: TempDir,
+    /// This filesystem is provided to the payload in memory with read-write support.
+    vfs_state: VfsState,
 
     /// A limited buffer for stderr.
     ///
@@ -75,15 +73,15 @@ struct WasmStateImpl {
 impl std::fmt::Debug for WasmStateImpl {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let Self {
-            root,
+            vfs_state,
             stderr,
             wasi_ctx: _,
             wasi_http_ctx: _,
             resource_table,
             http_validator,
         } = self;
         f.debug_struct("WasmStateImpl")
-            .field("root", root)
+            .field("vfs_state", vfs_state)
             .field("stderr", stderr)
             .field("wasi_ctx", &"<WASI_CTX>")
             .field("resource_table", resource_table)
@@ -136,6 +134,15 @@ impl WasiHttpView for WasmStateImpl {
     }
 }
 
+impl VfsView for WasmStateImpl {
+    fn vfs(&mut self) -> VfsCtxView<'_> {
+        VfsCtxView {
+            table: &mut self.resource_table,
+            vfs_state: &mut self.vfs_state,
+        }
+    }
+}
+
 /// Pre-compiled WASM component.
 ///
 /// The pre-compilation is stateless and can be used to [create](WasmScalarUdf::new) multiple instances that do not share
@@ -266,21 +273,13 @@ impl WasmScalarUdf {
     ) -> DataFusionResult<Vec<Self>> {
         let WasmComponentPrecompiled { engine, component } = component;
 
-        // TODO: we need an in-mem file system for this, see
-        //       - https://github.com/bytecodealliance/wasmtime/issues/8963
-        //       - https://github.com/Timmmm/wasmtime_fs_demo
-        let root = blocking_io(TempDir::new)
-            .await
-            .map_err(DataFusionError::IoError)?;
+        // Create in-memory VFS
+        let vfs_state = VfsState::new();
 
         let stderr = MemoryOutputPipe::new(1024);
-        let wasi_ctx = WasiCtx::builder()
-            .stderr(stderr.clone())
-            .preopened_dir(root.path(), "/", DirPerms::READ, FilePerms::READ)
-            .context("pre-open root dir", None)?
-            .build();
+        let wasi_ctx = WasiCtx::builder().stderr(stderr.clone()).build();
         let state = WasmStateImpl {
-            root,
+            vfs_state,
             stderr,
             wasi_ctx,
             wasi_http_ctx: WasiHttpCtx::new(),
@@ -291,7 +290,7 @@ impl WasmScalarUdf {
             .await
             .context("link WASM components", None)?;
 
-        // fill root FS
+        // Populate VFS from tar archive
         let root_data = bindings
             .datafusion_udf_wasm_udf_types()
             .call_root_fs_tar(&mut store)
@@ -301,13 +300,11 @@ impl WasmScalarUdf {
                 Some(&store.data().stderr.contents()),
             )?;
         if let Some(root_data) = root_data {
-            let root_path = store.data().root.path().to_owned();
-            blocking_io(move || {
-                let mut a = tar::Archive::new(Cursor::new(root_data));
-                a.unpack(root_path)
-            })
-            .await
-            .map_err(DataFusionError::IoError)?;
+            store
+                .data_mut()
+                .vfs_state
+                .populate_from_tar(&root_data)
+                .map_err(DataFusionError::IoError)?;
         }
 
         let udf_resources = bindings

host/src/lib.rs

@@ -9,7 +9,11 @@ use wasmtime::{
 };
 use wasmtime_wasi::{ResourceTable, WasiView};
 
-use crate::{WasmStateImpl, bindings::Datafusion};
+use crate::{
+    WasmStateImpl,
+    bindings::Datafusion,
+    vfs::{HasFs, VfsView},
+};
 
 /// Link everything.
 pub(crate) async fn link(
@@ -37,7 +41,6 @@ fn link_wasi_p2(linker: &mut Linker<WasmStateImpl>) -> Result<()> {
     use wasmtime_wasi::{
         cli::{WasiCli, WasiCliView},
         clocks::{WasiClocks, WasiClocksView},
-        filesystem::{WasiFilesystem, WasiFilesystemView},
         p2::bindings,
         random::WasiRandom,
         sockets::{WasiSockets, WasiSocketsView},
@@ -96,13 +99,10 @@ fn link_wasi_p2(linker: &mut Linker<WasmStateImpl>) -> Result<()> {
         linker,
         WasmStateImpl::cli,
     )?;
-    bindings::filesystem::types::add_to_linker::<WasmStateImpl, WasiFilesystem>(
-        linker,
-        WasmStateImpl::filesystem,
-    )?;
-    bindings::filesystem::preopens::add_to_linker::<WasmStateImpl, WasiFilesystem>(
+    bindings::filesystem::types::add_to_linker::<WasmStateImpl, HasFs>(linker, WasmStateImpl::vfs)?;
+    bindings::filesystem::preopens::add_to_linker::<WasmStateImpl, HasFs>(
         linker,
-        WasmStateImpl::filesystem,
+        WasmStateImpl::vfs,
     )?;
     bindings::random::random::add_to_linker::<WasmStateImpl, WasiRandom>(linker, |t| {
         t.ctx().ctx.random()

host/src/linker.rs

@@ -11,12 +11,3 @@ where
 {
     tokio::task::block_in_place(move || tokio::runtime::Handle::current().block_on(fut))
 }
-
-/// Spawn blocking I/O.
-pub(crate) async fn blocking_io<F, T>(f: F) -> std::io::Result<T>
-where
-    F: FnOnce() -> std::io::Result<T> + Send + 'static,
-    T: Send + 'static,
-{
-    tokio::task::spawn_blocking(f).await?
-}