main
use std::collections::VecDeque;
use tokio::sync::{mpsc, oneshot};
use crate::Frame;
use crate::engine::{self, Packet as EnginePacket};
use super::super::{EmitPacket, SocketPayload};
const MAX_POLL_QUEUE: usize = 1024;
pub(crate) enum Input {
Packets(Vec<EnginePacket>),
Poll(oneshot::Sender<Vec<EnginePacket>>),
Upgrade {
tx: mpsc::Sender<Vec<Frame>>,
result: oneshot::Sender<bool>,
},
Emit {
packet: EmitPacket,
ack: Option<(oneshot::Sender<SocketPayload>, tokio::time::Instant)>,
accepted: Option<oneshot::Sender<bool>>,
},
Ack {
namespace: String,
id: u64,
payload: SocketPayload,
},
Close,
}
pub(crate) enum Output {
Polling {
queue: VecDeque<EnginePacket>,
waiter: Option<oneshot::Sender<Vec<EnginePacket>>>,
},
WebSocket {
tx: mpsc::Sender<Vec<Frame>>,
},
}
impl Output {
pub(crate) fn polling() -> Self {
Self::Polling {
queue: VecDeque::new(),
waiter: None,
}
}
pub(crate) fn websocket(tx: mpsc::Sender<Vec<Frame>>) -> Self {
Self::WebSocket { tx }
}
pub(super) async fn send(&mut self, packets: Vec<EnginePacket>) -> bool {
if packets.is_empty() {
return true;
}
match self {
Self::Polling { queue, waiter } => {
if queue.len() + packets.len() > MAX_POLL_QUEUE {
return false;
}
queue.extend(packets);
if let Some(waiter) = waiter.take() {
let packets = queue.drain(..).collect();
if let Err(packets) = waiter.send(packets) {
queue.extend(packets);
}
}
true
}
Self::WebSocket { tx } => {
let frames = engine::encode_frames(&packets);
if frames.is_empty() {
return true;
}
tx.try_send(frames).is_ok()
}
}
}
pub(super) fn poll(&mut self, responder: oneshot::Sender<Vec<EnginePacket>>) -> bool {
match self {
Self::Polling { queue, waiter } => {
if queue.is_empty() {
if let Some(previous) = waiter.replace(responder) {
let _ = previous.send(Vec::new());
}
} else {
let packets = queue.drain(..).collect();
if let Err(packets) = responder.send(packets) {
queue.extend(packets);
}
}
true
}
Self::WebSocket { .. } => {
let _ = responder.send(vec![EnginePacket::Noop]);
true
}
}
}
pub(super) async fn upgrade(&mut self, tx: mpsc::Sender<Vec<Frame>>) -> bool {
let pending = match self {
Self::Polling { queue, waiter } => {
if let Some(waiter) = waiter.take() {
let _ = waiter.send(vec![EnginePacket::Noop]);
}
queue.drain(..).collect::<Vec<_>>()
}
Self::WebSocket { .. } => return false,
};
*self = Self::WebSocket { tx };
self.send(pending).await
}
pub(super) async fn close(&mut self) -> bool {
match self {
Self::Polling { queue, waiter } => {
queue.push_back(EnginePacket::Close);
if let Some(waiter) = waiter.take() {
let packets = queue.drain(..).collect();
if let Err(packets) = waiter.send(packets) {
queue.extend(packets);
}
}
true
}
Self::WebSocket { tx } => tx.try_send(vec![Frame::Text("1".to_string())]).is_ok(),
}
}
}