tauri-plugin-deeprespond-im 1.0.0-alpha

//! 跨进程边界的事务（《1.本地存储》第 5 节）：tx_begin 返回事务 ID，之后的读写带上它，在同一个连接上依次执行。
//! - 连接在专用线程中（Worker），命令只经通道把操作交过去、等回复，不阻塞 Tauri 的主线程和异步运行时；
//! - 写事务独占写连接（同一时刻一个，其他的在 tx_begin 中排队）；只读事务从 4 个只读连接中取一个，BEGIN 后先读一次，
//!   固定事务开始时的快照（WAL 下读写互不阻塞）；
//! - 两次操作之间空闲超过 10 秒的事务回滚（从上一个操作结束时算起，清空大表这类长操作的执行时间不算空闲）；
//! - 页面实例消失、数据库关闭或删除时，回滚它们的事务。
//!
//! 同一个连接上的操作按交给线程的顺序执行：结束事务时先把 ROLLBACK/COMMIT 交给线程、再放开连接，下一个事务的 BEGIN
//! 一定排在后面，所以不必等回滚执行完。

use std::collections::{HashMap, HashSet, VecDeque};
use std::panic::AssertUnwindSafe;
use std::sync::{mpsc, Arc, Mutex, Weak};
use std::time::{Duration, Instant};

use rusqlite::Connection;
use serde_json::Value;
use tokio::sync::{oneshot, OwnedSemaphorePermit};

use super::keys::{require_key, KeyRangeArg, Range};
use super::query::{self, Cursor, Query, ReadResult, WriteOp};
use super::{Caller, Core, DbFile};
use crate::error::{Error, Result};

/// 交给连接线程执行的操作。
pub(crate) type Job = Box<dyn FnOnce(&mut Connection) + Send + 'static>;

pub(crate) fn closed() -> Error {
    Error::invalid_state("db_closed")
}

/// 等连接线程的回复；线程已经结束（数据库关闭）时回复通道被丢弃。
pub(crate) async fn reply<T>(rx: oneshot::Receiver<T>) -> Result<T> {
    rx.await.map_err(|_| closed())
}

/// 持有一个连接的线程：依次执行交来的操作，通道关闭后关闭连接、结束。
pub(crate) struct Worker {
    sender: Mutex<Option<mpsc::Sender<Job>>>,
    done: Mutex<Option<mpsc::Receiver<()>>>,
}

impl Worker {
    pub fn spawn(name: String, conn: Connection) -> Result<Worker> {
        let (sender, jobs) = mpsc::channel::<Job>();
        let (done_tx, done) = mpsc::channel();
        std::thread::Builder::new()
            .name(name)
            .spawn(move || {
                let mut conn = conn;
                while let Ok(job) = jobs.recv() {
                    // 一个操作 panic 不能带走连接和后面排队的操作；它的回复通道随之丢弃，调用方得到 db_closed
                    if std::panic::catch_unwind(AssertUnwindSafe(|| job(&mut conn))).is_err() {
                        log::error!("数据库线程中的操作 panic，已跳过");
                    }
                }
                drop(conn);
                let _ = done_tx.send(());
            })
            .map_err(|e| Error::internal(format!("创建数据库线程失败：{e}")))?;
        Ok(Worker { sender: Mutex::new(Some(sender)), done: Mutex::new(Some(done)) })
    }

    pub fn send(&self, job: Job) -> Result<()> {
        match self.sender.lock().unwrap().as_ref() {
            Some(sender) => sender.send(job).map_err(|_| closed()),
            None => Err(closed()),
        }
    }

    pub fn call<T: Send + 'static>(&self, f: impl FnOnce(&mut Connection) -> T + Send + 'static) -> Result<oneshot::Receiver<T>> {
        let (tx, rx) = oneshot::channel();
        self.send(Box::new(move |conn| {
            let _ = tx.send(f(conn));
        }))?;
        Ok(rx)
    }

    /// 不再接受新的操作；已排队的（回滚）执行完后线程关闭连接。
    pub fn close(&self) {
        self.sender.lock().unwrap().take();
    }

    /// 等线程关闭连接（删除文件之前）；超时返回 false。
    pub fn wait(&self, timeout: Duration) -> bool {
        let Some(done) = self.done.lock().unwrap().take() else { return true };
        match done.recv_timeout(timeout) {
            Ok(()) | Err(mpsc::RecvTimeoutError::Disconnected) => true,
            Err(mpsc::RecvTimeoutError::Timeout) => {
                *self.done.lock().unwrap() = Some(done);
                false
            }
        }
    }
}

