mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-25 23:52:04 +00:00
fix: debounce inline generations and cap results
This commit is contained in:
@@ -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));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user