src/client.rs

@@ -15,7 +15,8 @@ impl Default for Args {
 }
 
 pub(crate) struct App {
-    client: Client,
+    outbox: VecDeque<ToServer>,
+    sequence: u64,
     stream: Framed<TcpStream, LengthDelimitedCodec>,
     timer: Interval,
 }
@@ -27,92 +28,69 @@ impl App {
         let mut timer = tokio::time::interval(Duration::from_millis(args.period_ms.into()));
         timer.set_missed_tick_behavior(MissedTickBehavior::Skip);
         Ok(Self {
-            client: Default::default(),
+            outbox: Default::default(),
+            sequence: 0,
             stream,
             timer,
         })
     }
 
     pub(crate) fn poll_run(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        match self.step(cx) {
-            Ok(()) => Poll::Pending,
-            Err(err) => Poll::Ready(Err(err)),
-        }
-    }
-
-    fn step(&mut self, cx: &mut Context<'_>) -> Result<()> {
-        let mut stream = pin!(&mut self.stream);
-        if let Poll::Ready(result) =
-            <_ as futures_sink::Sink<Bytes>>::poll_ready(stream.as_mut(), cx)
         {
-            result?;
+            let mut stream = pin!(&mut self.stream);
-            if let Some(frame) = self.client.poll_send() {
-                stream
-                    .as_mut()
-                    .start_send(frame)
-                    .context("stream.start_send")?;
-                tracing::debug!("Started send");
-            }
             match <_ as futures_sink::Sink<Bytes>>::poll_flush(stream.as_mut(), cx) {
                 Poll::Pending => {}
-                Poll::Ready(result) => {
+                Poll::Ready(Err(err)) => return Poll::Ready(Err(err).context("poll_flush")),
-                    result.context("poll_flush")?;
+                Poll::Ready(Ok(())) => {}
+            }
+        }
+
+        if !self.outbox.is_empty() {
+            let mut stream = pin!(&mut self.stream);
+            match <_ as futures_sink::Sink<Bytes>>::poll_ready(stream.as_mut(), cx) {
+                Poll::Pending => {}
+                Poll::Ready(Err(err)) => {
+                    return Poll::Ready(
+                        Err(err).context("Can't check network write half for readiness"),
+                    );
+                }
+                Poll::Ready(Ok(())) => {
+                    cx.waker().wake_by_ref();
+                    let frame = rmp_serde::to_vec(&self.outbox.pop_front())?;
+                    if let Err(err) = stream.start_send(Bytes::from(frame)) {
+                        return Poll::Ready(Err(err).context("start_send"));
+                    }
+                    tracing::debug!("Started send");
                 }
             }
         }
 
-        if let Poll::Ready(frame_opt) = stream.as_mut().poll_next(cx) {
+        let stream = pin!(&mut self.stream);
-            let frame = frame_opt.context("Server closed cxn")?;
+        match stream.poll_next(cx) {
+            Poll::Pending => {}
+            Poll::Ready(None) => return Poll::Ready(Err(anyhow!("Server closed cxn"))),
+            Poll::Ready(Some(frame)) => {
                 cx.waker().wake_by_ref();
                 let frame = frame.context("network framing decode")?;
-            self.client
+                match rmp_serde::from_slice(&frame)? {
-                .handle_frame(frame.into())
+                    ToClient::ChatLine { id, line } => {
-                .context("client.handle_frame")?;
+                        tracing::info!(?id, ?line);
+                    }
+                }
+            }
         }
 
         if self.timer.poll_tick(cx).is_ready() {
             cx.waker().wake_by_ref();
-            self.client
+            if self.outbox.is_empty() {
-                .handle_timeout()
+                self.outbox.push_back(ToServer::ChatLine {
-                .context("client.handle_timeout")?;
-        }
-
-        Ok(())
-    }
-}
-
-#[derive(Default)]
-struct Client {
-    outbox: VecDeque<Bytes>,
-    sequence: u64,
-}
-
-impl Client {
-    fn handle_frame(&self, frame: Bytes) -> Result<()> {
-        match rmp_serde::from_slice(&frame)? {
-            ToClient::ChatLine { id, line } => tracing::info!(?id, ?line),
-            ToClient::ClientConnected { id } => tracing::info!(?id, "Connected"),
-            ToClient::ClientDisconnected { id } => tracing::info!(?id, "Disconnected"),
-        }
-        Ok(())
-    }
-
-    fn handle_timeout(&mut self) -> Result<()> {
-        if !self.outbox.is_empty() {
-            bail!("Dropped message, outbox is full");
-        }
-
-        let msg = ToServer::ChatLine {
                     line: "There was a time, in the era of great chaos, when the Earth and the Moon were at war with each other. A daredevil from the Moon piloted a bizarre aircraft. It was feared, and because of its shape, called EINHANDER.".to_string(),
                     sequence: self.sequence,
-        };
+                });
-        let frame = rmp_serde::to_vec(&msg)?;
-        self.outbox.push_back(frame.into());
                 self.sequence += 1;
-        Ok(())
+            }
         }
 
-    fn poll_send(&mut self) -> Option<Bytes> {
+        Poll::Pending
-        self.outbox.pop_front()
     }
 }

src/messages.rs

@@ -1,11 +1,8 @@
-use crate::prelude::*;
 use serde::{Deserialize, Serialize};
 
 #[derive(Deserialize, Serialize)]
 pub(crate) enum ToClient {
-    ChatLine { id: Id, line: String },
+    ChatLine { id: u64, line: String },
-    ClientConnected { id: Id },
-    ClientDisconnected { id: Id },
 }
 
 #[derive(Deserialize, Serialize)]

src/prelude.rs

@@ -1,13 +1,14 @@
 pub(crate) use crate::messages::{ToClient, ToServer};
-pub use anyhow::{Context as _, Result, bail};
+pub use anyhow::{Context as _, Result, anyhow, bail};
 pub use bytes::Bytes;
 pub use futures_core::stream::Stream;
 pub use futures_sink::Sink;
 pub use std::{
-    collections::{BTreeMap, VecDeque},
+    collections::VecDeque,
     future::poll_fn,
     ops::ControlFlow,
     pin::pin,
+    rc::Rc,
     str::FromStr,
     task::{Context, Poll},
     time::Duration,
@@ -19,5 +20,3 @@ pub use tokio::{
 };
 // Don't use BytesCodec, it is _nonsense_
 pub use tokio_util::codec::{Framed, LengthDelimitedCodec};
-
-pub type Id = u64;

src/server.rs

@@ -11,9 +11,9 @@ impl Default for Args {
 }
 
 pub(crate) struct App {
-    client_streams: BTreeMap<Id, Framed<TcpStream, LengthDelimitedCodec>>,
+    clients: Vec<Client>,
     listener: TcpListener,
-    server: Server,
+    next_client_id: u64,
 }
 
 impl App {
@@ -21,170 +21,147 @@ impl App {
         let listener = TcpListener::bind(("0.0.0.0", args.port)).await?;
 
         Ok(Self {
-            client_streams: Default::default(),
+            clients: Default::default(),
             listener,
-            server: Default::default(),
+            next_client_id: 1,
         })
     }
 
     pub(crate) fn poll_run(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        match self.step(cx) {
+        let previous_client_count = self.clients.len();
-            Ok(()) => Poll::Pending,
+
-            Err(err) => Poll::Ready(Err(err)),
+        let mut new_messages = vec![];
+        for client in &mut self.clients {
+            {
+                let mut stream = pin!(&mut client.stream);
+                match <_ as futures_sink::Sink<Bytes>>::poll_flush(stream.as_mut(), cx) {
+                    Poll::Pending => {}
+                    Poll::Ready(Err(err)) => return Poll::Ready(Err(err).context("poll_flush")),
+                    Poll::Ready(Ok(())) => {}
                 }
             }
 
-    fn step(&mut self, cx: &mut Context<'_>) -> Result<()> {
+            if !client.outbox.is_empty() {
-        let mut clients_to_remove = vec![];
+                let mut stream = pin!(&mut client.stream);
-
+                match <_ as futures_sink::Sink<Bytes>>::poll_ready(stream.as_mut(), cx) {
-        // Pump the network streams
+                    Poll::Pending => {}
-        for (id, stream) in &mut self.client_streams {
+                    Poll::Ready(Err(err)) => {
-            let mut stream = pin!(stream);
+                        return Poll::Ready(
-
+                            Err(err).context("Can't check network write half for readiness"),
-            // Try to send from the client's outbox
+                        );
-            if let Poll::Ready(result) =
+                    }
-                <_ as futures_sink::Sink<Bytes>>::poll_ready(stream.as_mut(), cx)
+                    Poll::Ready(Ok(())) => {
-            {
+                        cx.waker().wake_by_ref();
-                result.context("Can't check network write half for readiness")?;
+                        let Some(msg) = client.outbox.pop_front() else {
-                if let Some(frame) = self.server.poll_send(*id)? {
+                            return Poll::Ready(Err(anyhow!(
-                    stream.as_mut().start_send(frame).context("start_send")?;
+                                "Can't pop from outbox even though we just checked it was not empty"
-                    if let Poll::Ready(result) =
+                            )));
-                        <_ as futures_sink::Sink<Bytes>>::poll_flush(stream.as_mut(), cx)
+                        };
-                    {
+                        let frame = rmp_serde::to_vec(msg.as_ref())?;
-                        result.context("poll_flush")?;
+                        if let Err(err) = stream.start_send(Bytes::from(frame)) {
+                            return Poll::Ready(Err(err).context("start_send"));
                         }
                         tracing::debug!("Started send");
                     }
                 }
+            }
 
-            // Try to read data in
+            match client.poll_next(cx) {
-            match stream.as_mut().poll_next(cx) {
                 Poll::Pending => {}
-                Poll::Ready(None) => clients_to_remove.push(*id),
+                Poll::Ready(Err(err)) => {
-                Poll::Ready(Some(result)) => {
+                    tracing::error!(?err, "client.poll_next error");
-                    let frame = result.context("stream.poll_next")?;
+                }
-                    cx.waker().wake_by_ref();
+                Poll::Ready(Ok(None)) => client.closed = true,
-                    self.server.handle_client_frame(*id, frame.into())?;
+                Poll::Ready(Ok(Some(ToServer::ChatLine { line, sequence }))) => {
+                    tracing::debug!(id = client.id, ?line);
+                    if sequence != client.sequence {
+                        tracing::error!(
+                            expected = client.sequence,
+                            actual = sequence,
+                            "Sequence mismatch"
+                        );
+                        return Poll::Ready(Err(anyhow!("Sequence mismatch")));
+                    }
+                    tracing::info!(?sequence, id = client.id);
+                    client.sequence += 1;
+                    new_messages.push(Rc::new(ToClient::ChatLine {
+                        id: client.id,
+                        line,
+                    }));
                 }
             }
         }
 
-        // Close out disconnected clients
+        self.clients.retain(|client| {
-        for id in clients_to_remove {
+            if client.closed {
-            cx.waker().wake_by_ref();
+                tracing::info!(id = client.id, "Closing client");
-            tracing::info!(?id, "Closing client");
+            }
-            self.client_streams.remove(&id);
+            !client.closed
-            self.server.handle_client_disconnected(id)?;
+        });
+
+        for client in &mut self.clients {
+            for msg in &new_messages {
+                client.outbox.push_back(Rc::clone(msg));
+            }
         }
 
-        if let Poll::Ready(result) = self.listener.poll_accept(cx) {
+        match self.listener.poll_accept(cx) {
-            let (stream, _addr) = result.context("listener.poll_accept")?;
+            Poll::Pending => {}
+            Poll::Ready(Err(err)) => return Poll::Ready(Err(err.into())),
+            Poll::Ready(Ok((stream, _addr))) => {
                 cx.waker().wake_by_ref();
                 let stream = Framed::new(stream, LengthDelimitedCodec::new());
-            let id = self.server.handle_new_client()?;
-            self.client_streams.insert(id, stream);
-        }
-
-        Ok(())
-    }
-}
-
-#[derive(Default)]
-struct Server {
-    clients: BTreeMap<Id, Client>,
-    next_client_id: Id,
-}
-
-impl Server {
-    fn broadcast(&mut self, msg: &ToClient) -> Result<()> {
-        for client in &mut self.clients.values_mut() {
-            client.handle_outgoing(msg)?;
-        }
-        Ok(())
-    }
-
-    fn handle_client_disconnected(&mut self, id: Id) -> Result<()> {
-        self.clients.remove(&id);
-        self.broadcast(&ToClient::ClientDisconnected { id })?;
-        Ok(())
-    }
-
-    fn handle_client_frame(&mut self, id: Id, frame: Bytes) -> Result<()> {
-        let msg = {
-            let client = self
-                .clients
-                .get_mut(&id)
-                .context("Logic error: Stream has no associated client")?;
-            let Some(msg) = client.handle_frame(frame).context("client.handle_frame")? else {
-                return Ok(());
-            };
-            msg
-        };
-        self.broadcast(&msg)?;
-        Ok(())
-    }
-
-    fn handle_new_client(&mut self) -> Result<Id> {
-        let id = self.next_client_id;
-        self.next_client_id += 1;
                 let client = Client {
-            id,
+                    closed: false,
+                    id: self.next_client_id,
                     outbox: Default::default(),
                     sequence: 0,
+                    stream,
                 };
-        tracing::info!(
+                self.next_client_id += 1;
-            id = client.id,
+                tracing::info!(id = client.id, "Accepted client");
-            total = self.clients.len(),
+                self.clients.push(client);
-            "Accepted client"
+            }
-        );
-        self.clients.insert(id, client);
-        self.broadcast(&ToClient::ClientConnected { id })?;
-        Ok(id)
         }
 
-    fn poll_send(&mut self, id: Id) -> Result<Option<Bytes>> {
+        if self.clients.len() != previous_client_count {
-        let client = self
+            tracing::info!(client_count = self.clients.len());
-            .clients
+        }
-            .get_mut(&id)
+
-            .context("Logic error: Stream has no associated client")?;
+        Poll::Pending
-        Ok(client.poll_send())
     }
 }
 
 struct Client {
+    closed: bool,
     id: u64,
-    outbox: VecDeque<Bytes>,
+    outbox: VecDeque<Rc<ToClient>>,
     sequence: u64,
+    stream: Framed<TcpStream, LengthDelimitedCodec>,
 }
 
 impl Client {
-    fn handle_frame(&mut self, frame: Bytes) -> Result<Option<ToClient>> {
+    fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<ToServer>>> {
-        match rmp_serde::from_slice(&frame)? {
+        let stream = pin!(&mut self.stream);
-            ToServer::ChatLine { line, sequence } => {
+        let Poll::Ready(frame_opt) = stream.poll_next(cx) else {
-                if sequence != self.sequence {
+            return Poll::Pending;
-                    tracing::error!(
+        };
-                        expected = self.sequence,
-                        actual = sequence,
-                        "Sequence mismatch"
-                    );
-                    bail!("Sequence mismatch");
-                }
-                self.sequence += 1;
-                Ok(Some(ToClient::ChatLine { id: self.id, line }))
-            }
-        }
-    }
 
-    fn handle_outgoing(&mut self, msg: &ToClient) -> Result<()> {
+        let Some(frame) = frame_opt else {
-        if !self.outbox.is_empty() {
+            return Poll::Ready(Ok(None));
-            bail!("Outbox full");
+        };
-        }
-        let bytes = rmp_serde::to_vec(msg)?.into();
-        self.outbox.push_back(bytes);
-        Ok(())
-    }
 
-    fn poll_send(&mut self) -> Option<Bytes> {
+        cx.waker().wake_by_ref();
-        self.outbox.pop_front()
+        tracing::debug!("Got network data");
+
+        let frame = match frame {
+            Err(err) => return Poll::Ready(Err(err).context("Network framing decode error")),
+            Ok(x) => x,
+        };
+        let msg = match rmp_serde::from_slice(&frame) {
+            Err(err) => return Poll::Ready(Err(err).context("MsgPack decode error")),
+            Ok(x) => x,
+        };
+        Poll::Ready(Ok(Some(msg)))
     }
 }