/// 事务占用的连接：写连接的许可，或一个只读连接（结束时还回池中，再放开许可）。
pub(crate) enum Lease {
    Write(#[allow(dead_code)] OwnedSemaphorePermit),
    Read { file: Arc<DbFile>, slot: usize, _permit: OwnedSemaphorePermit },
}

impl Drop for Lease {
    fn drop(&mut self) {
        if let Lease::Read { file, slot, .. } = self {
            file.return_reader(*slot);
        }
    }
}

struct TxState {
    /// 结束的原因（committed、aborted、timeout、released、closed）；None 为进行中。
    ended: Option<&'static str>,
    /// 已交给线程、还没有回复的操作数。
    active: usize,
    /// 上一个操作结束的时间：空闲从这里算起。
    last_end: Instant,
}

pub(crate) struct Tx {
    pub id: u64,
    pub db: u64,
    pub file: Arc<DbFile>,
    pub label: String,
    pub instance: u64,
    pub write: bool,
    pub stores: Vec<String>,
    slot: Option<usize>,
    lease: Mutex<Option<Lease>>,
    state: Mutex<TxState>,
}

fn ended_error(reason: &str) -> Error {
    match reason {
        "timeout" => Error::new("invalid_state", Some("tx_timeout"), "事务空闲超时，已回滚"),
        _ => Error::new("invalid_state", Some("tx_finished"), format!("事务已结束（{reason}）")),
    }
}

impl Tx {
    fn worker(&self) -> &Worker {
        match self.slot {
            None => &self.file.writer,
            Some(i) => &self.file.readers[i],
        }
    }

    /// 在事务的连接上执行一个操作。交给线程在状态锁内进行：结束事务时不会有操作插到 ROLLBACK/COMMIT 后面。
    pub async fn run<T: Send + 'static>(&self, f: impl FnOnce(&mut Connection) -> Result<T> + Send + 'static) -> Result<T> {
        let rx = {
            let mut state = self.state.lock().unwrap();
            if let Some(reason) = state.ended {
                return Err(ended_error(reason));
            }
            let rx = self.worker().call(move |conn| {
                // SQLite 在磁盘已满等错误时可能自行回滚整个事务；之后的写入不能落到自动提交模式
                if conn.is_autocommit() {
                    return Err(Error::invalid_state("tx_rolled_back"));
                }
                f(conn)
            })?;
            state.active += 1;
            rx
        };
        let _guard = OpGuard(self);
        reply(rx).await?
    }

    /// 标记结束；已经结束的返回 false。
    fn end(&self, reason: &'static str) -> bool {
        let mut state = self.state.lock().unwrap();
        if state.ended.is_some() {
            return false;
        }
        state.ended = Some(reason);
        true
    }

    fn try_expire(&self, idle: Duration) -> bool {
        let mut state = self.state.lock().unwrap();
        if state.ended.is_none() && state.active == 0 && state.last_end.elapsed() >= idle {
            state.ended = Some("timeout");
            return true;
        }
        false
    }

    /// 回滚并放开连接（不等回滚执行完，见文件开头的说明）。
    pub fn rollback(&self) {
        let _ = self.worker().send(Box::new(|conn| {
            let _ = end_sql(conn, false);
        }));
        self.lease.lock().unwrap().take();
    }

    fn release(&self) {
        self.lease.lock().unwrap().take();
    }
}

struct OpGuard<'a>(&'a Tx);

impl Drop for OpGuard<'_> {
    fn drop(&mut self) {
        let mut state = self.0.state.lock().unwrap();
        state.active -= 1;
        state.last_end = Instant::now();
    }
}

fn begin_sql(conn: &mut Connection, write: bool, strict: bool) -> Result<()> {
    // 上一个事务的结束没有执行到（不应发生）：先回滚，不能把新事务套在旧事务里
    if !conn.is_autocommit() {
        log::warn!("连接上有未结束的事务，先回滚");
        conn.execute_batch("ROLLBACK")?;
    }
    if write {
        // 发送队列、上传任务、reset_required 的清空在 FULL 下提交（与 Flutter 相同），其余 NORMAL
        conn.execute_batch(if strict { "PRAGMA synchronous = FULL" } else { "PRAGMA synchronous = NORMAL" })?;
        conn.execute_batch("BEGIN IMMEDIATE")?;
    } else {
        conn.execute_batch("BEGIN")?;
        // WAL 的快照在第一次读取时确定：现在就读一次，读到的是事务开始时已提交的数据
        if let Err(err) = conn.query_row("SELECT count(*) FROM _dr_schema", [], |_| Ok(())) {
            let _ = conn.execute_batch("ROLLBACK");
            return Err(err.into());
        }
    }
    Ok(())
}

