use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
use tokio_tungstenite::tungstenite::protocol::CloseFrame;
use tokio_tungstenite::tungstenite::Message as WsMessage;
pub(crate) enum GatewayEvent {
Dispatch(String, Value),
Reconnecting {
code: u16,
reason: String,
wait: Duration,
},
Fatal(String),
}
const FATAL: [u16; 1] = [4004];
pub(crate) struct Gateway {
url: String,
token: String,
intents: u64,
out: Mutex<Option<mpsc::UnboundedSender<WsMessage>>>,
stop: tokio::sync::watch::Sender<bool>,
}
pub(crate) fn backoff(retry: u32, jitter01: f64) -> Duration {
let base = 1000f64 * 2f64.powi(retry.min(16) as i32);
let ms = base.min(30_000.0) * (0.75 + jitter01 / 2.0);
Duration::from_millis(ms as u64)
}
fn jitter() -> f64 {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.subsec_nanos())
.unwrap_or(0);
(nanos % 1000) as f64 / 1000.0
}
async fn stopped(rx: &mut tokio::sync::watch::Receiver<bool>) {
let _ = rx.wait_for(|s| *s).await;
}
impl Gateway {
pub fn new(url: String, token: String, intents: u64) -> Arc<Self> {
let (stop, _) = tokio::sync::watch::channel(false);
Arc::new(Self {
url,
token,
intents,
out: Mutex::new(None),
stop,
})
}
pub fn send(&self, op: u8, d: Value) -> bool {
let out = self.out.lock().unwrap_or_else(|e| e.into_inner());
match out.as_ref() {
Some(tx) => tx
.send(WsMessage::Text(
json!({ "op": op, "d": d }).to_string().into(),
))
.is_ok(),
None => false,
}
}
pub fn close(&self) {
let _ = self.stop.send(true);
}
pub async fn run(self: Arc<Self>, events: mpsc::UnboundedSender<GatewayEvent>) {
let _ = self.stop.send(false);
let mut stop = self.stop.subscribe();
let mut retry: u32 = 0;
loop {
if *stop.borrow() {
return;
}
let (code, reason, ready) = self.session(&events, &mut stop).await;
*self.out.lock().unwrap_or_else(|e| e.into_inner()) = None;
if *stop.borrow() {
return;
}
if ready {
retry = 0;
}
if FATAL.contains(&code) {
let _ = events.send(GatewayEvent::Fatal(format!(
"Token recusado pelo gateway ({code} {reason})"
)));
return;
}
let wait = backoff(retry, jitter());
retry += 1;
let _ = events.send(GatewayEvent::Reconnecting { code, reason, wait });
tokio::select! {
_ = tokio::time::sleep(wait) => {}
_ = stopped(&mut stop) => return,
}
}
}
async fn session(
&self,
events: &mpsc::UnboundedSender<GatewayEvent>,
stop: &mut tokio::sync::watch::Receiver<bool>,
) -> (u16, String, bool) {
let ws = tokio::select! {
r = tokio_tungstenite::connect_async(self.url.as_str()) => r,
_ = stopped(stop) => return (1000, "fechado".into(), false),
};
let (mut sink, mut stream) = match ws {
Ok((ws, _)) => ws.split(),
Err(e) => return (1006, e.to_string(), false),
};
let (tx, mut rx) = mpsc::unbounded_channel::<WsMessage>();
*self.out.lock().unwrap_or_else(|e| e.into_inner()) = Some(tx.clone());
let mut ready = false;
let mut every = Duration::from_secs(30);
let mut beat: Option<tokio::time::Interval> = None;
let mut last_ack = Instant::now();
loop {
tokio::select! {
_ = stopped(stop) => {
let _ = sink.send(WsMessage::Close(Some(CloseFrame { code: CloseCode::Normal, reason: "".into() }))).await;
return (1000, "fechado".into(), ready);
}
out = rx.recv() => {
if let Some(m) = out {
if sink.send(m).await.is_err() {
return (1006, "falha ao enviar".into(), ready);
}
}
}
_ = async {
match beat.as_mut() {
Some(b) => {
b.tick().await;
}
None => std::future::pending::<()>().await,
}
} => {
if last_ack.elapsed() > every * 2 + Duration::from_secs(5) {
let _ = sink.send(WsMessage::Close(Some(CloseFrame { code: CloseCode::from(4900), reason: "heartbeat timeout".into() }))).await;
return (4900, "heartbeat timeout".into(), ready);
}
let _ = sink.send(WsMessage::Text(json!({ "op": 1, "d": null }).to_string().into())).await;
}
msg = stream.next() => {
let text = match msg {
Some(Ok(WsMessage::Text(t))) => t.to_string(),
Some(Ok(WsMessage::Binary(b))) => String::from_utf8_lossy(&b).into_owned(),
Some(Ok(WsMessage::Close(frame))) => {
let (code, reason) = frame.map(|f| (u16::from(f.code), f.reason.to_string())).unwrap_or((1005, String::new()));
return (code, reason, ready);
}
Some(Ok(_)) => continue,
Some(Err(e)) => return (1006, e.to_string(), ready),
None => return (1006, "conexão encerrada".into(), ready),
};
let Ok(msg) = serde_json::from_str::<Value>(&text) else { continue };
match msg["op"].as_u64() {
Some(10) => {
let identify = json!({ "op": 2, "d": { "token": self.token, "intents": self.intents } });
let _ = sink.send(WsMessage::Text(identify.to_string().into())).await;
every = Duration::from_millis(msg["d"]["heartbeat_interval"].as_u64().unwrap_or(30_000));
last_ack = Instant::now();
let mut b = tokio::time::interval_at(tokio::time::Instant::now() + every, every);
b.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
beat = Some(b);
}
Some(11) => last_ack = Instant::now(),
Some(0) => {
last_ack = Instant::now();
let Some(t) = msg["t"].as_str().map(str::to_string) else { continue };
if t == "READY" {
ready = true;
}
let mut msg = msg;
let d = msg["d"].take();
let _ = events.send(GatewayEvent::Dispatch(t, d));
}
_ => {}
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn backoff_doubles_up_to_thirty_seconds() {
assert_eq!(backoff(0, 0.5), Duration::from_secs(1));
assert_eq!(backoff(1, 0.5), Duration::from_secs(2));
assert_eq!(backoff(2, 0.5), Duration::from_secs(4));
assert_eq!(backoff(10, 0.5), Duration::from_secs(30));
assert_eq!(backoff(100, 0.5), Duration::from_secs(30));
assert!(backoff(0, 0.0) >= Duration::from_millis(750));
assert!(backoff(0, 1.0) <= Duration::from_millis(1250));
}
}