fix: run HTTP in I/O runtime

Closes #140.

Author: crepererum

host/src/lib.rs

@@ -8,7 +8,7 @@ use ::http::HeaderName;
 use arrow::datatypes::DataType;
 use datafusion_common::{DataFusionError, Result as DataFusionResult};
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
-use tokio::sync::Mutex;
+use tokio::{runtime::Handle, sync::Mutex};
 use wasmtime::{
     Engine, Store,
     component::{Component, ResourceAny},
@@ -72,6 +72,9 @@ struct WasmStateImpl {
 
     /// HTTP request validator.
     http_validator: Arc<dyn HttpRequestValidator>,
+
+    /// Handle to tokio I/O runtime.
+    io_rt: Handle,
 }
 
 impl std::fmt::Debug for WasmStateImpl {
@@ -83,13 +86,15 @@ impl std::fmt::Debug for WasmStateImpl {
             wasi_http_ctx: _,
             resource_table,
             http_validator,
+            io_rt,
         } = self;
         f.debug_struct("WasmStateImpl")
             .field("vfs_state", vfs_state)
             .field("stderr", stderr)
             .field("wasi_ctx", &"<WASI_CTX>")
             .field("resource_table", resource_table)
             .field("http_validator", http_validator)
+            .field("io_rt", io_rt)
             .finish()
     }
 }
@@ -117,6 +122,8 @@ impl WasiHttpView for WasmStateImpl {
         mut request: hyper::Request<HyperOutgoingBody>,
         config: OutgoingRequestConfig,
     ) -> HttpResult<HostFutureIncomingResponse> {
+        let _guard = self.io_rt.enter();
+
         // Python `requests` sends this so we allow it but later drop it from the actual request.
         request.headers_mut().remove(hyper::header::CONNECTION);
 
@@ -298,6 +305,7 @@ impl WasmScalarUdf {
     pub async fn new(
         component: &WasmComponentPrecompiled,
         permissions: &WasmPermissions,
+        io_rt: Handle,
         source: String,
     ) -> DataFusionResult<Vec<Self>> {
         let WasmComponentPrecompiled { engine, component } = component;
@@ -314,6 +322,7 @@ impl WasmScalarUdf {
             wasi_http_ctx: WasiHttpCtx::new(),
             resource_table: ResourceTable::new(),
             http_validator: Arc::clone(&permissions.http),
+            io_rt,
         };
         let (bindings, mut store) = link(engine, component, state)
             .await

host/tests/integration_tests/python/runtime/fs.rs

@@ -6,6 +6,7 @@ use arrow::{
 };
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
 use datafusion_udf_wasm_host::{WasmPermissions, WasmScalarUdf, vfs::VfsLimits};
+use tokio::runtime::Handle;
 
 use crate::integration_tests::{
     python::test_utils::{python_component, python_scalar_udf},
@@ -239,6 +240,7 @@ async fn test_limit_inodes() {
             inodes: 42,
             ..Default::default()
         }),
+        Handle::current(),
         "".to_owned(),
     )
     .await
@@ -260,6 +262,7 @@ async fn test_limit_bytes() {
             bytes: 1337,
             ..Default::default()
         }),
+        Handle::current(),
         "".to_owned(),
     )
     .await

host/tests/integration_tests/python/runtime/http.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use arrow::{
     array::{Array, StringArray, StringBuilder},
@@ -10,6 +10,7 @@ use datafusion_udf_wasm_host::{
     WasmPermissions, WasmScalarUdf,
     http::{AllowCertainHttpRequests, HttpRequestValidator, Matcher},
 };
+use tokio::runtime::Handle;
 use wasmtime_wasi_http::types::DEFAULT_FORBIDDEN_HEADERS;
 use wiremock::{Mock, MockServer, ResponseTemplate, matchers};
 
@@ -568,6 +569,7 @@ where
     WasmScalarUdf::new(
         python_component().await,
         &WasmPermissions::new().with_http(permissions),
+        Handle::current(),
         code.to_owned(),
     )
     .await
@@ -582,3 +584,79 @@ where
     assert_eq!(udfs.len(), 1);
     udfs.into_iter().next().unwrap()
 }