fn end_sql(conn: &mut Connection, commit: bool) -> Result<()> {
    let result = if commit && !conn.is_autocommit() { conn.execute_batch("COMMIT").map_err(Error::from) } else { Ok(()) };
    if !conn.is_autocommit() {
        let _ = conn.execute_batch("ROLLBACK");
    }
    result
}

#[derive(Default)]
struct TableInner {
    live: HashMap<u64, Arc<Tx>>,
    /// 最近结束的事务及原因：之后的操作给出“超时”还是“已结束”。
    ended: VecDeque<(u64, &'static str)>,
    /// 已经不在的页面实例（实例 ID 全局递增，不会重复）。
    released: HashSet<u64>,
}

const ENDED_KEEP: usize = 256;

/// 进行中的事务。
#[derive(Default)]
pub(crate) struct TxTable {
    inner: Mutex<TableInner>,
}

impl TxTable {
    /// 登记；页面实例已经不在或数据库已关闭时返回 false（与 release_instance、关闭在同一个锁内判断）。
    fn insert(&self, tx: Arc<Tx>) -> bool {
        let mut inner = self.inner.lock().unwrap();
        if inner.released.contains(&tx.instance) || tx.file.is_closed() {
            return false;
        }
        inner.live.insert(tx.id, tx);
        true
    }

    fn get(&self, id: u64, caller: &Caller) -> Result<Arc<Tx>> {
        let inner = self.inner.lock().unwrap();
        match inner.live.get(&id) {
            Some(tx) if tx.instance == caller.instance && tx.label == caller.label => Ok(tx.clone()),
            Some(_) => Err(Error::invalid_state("tx_not_found")),
            None => match inner.ended.iter().find(|(i, _)| *i == id) {
                Some((_, reason)) => Err(ended_error(reason)),
                None => Err(Error::invalid_state("tx_not_found")),
            },
        }
    }

    fn remove(&self, id: u64, reason: &'static str) {
        let mut inner = self.inner.lock().unwrap();
        inner.live.remove(&id);
        push_ended(&mut inner, id, reason);
    }

    /// 取出满足条件的事务（条件中标记结束），记下结束的原因。
    fn take_where(&self, reason: &'static str, pred: impl Fn(&Tx) -> bool) -> Vec<Arc<Tx>> {
        let mut inner = self.inner.lock().unwrap();
        let ids: Vec<u64> = inner.live.values().filter(|tx| pred(tx)).map(|tx| tx.id).collect();
        let mut out = Vec::with_capacity(ids.len());
        for id in ids {
            if let Some(tx) = inner.live.remove(&id) {
                out.push(tx);
            }
            push_ended(&mut inner, id, reason);
        }
        out
    }

    pub fn mark_released(&self, instance: u64) {
        self.inner.lock().unwrap().released.insert(instance);
    }

    pub fn is_released(&self, instance: u64) -> bool {
        self.inner.lock().unwrap().released.contains(&instance)
    }

