kiromi-ai-memory 0.2.2

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! Per-tenant single-writer mutex registry.

use std::sync::Arc;

use parking_lot::Mutex as SyncMutex;
use tokio::sync::Mutex as AsyncMutex;
use tokio::sync::OwnedMutexGuard;

use std::collections::HashMap;

use crate::tenant::TenantId;

/// Process-local registry: one async mutex per `TenantId`.
#[derive(Debug, Default, Clone)]
pub struct TenantLocks {
    inner: Arc<SyncMutex<HashMap<TenantId, Arc<AsyncMutex<()>>>>>,
}

impl TenantLocks {
    /// Acquire (or create-then-acquire) the write lock for a tenant.
    pub async fn lock(&self, tenant: &TenantId) -> OwnedMutexGuard<()> {
        let mu = {
            let mut g = self.inner.lock();
            g.entry(tenant.clone())
                .or_insert_with(|| Arc::new(AsyncMutex::new(())))
                .clone()
        };
        mu.lock_owned().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn lock_serialises_writers() {
        let t = TenantId::new("a").unwrap();
        let locks = TenantLocks::default();
        let g = locks.lock(&t).await;
        let waiter = {
            let locks = locks.clone();
            let t = t.clone();
            tokio::spawn(async move { locks.lock(&t).await })
        };
        // Give the spawn a moment to start.
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        assert!(!waiter.is_finished());
        drop(g);
        let _g2 = waiter.await.unwrap();
    }
}