use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use crate::config::Config;
#[cfg(test)]
static CLOCK_OFFSET_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
#[cfg(not(test))]
fn now() -> Instant {
Instant::now()
}
#[cfg(test)]
fn now() -> Instant {
use std::sync::atomic::Ordering;
let off = CLOCK_OFFSET_MS.load(Ordering::Relaxed);
Instant::now() + Duration::from_millis(off)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Surface {
Send,
Final,
Chrome,
}
impl Surface {
pub(crate) fn is_content(self) -> bool {
matches!(self, Surface::Send | Surface::Final)
}
}
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum Decision {
Admit,
Hold(Duration),
Drop,
FailOpen,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Limits {
pub enabled: bool,
pub ceiling_per_min: u32,
pub burst: u32,
pub chrome_spacing: Duration,
pub content_max_hold: Duration,
pub pause_fallback: Duration,
pub pause_ceiling: Duration,
}
impl Limits {
fn from_config() -> Self {
let g = &Config::current().channels.discord.governor;
Self {
enabled: g.enabled,
ceiling_per_min: g.writes_per_minute,
burst: g.burst.max(1),
chrome_spacing: Duration::from_millis(g.chrome_min_spacing_ms),
content_max_hold: Duration::from_secs(25 * 60),
pause_fallback: Duration::from_secs(g.pause_secs.max(1)),
pause_ceiling: Duration::from_secs(g.pause_secs.max(1) * 4),
}
}
}
#[derive(Debug)]
pub(crate) struct ChannelBudget {
tokens: f64,
last_refill: Option<Instant>,
last_chrome: Option<Instant>,
pause_until: Option<Instant>,
}
impl Default for ChannelBudget {
fn default() -> Self {
Self {
tokens: f64::INFINITY,
last_refill: None,
last_chrome: None,
pause_until: None,
}
}
}
impl ChannelBudget {
fn refresh(&mut self, l: &Limits, now: Instant) {
if l.ceiling_per_min == 0 {
self.tokens = f64::INFINITY;
self.last_refill = Some(now);
} else {
let capacity = l.burst.max(1) as f64;
let rate = l.ceiling_per_min as f64 / 60.0;
self.tokens = match self.last_refill {
Some(last) => {
let dt = now.saturating_duration_since(last).as_secs_f64();
(self.tokens + dt * rate).min(capacity)
}
None => capacity,
};
self.last_refill = Some(now);
}
if self.pause_until.is_some_and(|until| until <= now) {
self.pause_until = None;
}
}
fn is_paused(&self, now: Instant) -> bool {
self.pause_until.is_some_and(|until| until > now)
}
}
pub(crate) fn admit(
b: &mut ChannelBudget,
l: &Limits,
surface: Surface,
held: Duration,
now: Instant,
) -> Decision {
if !l.enabled {
return Decision::Admit;
}
b.refresh(l, now);
if b.is_paused(now) {
let until = b.pause_until.unwrap_or(now);
let wait = until.saturating_duration_since(now);
return if surface.is_content() {
if held + wait > l.content_max_hold {
Decision::FailOpen
} else {
Decision::Hold(wait)
}
} else {
Decision::Drop
};
}
if surface == Surface::Chrome
&& !l.chrome_spacing.is_zero()
&& let Some(last) = b.last_chrome
&& now.saturating_duration_since(last) < l.chrome_spacing
{
return Decision::Drop;
}
if l.ceiling_per_min > 0 && b.tokens < 1.0 {
let need = Duration::from_secs_f64((1.0 - b.tokens) * 60.0 / l.ceiling_per_min as f64);
if !surface.is_content() {
return Decision::Drop;
}
if held + need > l.content_max_hold {
tracing::warn!(
"Discord: governor write ceiling held a content write past its ceiling, failing open"
);
return Decision::FailOpen;
}
return Decision::Hold(need);
}
if l.ceiling_per_min > 0 {
b.tokens -= 1.0;
}
if surface == Surface::Chrome {
b.last_chrome = Some(now);
}
Decision::Admit
}
fn arm_park(b: &mut ChannelBudget, until: Instant, now: Instant) -> Duration {
match b.pause_until {
Some(existing) if existing >= until => existing.saturating_duration_since(now),
_ => {
b.pause_until = Some(until);
until.saturating_duration_since(now)
}
}
}
pub(crate) fn note_429(
b: &mut ChannelBudget,
l: &Limits,
retry_after: Option<Duration>,
now: Instant,
) -> Duration {
let requested = retry_after.unwrap_or(l.pause_fallback);
let clamped = requested.min(l.pause_ceiling);
arm_park(b, now + clamped, now)
}
#[derive(Default)]
pub(crate) struct Governor {
channels: Mutex<HashMap<u64, ChannelBudget>>,
}
impl Governor {
pub(crate) async fn gate(&self, channel: u64, surface: Surface) {
let mut held = Duration::ZERO;
loop {
let decision = {
let limits = Limits::from_config();
let mut channels = self.channels.lock().expect("discord governor lock");
admit(
channels.entry(channel).or_default(),
&limits,
surface,
held,
now(),
)
};
match decision {
Decision::Admit | Decision::Drop => return,
Decision::FailOpen => {
tracing::warn!(
"Discord: governor held a {surface:?} write past its ceiling on \
channel {channel} and let it through unpaced (#1910)"
);
return;
}
Decision::Hold(wait) => {
tokio::time::sleep(wait).await;
held += wait;
}
}
}
}
pub(crate) fn chrome_admits(&self, channel: u64) -> bool {
let limits = Limits::from_config();
let mut channels = self.channels.lock().expect("discord governor lock");
matches!(
admit(
channels.entry(channel).or_default(),
&limits,
Surface::Chrome,
Duration::ZERO,
now(),
),
Decision::Admit
)
}
pub(crate) fn note_rate_limited(&self, channel: u64, retry_after: Option<Duration>) {
let limits = Limits::from_config();
let pause = {
let mut channels = self.channels.lock().expect("discord governor lock");
note_429(
channels.entry(channel).or_default(),
&limits,
retry_after,
now(),
)
};
tracing::warn!(
channel = channel,
pause_ms = pause.as_millis() as u64,
named = retry_after.map(|d| d.as_millis() as u64),
"Discord: 429 on channel {channel}, parking outbound writes for {} ms{}",
pause.as_millis(),
match retry_after {
Some(d) => format!(" (API named {d:?}, clamped to the ceiling)"),
None => " (response named no retry_after, fallback window)".to_string(),
}
);
}
pub(crate) fn is_rate_limited(err: &serenity::Error) -> bool {
matches!(
err,
serenity::Error::Http(http)
if http.status_code().is_some_and(|code| code.as_u16() == 429)
)
}
pub(crate) fn retry_after(err: &serenity::Error) -> Option<Duration> {
let text = err.to_string().to_lowercase();
let at = text
.find("retry_after")
.or_else(|| text.find("retry after"))?;
let rest = &text[at..];
let digits = rest.find(char::is_numeric)?;
let tail = &rest[digits..];
let end = tail
.find(|c: char| !c.is_ascii_digit() && c != '.')
.unwrap_or(tail.len());
let secs: f64 = tail[..end].parse().ok()?;
if secs > 0.0 {
Some(Duration::from_secs_f64(secs))
} else {
None
}
}
}
pub(crate) fn note_if_429(governor: &Governor, channel: u64, err: Option<&serenity::Error>) {
if let Some(e) = err.filter(|e| Governor::is_rate_limited(e)) {
governor.note_rate_limited(channel, Governor::retry_after(e));
}
}
#[allow(clippy::result_large_err)]
pub(crate) async fn send_content(
governor: &Governor,
channel: serenity::model::id::ChannelId,
http: &serenity::http::Http,
surface: Surface,
builder: serenity::builder::CreateMessage,
) -> Result<serenity::model::channel::Message, serenity::Error> {
governor.gate(channel.get(), surface).await;
let res = channel.send_message(http, builder).await;
note_if_429(governor, channel.get(), res.as_ref().err());
res
}
#[allow(clippy::result_large_err)]
pub(crate) async fn edit_content(
governor: &Governor,
channel: serenity::model::id::ChannelId,
http: &serenity::http::Http,
message: serenity::model::id::MessageId,
surface: Surface,
edit: serenity::builder::EditMessage,
) -> Result<serenity::model::channel::Message, serenity::Error> {
governor.gate(channel.get(), surface).await;
let res = channel.edit_message(http, message, edit).await;
note_if_429(governor, channel.get(), res.as_ref().err());
res
}
#[allow(clippy::result_large_err)]
pub(crate) async fn say(
governor: &Governor,
channel: serenity::model::id::ChannelId,
http: &serenity::http::Http,
surface: Surface,
content: impl std::fmt::Display,
) -> Result<serenity::model::channel::Message, serenity::Error> {
governor.gate(channel.get(), surface).await;
let res = channel.say(http, content.to_string()).await;
note_if_429(governor, channel.get(), res.as_ref().err());
res
}
pub(crate) async fn edit_chrome(
governor: &Governor,
channel: serenity::model::id::ChannelId,
http: &serenity::http::Http,
message: serenity::model::id::MessageId,
edit: serenity::builder::EditMessage,
) -> Option<Result<serenity::model::channel::Message, serenity::Error>> {
if !governor.chrome_admits(channel.get()) {
tracing::debug!(
"Discord: governor dropped chrome edit on channel {}",
channel.get()
);
return None;
}
let res = channel.edit_message(http, message, edit).await;
note_if_429(governor, channel.get(), res.as_ref().err());
Some(res)
}
#[cfg(test)]
pub(crate) mod test_support {
use super::*;
use std::sync::atomic::Ordering;
pub(crate) async fn registry_guard() -> tokio::sync::MutexGuard<'static, ()> {
static GUARD: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
GUARD.lock().await
}
pub(crate) fn reset(offset_ms: u64) {
CLOCK_OFFSET_MS.store(offset_ms, Ordering::Relaxed);
}
pub(crate) fn advance(ms: u64) {
CLOCK_OFFSET_MS.fetch_add(ms, Ordering::Relaxed);
}
}