main
_ 2025-02-23 21:11:08 -06:00
parent 7119603145
commit 2606946240
1 changed files with 83 additions and 65 deletions

View File

@ -14,9 +14,8 @@ impl Default for Args {
pub(crate) struct App { pub(crate) struct App {
client_streams: BTreeMap<Id, Framed<TcpStream, LengthDelimitedCodec>>, client_streams: BTreeMap<Id, Framed<TcpStream, LengthDelimitedCodec>>,
clients: BTreeMap<Id, Client>,
listener: TcpListener, listener: TcpListener,
next_client_id: Id, server: Server,
} }
impl App { impl App {
@ -25,75 +24,81 @@ impl App {
Ok(Self { Ok(Self {
client_streams: Default::default(), client_streams: Default::default(),
clients: Default::default(),
listener, listener,
next_client_id: 1, server: Default::default(),
}) })
} }
pub(crate) fn poll_run(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> { pub(crate) fn poll_run(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
let previous_client_count = self.clients.len(); match self.step(cx) {
Ok(()) => Poll::Pending,
Err(err) => Poll::Ready(Err(err)),
}
}
let mut new_messages = vec![]; fn step(&mut self, cx: &mut Context<'_>) -> Result<()> {
let mut clients_to_delete = vec![]; let mut clients_to_remove = vec![];
for (id, client) in &mut self.clients {
let id = *id; // Pump the network streams
let stream = self for (id, stream) in &mut self.client_streams {
.client_streams
.get_mut(&id)
.context("Logic error: Client has no stream")?;
let mut stream = pin!(stream); let mut stream = pin!(stream);
match <_ as futures_sink::Sink<Bytes>>::poll_ready(stream.as_mut(), cx) { // Try to send from the client's outbox
Poll::Pending => {} if let Poll::Ready(result) =
Poll::Ready(Err(err)) => { <_ as futures_sink::Sink<Bytes>>::poll_ready(stream.as_mut(), cx)
return Poll::Ready( {
Err(err).context("Can't check network write half for readiness"), result.context("Can't check network write half for readiness")?;
); let client = self
} .server
Poll::Ready(Ok(())) => { .clients
if let Some(frame) = client.poll_send() { .get_mut(id)
if let Err(err) = stream.as_mut().start_send(frame) { .context("Logic error: Stream has no associated client")?;
return Poll::Ready(Err(err).context("start_send")); if let Some(frame) = client.poll_send() {
} stream.as_mut().start_send(frame).context("start_send")?;
match <_ as futures_sink::Sink<Bytes>>::poll_flush(stream.as_mut(), cx) { match <_ as futures_sink::Sink<Bytes>>::poll_flush(stream.as_mut(), cx) {
Poll::Pending => {} Poll::Pending => {}
Poll::Ready(Err(err)) => { Poll::Ready(result) => result.context("poll_flush")?,
return Poll::Ready(Err(err).context("poll_flush"));
}
Poll::Ready(Ok(())) => {}
}
tracing::debug!("Started send");
} }
tracing::debug!("Started send");
} }
} }
if let Some(ChatLine { id, line }) = client.poll_inbox() { // Try to read into the client's inbox
cx.waker().wake_by_ref();
new_messages.push(ToClient::ChatLine { id, line });
}
match stream.as_mut().poll_next(cx) { match stream.as_mut().poll_next(cx) {
Poll::Pending => {} Poll::Pending => {}
Poll::Ready(None) => clients_to_delete.push(id), Poll::Ready(None) => clients_to_remove.push(*id),
Poll::Ready(Some(Err(err))) => { Poll::Ready(Some(result)) => {
tracing::error!(?err, "stream.poll_next error"); let frame = result.context("stream.poll_next")?;
}
Poll::Ready(Some(Ok(frame))) => {
cx.waker().wake_by_ref(); cx.waker().wake_by_ref();
if let Err(err) = client.handle_frame(frame.into()) { let client = self
return Poll::Ready(Err(err).context("client.handle_frame")); .server
} .clients
.get_mut(id)
.context("Logic error: Stream has no associated client")?;
client
.handle_frame(frame.into())
.context("client.handle_frame")?
} }
} }
} }
for id in clients_to_delete { // Close out disconnected clients
for id in clients_to_remove {
tracing::info!(?id, "Closing client"); tracing::info!(?id, "Closing client");
self.clients.remove(&id); self.client_streams.remove(&id);
self.server.clients.remove(&id);
} }
for client in &mut self.clients.values_mut() { // Broadcast chat lines across all clients
let mut new_messages = vec![];
for client in &mut self.server.clients.values_mut() {
if let Some(ChatLine { id, line }) = client.poll_inbox() {
new_messages.push(ToClient::ChatLine { id, line });
}
}
for client in &mut self.server.clients.values_mut() {
for msg in &new_messages { for msg in &new_messages {
client.handle_outgoing(msg)?; client.handle_outgoing(msg)?;
} }
@ -101,29 +106,42 @@ impl App {
match self.listener.poll_accept(cx) { match self.listener.poll_accept(cx) {
Poll::Pending => {} Poll::Pending => {}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err.into())), Poll::Ready(result) => {
Poll::Ready(Ok((stream, _addr))) => { let (stream, _addr) = result.context("listener.poll_accept")?;
cx.waker().wake_by_ref(); cx.waker().wake_by_ref();
let id = self.next_client_id;
self.next_client_id += 1;
let client = Client {
id,
inbox: Default::default(),
outbox: Default::default(),
sequence: 0,
};
let stream = Framed::new(stream, LengthDelimitedCodec::new()); let stream = Framed::new(stream, LengthDelimitedCodec::new());
tracing::info!(id = client.id, "Accepted client"); let id = self.server.handle_new_client();
self.clients.insert(id, client);
self.client_streams.insert(id, stream); self.client_streams.insert(id, stream);
} }
} }
if self.clients.len() != previous_client_count { Ok(())
tracing::info!(client_count = self.clients.len()); }
} }
Poll::Pending #[derive(Default)]
struct Server {
clients: BTreeMap<Id, Client>,
next_client_id: Id,
}
impl Server {
fn handle_new_client(&mut self) -> Id {
let id = self.next_client_id;
self.next_client_id += 1;
let client = Client {
id,
inbox: Default::default(),
outbox: Default::default(),
sequence: 0,
};
tracing::info!(
id = client.id,
total = self.clients.len(),
"Accepted client"
);
self.clients.insert(id, client);
id
} }
} }