JKHKQKJV/src/server.rs

use crate::prelude::*;

type Id = u64;

pub(crate) struct Args {
    pub(crate) port: u16,
}

impl Default for Args {
    fn default() -> Self {
        Self { port: 9000 }
    }
}

pub(crate) struct App {
    client_streams: BTreeMap<Id, Framed<TcpStream, LengthDelimitedCodec>>,
    clients: BTreeMap<Id, Client>,
    listener: TcpListener,
    next_client_id: Id,
}

impl App {
    pub(crate) async fn new(args: Args) -> Result<Self> {
        let listener = TcpListener::bind(("0.0.0.0", args.port)).await?;

        Ok(Self {
            client_streams: Default::default(),
            clients: Default::default(),
            listener,
            next_client_id: 1,
        })
    }

    pub(crate) fn poll_run(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
        let previous_client_count = self.clients.len();

        let mut new_messages = vec![];
        let mut clients_to_delete = vec![];
        for (id, client) in &mut self.clients {
            let id = *id;
            let stream = self
                .client_streams
                .get_mut(&id)
                .context("Logic error: Client has no stream")?;
            let mut stream = pin!(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(())) => {
                    if let Some(frame) = client.poll_send() {
                        if let Err(err) = stream.as_mut().start_send(frame) {
                            return Poll::Ready(Err(err).context("start_send"));
                        }
                        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(())) => {}
                        }
                        tracing::debug!("Started send");
                    }
                }
            }

            if let Some(ChatLine { id, line }) = client.poll_inbox() {
                cx.waker().wake_by_ref();
                new_messages.push(ToClient::ChatLine { id, line });
            }

            match stream.as_mut().poll_next(cx) {
                Poll::Pending => {}
                Poll::Ready(None) => clients_to_delete.push(id),
                Poll::Ready(Some(Err(err))) => {
                    tracing::error!(?err, "stream.poll_next error");
                }
                Poll::Ready(Some(Ok(frame))) => {
                    cx.waker().wake_by_ref();
                    if let Err(err) = client.handle_frame(frame.into()) {
                        return Poll::Ready(Err(err).context("client.handle_frame"));
                    }
                }
            }
        }

        for id in clients_to_delete {
            tracing::info!(?id, "Closing client");
            self.clients.remove(&id);
        }

        for client in &mut self.clients.values_mut() {
            for msg in &new_messages {
                client.handle_outgoing(msg)?;
            }
        }

        match self.listener.poll_accept(cx) {
            Poll::Pending => {}
            Poll::Ready(Err(err)) => return Poll::Ready(Err(err.into())),
            Poll::Ready(Ok((stream, _addr))) => {
                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());
                tracing::info!(id = client.id, "Accepted client");
                self.clients.insert(id, client);
                self.client_streams.insert(id, stream);
            }
        }

        if self.clients.len() != previous_client_count {
            tracing::info!(client_count = self.clients.len());
        }

        Poll::Pending
    }
}

struct ChatLine {
    id: Id,
    line: String,
}

struct Client {
    id: u64,
    inbox: VecDeque<ChatLine>,
    outbox: VecDeque<Bytes>,
    sequence: u64,
}

impl Client {
    fn handle_frame(&mut self, frame: Bytes) -> Result<()> {
        match rmp_serde::from_slice(&frame)? {
            ToServer::ChatLine { line, sequence } => {
                if sequence != self.sequence {
                    tracing::error!(
                        expected = self.sequence,
                        actual = sequence,
                        "Sequence mismatch"
                    );
                    bail!("Sequence mismatch");
                }
                self.sequence += 1;
                self.inbox.push_back(ChatLine { id: self.id, line });
            }
        }
        Ok(())
    }

    fn handle_outgoing(&mut self, msg: &ToClient) -> Result<()> {
        if !self.outbox.is_empty() {
            bail!("Outbox full");
        }
        let bytes = rmp_serde::to_vec(msg)?.into();
        self.outbox.push_back(bytes);
        Ok(())
    }

    fn poll_inbox(&mut self) -> Option<ChatLine> {
        self.inbox.pop_front()
    }

    fn poll_send(&mut self) -> Option<Bytes> {
        self.outbox.pop_front()
    }
}