zynko 0.1.0

Biblioteca oficial para bots do Zynko: API, gateway e voz.
Documentation
//! The realtime link (docs.zynko.org/api/gateway): HELLO, IDENTIFY with the token and the
//! intents, heartbeats, and a fresh IDENTIFY after any drop (there is no RESUME: READY always
//! brings the whole state again).

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;

/// What the gateway tells its owner.
pub(crate) enum GatewayEvent {
    Dispatch(String, Value),
    /// The link dropped; trying again after `wait`.
    Reconnecting {
        code: u16,
        reason: String,
        wait: Duration,
    },
    /// The token was refused: no retry.
    Fatal(String),
}

/// Close codes after which trying again cannot help.
const FATAL: [u16; 1] = [4004];

pub(crate) struct Gateway {
    url: String,
    token: String,
    intents: u64,
    /// The open link's writer, while there is one.
    out: Mutex<Option<mpsc::UnboundedSender<WsMessage>>>,
    stop: tokio::sync::watch::Sender<bool>,
}

/// 1s, 2s, 4s... up to 30s, with some jitter so a server restart is not hit all at once.
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,
        })
    }

    /// Sends an op (4: voice state). Returns whether the link was open.
    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);
    }

    /// Runs until `close` or a fatal close code, reconnecting in between.
    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,
            }
        }
    }

    /// One connection, from open to close. Returns the close code and reason, and whether READY
    /// came.
    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,
                    }
                } => {
                    // No answer to two beats in a row: the link is dead even if the socket says otherwise.
                    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() {
        // With jitter 0.5 the factor is exactly 1.
        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));
    }
}