add more stream/future/resource tests

This addresses some gaps in the test coverage and fixes the issued uncovered
(e.g. not handling drops of exported resources and not restoring payload
ownership after a failed `future.write`).

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

bundled/componentize_py_async_support/__init__.py

@@ -26,6 +26,11 @@ class _FutureState:
     handles: list[asyncio.Handle]
     pending_count: int
 
+class _ReturnCode:
+    COMPLETED = 0
+    DROPPED = 1
+    CANCELLED = 2
+
 class _CallbackCode:
     EXIT = 0
     YIELD = 1

bundled/componentize_py_async_support/futures.py

@@ -4,6 +4,7 @@
 
 from typing import TypeVar, Generic, cast, Self, Any, Callable
 from types import TracebackType
+from componentize_py_async_support import _ReturnCode
 
 T = TypeVar('T')
 
@@ -13,7 +14,7 @@ def __init__(self, type_: int, handle: int):
         self.handle: int | None = handle
         self.finalizer = weakref.finalize(self, componentize_py_runtime.future_drop_readable, type_, handle)
 
-    async def read(self) -> T | None:
+    async def read(self) -> T:
         self.finalizer.detach()
         handle = self.handle
         self.handle = None
@@ -57,15 +58,25 @@ def __init__(self, type_: int, handle: int, default: Callable[[], T]):
         self.default = default
         self.finalizer = weakref.finalize(self, write_default, type_, handle, default)
 
-    async def write(self, value: T) -> None:
+    async def write(self, value: T) -> bool:
         self.finalizer.detach()
         handle = self.handle
         self.handle = None
         if handle is not None:
-            await componentize_py_async_support.await_result(
+            code, _ = await componentize_py_async_support.await_result(
                 componentize_py_runtime.future_write(self.type_, handle, value)
             )
             componentize_py_runtime.future_drop_writable(self.type_, handle)
+            match code:
+                case _ReturnCode.COMPLETED:
+                    return True
+                case _ReturnCode.DROPPED:
+                    return False
+                case _ReturnCode.CANCELLED:
+                    # todo
+                    raise NotImplementedError
+                case _:
+                    raise AssertionError
         else:
             raise AssertionError
 

bundled/componentize_py_async_support/streams.py

@@ -4,11 +4,7 @@
 
 from typing import TypeVar, Generic, Self, cast
 from types import TracebackType
-
-class _ReturnCode:
-    COMPLETED = 0
-    DROPPED = 1
-    CANCELLED = 2
+from componentize_py_async_support import _ReturnCode
 
 class ByteStreamReader:
     def __init__(self, type_: int, handle: int):

bundled/componentize_py_runtime.pyi

@@ -19,7 +19,7 @@ def promise_get_result(event: int, promise: int) -> Any: ...
 
 def future_read(ty: int, future: int) -> Result[Any, tuple[int, int]]: ...
 
-def future_write(ty: int, future: int, value: Any) -> Result[None, tuple[int, int]]: ...
+def future_write(ty: int, future: int, value: Any) -> Result[tuple[int, int], tuple[int, int]]: ...
 
 def future_drop_readable(ty: int, future: int) -> None: ...
 

runtime/src/lib.rs

@@ -310,6 +310,7 @@ mod async_ {
         },
         FutureWrite {
             _call: MyCall<'static>,
+            resources: Option<Vec<EmptyResource>>,
         },
     }
 
@@ -391,9 +392,16 @@ mod async_ {
                 }
                 call.stack.pop().unwrap_or(py.None()).into_bound(py)
             }
-            Promise::FutureWrite { .. } => {
+            Promise::FutureWrite { ref resources, .. } => {
                 let count = event >> 4;
                 let code = event & 0xF;
+
+                if let (RETURN_CODE_DROPPED, Some(resources)) = (code, resources) {
+                    for resource in resources {
+                        resource.restore(py)
+                    }
+                }
+
                 PyTuple::new(py, [code, count]).unwrap().into_any()
             }
         }
