Skip to content

Commit 36b48a4

Browse files
committed
chore: manual WASI linking
1 parent 6d6ce5f commit 36b48a4

File tree

5 files changed

+165
-13
lines changed

5 files changed

+165
-13
lines changed

Cargo.lock

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ members = [
1515
resolver = "3"
1616

1717
[workspace.dependencies]
18+
anyhow = { version = "1.0.100", default-features = false }
1819
arrow = { version = "55.2.0", default-features = false, features = ["ipc"] }
1920
chrono = { version = "0.4.42", default-features = false }
2021
datafusion-common = { version = "49.0.1", default-features = false }
@@ -33,6 +34,7 @@ wasip2 = { version = "1" }
3334
wasmtime = { version = "38.0.3", default-features = false, features = ["async", "cranelift"] }
3435
wasmtime-wasi = { version = "38.0.3", default-features = false }
3536
wasmtime-wasi-http = { version = "38.0.3", default-features = false, features = ["default-send-request"] }
37+
wasmtime-wasi-io = { version = "38.0.3", default-features = false }
3638
wit-bindgen = { version = "0.47", default-features = false, features = ["macros"] }
3739

3840
[workspace.lints.rust]

host/Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ license.workspace = true
99
workspace = true
1010

1111
[dependencies]
12+
anyhow.workspace = true
1213
arrow.workspace = true
1314
datafusion-common.workspace = true
1415
datafusion-expr.workspace = true
@@ -21,6 +22,7 @@ tokio = { workspace = true, features = ["rt", "rt-multi-thread", "sync"] }
2122
wasmtime.workspace = true
2223
wasmtime-wasi.workspace = true
2324
wasmtime-wasi-http.workspace = true
25+
wasmtime-wasi-io.workspace = true
2426

2527
[dev-dependencies]
2628
datafusion-udf-wasm-bundle = { workspace = true, features = ["example", "python"] }

host/src/lib.rs

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use tempfile::TempDir;
1111
use tokio::sync::Mutex;
1212
use wasmtime::{
1313
Engine, Store,
14-
component::{Component, Linker, ResourceAny},
14+
component::{Component, ResourceAny},
1515
};
1616
use wasmtime_wasi::{
1717
DirPerms, FilePerms, ResourceTable, WasiCtx, WasiCtxView, WasiView, p2::pipe::MemoryOutputPipe,
@@ -27,6 +27,7 @@ use crate::{
2727
bindings::exports::datafusion_udf_wasm::udf::types as wit_types,
2828
error::DataFusionResultExt,
2929
http::{HttpRequestValidator, RejectAllHttpRequests},
30+
linker::link,
3031
tokio_helpers::async_in_sync_context,
3132
};
3233
use crate::{error::WasmToDataFusionResultExt, tokio_helpers::blocking_io};
@@ -43,6 +44,7 @@ mod bindings;
4344
mod conversion;
4445
mod error;
4546
pub mod http;
47+
mod linker;
4648
mod tokio_helpers;
4749

4850
/// State of the WASM payload.
@@ -285,18 +287,9 @@ impl WasmScalarUdf {
285287
resource_table: ResourceTable::new(),
286288
http_validator: Arc::clone(&permissions.http),
287289
};
288-
let mut store = Store::new(engine, state);
289-
290-
let mut linker = Linker::new(engine);
291-
wasmtime_wasi::p2::add_to_linker_async(&mut linker).context("link WASI p2", None)?;
292-
wasmtime_wasi_http::add_only_http_to_linker_async(&mut linker)
293-
.context("link WASI p2 HTTP", None)?;
294-
295-
let bindings = Arc::new(
296-
bindings::Datafusion::instantiate_async(&mut store, component, &linker)
297-
.await
298-
.context("initialize bindings", Some(&store.data().stderr.contents()))?,
299-
);
290+
let (bindings, mut store) = link(engine, component, state)
291+
.await
292+
.context("link WASM components", None)?;
300293

301294
// fill root FS
302295
let root_data = bindings

host/src/linker.rs

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
//! WebAssembly linker code.
2+
3+
use std::sync::Arc;
4+
5+
use anyhow::{Context, Result};
6+
use wasmtime::{
7+
Engine, Store,
8+
component::{Component, HasData, Linker},
9+
};
10+
use wasmtime_wasi::{ResourceTable, WasiView};
11+
12+
use crate::{WasmStateImpl, bindings::Datafusion};
13+
14+
/// Link everything.
15+
pub(crate) async fn link(
16+
engine: &Engine,
17+
component: &Component,
18+
state: WasmStateImpl,
19+
) -> Result<(Arc<Datafusion>, Store<WasmStateImpl>)> {
20+
let mut store = Store::new(engine, state);
21+
22+
let mut linker = Linker::new(engine);
23+
link_wasi_p2(&mut linker).context("link WASI p2")?;
24+
wasmtime_wasi_http::add_only_http_to_linker_async(&mut linker)
25+
.context("link WASI p2 HWasmStateImplWasmStateImplP")?;
26+
27+
let bindings = Arc::new(
28+
Datafusion::instantiate_async(&mut store, component, &linker)
29+
.await
30+
.context("initialize bindings")?,
31+
);
32+
Ok((bindings, store))
33+
}
34+
35+
/// Link WASIp2 interfaces.
36+
fn link_wasi_p2(linker: &mut Linker<WasmStateImpl>) -> Result<()> {
37+
use wasmtime_wasi::{
38+
cli::{WasiCli, WasiCliView},
39+
clocks::{WasiClocks, WasiClocksView},
40+
filesystem::{WasiFilesystem, WasiFilesystemView},
41+
p2::bindings,
42+
random::WasiRandom,
43+
sockets::{WasiSockets, WasiSocketsView},
44+
};
45+
46+
let options = bindings::LinkOptions::default();
47+
wasmtime_wasi_io::bindings::wasi::io::error::add_to_linker::<WasmStateImpl, HasIo>(
48+
linker,
49+
|t| t.ctx().table,
50+
)?;
51+
wasmtime_wasi_io::bindings::wasi::io::poll::add_to_linker::<WasmStateImpl, HasIo>(
52+
linker,
53+
|t| t.ctx().table,
54+
)?;
55+
wasmtime_wasi_io::bindings::wasi::io::streams::add_to_linker::<WasmStateImpl, HasIo>(
56+
linker,
57+
|t| t.ctx().table,
58+
)?;
59+
bindings::clocks::wall_clock::add_to_linker::<WasmStateImpl, WasiClocks>(
60+
linker,
61+
WasmStateImpl::clocks,
62+
)?;
63+
bindings::clocks::monotonic_clock::add_to_linker::<WasmStateImpl, WasiClocks>(
64+
linker,
65+
WasmStateImpl::clocks,
66+
)?;
67+
bindings::cli::exit::add_to_linker::<WasmStateImpl, WasiCli>(
68+
linker,
69+
&(&options).into(),
70+
WasmStateImpl::cli,
71+
)?;
72+
bindings::cli::environment::add_to_linker::<WasmStateImpl, WasiCli>(
73+
linker,
74+
WasmStateImpl::cli,
75+
)?;
76+
bindings::cli::stdin::add_to_linker::<WasmStateImpl, WasiCli>(linker, WasmStateImpl::cli)?;
77+
bindings::cli::stdout::add_to_linker::<WasmStateImpl, WasiCli>(linker, WasmStateImpl::cli)?;
78+
bindings::cli::stderr::add_to_linker::<WasmStateImpl, WasiCli>(linker, WasmStateImpl::cli)?;
79+
bindings::cli::terminal_input::add_to_linker::<WasmStateImpl, WasiCli>(
80+
linker,
81+
WasmStateImpl::cli,
82+
)?;
83+
bindings::cli::terminal_output::add_to_linker::<WasmStateImpl, WasiCli>(
84+
linker,
85+
WasmStateImpl::cli,
86+
)?;
87+
bindings::cli::terminal_stdin::add_to_linker::<WasmStateImpl, WasiCli>(
88+
linker,
89+
WasmStateImpl::cli,
90+
)?;
91+
bindings::cli::terminal_stdout::add_to_linker::<WasmStateImpl, WasiCli>(
92+
linker,
93+
WasmStateImpl::cli,
94+
)?;
95+
bindings::cli::terminal_stderr::add_to_linker::<WasmStateImpl, WasiCli>(
96+
linker,
97+
WasmStateImpl::cli,
98+
)?;
99+
bindings::filesystem::types::add_to_linker::<WasmStateImpl, WasiFilesystem>(
100+
linker,
101+
WasmStateImpl::filesystem,
102+
)?;
103+
bindings::filesystem::preopens::add_to_linker::<WasmStateImpl, WasiFilesystem>(
104+
linker,
105+
WasmStateImpl::filesystem,
106+
)?;
107+
bindings::random::random::add_to_linker::<WasmStateImpl, WasiRandom>(linker, |t| {
108+
t.ctx().ctx.random()
109+
})?;
110+
bindings::random::insecure::add_to_linker::<WasmStateImpl, WasiRandom>(linker, |t| {
111+
t.ctx().ctx.random()
112+
})?;
113+
bindings::random::insecure_seed::add_to_linker::<WasmStateImpl, WasiRandom>(linker, |t| {
114+
t.ctx().ctx.random()
115+
})?;
116+
bindings::sockets::instance_network::add_to_linker::<WasmStateImpl, WasiSockets>(
117+
linker,
118+
WasmStateImpl::sockets,
119+
)?;
120+
bindings::sockets::network::add_to_linker::<WasmStateImpl, WasiSockets>(
121+
linker,
122+
&(&options).into(),
123+
WasmStateImpl::sockets,
124+
)?;
125+
bindings::sockets::ip_name_lookup::add_to_linker::<WasmStateImpl, WasiSockets>(
126+
linker,
127+
WasmStateImpl::sockets,
128+
)?;
129+
bindings::sockets::tcp::add_to_linker::<WasmStateImpl, WasiSockets>(
130+
linker,
131+
WasmStateImpl::sockets,
132+
)?;
133+
bindings::sockets::tcp_create_socket::add_to_linker::<WasmStateImpl, WasiSockets>(
134+
linker,
135+
WasmStateImpl::sockets,
136+
)?;
137+
bindings::sockets::udp::add_to_linker::<WasmStateImpl, WasiSockets>(
138+
linker,
139+
WasmStateImpl::sockets,
140+
)?;
141+
bindings::sockets::udp_create_socket::add_to_linker::<WasmStateImpl, WasiSockets>(
142+
linker,
143+
WasmStateImpl::sockets,
144+
)?;
145+
Ok(())
146+
}
147+
148+
/// Marker struct to tell linker that we do in fact provide IO-related resource tables.
149+
struct HasIo;
150+
151+
impl HasData for HasIo {
152+
type Data<'a> = &'a mut ResourceTable;
153+
}

0 commit comments

Comments
 (0)