    #[cfg(test)]
    pub fn len(&self) -> usize {
        self.inner.lock().unwrap().live.len()
    }
}

fn push_ended(inner: &mut TableInner, id: u64, reason: &'static str) {
    if inner.ended.len() >= ENDED_KEEP {
        inner.ended.pop_front();
    }
    inner.ended.push_back((id, reason));
}

/// 空闲检查的线程：只持有弱引用，StoreManager 丢弃后自行结束。
pub(crate) fn start_reaper(core: Weak<Core>, idle: Duration) {
    let tick = (idle / 5).clamp(Duration::from_millis(10), Duration::from_secs(1));
    let spawned = std::thread::Builder::new().name("drim-db-idle".into()).spawn(move || loop {
        std::thread::sleep(tick);
        let Some(core) = core.upgrade() else { return };
        core.expire_idle();
    });
    if let Err(err) = spawned {
        log::error!("创建事务空闲检查线程失败：{err}");
    }
}

/// tx_read 的参数。
pub struct ReadRequest {
    pub op: String,
    pub store: String,
    pub key: Option<Value>,
    pub index: Option<String>,
    pub range: Option<KeyRangeArg>,
    pub limit: Option<u64>,
    pub direction: Option<String>,
    pub cursor: Option<String>,
}

impl Core {
    pub(crate) async fn begin(&self, caller: &Caller, db: u64, stores: Vec<String>, mode: &str, durability: Option<&str>) -> Result<u64> {
        let write = match mode {
            "readwrite" => true,
            "readonly" => false,
            _ => return Err(Error::validation("invalid_mode", format!("不认识的事务模式 {mode}"))),
        };
        // 只有引擎写入（即使 App 误把引擎的权限授予了界面窗口，《3.接口设计》第 2 节）
        if write && !caller.engine {
            return Err(Error::not_engine());
        }
        let file = self.handle_file(db, caller)?;
        if stores.is_empty() {
            return Err(Error::validation("unknown_store", "事务没有对象仓库"));
        }
        {
            let schema = file.schema();
            if let Some(name) = stores.iter().find(|s| schema.store(s).is_none()) {
                return Err(Error::validation("unknown_store", format!("没有对象仓库 {name}")));
            }
        }
        let strict = write && durability == Some("strict");
        let (lease, slot) = if write {
            // 同一时刻一个写事务：其他的在这里排队
            let permit = file.write_gate.clone().acquire_owned().await.map_err(|_| closed())?;
            (Lease::Write(permit), None)
        } else {
            let permit = file.read_gate.clone().acquire_owned().await.map_err(|_| closed())?;
            let slot = file.take_reader().ok_or_else(closed)?;
            (Lease::Read { file: file.clone(), slot, _permit: permit }, Some(slot))
        };
        let tx = Arc::new(Tx {
            id: self.next_id(),
            db,
            file,
            label: caller.label.clone(),
            instance: caller.instance,
            write,
            stores,
            slot,
            lease: Mutex::new(Some(lease)),
            state: Mutex::new(TxState { ended: None, active: 0, last_end: Instant::now() }),
        });
        // 还没有登记，不会有其他操作和空闲检查：直接交给线程
        let rx = tx.worker().call(move |conn| begin_sql(conn, write, strict))?;
        reply(rx).await??;
        tx.state.lock().unwrap().last_end = Instant::now();
        if !self.txs.insert(tx.clone()) {
            // 排队期间页面已经不在了，或数据库被关闭、删除
            tx.end("released");
            tx.rollback();
            return Err(Error::aborted());
        }
        Ok(tx.id)
    }

    pub(crate) async fn read(&self, caller: &Caller, id: u64, req: ReadRequest) -> Result<ReadResult> {
        let tx = self.txs.get(id, caller)?;
        let schema = tx.file.schema();
        let def = query::store_in_scope(&schema, &tx.stores, &req.store)?;
        let column = match &req.index {
            Some(name) => Some(def.index(name).ok_or_else(|| Error::validation("unknown_index", format!("{} 上没有索引 {name}", req.store)))?.column.clone()),
            None => None,
        };
        let range = match (&req.range, &req.key) {
            (Some(range), _) => range.encode()?,
            // get_all、count、iterate 带 key 时只取这一个键
            (None, Some(key)) if req.op != "get" => {
                let key = require_key(key)?;
                Range { lower: Some((key.clone(), false)), upper: Some((key, false)) }
            }
            _ => Range::default(),
        };
        let reverse = match req.direction.as_deref() {
            None | Some("next") => false,
            Some("prev") => true,
            Some(other) => return Err(Error::validation("invalid_direction", format!("不认识的方向 {other}"))),
        };
        let store = req.store.clone();
        match req.op.as_str() {
            "get" => {
                let key = require_key(req.key.as_ref().unwrap_or(&Value::Null))?;
                tx.run(move |conn| query::get(conn, &store, &key).map(ReadResult::One)).await
            }
            "get_all" | "count" | "iterate" => {
                let after = match (&req.cursor, req.op.as_str()) {
                    (Some(text), "iterate") => Some(Cursor::parse(text, column.is_some())?),
                    _ => None,
                };
                let q = Query { store, column, range, reverse, limit: req.limit };
                match req.op.as_str() {
                    "get_all" => tx.run(move |conn| query::get_all(conn, &q).map(ReadResult::Many)).await,
                    "count" => tx.run(move |conn| query::count(conn, &q).map(ReadResult::Count)).await,
                    _ => tx.run(move |conn| query::iterate(conn, &q, after.as_ref())).await,
                }
            }
            other => Err(Error::validation("invalid_op", format!("不认识的读取 {other}"))),
        }
    }

