use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock, Weak};
use std::time::{Duration, Instant};
use serde::Serialize;
use tokio::sync::{oneshot, Semaphore};
use crate::error::{Error, Result};
use crate::paths;
use crate::secure::Secure;
pub mod commands;
mod disk;
mod keys;
mod query;
mod schema;
#[cfg(test)]
mod tests;
mod tx;
pub use query::{ReadResult, WriteOp};
pub use schema::StoreSchema;
pub use tx::ReadRequest;
pub type OnFull = Arc<dyn Fn(String, usize) + Send + Sync>;
const IDLE_TIMEOUT: Duration = Duration::from_secs(10);
const READERS: usize = 4;
const FULL_INTERVAL: Duration = Duration::from_secs(600);
#[derive(Clone, Debug)]
pub struct Caller {
pub label: String,
pub instance: u64,
pub engine: bool,
}
pub struct OpenRequest {
pub app_key: String,
pub username: String,
pub version: u32,
pub schema: Vec<StoreSchema>,
pub encrypt: bool,
}
#[derive(Debug, Serialize)]
pub struct Opened {
pub db: u64,
pub old_version: u32,
pub reset: Option<&'static str>,
}
pub(crate) struct DbFile {
pub app_key: String,
pub username: String,
pub path: PathBuf,
pub encrypted: bool,
schema: RwLock<(Arc<schema::Schema>, u32)>,
pub writer: tx::Worker,
pub readers: Vec<tx::Worker>,
free: Mutex<Vec<usize>>,
pub write_gate: Arc<Semaphore>,
pub read_gate: Arc<Semaphore>,
closed: AtomicBool,
}
impl DbFile {
#[allow(clippy::too_many_arguments)]
fn new(app_key: String, username: String, path: PathBuf, encrypted: bool, schema: Arc<schema::Schema>, version: u32, writer: tx::Worker, readers: Vec<tx::Worker>) -> Self {
let n = readers.len();
Self {
app_key,
username,
path,
encrypted,
schema: RwLock::new((schema, version)),
writer,
readers,
free: Mutex::new((0..n).collect()),
write_gate: Arc::new(Semaphore::new(1)),
read_gate: Arc::new(Semaphore::new(n)),
closed: AtomicBool::new(false),
}
}
pub fn schema(&self) -> Arc<schema::Schema> {
self.schema.read().unwrap().0.clone()
}
pub fn version(&self) -> u32 {
self.schema.read().unwrap().1
}
pub fn is_closed(&self) -> bool {
self.closed.load(Ordering::SeqCst)
}
pub fn take_reader(&self) -> Option<usize> {
self.free.lock().unwrap().pop()
}
pub fn return_reader(&self, slot: usize) {
self.free.lock().unwrap().push(slot);
}
fn is_done(&self) -> bool {
self.writer.wait(Duration::ZERO) && self.readers.iter().all(|r| r.wait(Duration::ZERO))
}
fn wait_closed(&self, timeout: Duration) -> bool {
let deadline = Instant::now() + timeout;
let mut ok = self.writer.wait(timeout);
for reader in &self.readers {
ok &= reader.wait(deadline.saturating_duration_since(Instant::now()));
}
ok
}
}
struct Handle {
file: Arc<DbFile>,
label: String,
instance: u64,
}
#[derive(Default)]
struct State {
files: HashMap<PathBuf, Arc<DbFile>>,
handles: HashMap<u64, Handle>,
}
pub(crate) struct Core {
me: Weak<Core>,
data_dir: PathBuf,
secure: Arc<Secure>,
on_full: OnFull,
idle: Duration,
next: AtomicU64,
open_lock: tokio::sync::Mutex<()>,
state: Mutex<State>,
closing: Mutex<Vec<Arc<DbFile>>>,
txs: tx::TxTable,
full_at: Mutex<HashMap<String, Instant>>,
}
pub struct StoreManager {
core: Arc<Core>,
}
impl StoreManager {
pub fn new(data_dir: PathBuf, secure: Arc<Secure>, on_full: OnFull) -> Self {
Self::with_idle(data_dir, secure, on_full, IDLE_TIMEOUT)
}
fn with_idle(data_dir: PathBuf, secure: Arc<Secure>, on_full: OnFull, idle: Duration) -> Self {
let core = Arc::new_cyclic(|me| Core {
me: me.clone(),
data_dir,
secure,
on_full,
idle,
next: AtomicU64::new(0),
open_lock: tokio::sync::Mutex::new(()),
state: Mutex::new(State::default()),
closing: Mutex::new(Vec::new()),
txs: tx::TxTable::default(),
full_at: Mutex::new(HashMap::new()),
});
tx::start_reaper(Arc::downgrade(&core), idle);
Self { core }
}
pub async fn open(&self, caller: &Caller, req: OpenRequest) -> Result<Opened> {
self.core.open(caller, req).await
}
pub fn close(&self, caller: &Caller, db: u64) {
self.core.close_handle(caller, db);
}
pub async fn delete(&self, app_key: &str, username: &str) -> Result<()> {
self.core.delete(app_key, username).await
}
pub async fn begin(&self, caller: &Caller, db: u64, stores: Vec<String>, mode: &str, durability: Option<&str>) -> Result<u64> {
self.core.begin(caller, db, stores, mode, durability).await
}
pub async fn read(&self, caller: &Caller, tx: u64, req: ReadRequest) -> Result<ReadResult> {
self.core.read(caller, tx, req).await
}
pub async fn write(&self, caller: &Caller, tx: u64, ops: Vec<WriteOp>) -> Result<()> {
self.core.write(caller, tx, ops).await
}
pub async fn commit(&self, caller: &Caller, tx: u64, ops: Vec<WriteOp>) -> Result<()> {
self.core.commit(caller, tx, ops).await
}
pub fn abort(&self, caller: &Caller, tx: u64) {
self.core.abort(caller, tx);
}
pub fn release_instance(&self, label: &str, instance: u64) {
let core = &self.core;
core.txs.mark_released(instance);
core.rollback_where("released", |tx| tx.instance == instance && tx.label == label);
let files = {
let mut state = core.state.lock().unwrap();
let ids: Vec<u64> = state.handles.iter().filter(|(_, h)| h.instance == instance && h.label == label).map(|(id, _)| *id).collect();
let mut files: Vec<Arc<DbFile>> = Vec::new();
for id in ids {
if let Some(h) = state.handles.remove(&id) {
if !files.iter().any(|f| Arc::ptr_eq(f, &h.file)) {
files.push(h.file);
}
}
}
files.retain(|f| detach_unused(&mut state, f));
files
};
for file in files {
core.shutdown_file(&file);
}
}
pub fn close_all(&self) {
let core = &self.core;
let files: Vec<Arc<DbFile>> = {
let mut state = core.state.lock().unwrap();
state.handles.clear();
let files: Vec<Arc<DbFile>> = state.files.drain().map(|(_, f)| f).collect();
for f in &files {
f.closed.store(true, Ordering::SeqCst);
}
files
};
for file in &files {
core.shutdown_file(file);
}
let closing: Vec<Arc<DbFile>> = core.closing.lock().unwrap().drain(..).collect();
let deadline = Instant::now() + Duration::from_secs(2);
for file in &closing {
if !file.wait_closed(deadline.saturating_duration_since(Instant::now())) {
log::warn!("退出时关闭本地数据库超时");
}
}
}
pub fn sizes(&self) -> serde_json::Value {
let open: HashSet<PathBuf> = self.core.state.lock().unwrap().files.keys().cloned().collect();
disk::sizes(&self.core.data_dir, &open)
}
}
fn detach_unused(state: &mut State, file: &Arc<DbFile>) -> bool {
if state.handles.values().any(|h| Arc::ptr_eq(&h.file, file)) {
return false;
}
if state.files.get(&file.path).is_some_and(|f| Arc::ptr_eq(f, file)) {
state.files.remove(&file.path);
}
!file.closed.swap(true, Ordering::SeqCst)
}
async fn blocking<T: Send + 'static>(f: impl FnOnce() -> T + Send + 'static) -> Result<T> {
let (tx, rx) = oneshot::channel();
std::thread::Builder::new()
.name("drim-db-io".into())
.spawn(move || {
let _ = tx.send(f());
})
.map_err(|e| Error::internal(format!("创建线程失败:{e}")))?;
rx.await.map_err(|_| Error::internal("数据库的后台线程意外结束"))
}
impl Core {
fn next_id(&self) -> u64 {
self.next.fetch_add(1, Ordering::Relaxed) + 1
}
fn handle_file(&self, db: u64, caller: &Caller) -> Result<Arc<DbFile>> {
let state = self.state.lock().unwrap();
match state.handles.get(&db) {
Some(h) if h.instance == caller.instance && h.label == caller.label => {
if h.file.is_closed() {
Err(tx::closed())
} else {
Ok(h.file.clone())
}
}
_ => Err(Error::invalid_state("db_not_found")),
}
}
async fn open(&self, caller: &Caller, req: OpenRequest) -> Result<Opened> {
if req.encrypt && !cfg!(feature = "sqlcipher") {
return Err(Error::unsupported("sqlcipher_required"));
}
if req.app_key.is_empty() || req.username.is_empty() {
return Err(Error::validation("invalid_argument", "app_key、username 不能为空"));
}
if req.version == 0 {
return Err(Error::validation("invalid_version", "结构版本从 1 开始"));
}
let schema = Arc::new(schema::Schema::compile(req.schema)?);
let username = req.username.to_lowercase();
let _guard = self.open_lock.lock().await;
let path = paths::db_path(&self.data_dir, &req.app_key, &username);
loop {
let existing = self.state.lock().unwrap().files.get(&path).cloned();
let (file, old_version, reset) = match existing {
Some(file) => {
let old = reuse(&file, schema.clone(), req.version, req.encrypt).await?;
(file, old, None)
}
None => {
let core = self.me.upgrade().ok_or_else(tx::closed)?;
let params = disk::OpenParams {
app_key: req.app_key.clone(),
username: username.clone(),
version: req.version,
schema: schema.clone(),
encrypt: req.encrypt,
touch: caller.engine,
};
let opened = blocking(move || disk::open_file(&core, params)).await??;
(Arc::new(opened.file), opened.old_version, opened.reset)
}
};
let db = self.next_id();
{
let mut state = self.state.lock().unwrap();
if file.is_closed() {
continue;
}
if !self.txs.is_released(caller.instance) {
state.files.entry(path.clone()).or_insert_with(|| file.clone());
state.handles.insert(db, Handle { file: file.clone(), label: caller.label.clone(), instance: caller.instance });
return Ok(Opened { db, old_version, reset });
}
state.files.entry(path.clone()).or_insert_with(|| file.clone());
if !detach_unused(&mut state, &file) {
return Err(Error::aborted());
}
}
self.shutdown_file(&file);
return Err(Error::aborted());
}
}
fn close_handle(&self, caller: &Caller, db: u64) {
let detached = {
let mut state = self.state.lock().unwrap();
match state.handles.get(&db) {
Some(h) if h.instance == caller.instance && h.label == caller.label => {}
_ => return,
}
let handle = state.handles.remove(&db).expect("刚查到的句柄");
detach_unused(&mut state, &handle.file).then_some(handle.file)
};
self.rollback_where("closed", |tx| tx.db == db);
if let Some(file) = detached {
self.shutdown_file(&file);
}
}
fn shutdown_file(&self, file: &Arc<DbFile>) {
file.write_gate.close();
file.read_gate.close();
self.rollback_where("closed", |tx| Arc::ptr_eq(&tx.file, file));
file.writer.close();
for reader in &file.readers {
reader.close();
}
let mut closing = self.closing.lock().unwrap();
closing.retain(|f| !f.is_done());
closing.push(file.clone());
}
fn wait_closing(&self, path: &std::path::Path) {
let same: Vec<Arc<DbFile>> = {
let mut closing = self.closing.lock().unwrap();
let (same, rest) = closing.drain(..).partition(|f| f.path == path);
*closing = rest;
same
};
for file in same {
if !file.wait_closed(Duration::from_secs(5)) {
log::warn!("等待数据库连接关闭超时");
}
}
}
async fn delete(&self, app_key: &str, username: &str) -> Result<()> {
if app_key.is_empty() || username.is_empty() {
return Err(Error::validation("invalid_argument", "app_key、username 不能为空"));
}
let _guard = self.open_lock.lock().await;
let username = username.to_lowercase();
let path = paths::db_path(&self.data_dir, app_key, &username);
let file = {
let mut state = self.state.lock().unwrap();
let file = state.files.remove(&path);
if let Some(f) = &file {
f.closed.store(true, Ordering::SeqCst);
}
file
};
if let Some(file) = &file {
self.shutdown_file(file);
}
let core = self.me.upgrade().ok_or_else(tx::closed)?;
let app_key = app_key.to_string();
drop(file);
blocking(move || {
core.wait_closing(&path);
disk::delete_user(&core, &app_key, &username)
})
.await?
}
}
async fn reuse(file: &Arc<DbFile>, schema: Arc<schema::Schema>, version: u32, encrypt: bool) -> Result<u32> {
let current = file.version();
if encrypt != file.encrypted {
log::warn!("同一个数据库以不同的加密设置打开,沿用第一次打开时的设置");
}
if version < current {
return Err(Error::new("invalid_state", Some("version_too_old"), format!("数据库已按结构版本 {current} 打开,请求的是 {version}")));
}
if version > current {
let _permit = file.write_gate.clone().acquire_owned().await.map_err(|_| tx::closed())?;
let next = schema.clone();
let rx = file.writer.call(move |conn| schema::upgrade(conn, &next, version).map_err(Error::from))?;
tx::reply(rx).await??;
*file.schema.write().unwrap() = (schema, version);
}
Ok(current)
}