+
+#[test]
+fn test_io_runtime() {
+    const CODE: &str = r#"
+import urllib3
+
+def perform_request(url: str) -> str:
+    resp = urllib3.request("GET", url)
+    return resp.data.decode("utf-8")
+"#;
+
+    let rt_tmp = tokio::runtime::Builder::new_current_thread()
+        .build()
+        .unwrap();
+    let rt_cpu = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(1)
+        // It would be nice if all the timeouts-related timers would also run within the within the I/O runtime, but
+        // that requires some larger intervention (either upstream or with a custom WASI HTTP implementation).
+        // Hence, we don't do that yet.
+        .enable_time()
+        .build()
+        .unwrap();
+    let rt_io = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(1)
+        .enable_all()
+        .build()
+        .unwrap();
+
+    let server = rt_io.block_on(async {
+        let server = MockServer::start().await;
+        Mock::given(matchers::any())
+            .respond_with(ResponseTemplate::new(200).set_body_string("hello world!"))
+            .expect(1)
+            .mount(&server)
+            .await;
+        server
+    });
+
+    // deliberately use a runtime what we are going to throw away later to prevent tricks like `Handle::current`
+    let udf = rt_tmp.block_on(async {
+        let mut permissions = AllowCertainHttpRequests::new();
+        permissions.allow(Matcher {
+            method: http::Method::GET,
+            host: server.address().ip().to_string().into(),
+            port: server.address().port(),
+        });
+
+        let udfs = WasmScalarUdf::new(
+            python_component().await,
+            &WasmPermissions::new().with_http(permissions),
+            rt_io.handle().clone(),
+            CODE.to_owned(),
+        )
+        .await
+        .unwrap();
+        assert_eq!(udfs.len(), 1);
+        udfs.into_iter().next().unwrap()
+    });
+    rt_tmp.shutdown_timeout(Duration::from_secs(1));
+
+    let array = rt_cpu.block_on(async {
+        udf.invoke_with_args(ScalarFunctionArgs {
+            args: vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(server.uri())))],
+            arg_fields: vec![Arc::new(Field::new("uri", DataType::Utf8, true))],
+            number_rows: 1,
+            return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
+        })
+        .unwrap()
+        .unwrap_array()
+    });
+
+    assert_eq!(
+        array.as_ref(),
+        &StringArray::from_iter([Some("hello world!".to_owned()),]) as &dyn Array,
+    );
+}

host/tests/integration_tests/python/test_utils.rs

@@ -1,6 +1,6 @@
 use datafusion_common::DataFusionError;
 use datafusion_udf_wasm_host::{WasmComponentPrecompiled, WasmScalarUdf};
-use tokio::sync::OnceCell;
+use tokio::{runtime::Handle, sync::OnceCell};
 
 /// Static precompiled Python WASM component for tests
 static COMPONENT: OnceCell<WasmComponentPrecompiled> = OnceCell::const_new();
@@ -20,7 +20,13 @@ pub(crate) async fn python_component() -> &'static WasmComponentPrecompiled {
 pub(crate) async fn python_scalar_udfs(code: &str) -> Result<Vec<WasmScalarUdf>, DataFusionError> {
     let component = python_component().await;
 
-    WasmScalarUdf::new(component, &Default::default(), code.to_owned()).await
+    WasmScalarUdf::new(
+        component,
+        &Default::default(),
+        Handle::current(),
+        code.to_owned(),
+    )
+    .await
 }
 
 /// Compiles the provided Python UDF code into a single WasmScalarUdf instance.

host/tests/integration_tests/rust.rs

