feat: fix write concurrency with lock

This commit is contained in:
kaichaosun 2026-07-10 18:03:49 +08:00
parent da11c6f51b
commit d891cdbfa4
No known key found for this signature in database
GPG Key ID: 223E0F992F4F03BF
2 changed files with 27 additions and 2 deletions

View File

@ -363,7 +363,7 @@ fn start_local_server(config: &Config) -> Result<LocalServer> {
.args(["--bind", &address.to_string(), "--db"])
.arg(&db_path)
.stdout(Stdio::null())
.stderr(Stdio::null())
.stderr(Stdio::inherit())
.spawn()
.with_context(|| format!("start local server {}", config.server_bin.display()))?;

View File

@ -4,9 +4,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result};
use sqlx::SqlitePool;
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions};
use tokio::sync::Mutex;
pub struct Store {
pool: SqlitePool,
write_lock: Mutex<()>,
}
#[derive(Debug, Clone, sqlx::FromRow)]
@ -72,10 +74,14 @@ impl Store {
.await
.context("run migrations")?;
Ok(Self { pool })
Ok(Self {
pool,
write_lock: Mutex::new(()),
})
}
pub async fn insert(&self, device_id: &str, bundle: &StoredKeyPackageBundle) -> Result<()> {
let _write_guard = self.write_lock.lock().await;
let received_at = now_ms() as i64;
sqlx::query(
"INSERT INTO keypackages (device_id, received_at, payload, signature)
@ -127,6 +133,10 @@ impl Store {
lamport: u64,
bundle: &StoredAccountBundle,
) -> Result<bool> {
// SQLite allows only one writer. This transaction reads first to enforce
// Lamport ordering, so concurrent mutations can otherwise race to
// upgrade their shared read locks and fail with SQLITE_BUSY.
let _write_guard = self.write_lock.lock().await;
let updated_at = now_ms() as i64;
let mut tx = self.pool.begin().await?;
@ -175,6 +185,7 @@ impl Store {
/// Drops account bundles that have not been refreshed within `retention`.
pub async fn prune_accounts(&self, retention: Duration) -> Result<()> {
let _write_guard = self.write_lock.lock().await;
let cutoff_ms = now_ms().saturating_sub(retention.as_millis() as u64) as i64;
sqlx::query("DELETE FROM account_bundles WHERE updated_at < ?")
.bind(cutoff_ms)
@ -191,6 +202,7 @@ impl Store {
max_per_identity: usize,
retention: Duration,
) -> Result<()> {
let _write_guard = self.write_lock.lock().await;
let cutoff_ms = now_ms().saturating_sub(retention.as_millis() as u64) as i64;
sqlx::query("DELETE FROM keypackages WHERE received_at < ?")
.bind(cutoff_ms)
@ -303,6 +315,19 @@ mod tests {
assert_eq!(after_first, after_replay);
}
#[tokio::test]
async fn concurrent_account_publishes_succeed() {
let store = std::sync::Arc::new(Store::open(Path::new(":memory:")).await.unwrap());
let mut tasks = tokio::task::JoinSet::new();
for index in 0..16 {
let store = store.clone();
tasks.spawn(async move { upsert(&store, &format!("acct-{index}"), 1).await });
}
while let Some(result) = tasks.join_next().await {
assert!(result.unwrap());
}
}
#[test]
fn payload_lamport_requires_domain_and_full_header() {
assert_eq!(payload_lamport(&payload_with_lamport(42)), Some(42));