Fix issue when another pollable cancels

Signed-off-by: James Sturtevant <jstur@microsoft.com>

Author: jsturtevant

Cargo.lock

@@ -587,4 +587,4 @@ wasm-mutate =  { git = "https://github.com/jsturtevant/wasm-tools.git", branch="
 wit-parser =  { git = "https://github.com/jsturtevant/wasm-tools.git", branch="unstable-refrence-stable" }
 wit-component =  { git = "https://github.com/jsturtevant/wasm-tools.git", branch="unstable-refrence-stable" }
 wasm-wave =  { git = "https://github.com/jsturtevant/wasm-tools.git", branch="unstable-refrence-stable" }
-wasm-metadata =  { git = "https://github.com/jsturtevant/wasm-tools.git", branch="unstable-refrence-stable" }
+wasm-metadata =  { git = "https://github.com/jsturtevant/wasm-tools.git", branch="unstable-refrence-stable" }

Cargo.toml

@@ -33,3 +33,4 @@ webpki-roots = { workspace = true }
 test-programs-artifacts = { workspace = true }
 wasmtime-wasi = { workspace = true }
 tokio = { workspace = true, features = ["macros"] }
+futures = { workspace = true }

crates/wasi-tls/Cargo.toml

@@ -186,7 +186,7 @@ impl<'a> generated::types::HostClientHandshake for WasiTlsCtx<'a> {
 
         Ok(self
             .table
-            .push(FutureClientStreams(StreamState::Pending(Box::pin(
+            .push(FutureStreams(StreamState::Pending(Box::pin(
                 async move {
                     let connector = tokio_rustls::TlsConnector::from(default_client_config());
                     connector
@@ -206,21 +206,20 @@ impl<'a> generated::types::HostClientHandshake for WasiTlsCtx<'a> {
     }
 }
 
-/// Future TLS connection after the handshake is completed.
-pub struct FutureClientStreams(StreamState<Result<TlsStream<WasiStreams>>>);
+struct FutureStreams<T>(StreamState<Result<T>>);
+
+/// Library specific version of TLS connection after the handshake is completed.
+/// This alias allows it to use with wit-bindgen component generator which won't take generic types
+#[allow(private_interfaces)]
+pub type FutureClientStreams = FutureStreams<TlsStream<WasiStreams>>;
 
 #[async_trait]
-impl Pollable for FutureClientStreams {
+impl<T: Send + 'static> Pollable for FutureStreams<T> {
     async fn ready(&mut self) {
-        match &self.0 {
-            StreamState::Pending(_) => (),
+        match &mut self.0 {
             StreamState::Ready(_) | StreamState::Closed => return,
+            StreamState::Pending(task) => self.0 = StreamState::Ready(task.as_mut().await),
         }
-
-        let StreamState::Pending(future) = mem::replace(&mut self.0, StreamState::Closed) else {
-            unreachable!()
-        };
-        self.0 = StreamState::Ready(future.await);
     }
 }
 
@@ -635,3 +634,41 @@ fn try_lock_for_stream<TlsWriter>(
         .try_lock()
         .map_err(|_| StreamError::trap("concurrent access to resource not supported"))
 }
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::sync::oneshot;
+
+    #[tokio::test]
+    async fn test_future_client_streams_ready_can_be_canceled() {
+        let (tx1, rx1) = oneshot::channel::<()>();
+
+        let mut future_streams = FutureStreams(StreamState::Pending(Box::pin(async move {
+            rx1.await.map_err(|_| anyhow::anyhow!("oneshot canceled"))
+        })));
+
+        let mut fut = future_streams.ready();
+
+        let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
+        assert!(fut.as_mut().poll(&mut cx).is_pending());
+
+        //cancel the readiness check
+        drop(fut);
+
+        match future_streams.0 {
+            StreamState::Closed => panic!("First future should be in Pending/ready state"),
+            _ => (),
+        }
+
+        // make it ready and wait for it to progress
+        tx1.send(()).unwrap();
+        future_streams.ready().await;
+
+        match future_streams.0 {
+            StreamState::Ready(Ok(())) => (),
+            _ => panic!("First future should be in Ready(Err) state"),
+        }
+    }
+}