@@ -7,6 +7,7 @@ use arrow::{
 use datafusion_common::ScalarValue;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
 use datafusion_udf_wasm_host::{WasmComponentPrecompiled, WasmScalarUdf};
+use tokio::runtime::Handle;
 
 use crate::integration_tests::test_utils::ColumnarValueExt;
 
@@ -15,9 +16,14 @@ async fn test_add_one() {
     let component = WasmComponentPrecompiled::new(datafusion_udf_wasm_bundle::BIN_EXAMPLE.into())
         .await
         .unwrap();
-    let mut udfs = WasmScalarUdf::new(&component, &Default::default(), "".to_owned())
-        .await
-        .unwrap();
+    let mut udfs = WasmScalarUdf::new(
+        &component,
+        &Default::default(),
+        Handle::current(),
+        "".to_owned(),
+    )
+    .await
+    .unwrap();
     assert_eq!(udfs.len(), 1);
     let udf = udfs.pop().unwrap();
 

query/src/lib.rs

@@ -10,6 +10,7 @@ use sqlparser::ast::{CreateFunctionBody, Expr, Statement as SqlStatement, Value}
 use sqlparser::dialect::dialect_from_str;
 
 use datafusion_udf_wasm_host::{WasmComponentPrecompiled, WasmPermissions, WasmScalarUdf};
+use tokio::runtime::Handle;
 
 /// A [ParsedQuery] contains the extracted UDFs and SQL query string
 #[derive(Debug)]
@@ -48,6 +49,7 @@ impl<'a> UdfQueryParser<'a> {
         &self,
         udf_query: &str,
         permissions: &WasmPermissions,
+        io_rt: Handle,
         task_ctx: &TaskContext,
     ) -> DataFusionResult<ParsedQuery> {
         let (code, sql) = self.parse_inner(udf_query, task_ctx)?;
@@ -62,7 +64,7 @@ impl<'a> UdfQueryParser<'a> {
             })?;
 
             for code in blocks {
-                udfs.extend(WasmScalarUdf::new(component, permissions, code).await?);
+                udfs.extend(WasmScalarUdf::new(component, permissions, io_rt.clone(), code).await?);
             }
         }
 

query/tests/integration.rs

@@ -14,6 +14,7 @@ use datafusion::{
 use datafusion_common::{Result as DataFusionResult, test_util::batches_to_string};
 use datafusion_udf_wasm_host::WasmPermissions;
 use datafusion_udf_wasm_query::{ParsedQuery, UdfQueryParser};
+use tokio::runtime::Handle;
 
 mod integration_tests;
 
@@ -54,7 +55,12 @@ SELECT add_one(1);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([("python".to_string(), component)]));
     let parsed_query = parser
-        .parse(query, &WasmPermissions::new(), ctx.task_ctx().as_ref())
+        .parse(
+            query,
+            &WasmPermissions::new(),
+            Handle::current(),
+            ctx.task_ctx().as_ref(),
+        )
         .await
         .unwrap();
 
@@ -98,7 +104,12 @@ SELECT add_one(1), multiply_two(3);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([("python".to_string(), component)]));
     let parsed_query = parser
-        .parse(query, &WasmPermissions::new(), ctx.task_ctx().as_ref())
+        .parse(
+            query,
+            &WasmPermissions::new(),
+            Handle::current(),
+            ctx.task_ctx().as_ref(),
+        )
         .await
         .unwrap();
 
@@ -138,7 +149,12 @@ SELECT add_one(1), multiply_two(3);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([("python".to_string(), component)]));
     let parsed_query = parser
-        .parse(query, &WasmPermissions::new(), ctx.task_ctx().as_ref())
+        .parse(
+            query,
+            &WasmPermissions::new(),
+            Handle::current(),
+            ctx.task_ctx().as_ref(),
+        )
         .await
         .unwrap();
 
@@ -172,7 +188,12 @@ SELECT add_one(1)
 
     let parser = UdfQueryParser::new(HashMap::from_iter([("python".to_string(), component)]));
     let parsed_query = parser
-        .parse(query, &WasmPermissions::new(), ctx.task_ctx().as_ref())
+        .parse(
+            query,
+            &WasmPermissions::new(),
+            Handle::current(),
+            ctx.task_ctx().as_ref(),
+        )
         .await
         .unwrap();
 
@@ -201,7 +222,12 @@ EXPLAIN SELECT add_one(1);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([("python".to_string(), component)]));
     let parsed_query = parser
-        .parse(query, &WasmPermissions::new(), ctx.task_ctx().as_ref())
+        .parse(
+            query,
+            &WasmPermissions::new(),
+            Handle::current(),
+            ctx.task_ctx().as_ref(),
+        )
         .await
         .unwrap();