libasync: add TcpSocket::{close, abort, keep_alive, timeout}()

experiments/src/main.rs

@@ -241,17 +241,27 @@ pub fn main_core0() {
         Ok(())
     }
 
-    TcpStream::listen(TCP_PORT, 2048, 2048, 8, |stream| async {
+    let mut counter = alloc::rc::Rc::new(core::cell::RefCell::new(0));
+    task::spawn(async move {
+        while let stream = TcpStream::accept(TCP_PORT, 2048, 2408).await.unwrap() {
+            let counter = counter.clone();
+            task::spawn(async move {
+                *counter.borrow_mut() += 1;
+                println!("Serving {} connections", *counter.borrow());
                 handle_connection(stream)
                     .await
                     .map_err(|e| println!("Connection: {:?}", e));
+                *counter.borrow_mut() -= 1;
+                println!("Now serving {} connections", *counter.borrow());
+            });
+        }
     });
 
     let mut time = 0u32;
     Sockets::run(&mut iface, || {
         time += 1;
         Instant::from_millis(time)
-    });
+    })
 }
 
 static CORE1_REQ: Mutex<Option<sync_channel::Receiver<usize>>> = Mutex::new(None);

libasync/src/smoltcp/mod.rs

@@ -44,14 +44,14 @@ impl Sockets {
     pub fn run<'b, 'c, 'e, D: for<'d> Device<'d>>(
         iface: &mut EthernetInterface<'b, 'c, 'e, D>,
         mut get_time: impl FnMut() -> Instant,
-    ) {
+    ) -> ! {
         task::block_on(async {
             loop {
                 let instant = get_time();
                 Self::instance().poll(iface, instant);
                 task::r#yield().await;
             }
-        });
+        })
     }
 
     pub(crate) fn instance() -> &'static Self {

libasync/src/smoltcp/tcp_stream.rs

@@ -9,10 +9,12 @@ use core::{
 };
 use alloc::vec::Vec;
 use smoltcp::{
+    Error, Result,
     socket::{
         SocketHandle, SocketRef,
-        TcpSocketBuffer, TcpSocket,
+        TcpSocketBuffer, TcpSocket, TcpState,
     },
+    time::Duration,
 };
 use crate::task;
 use super::Sockets;
@@ -77,47 +79,26 @@ impl TcpStream {
         f(socket_ref)
     }
 
-    /// Spawns `backlog` tasks with listening sockets so that more
-    /// connections can be accepted while some are still
-    /// handshaking. Spawns additional tasks for each connection.
-    pub fn listen<F, R, T>(port: u16, rx_bufsize: usize, tx_bufsize: usize, backlog: usize, f: F)
-    where
-        F: Fn(Self) -> R + Clone + 'static,
-        R: Future<Output = T> + 'static,
-    {
-        for _ in 0..backlog {
-            let f = f.clone();
-            task::spawn(async move {
-                loop {
-                    // Wait for new connection
-                    let stream = TcpStream::accept(port, rx_bufsize, tx_bufsize).await;
-                    // Spawn async task for new connection
-                    task::spawn(f(stream));
-                }
-            });
-        }
-    }
-
     /// Listen for the next incoming connection on a TCP
     /// port. Succeeds on connection attempt.
     ///
     /// Calling this serially in a loop will cause slow/botched
     /// connection attempts stall any more new connections. Use
     /// `listen()` with a backlog instead.
