fix: debounce inline generations and cap results

This commit is contained in:
2026-09-24 20:40:27 +08:00
parent 85c3be5ef5
commit b2b01721b2
+73 -61
View File
@@ -36,54 +36,55 @@ const INLINE_STATE_TTL: std::time::Duration = std::time::Duration::from_secs(300
/// different user's query) cancel another user's pending answer. /// different user's query) cancel another user's pending answer.
struct InlineDebounceState { struct InlineDebounceState {
query: String, query: String,
generation: u64,
answered: bool, answered: bool,
/// When a query last touched this entry, so the periodic sweep can drop
/// one per user who ever used inline mode (the map had no eviction at all,
/// unlike the rate limiter's buckets and the chat store).
last_seen: std::time::Instant, last_seen: std::time::Instant,
} }
#[derive(Default)] #[derive(Default)]
struct DebounceStates(HashMap<u64, InlineDebounceState>); struct DebounceStates {
entries: HashMap<u64, InlineDebounceState>,
generation: u64,
}
impl DebounceStates { impl DebounceStates {
/// Records `query` as the user's newest query. Returns false when it is a /// Records `query` as the user's newest query. Returns false when it is a
/// repeat whose answer already went out (Telegram's inline cache serves /// repeat whose answer already went out (Telegram's inline cache serves
/// it; re-fetching would only hit the source site again). /// it; re-fetching would only hit the source site again).
fn note(&mut self, user_id: u64, query: &str) -> bool { fn note(&mut self, user_id: u64, query: &str) -> (bool, u64) {
if let Some(prev) = self.0.get(&user_id) if let Some(prev) = self.entries.get(&user_id)
&& prev.query == query && prev.query == query
&& prev.answered && prev.answered
{ {
return false; return (false, prev.generation);
} }
self.0.insert( self.generation = self.generation.wrapping_add(1);
let generation = self.generation;
self.entries.insert(
user_id, user_id,
InlineDebounceState { InlineDebounceState {
query: query.to_string(), query: query.to_string(),
generation,
answered: false, answered: false,
last_seen: std::time::Instant::now(), last_seen: std::time::Instant::now(),
}, },
); );
true (true, generation)
} }
/// Drops entries no query has touched for `idle_for`. Split from the clock /// Drops entries no query has touched for `idle_for`. Split from the clock
/// so the boundary is testable without ageing a monotonic instant. /// so the boundary is testable without ageing a monotonic instant.
fn prune_idle_at(&mut self, now: std::time::Instant, idle_for: std::time::Duration) -> usize { fn prune_idle_at(&mut self, now: std::time::Instant, idle_for: std::time::Duration) -> usize {
let before = self.0.len(); let before = self.entries.len();
self.0 self.entries
.retain(|_, state| now.saturating_duration_since(state.last_seen) < idle_for); .retain(|_, state| now.saturating_duration_since(state.last_seen) < idle_for);
before - self.0.len() before - self.entries.len()
} }
fn claim(&mut self, user_id: u64, query: &str, generation: u64) -> bool {
/// Claims the answer for the user's newest query; false when a newer query let Some(state) = self.entries.get_mut(&user_id) else {
/// superseded it or the answer was already claimed.
fn claim(&mut self, user_id: u64, query: &str) -> bool {
let Some(state) = self.0.get_mut(&user_id) else {
return false; return false;
}; };
if state.query != query || state.answered { if state.query != query || state.generation != generation || state.answered {
return false; return false;
} }
state.answered = true; state.answered = true;
@@ -91,10 +92,10 @@ impl DebounceStates {
true true
} }
/// Releases a claimed-but-unsent answer so a repeat can retry the fetch. fn release(&mut self, user_id: u64, query: &str, generation: u64) {
fn release(&mut self, user_id: u64, query: &str) { if let Some(state) = self.entries.get_mut(&user_id)
if let Some(state) = self.0.get_mut(&user_id)
&& state.query == query && state.query == query
&& state.generation == generation
{ {
state.answered = false; state.answered = false;
state.last_seen = std::time::Instant::now(); state.last_seen = std::time::Instant::now();
@@ -119,20 +120,26 @@ pub async fn inline_query_handler(bot: Bot, query: InlineQuery) -> Result<(), Re
return answer_inline_query(&ctx, query).await.map(|_| ()); return answer_inline_query(&ctx, query).await.map(|_| ());
} }
let user_id = query.from.id.0; let user_id = query.from.id.0;
if !INLINE_DEBOUNCE_STATE.lock().note(user_id, &query.query) { let (should_answer, generation) = INLINE_DEBOUNCE_STATE.lock().note(user_id, &query.query);
if !should_answer {
return respond(()); return respond(());
} }
let query_text = query.query.clone(); let query_text = query.query.clone();
tokio::spawn(async move { tokio::spawn(async move {
tokio::time::sleep(INLINE_DEBOUNCE).await; tokio::time::sleep(INLINE_DEBOUNCE).await;
if !INLINE_DEBOUNCE_STATE.lock().claim(user_id, &query_text) { if !INLINE_DEBOUNCE_STATE
.lock()
.claim(user_id, &query_text, generation)
{
return; return;
} }
let ctx = AppContext::from_statics(&bot); let ctx = AppContext::from_statics(&bot);
match answer_inline_query(&ctx, query).await { match answer_inline_query(&ctx, query).await {
Ok(true) => {} Ok(true) => {}
Ok(false) | Err(_) => { Ok(false) | Err(_) => {
INLINE_DEBOUNCE_STATE.lock().release(user_id, &query_text); INLINE_DEBOUNCE_STATE
.lock()
.release(user_id, &query_text, generation);
} }
} }
}); });
@@ -185,7 +192,7 @@ async fn answer_inline_query(
ctx.config.caption_quote_text_chars, ctx.config.caption_quote_text_chars,
); );
let mut results: Vec<InlineQueryResult> = Vec::new(); let mut results: Vec<InlineQueryResult> = Vec::new();
for (i, media) in fetched.media.iter().enumerate() { for (i, media) in fetched.media.iter().enumerate().take(50) {
// Telegram fetches an inline result's URL itself and cannot // Telegram fetches an inline result's URL itself and cannot
// send site-specific headers, so hotlink-protected media // send site-specific headers, so hotlink-protected media
// (pixiv's pximg.net) would render as a broken file there. // (pixiv's pximg.net) would render as a broken file there.
@@ -328,6 +335,7 @@ fn cached_inline_results(cached: &CachedPost, caption: &str) -> Vec<InlineQueryR
.media .media
.iter() .iter()
.enumerate() .enumerate()
.take(50)
.filter_map(|(i, media)| { .filter_map(|(i, media)| {
let id = i.to_string(); let id = i.to_string();
let caption = || caption.to_string(); let caption = || caption.to_string();
@@ -381,7 +389,7 @@ async fn answer(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{DebounceStates, INLINE_STATE_TTL, answer_inline_query}; use super::{DebounceStates, INLINE_STATE_TTL, answer_inline_query, cached_inline_results};
use crate::ctx::test_support::{TestStores, api_error, cached_photo}; use crate::ctx::test_support::{TestStores, api_error, cached_photo};
use crate::link_cache::{CachedMedia, CachedMediaKind}; use crate::link_cache::{CachedMedia, CachedMediaKind};
use crate::media_sender::test_support::MockSender; use crate::media_sender::test_support::MockSender;
@@ -489,64 +497,68 @@ mod tests {
); );
} }
#[test]
fn inline_results_are_capped_at_telegram_limit() {
let mut entry = cached_photo();
entry.media = (0..51)
.map(|i| CachedMedia {
kind: CachedMediaKind::Photo,
file_id: format!("id-{i}"),
url: format!("https://p/{i}.jpg"),
})
.collect();
assert_eq!(cached_inline_results(&entry, "caption").len(), 50);
}
#[test] #[test]
fn debounce_state_is_per_user() { fn debounce_state_is_per_user() {
let mut states = DebounceStates::default(); let mut states = DebounceStates::default();
// Two users query different links: both proceed, and neither timer assert!(states.note(1, URL_A).0);
// cancels the other (a single shared slot dropped one of them). assert!(states.note(2, URL_B).0);
assert!(states.note(1, URL_A)); let (_, generation_a) = states.note(1, URL_A);
assert!(states.note(2, URL_B)); let (_, generation_b) = states.note(2, URL_B);
assert!(states.claim(1, URL_A), "user 1's answer was cancelled"); assert!(states.claim(1, URL_A, generation_a));
assert!(states.claim(2, URL_B), "user 2's answer was cancelled"); assert!(states.claim(2, URL_B, generation_b));
} }
#[test] #[test]
fn answered_query_is_suppressed_per_user_only() { fn answered_query_is_suppressed_per_user_only() {
let mut states = DebounceStates::default(); let mut states = DebounceStates::default();
assert!(states.note(1, URL_A)); assert!(states.note(1, URL_A).0);
assert!(states.claim(1, URL_A)); let (_, generation) = states.note(1, URL_A);
// A repeat of the answered query by the same user is left to assert!(states.claim(1, URL_A, generation));
// Telegram's inline cache. assert!(!states.note(1, URL_A).0);
assert!(!states.note(1, URL_A)); assert!(states.note(2, URL_A).0);
// Another user pasting the same link still gets an answer. let (_, generation) = states.note(2, URL_A);
assert!(states.note(2, URL_A)); assert!(states.claim(2, URL_A, generation));
assert!(states.claim(2, URL_A));
} }
#[test] #[test]
fn idle_states_are_pruned_and_live_ones_kept() { fn idle_states_are_pruned_and_live_ones_kept() {
let mut states = DebounceStates::default(); let mut states = DebounceStates::default();
assert!(states.note(1, URL_A)); assert!(states.note(1, URL_A).0);
let first = states.0[&1].last_seen; let first = states.entries[&1].last_seen;
// Entry 2 is strictly newer, so one timestamp can sit exactly on the
// window's edge for one and comfortably inside it for the other.
std::thread::sleep(std::time::Duration::from_millis(2)); std::thread::sleep(std::time::Duration::from_millis(2));
assert!(states.note(2, URL_B)); assert!(states.note(2, URL_B).0);
assert_eq!( assert_eq!(
states.prune_idle_at(first + INLINE_STATE_TTL, INLINE_STATE_TTL), states.prune_idle_at(first + INLINE_STATE_TTL, INLINE_STATE_TTL),
1 1
); );
assert!( assert!(!states.entries.contains_key(&1));
!states.0.contains_key(&1), assert!(states.entries.contains_key(&2));
"the entry past the window must go" assert!(states.note(1, URL_A).0);
);
assert!(states.0.contains_key(&2), "the live entry must stay");
// A pruned user's repeat is answered fresh instead of suppressed.
assert!(states.note(1, URL_A));
} }
#[test] #[test]
fn newer_query_supersedes_and_failed_answer_is_released() { fn newer_query_supersedes_and_failed_answer_is_released() {
let mut states = DebounceStates::default(); let mut states = DebounceStates::default();
assert!(states.note(1, URL_A)); assert!(states.note(1, URL_A).0);
assert!(states.note(1, URL_B)); let (_, generation_a) = states.note(1, URL_B);
// The stale timer for the half-typed query gives up… assert!(!states.claim(1, URL_A, generation_a));
assert!(!states.claim(1, URL_A)); let (_, generation_b) = states.note(1, URL_B);
// …and the newest one answers. assert!(states.claim(1, URL_B, generation_b));
assert!(states.claim(1, URL_B)); states.release(1, URL_B, generation_b);
// No results → release so a repeat may retry the fetch. assert!(states.claim(1, URL_B, generation_b));
states.release(1, URL_B);
assert!(states.claim(1, URL_B));
} }
} }