feat(go): add support for future/stream task cancellation

crates/go/src/lib.rs

@@ -599,6 +599,12 @@ func wasm_{kind}_drop_readable_{snake}(handle int32)
 //go:wasmimport {module} [{kind}-drop-writable-{index}]{func_name}
 func wasm_{kind}_drop_writable_{snake}(handle int32)
 
+//go:wasmimport {module} [{kind}-cancel-read-{index}]{func_name}
+func wasm_{kind}_cancel_read_{snake}(handle int32) uint32
+
+//go:wasmimport {module} [{kind}-cancel-write-{index}]{func_name}
+func wasm_{kind}_cancel_write_{snake}(handle int32) uint32
+
 {lift}
 
 {lower}
@@ -608,8 +614,8 @@ var wasm_{kind}_vtable_{snake} = witTypes.{upper_kind}Vtable[{payload}]{{
 	Align: {align},
 	Read: wasm_{kind}_read_{snake},
 	Write: wasm_{kind}_write_{snake},
-	CancelRead: nil,
-	CancelWrite: nil,
+	CancelRead: wasm_{kind}_cancel_read_{snake},
+	CancelWrite: wasm_{kind}_cancel_write_{snake},
 	DropReadable: wasm_{kind}_drop_readable_{snake},
 	DropWritable: wasm_{kind}_drop_writable_{snake},
 	Lift: {lift_name},

tests/runtime/cancel-future-stream/runner.rs

@@ -0,0 +1,57 @@
+//@ wasmtime-flags = '-Wcomponent-model-async'
+
+include!(env!("BINDINGS"));
+
+use crate::my::test::i::*;
+use futures::task::noop_waker_ref;
+use std::future::Future;
+use std::task::Context;
+
+struct Component;
+
+export!(Component);
+
+impl Guest for Component {
+    // Start each import call and poll it once so it's pending, then drop
+    // it to cancel the subtask.
+    async fn run() {
+        // Cancel a call blocked reading a future. The callee must drop its
+        // reader, so our write fails instead of hanging.
+        let (tx, rx) = wit_future::new(|| unreachable!());
+        let mut import = Box::pin(pending_future(rx));
+        assert!(
+            import
+                .as_mut()
+                .poll(&mut Context::from_waker(noop_waker_ref()))
+                .is_pending()
+        );
+        drop(import);
+        assert!(tx.write(()).await.is_err());
+
+        // Cancel a call blocked reading a stream. The callee must drop its
+        // reader, so our writer sees the stream as dropped instead of hanging.
+        let (mut tx, rx) = wit_stream::new::<()>();
+        let mut import = Box::pin(pending_stream(rx));
+        assert!(
+            import
+                .as_mut()
+                .poll(&mut Context::from_waker(noop_waker_ref()))
+                .is_pending()
+        );
+        drop(import);
+        assert!(tx.write_one(()).await.is_some());
+
+        // Cancel a call blocked writing to a stream. The callee must drop its
+        // writer, so our reader sees the stream as dropped instead of hanging.
+        let mut rx = make_stream();
+        let mut import = Box::pin(pending_stream_write());
+        assert!(
+            import
+                .as_mut()
+                .poll(&mut Context::from_waker(noop_waker_ref()))
+                .is_pending()
+        );
+        drop(import);
+        assert!(rx.next().await.is_none());
+    }
+}

tests/runtime/cancel-future-stream/test.go

@@ -0,0 +1,27 @@
+package export_my_test_i
+
+import (
+	"wit_component/my_test_i"
+
+	witTypes "go.bytecodealliance.org/pkg/wit/types"
+)
+
+var writer *witTypes.StreamWriter[witTypes.Unit]
+
+func PendingFuture(x *witTypes.FutureReader[witTypes.Unit]) {
+	x.Read()
+}
+
+func PendingStream(x *witTypes.StreamReader[witTypes.Unit]) {
+	x.Read(make([]witTypes.Unit, 1))
+}
+
+func MakeStream() *witTypes.StreamReader[witTypes.Unit] {
+	tx, rx := my_test_i.MakeStreamUnit()
+	writer = tx
+	return rx
+}
+
+func PendingStreamWrite() {
+	writer.Write(make([]witTypes.Unit, 1))
+}

tests/runtime/cancel-future-stream/test.rs

@@ -0,0 +1,31 @@
+use std::sync::Mutex;
+use wit_bindgen::{FutureReader, StreamReader, StreamWriter};
+
+include!(env!("BINDINGS"));
+
+struct Component;
+
+export!(Component);
+
+static WRITER: Mutex<Option<StreamWriter<()>>> = Mutex::new(None);
+
+impl crate::exports::my::test::i::Guest for Component {
+    async fn pending_future(x: FutureReader<()>) {
+        x.await
+    }
+
+    async fn pending_stream(mut x: StreamReader<()>) {
+        x.next().await;
+    }
+
+    fn make_stream() -> StreamReader<()> {
+        let (tx, rx) = wit_stream::new();
+        *WRITER.lock().unwrap() = Some(tx);
+        rx
+    }
+
+    async fn pending_stream_write() {
+        let mut tx = WRITER.lock().unwrap().take().unwrap();
+        tx.write_one(()).await;
+    }
+}

tests/runtime/cancel-future-stream/test.wit

@@ -0,0 +1,22 @@
+//@ async = true
+package my:test;
+
+interface i {
+  pending-future: async func(x: future);
+  pending-stream: async func(x: stream);
+
+  // Creates a stream and stashes its writer for `pending-stream-write`.
+  make-stream: func() -> stream;
+  // Writes to the stashed writer; blocks until the reader reads or drops.
+  pending-stream-write: async func();
+}
+
+world test {
+  export i;
+}
+
+world runner {
+  import i;
+
+  export run: async func();
+}