-    pub async fn accept(port: u16, rx_bufsize: usize, tx_bufsize: usize) -> Self {
+    pub async fn accept(port: u16, rx_bufsize: usize, tx_bufsize: usize) -> Result<Self> {
         let stream = Self::new(rx_bufsize, tx_bufsize);
         // Set socket to listen
-        stream.with_socket(|mut s| s.listen(port))
+        stream.with_socket(|mut s| s.listen(port))?;
-            .expect("listen");
         // Wait for a connection
         poll_stream!(&stream, (), |socket| {
-            if socket.is_active() {
+            if socket.state() != TcpState::Listen {
                 Poll::Ready(())
             } else {
                 Poll::Pending
             }
         }).await;
-        stream
+
+        Ok(stream)
     }
 
     /// Probe the receive buffer
@@ -127,7 +108,7 @@ impl TcpStream {
     /// just return `Poll::Pending` if there is not enough data
     /// yet. Likewise, return the amount of bytes consumed from the
     /// buffer in the `Poll::Ready` result.
-    pub async fn recv<F, R>(&self, f: F) -> smoltcp::Result<R>
+    pub async fn recv<F, R>(&self, f: F) -> Result<R>
     where
         F: Fn(&[u8]) -> Poll<(usize, R)>,
     {
@@ -137,10 +118,14 @@ impl TcpStream {
         }
 
         impl<'a, F: Fn(&[u8]) -> Poll<(usize, R)>, R> Future for Recv<'a, F, R> {
-            type Output = smoltcp::Result<R>;
+            type Output = Result<R>;
 
             fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                 let result = self.stream.with_socket(|mut socket| {
+                    if socket_is_handhshaking(&socket) {
+                        return Ok(Poll::Pending);
+                    }
+
                     socket.recv(|buf| {
                         if buf.len() > 0 {
                             match (self.f)(buf) {
@@ -156,10 +141,11 @@ impl TcpStream {
                     })
                 });
                 match result {
-                    Ok(result) => {
+                    Ok(Poll::Pending) => {
-                        if !result.is_ready() {
                         Sockets::register_waker(cx.waker().clone());
+                        Poll::Pending
                     }
+                    Ok(result) => {
                         result
                     }
                     Err(e) =>
@@ -175,12 +161,14 @@ impl TcpStream {
     }
 
     /// Wait until there is any space in the socket's send queue
-    async fn wait_can_send(&self) -> smoltcp::Result<()> {
+    async fn wait_can_send(&self) -> Result<()> {
-        poll_stream!(self, smoltcp::Result<()>, |socket| {
+        poll_stream!(self, Result<()>, |socket| {
-            if !socket.is_active() {
+            if socket_is_handhshaking(&socket) {
-                Poll::Ready(Err(smoltcp::Error::Illegal))
+                Poll::Pending
             } else if socket.can_send() {
                 Poll::Ready(Ok(()))
+            } else if ! socket.may_send() {
+                Poll::Ready(Err(Error::Truncated))
             } else {
                 Poll::Pending
             }
@@ -188,7 +176,7 @@ impl TcpStream {
     }
 
     /// Yields to wait for more buffer space
-    pub async fn send<I: IntoIterator<Item = u8>>(&self, data: I) -> Result<(), smoltcp::Error> {
+    pub async fn send<I: IntoIterator<Item = u8>>(&self, data: I) -> Result<()> {
         let mut data = data.into_iter();
         let mut done = false;
         while !done {
@@ -217,15 +205,51 @@ impl TcpStream {
     /// **Warning:** this may not work as immediately as expected! The
     /// other side may wait until it sends packets to you for
     /// piggybacking the ACKs.
-    pub async fn flush(&self) {
+    pub async fn flush(&self) -> Result<()> {
-        poll_stream!(self, (), |socket| {
+        poll_stream!(self, Result<()>, |socket| {
-            if socket.may_send() && socket.send_queue() > 0 {
+            if socket_is_handhshaking(&socket) {
                 Poll::Pending
+            } else if socket.may_send() && socket.send_queue() > 0 {
+                Poll::Pending
+            } else if socket.may_send() {
+                Poll::Ready(Ok(()))
             } else {
-                Poll::Ready(())
+                Poll::Ready(Err(Error::Truncated))
             }
         }).await
     }
+
+    /// Close the transmit half of the connection
+    pub async fn close(&self) {
+        self.with_socket(|mut socket| socket.close());
+
+        // Yield for one iface.poll() to send the packet
+        task::r#yield().await;
+    }
+
+    /// Destroy the socket, sending the RST
+    pub async fn abort(self) {
+        self.with_socket(|mut socket| socket.abort());
+
+        // Yield for one iface.poll() to send the packet
+        task::r#yield().await;
+    }
+
+    pub fn keep_alive(&self) -> Option<Duration> {
+        self.with_socket(|socket| socket.keep_alive())
+    }
+
+    pub fn set_keep_alive(&mut self, interval: Option<Duration>) {
+        self.with_socket(|mut socket| socket.set_keep_alive(interval));
+    }
+
+    pub fn timeout(&self) -> Option<Duration> {
+        self.with_socket(|socket| socket.timeout())
+    }
+
+    pub fn set_timeout(&mut self, duration: Option<Duration>) {
+        self.with_socket(|mut socket| socket.set_timeout(duration));
+    }
 }
 
 impl Drop for TcpStream {
@@ -236,3 +260,12 @@ impl Drop for TcpStream {
             .remove(self.handle);
     }
 }
+
+fn socket_is_handhshaking(socket: &SocketRef<TcpSocket>) -> bool {
+    match socket.state() {
+        TcpState::SynSent | TcpState::SynReceived =>
+            true,
+        _ =>
+            false,
+    }
+}