cubic_protocol_server/
write.rs

1use std::sync::Arc;
2use tokio::*;
3use tokio::io::AsyncWriteExt;
4use tokio::net::tcp::OwnedWriteHalf;
5use tokio::sync::mpsc::Receiver;
6use crate::connection::Connection;
7use crate::handler::{ConnectionHandler, ReadHandler};
8use crate::server::ProtocolServerDeclare;
9
10pub enum WriteMessage {
11    Close,
12    Bytes(Vec<u8>),
13}
14
15pub(crate) struct WriteStreamQueue {
16    pub write_half: OwnedWriteHalf,
17    pub receiver: Receiver<WriteMessage>,
18}
19
20impl WriteStreamQueue {
21    pub async fn run<
22        H: ReadHandler + Sized + Send + Sync + 'static,
23        C: ConnectionHandler + Sized + Send + Sync + 'static
24    >(mut self, connection: Arc<Connection>, declare: Arc<ProtocolServerDeclare<H, C>>) -> io::Result<()> {
25        while let Some(bytes) = self.receiver.recv().await {
26            match bytes {
27                WriteMessage::Close => {
28                    declare.connection_handler.handle_disconnect(connection);
29                    return self.write_half.shutdown().await;
30                }
31                WriteMessage::Bytes(bytes) =>
32                    self.write_half.write_all(bytes.as_slice()).await?,
33            }
34        }
35        Ok(())
36    }
37}