From d891cdbfa45c531923846f90433b6792f3394f20 Mon Sep 17 00:00:00 2001
From: kaichaosun <kaichaosuna@gmail.com>
Date: Fri, 10 Jul 2026 18:03:49 +0800
Subject: [PATCH] feat: fix write concurrency with lock

---
 examples/benchmark.rs |  2 +-
 src/store.rs          | 27 ++++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/examples/benchmark.rs b/examples/benchmark.rs
index d54708b..1d1f090 100644
--- a/examples/benchmark.rs
+++ b/examples/benchmark.rs
@@ -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()))?;
 
diff --git a/src/store.rs b/src/store.rs
index 77607c9..e20b55d 100644
--- a/src/store.rs
+++ b/src/store.rs
@@ -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));