Fix wake_up_after being ignored in the wasm-node streams (#1023)

* Fix `wake_up_after` being ignored in the wasm-node streams * PR link
smol-dot · Aug 11, 2023 · 48864e7 · 48864e7
1 parent 0d726b8
commit 48864e7
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 1 deletion.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/wasm-node/CHANGELOG.md b/wasm-node/CHANGELOG.md
@@ -12,6 +12,7 @@
 - The block announces substream handshake of a parachain peer-to-peer network now properly contains the block that smoldot thinks is the best. The genesis block was previously always reported. ([#1012](https://github.com/smol-dot/smoldot/pull/1012))
 - Fix panic when removing a chain while a networking connection is being opened. ([#1011](https://github.com/smol-dot/smoldot/pull/1011))
 - Fix epoch start slot calculation when epochs have been skipped. ([#1015](https://github.com/smol-dot/smoldot/pull/1015))
+- Fix timeouts not working properly in the networking. ([#1023](https://github.com/smol-dot/smoldot/pull/1023))
 
 ## 1.0.15 - 2023-08-08
 

diff --git a/wasm-node/rust/Cargo.toml b/wasm-node/rust/Cargo.toml
@@ -18,6 +18,7 @@ async-task = { version = "4.4.0", default-features = false }
 derive_more = "0.99.17"
 event-listener = { version = "2.5.3" }
 fnv = { version = "1.0.7", default-features = false }
+futures-lite = { version = "1.13.0", default-features = false, features = ["alloc"] }
 futures-util = { version = "0.3.27", default-features = false }
 hashbrown = { version = "0.14.0", default-features = false }
 log = { version = "0.4.18", features = ["std"] }

diff --git a/wasm-node/rust/src/platform.rs b/wasm-node/rust/src/platform.rs
@@ -17,6 +17,8 @@
 
 use crate::{bindings, timers::Delay};
 
+use futures_lite::future::FutureExt as _;
+
 use smoldot_light::platform::{read_write, ConnectError, SubstreamDirection};
 
 use core::{future, iter, mem, ops, pin, str, task, time::Duration};
@@ -293,6 +295,7 @@ impl smoldot_light::platform::PlatformRef for PlatformRef {
                         writable_bytes: 0,
                         write_closable,
                         write_closed: false,
+                        when_wake_up: None,
                     })
                 }
                 ConnectionInner::Reset {
@@ -458,6 +461,7 @@ impl smoldot_light::platform::PlatformRef for PlatformRef {
                     writable_bytes: usize::try_from(initial_writable_bytes).unwrap(),
                     write_closable: false, // Note: this is currently hardcoded for WebRTC.
                     write_closed: false,
+                    when_wake_up: None,
                 },
                 direction,
             ))
@@ -540,7 +544,16 @@ impl smoldot_light::platform::PlatformRef for PlatformRef {
                     stream_inner.something_happened.listen()
                 };
 
-                listener.await
+                listener
+                    .or(async {
+                        if let Some(when_wake_up) = stream.when_wake_up.as_mut() {
+                            when_wake_up.await;
+                            stream.when_wake_up = None;
+                        } else {
+                            future::pending().await
+                        }
+                    })
+                    .await
             }
         })
     }
@@ -603,6 +616,11 @@ impl<'a> Drop for ReadWriteAccess<'a> {
             .get_mut(&(self.stream.connection_id, self.stream.stream_id))
             .unwrap();
 
+        self.stream.when_wake_up = self
+            .read_write
+            .wake_up_after
+            .map(|when| Delay::new_at(when));
+
         self.stream.read_buffer = mem::take(&mut self.read_write.incoming_buffer);
 
         for buffer in self.read_write.write_buffers.drain(..) {
@@ -648,6 +666,8 @@ pub(crate) struct StreamWrapper {
     writable_bytes: usize,
     write_closable: bool,
     write_closed: bool,
+    /// The stream should wake up after this delay.
+    when_wake_up: Option<Delay>,
 }
 
 impl Drop for StreamWrapper {