TelegramTwitterMediaBot/crates/xmedia-bot/src/state.rs

//! 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::OptionalExtension;
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>,
}

impl ChatData {
    /// The chat's caption format for `site`, empty when it has none — the
    /// built-in caption then applies (`caption_from_fields`).
    pub fn format_for(&self, site: &str) -> String {
        self.message_format.get(site).cloned().unwrap_or_default()
    }
}

#[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,
        }
    }

    /// Loads persisted state, distinguishing a missing row and a failed read
    /// from valid default settings. Only the ordinary read-only `get` path is
    /// allowed to degrade to defaults; mutations must not write those defaults
    /// back over a real row.
    async fn load(&self, chat_id: i64) -> Result<ChatData, ()> {
        if let Some(data) = self.cache.lock().get(&chat_id) {
            return Ok(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.
                conn.query_row(
                    "SELECT payload FROM chat_state WHERE chat_id = ?1",
                    params![chat_key],
                    |row| row.get::<_, String>(0),
                )
                .optional()
            })
            .await
            .map_err(|e| {
                log::warn!("chat_state read failed: {e}");
            })?;
        let payload = payload.unwrap_or_default();
        let data = if payload.is_empty() {
            ChatData::default()
        } else {
            serde_json::from_str(&payload).map_err(|e| {
                log::warn!("chat_state payload is invalid: {e}");
            })?
        };
        // Only fill a miss: an unconditional insert would let this (possibly
        // stale) snapshot overwrite what a concurrent set just wrote.
        self.cache
            .lock()
            .entry(chat_id)
            .or_insert_with(|| data.clone());
        Ok(data)
    }

    /// Read-only access may degrade to defaults for display and control flow.
    /// Mutating callers use [`Self::update`], which refuses a failed load.
    pub async fn get(&self, chat_id: i64) -> ChatData {
        self.load(chat_id).await.unwrap_or_default()
    }

    /// Strict read for callers whose *behavior* depends on the settings. A
    /// send that reads them via [`Self::get`] would silently use the defaults
    /// on a failed read — no channel forward, no edit prompt — and persist
    /// those defaults into the retry payload it builds.
    pub async fn try_get(&self, chat_id: i64) -> Result<ChatData, ()> {
        self.load(chat_id).await
    }

    /// Write-through: update the cache and the DB. Returns whether the DB
    /// write landed: the cache is updated either way, so `false` means the
    /// change lives only until the next restart and the caller has to say so
    /// instead of reporting a save that did not happen.
    pub async fn set(&self, chat_id: i64, data: &ChatData) -> bool {
        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();
        self.pool
            .with_conn_or(
                log::Level::Warn,
                "chat_state write failed",
                false,
                move |conn| {
                    conn.execute(
                        "INSERT OR REPLACE INTO chat_state (chat_id, payload) VALUES (?1, ?2)",
                        params![chat_id, payload],
                    )?;
                    Ok(true)
                },
            )
            .await
    }

    /// 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 plus whether the DB write
    /// landed (see [`Self::set`]); callers that do not care ignore the flag.
    pub async fn update<R>(
        &self,
        chat_id: i64,
        f: impl FnOnce(&mut ChatData) -> R,
    ) -> Result<(R, bool), ()> {
        let lock = self.lock_for(chat_id);
        let _guard = lock.lock().await;
        let mut data = self.load(chat_id).await?;
        let r = f(&mut data);
        let saved = self.set(chat_id, &data).await;
        Ok((r, saved))
    }

    /// 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 mut 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 persisted = self
            .pool
            .with_conn_or(
                log::Level::Warn,
                "expired prompt scan failed",
                Vec::<i64>::new(),
                move |conn| {
                    let mut stmt = conn.prepare("SELECT chat_id, payload FROM chat_state")?;
                    let rows = stmt.query_map([], |row| {
                        let id: String = row.get(0)?;
                        let id: i64 = id.parse().map_err(|e| {
                            rusqlite::Error::FromSqlConversionFailure(
                                0,
                                rusqlite::types::Type::Text,
                                Box::new(e),
                            )
                        })?;
                        let payload: String = row.get(1)?;
                        let data: ChatData = serde_json::from_str(&payload).map_err(|e| {
                            rusqlite::Error::FromSqlConversionFailure(
                                0,
                                rusqlite::types::Type::Text,
                                Box::new(e),
                            )
                        })?;
                        Ok((id, data))
                    })?;
                    Ok(rows
                        .filter_map(Result::ok)
                        .filter(|(_, data)| {
                            data.edit_message
                                .values()
                                .any(|entry| entry.created_at + ttl_secs <= now)
                        })
                        .map(|(id, _)| id)
                        .collect())
                },
            )
            .await;
        candidates.extend(persisted);
        candidates.sort_unstable();
        candidates.dedup();
        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();
            for chat_id in &evicted_chats {
                cache.remove(chat_id);
            }
        }
        // Per-chat locks go only while uncontended (the same rule as
        // rate_limit's prune): pulling a lock out from under an in-flight
        // update — between its `lock_for` clone and its `lock().await` —
        // would let a second writer `lock_for` a fresh one and enter the
        // critical section concurrently. A contended lock stays until a later
        // sweep, and dropping the uncontended ones also catches chats an
        // earlier sweep had to skip, so the map stays bounded.
        self.locks
            .lock()
            .retain(|_, lock| Arc::strong_count(lock) > 1);
        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
                    .unwrap();
            }));
        }
        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
            .unwrap();

        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 persisted_expired_prompts_are_pruned_after_restart() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("cold.db");
        let pool = crate::db::open_store(path.to_str().unwrap()).unwrap();
        let raw = rusqlite::Connection::open(&path).unwrap();
        let data = ChatData {
            edit_message: [(1, edit_entry(7, unix_now() - 3600))]
                .into_iter()
                .collect(),
            ..ChatData::default()
        };
        raw.execute(
            "INSERT INTO chat_state (chat_id, payload) VALUES ('7', ?1)",
            rusqlite::params![serde_json::to_string(&data).unwrap()],
        )
        .unwrap();
        let store = ChatStore::new(pool);
        assert!(store.cache.lock().get(&7).is_none());
        assert_eq!(
            store.prune_expired(Duration::from_secs(60)).await,
            vec![(7, 1)]
        );
    }

    #[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
            .unwrap();
        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
            .unwrap();

        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
            .unwrap();

        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"
        );
    }

    #[tokio::test]
    async fn a_failed_read_is_not_cached() {
        // A read that errors (busy, IO, a missing table) answers the default;
        // caching that answer would make the next get return it blind and the
        // next update write it back over the chat's real settings.
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("f.db");
        let pool = crate::db::open_store(path.to_str().unwrap()).unwrap();
        let raw = rusqlite::Connection::open(&path).unwrap();
        raw.execute_batch("DROP TABLE chat_state").unwrap();
        let store = ChatStore::new(pool);

        let first = store.get(7).await;
        assert!(first.forward_channel_id.is_none());
        assert!(
            !store.cache.lock().contains_key(&7),
            "a failed read must not poison the cache"
        );

        // The next get retries the DB and sees the real row.
        raw.execute_batch(
            "CREATE TABLE chat_state (chat_id TEXT PRIMARY KEY, payload TEXT NOT NULL)",
        )
        .unwrap();
        let real = ChatData {
            forward_channel_id: Some(42),
            ..ChatData::default()
        };
        raw.execute(
            "INSERT INTO chat_state (chat_id, payload) VALUES ('7', ?1)",
            rusqlite::params![serde_json::to_string(&real).unwrap()],
        )
        .unwrap();
        assert_eq!(store.get(7).await.forward_channel_id, Some(42));
    }

    #[tokio::test]
    async fn a_failed_read_does_not_overwrite_existing_state() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("update.db");
        let pool = crate::db::open_store(path.to_str().unwrap()).unwrap();
        let raw = rusqlite::Connection::open(&path).unwrap();
        let stored = "{\"forward_channel_id\":";
        raw.execute(
            "INSERT INTO chat_state (chat_id, payload) VALUES ('7', ?1)",
            rusqlite::params![stored],
        )
        .unwrap();
        let store = ChatStore::new(pool);
        let mut called = false;
        let result = store
            .update(7, |data| {
                called = true;
                data.message_format.insert("twitter".into(), "{url}".into());
            })
            .await;

        assert!(result.is_err());
        assert!(!called, "a failed load must not run a destructive mutation");
        assert!(!store.cache.lock().contains_key(&7));
        let payload: String = raw
            .query_row(
                "SELECT payload FROM chat_state WHERE chat_id='7'",
                [],
                |row| row.get(0),
            )
            .unwrap();
        assert_eq!(payload, stored, "the original row must remain unchanged");
    }

    #[tokio::test]
    async fn prune_spares_a_lock_someone_still_holds() {
        // The sweep evicts uncontended locks only: removing one an update
        // still holds (its `lock_for` clone alive) would let a second writer
        // create a fresh lock and enter the critical section concurrently.
        let dir = tempfile::tempdir().unwrap();
        let pool = crate::db::open_store(dir.path().join("l.db").to_str().unwrap()).unwrap();
        let store = ChatStore::new(pool);
        store.set(1, &ChatData::default()).await;

        let held = store.lock_for(1); // an update between lock_for and lock().await
        store.prune_expired(Duration::from_secs(60)).await;
        assert!(
            store.locks.lock().contains_key(&1),
            "a contended lock must survive the sweep"
        );

        drop(held);
        store.prune_expired(Duration::from_secs(60)).await;
        assert!(
            !store.locks.lock().contains_key(&1),
            "the next sweep drops it once uncontended"
        );
    }
}