@@ -804,11 +812,16 @@ mod async_ {
         } else {
             unsafe { call.defer_deallocate(buffer, layout) };
 
+            call.resources = Some(Vec::new());
             let code = unsafe {
                 ty.lower(&mut call, buffer);
 
                 ty.write()(handle, buffer.cast())
             };
+            let resources = call
+                .resources
+                .take()
+                .and_then(|v| if v.is_empty() { None } else { Some(v) });
 
             Ok(if code == RETURN_CODE_BLOCKED {
                 ERR_CONSTRUCTOR
@@ -822,6 +835,7 @@ mod async_ {
                                 usize::try_from(handle).unwrap(),
                                 Box::into_raw(Box::new(async_::Promise::FutureWrite {
                                     _call: call,
+                                    resources,
                                 })) as usize,
                             ],
                         )
@@ -831,6 +845,13 @@ mod async_ {
             } else {
                 let count = code >> 4;
                 let code = code & 0xF;
+
+                if let (RETURN_CODE_DROPPED, Some(resources)) = (code, &resources) {
+                    for resource in resources {
+                        resource.restore(py)
+                    }
+                }
+
                 OK_CONSTRUCTOR
                     .get()
                     .unwrap()
@@ -1330,8 +1351,10 @@ impl Interpreter for MyInterpreter {
     }
 
     fn resource_dtor(ty: wit::Resource, handle: usize) {
+        // We don't currently include a `drop` function as part of the abstract
+        // base class we generate for an exported resource, so there's nothing
+        // to do here.  If/when that changes, we'll want to call `drop` here.
         _ = (ty, handle);
-        todo!()
     }
 }
 

src/test/python_source/app.py

@@ -5,12 +5,14 @@
 import resource_alias1
 import resource_borrow_in_record
 import componentize_py_async_support
+import streams_and_futures as my_streams_and_futures
 
 from componentize_py_types import Result, Ok, Err
 from tests import exports, imports
 from tests.imports import resource_borrow_import
 from tests.imports import simple_import_and_export
 from tests.imports import simple_async_import_and_export
+from tests.imports import host_thing_interface
 from tests.exports import resource_alias2
 from tests.exports import streams_and_futures
 from typing import Tuple, List, Optional
@@ -147,6 +149,42 @@ async def pipe_things(rx: StreamReader[streams_and_futures.Thing], tx: StreamWri
     # Write the things all at once.  The host will read them only one at a time,
     # forcing us to re-take ownership of any unwritten items between writes.
     await tx.write_all(things)
+
+async def pipe_host_things(rx: StreamReader[host_thing_interface.HostThing], tx: StreamWriter[host_thing_interface.HostThing]):
+    # Read the things one at a time, forcing the host to re-take ownership of
+    # any unwritten items between writes.
+    things = []
+    while not rx.writer_dropped:
+        things += await rx.read(1)
+
+    # Write the things all at once.  The host will read them only one at a time,
+    # forcing us to re-take ownership of any unwritten items between writes.
+    await tx.write_all(things)
+
+async def write_thing(thing: my_streams_and_futures.Thing,
+                      tx1: FutureWriter[streams_and_futures.Thing],
+                      tx2: FutureWriter[streams_and_futures.Thing]):
+    # The host will drop the first reader without reading, which should give us
+    # back ownership of `thing`.
+    wrote = await tx1.write(thing)
+    assert not wrote
+    # The host will read from the second reader, though.
+    wrote = await tx2.write(thing)
+    assert wrote
+
+async def write_host_thing(thing: host_thing_interface.HostThing,
+                      tx1: FutureWriter[host_thing_interface.HostThing],
+                      tx2: FutureWriter[host_thing_interface.HostThing]):
+    # The host will drop the first reader without reading, which should give us
+    # back ownership of `thing`.
+    wrote = await tx1.write(thing)
+    assert not wrote
+    # The host will read from the second reader, though.
+    wrote = await tx2.write(thing)
+    assert wrote
+
+def unreachable() -> str:
+    raise AssertionError
 
 class StreamsAndFutures(exports.StreamsAndFutures):
     async def echo_stream_u8(self, stream: ByteStreamReader) -> ByteStreamReader:
@@ -155,8 +193,6 @@ async def echo_stream_u8(self, stream: ByteStreamReader) -> ByteStreamReader:
         return rx
 
     async def echo_future_string(self, future: FutureReader[str]) -> FutureReader[str]:
-        def unreachable() -> str:
-            raise AssertionError
         tx, rx = tests.string_future(unreachable)
         componentize_py_async_support.spawn(pipe_strings(future, tx))
         return rx
@@ -166,6 +202,23 @@ async def short_reads(self, stream: StreamReader[streams_and_futures.Thing]) ->
         componentize_py_async_support.spawn(pipe_things(stream, tx))
         return rx
 
+    async def short_reads_host(self, stream: StreamReader[host_thing_interface.HostThing]) -> StreamReader[host_thing_interface.HostThing]:
+        tx, rx = tests.host_thing_interface_host_thing_stream()
+        componentize_py_async_support.spawn(pipe_host_things(stream, tx))
+        return rx
+
+    async def dropped_future_reader(self, value: str) -> tuple[FutureReader[streams_and_futures.Thing], FutureReader[streams_and_futures.Thing]]:
+        tx1, rx1 = tests.streams_and_futures_thing_future(unreachable)
+        tx2, rx2 = tests.streams_and_futures_thing_future(unreachable)
+        componentize_py_async_support.spawn(write_thing(my_streams_and_futures.Thing(value), tx1, tx2))
+        return (rx1, rx2)
+
+    async def dropped_future_reader_host(self, value: str) -> tuple[FutureReader[host_thing_interface.HostThing], FutureReader[host_thing_interface.HostThing]]:
+        tx1, rx1 = tests.host_thing_interface_host_thing_future(unreachable)
+        tx2, rx2 = tests.host_thing_interface_host_thing_future(unreachable)
+        componentize_py_async_support.spawn(write_host_thing(host_thing_interface.HostThing(value), tx1, tx2))
+        return (rx1, rx2)
+
 class Tests(tests.Tests):
     def test_resource_borrow_import(self, v: int) -> int:
         return resource_borrow_import.foo(resource_borrow_import.Thing(v + 1)) + 4