mirror of
https://github.com/TheFunny/TelegramTwitterMediaBot.git
synced 2026-09-23 23:32:05 +00:00
YoursFunny
507c8ac317
Two things the 300 s sweep did the hard way: - The link cache is pruned by `created_at` (`DELETE FROM link_cache WHERE created_at < ?`) and had no index on it, so every sweep scanned the whole table — every post sent inside the TTL window, which is up to a week of them — while the `url` primary key served none of it. A new migration (appended; migration 1 is frozen and already shipped) creates the index, and the upgrade test now asserts it exists after an upgrade. - `ChatStore::prune_expired` only ever *looked* at chats that had an expired edit-before-forward record, so a chat with no prompt at all — the common case: every chat that ever sent a message or ran a command — stayed in the cache and in the per-chat lock map for the process lifetime. The candidate set now includes chats holding no records, which is what the eviction below was written for; the DB keeps the row, so the next use costs one SELECT (pinned by a new test that also shows the durable settings come back). Deliberately *not* done: skipping the write in `ChatStore::set` when the state is unchanged. Comparing against the cached copy would skip a serialize plus a blocking DB round trip for a no-op update — but every one of the 13 `update` callers mutates something, so the no-op case is a user repeating an identical command, and the same comparison would also skip the write that repairs a row whose earlier write failed. A rare saving against a rare repair, and the write is what makes the cache a cache rather than a source of truth. Verified: the new eviction test fails without the candidate change (checked by reverting it) and passes with it; 118 bot tests and 91 x-media tests pass. `cargo fmt --check`, `cargo clippy --workspace --all-targets --locked -- -D warnings` and `cargo test --workspace --locked` clean.
352 lines
13 KiB
Rust
352 lines
13 KiB
Rust
//! Per-chat state with SQLite persistence (table `chat_state` in
|
|
//! `data/task_queue.db`, shared with the task queue).
|
|
|
|
use crate::db::unix_now;
|
|
use parking_lot::Mutex;
|
|
use rusqlite::params;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
#[derive(Serialize, Deserialize, Default, Clone, Debug)]
|
|
pub struct ChatData {
|
|
pub forward_channel_id: Option<i64>,
|
|
pub edit_before_forward: bool,
|
|
/// Key: prompt message id.
|
|
pub edit_message: HashMap<i64, EditMessage>,
|
|
/// name -> HTML template containing "[]"
|
|
pub template: HashMap<String, String>,
|
|
/// site name (twitter/bsky/misskey/pixiv/bilibili) -> user-supplied caption format
|
|
/// with {url} {author} {author_url} {title} {content} {tags} placeholders.
|
|
pub message_format: HashMap<String, String>,
|
|
}
|
|
|
|
#[derive(Serialize, Deserialize, Clone, Debug, Default)]
|
|
pub struct EditMessage {
|
|
pub url: String,
|
|
pub chat_id: i64,
|
|
pub forward_message_ids: Vec<i64>,
|
|
pub template: String,
|
|
/// Unix seconds at registration; expiry = created_at + ttl.
|
|
pub created_at: i64,
|
|
}
|
|
|
|
pub struct ChatStore {
|
|
/// In-memory cache; the DB is the source of truth on first access.
|
|
cache: Mutex<HashMap<i64, ChatData>>,
|
|
/// Per-chat async locks serializing get→mutate→set so concurrent handler
|
|
/// tasks (batch-forwards, callbacks) cannot clobber each other's writes.
|
|
locks: Mutex<HashMap<i64, Arc<tokio::sync::Mutex<()>>>>,
|
|
pool: Arc<crate::db::DbPool>,
|
|
}
|
|
|
|
impl ChatStore {
|
|
/// Wraps the shared DB pool (schema initialized once by
|
|
/// [`crate::db::open_store`]; the `chat_state` table lives in the merged
|
|
/// schema alongside `tasks` and `link_cache`).
|
|
pub fn new(pool: Arc<crate::db::DbPool>) -> Self {
|
|
ChatStore {
|
|
cache: Mutex::new(HashMap::new()),
|
|
locks: Mutex::new(HashMap::new()),
|
|
pool,
|
|
}
|
|
}
|
|
|
|
pub async fn get(&self, chat_id: i64) -> ChatData {
|
|
if let Some(data) = self.cache.lock().get(&chat_id) {
|
|
return data.clone();
|
|
}
|
|
let chat_key = chat_id.to_string();
|
|
let payload = self
|
|
.pool
|
|
.with_conn(move |conn| {
|
|
// Concurrent handler tasks (batch-forwards) may write chat_state
|
|
// while this read runs; the shared busy timeout handles the
|
|
// write-lock collision instead of failing the query.
|
|
let mut stmt = conn.prepare("SELECT payload FROM chat_state WHERE chat_id = ?1")?;
|
|
let mut rows = stmt.query(params![chat_key])?;
|
|
match rows.next()? {
|
|
Some(row) => Ok(Some(row.get::<_, String>(0)?)),
|
|
None => Ok(None),
|
|
}
|
|
})
|
|
.await
|
|
.unwrap_or_else(|e| {
|
|
log::warn!("chat_state read failed: {e}");
|
|
None
|
|
})
|
|
.unwrap_or_default();
|
|
let data: ChatData = serde_json::from_str(&payload).unwrap_or_default();
|
|
self.cache.lock().insert(chat_id, data.clone());
|
|
data
|
|
}
|
|
|
|
/// Write-through: update the cache and the DB.
|
|
pub async fn set(&self, chat_id: i64, data: &ChatData) {
|
|
self.cache.lock().insert(chat_id, data.clone());
|
|
let payload = serde_json::to_string(data).expect("chat state serializes");
|
|
let chat_id = chat_id.to_string();
|
|
let result = self
|
|
.pool
|
|
.with_conn(move |conn| {
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO chat_state (chat_id, payload) VALUES (?1, ?2)",
|
|
params![chat_id, payload],
|
|
)?;
|
|
Ok(())
|
|
})
|
|
.await;
|
|
if let Err(e) = result {
|
|
log::warn!("chat_state write failed: {e}");
|
|
}
|
|
}
|
|
|
|
/// The per-chat async lock serializing get→mutate→set cycles.
|
|
fn lock_for(&self, chat_id: i64) -> Arc<tokio::sync::Mutex<()>> {
|
|
self.locks
|
|
.lock()
|
|
.entry(chat_id)
|
|
.or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
|
|
.clone()
|
|
}
|
|
|
|
/// Serializes a get→mutate→set cycle per chat: concurrent handler tasks
|
|
/// (the batch-forward design spawns several per chat) each snapshot the
|
|
/// same `ChatData` and last-writer-wins would silently drop mutations,
|
|
/// e.g. a second `edit_message` record. The per-chat lock makes the
|
|
/// cycle atomic. Returns the closure's result.
|
|
pub async fn update<R>(&self, chat_id: i64, f: impl FnOnce(&mut ChatData) -> R) -> R {
|
|
let lock = self.lock_for(chat_id);
|
|
let _guard = lock.lock().await;
|
|
let mut data = self.get(chat_id).await;
|
|
let r = f(&mut data);
|
|
self.set(chat_id, &data).await;
|
|
r
|
|
}
|
|
|
|
/// Removes edit-before-forward records whose `created_at + ttl` is in the
|
|
/// past. Returns the removed `(chat_id, prompt_message_id)` pairs so the
|
|
/// caller can clear the prompt's buttons.
|
|
pub async fn prune_expired(&self, ttl: Duration) -> Vec<(i64, i64)> {
|
|
let now = unix_now();
|
|
let ttl_secs = ttl.as_secs() as i64;
|
|
// Chats worth looking at, from a cache snapshot: the ones with an
|
|
// expired record, plus the ones holding no record at all. The latter
|
|
// used to be left alone for the process lifetime — every chat that ever
|
|
// sent a message or ran a command stayed in the cache and in the
|
|
// per-chat lock map — even though a chat with no live prompt is exactly
|
|
// what the eviction below is for. The pruning itself re-reads and
|
|
// writes under the per-chat lock below; taking no lock here means a
|
|
// chat appearing later is simply picked up by the next sweep.
|
|
let candidates: Vec<i64> = {
|
|
let cache = self.cache.lock();
|
|
cache
|
|
.iter()
|
|
.filter(|(_, data)| {
|
|
data.edit_message.is_empty()
|
|
|| data
|
|
.edit_message
|
|
.values()
|
|
.any(|entry| entry.created_at + ttl_secs <= now)
|
|
})
|
|
.map(|(chat_id, _)| *chat_id)
|
|
.collect()
|
|
};
|
|
let mut removed = Vec::new();
|
|
let mut evicted_chats = Vec::new();
|
|
for chat_id in candidates {
|
|
let lock = self.lock_for(chat_id);
|
|
let _guard = lock.lock().await;
|
|
let mut data = self.get(chat_id).await;
|
|
let before = data.edit_message.len();
|
|
data.edit_message.retain(|key, entry| {
|
|
if entry.created_at + ttl_secs > now {
|
|
return true;
|
|
}
|
|
removed.push((chat_id, *key));
|
|
false
|
|
});
|
|
if data.edit_message.len() != before {
|
|
self.set(chat_id, &data).await;
|
|
}
|
|
// Chats with no live edit records: evicted from the cache (and
|
|
// their per-chat lock) so the cache stays bounded to active
|
|
// prompts. The DB keeps the row; the next get() reloads it.
|
|
if data.edit_message.is_empty() {
|
|
evicted_chats.push(chat_id);
|
|
}
|
|
}
|
|
if !evicted_chats.is_empty() {
|
|
let mut cache = self.cache.lock();
|
|
let mut locks = self.locks.lock();
|
|
for chat_id in &evicted_chats {
|
|
cache.remove(chat_id);
|
|
locks.remove(chat_id);
|
|
}
|
|
}
|
|
if !removed.is_empty() {
|
|
log::info!(
|
|
"pruned {} expired edit-before-forward record(s)",
|
|
removed.len()
|
|
);
|
|
}
|
|
removed
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn concurrent_updates_do_not_lose_edit_records() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let pool = crate::db::open_store(dir.path().join("s.db").to_str().unwrap()).unwrap();
|
|
let store = std::sync::Arc::new(ChatStore::new(pool));
|
|
let mut handles = Vec::new();
|
|
for i in 0..4 {
|
|
let store = Arc::clone(&store);
|
|
handles.push(tokio::spawn(async move {
|
|
store
|
|
.update(1001, |data| {
|
|
data.edit_message.insert(
|
|
i,
|
|
EditMessage {
|
|
url: format!("https://x.com/u/status/{i}"),
|
|
chat_id: 1001,
|
|
forward_message_ids: vec![i],
|
|
template: String::new(),
|
|
created_at: 0,
|
|
},
|
|
);
|
|
})
|
|
.await;
|
|
}));
|
|
}
|
|
for h in handles {
|
|
h.await.unwrap();
|
|
}
|
|
let data = store.get(1001).await;
|
|
assert_eq!(
|
|
data.edit_message.len(),
|
|
4,
|
|
"concurrent get→mutate→set must not drop records"
|
|
);
|
|
}
|
|
|
|
fn edit_entry(chat_id: i64, created_at: i64) -> EditMessage {
|
|
EditMessage {
|
|
url: "https://x.com/u/status/1".into(),
|
|
chat_id,
|
|
forward_message_ids: vec![9],
|
|
template: String::new(),
|
|
created_at,
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prune_removes_only_expired_records() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let pool = crate::db::open_store(dir.path().join("p.db").to_str().unwrap()).unwrap();
|
|
let store = ChatStore::new(pool);
|
|
let now = unix_now();
|
|
store
|
|
.update(7, |data| {
|
|
data.template.insert("t".into(), "[]".into());
|
|
data.edit_message.insert(1, edit_entry(7, now - 3600));
|
|
data.edit_message.insert(2, edit_entry(7, now));
|
|
})
|
|
.await;
|
|
|
|
let removed = store.prune_expired(Duration::from_secs(60)).await;
|
|
|
|
assert_eq!(removed, vec![(7, 1)]);
|
|
let data = store.get(7).await;
|
|
assert!(data.edit_message.contains_key(&2), "live record pruned");
|
|
assert_eq!(
|
|
data.template.get("t").map(String::as_str),
|
|
Some("[]"),
|
|
"unrelated state lost by the prune"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_idle_chat_is_evicted_and_its_state_reloads() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let pool = crate::db::open_store(dir.path().join("e.db").to_str().unwrap()).unwrap();
|
|
let store = ChatStore::new(pool);
|
|
// Durable settings and no prompt at all: this chat used to sit in the
|
|
// cache (and in the per-chat lock map) for the process lifetime,
|
|
// because the sweep only ever looked at chats with an *expired* record.
|
|
store
|
|
.update(9, |data| {
|
|
data.forward_channel_id = Some(-100);
|
|
data.message_format.insert("twitter".into(), "{url}".into());
|
|
})
|
|
.await;
|
|
assert!(store.cache.lock().contains_key(&9));
|
|
|
|
let removed = store.prune_expired(Duration::from_secs(60)).await;
|
|
|
|
assert!(removed.is_empty(), "nothing had expired");
|
|
assert!(
|
|
!store.cache.lock().contains_key(&9),
|
|
"a chat with no live prompt must leave the cache"
|
|
);
|
|
assert!(!store.locks.lock().contains_key(&9), "…and its lock");
|
|
// The DB kept the row, so the next use reloads everything it held.
|
|
let data = store.get(9).await;
|
|
assert_eq!(data.forward_channel_id, Some(-100));
|
|
assert_eq!(
|
|
data.message_format.get("twitter").map(String::as_str),
|
|
Some("{url}")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_live_prompt_keeps_its_chat_cached() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let pool = crate::db::open_store(dir.path().join("k.db").to_str().unwrap()).unwrap();
|
|
let store = ChatStore::new(pool);
|
|
store
|
|
.update(10, |data| {
|
|
data.edit_message.insert(1, edit_entry(10, unix_now()));
|
|
})
|
|
.await;
|
|
|
|
store.prune_expired(Duration::from_secs(3600)).await;
|
|
|
|
assert!(
|
|
store.cache.lock().contains_key(&10),
|
|
"a live prompt holds its chat in the cache"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prune_eviction_keeps_the_persisted_state() {
|
|
// Every record expires → the chat is evicted from the cache; the
|
|
// pruned state must already be in the DB when that happens.
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let pool = crate::db::open_store(dir.path().join("p.db").to_str().unwrap()).unwrap();
|
|
let store = ChatStore::new(pool);
|
|
store
|
|
.update(8, |data| {
|
|
data.template.insert("keep".into(), "[]".into());
|
|
data.edit_message.insert(1, edit_entry(8, 0));
|
|
})
|
|
.await;
|
|
|
|
let removed = store.prune_expired(Duration::from_secs(60)).await;
|
|
|
|
assert_eq!(removed, vec![(8, 1)]);
|
|
let data = store.get(8).await;
|
|
assert!(data.edit_message.is_empty());
|
|
assert_eq!(
|
|
data.template.get("keep").map(String::as_str),
|
|
Some("[]"),
|
|
"eviction dropped state the DB never received"
|
|
);
|
|
}
|
|
}
|