    /// tx_write：一批写入要么全部生效、要么全部撤销；出错不结束事务，由 JavaScript 决定提交还是中止。
    pub(crate) async fn write(&self, caller: &Caller, id: u64, ops: Vec<WriteOp>) -> Result<()> {
        let tx = self.txs.get(id, caller)?;
        if !tx.write {
            return Err(Error::invalid_state("readonly_transaction"));
        }
        let prepared = query::prepare(&tx.file.schema(), &tx.stores, ops)?;
        let result = tx.run(move |conn| query::apply_batch(conn, &prepared)).await;
        self.check_full(&tx, &result);
        result
    }

    /// tx_commit：先执行 ops 再提交；任何一步失败都回滚整个事务。
    pub(crate) async fn commit(&self, caller: &Caller, id: u64, ops: Vec<WriteOp>) -> Result<()> {
        let tx = self.txs.get(id, caller)?;
        if !tx.write && !ops.is_empty() {
            return Err(Error::invalid_state("readonly_transaction"));
        }
        let prepared = query::prepare(&tx.file.schema(), &tx.stores, ops);
        {
            // 与 run 相同：在状态锁内交给线程，之后不会再有操作排进来
            let mut state = tx.state.lock().unwrap();
            if let Some(reason) = state.ended {
                return Err(ended_error(reason));
            }
            state.ended = Some("committed");
        }
        self.txs.remove(tx.id, "committed");
        let prepared = match prepared {
            Ok(p) => p,
            Err(err) => {
                tx.rollback();
                return Err(err);
            }
        };
        let rx = tx.worker().call(move |conn| match query::apply(conn, &prepared) {
            Ok(()) => end_sql(conn, true),
            Err(err) => {
                let _ = end_sql(conn, false);
                Err(err)
            }
        });
        tx.release();
        let result = match rx {
            Ok(rx) => reply(rx).await.and_then(|r| r),
            Err(err) => Err(err),
        };
        self.check_full(&tx, &result);
        result
    }

    /// tx_abort：回滚；已经结束的（超时、页面消失后的重复调用）直接返回。
    pub(crate) fn abort(&self, caller: &Caller, id: u64) {
        let Ok(tx) = self.txs.get(id, caller) else { return };
        if tx.end("aborted") {
            self.txs.remove(tx.id, "aborted");
            tx.rollback();
        }
    }

    /// 空闲超时的事务回滚（空闲检查线程调用）。
    pub(crate) fn expire_idle(&self) {
        let idle = self.idle;
        for tx in self.txs.take_where("timeout", |tx| tx.try_expire(idle)) {
            if tx.write {
                log::error!("写事务 {} 两次操作之间空闲超过 {} 毫秒，已回滚（引擎页可能卡住）", tx.id, idle.as_millis());
            } else {
                log::warn!("只读事务 {} 空闲超过 {} 毫秒，已结束并放开连接", tx.id, idle.as_millis());
            }
            tx.rollback();
        }
    }

    /// 回滚满足条件的事务（页面实例消失、句柄或数据库关闭）。
    pub(crate) fn rollback_where(&self, reason: &'static str, pred: impl Fn(&Tx) -> bool) {
        for tx in self.txs.take_where(reason, |tx| pred(tx) && tx.end(reason)) {
            log::debug!("回滚事务 {}（{reason}）", tx.id);
            tx.rollback();
        }
    }

    fn check_full(&self, tx: &Tx, result: &Result<()>) {
        if let Err(err) = result {
            if err.code == "storage_error" && err.reason.as_deref() == Some("quota_exceeded") {
                self.storage_full(&tx.file.app_key, &tx.file.username);
            }
        }
    }
}

#[cfg(test)]
impl Core {
    /// 测试用：在事务的连接上直接执行（模拟执行很久的操作）。
    pub(crate) async fn run_in_tx(&self, caller: &Caller, id: u64, f: impl FnOnce(&mut Connection) -> Result<()> + Send + 'static) -> Result<()> {
        self.txs.get(id, caller)?.run(f).await
    }
}