Skip to main content

sqlite_graphrag/
lock.rs

1//! Counting semaphore via lock files to limit parallel CLI invocations.
2//!
3//! `acquire_cli_slot` tries to acquire one of `N` available slots by opening the file
4//! `cli-slot-{N}.lock` in the OS cache directory and obtaining an exclusive `flock`.
5//! The returned [`std::fs::File`] MUST be kept alive for the entire duration of `main`;
6//! dropping it releases the slot automatically for the next invocation.
7//!
8//! When `wait_seconds` is `Some(n) > 0`, the function polls every
9//! [`crate::constants::CLI_LOCK_POLL_INTERVAL_MS`] milliseconds until the deadline. When it
10//! is `None` or `Some(0)`, a single attempt is made and `Err(AppError::AllSlotsFull)` is
11//! returned immediately if all slots are occupied.
12//!
13//! ## Job-type singleton (G28-B, v1.0.68)
14//!
15//! Heavy long-running jobs (`enrich`, `ingest --mode claude-code`,
16//! `ingest --mode codex`) also acquire a *singleton* lock per `(job_type,
17//! namespace)` via `acquire_job_singleton`.  This guarantees at most one
18//! heavy job per namespace runs at any time, which was the root cause
19//! of the 2026-06-03 process-proliferation incident (4 parallel `enrich`
20//! instances × N workers × 10 MCP servers = ~192 spawned processes).
21// Workload: I/O-bound (flock polling with exponential backoff sleep)
22
23use std::fs::{File, OpenOptions};
24use std::path::{Path, PathBuf};
25use std::thread;
26use std::time::{Duration, Instant};
27
28use fs4::fs_std::FileExt;
29
30use crate::constants::{
31    CLI_LOCK_POLL_INTERVAL_MS, JOB_SINGLETON_POLL_INTERVAL_MS, LLM_WORKER_RSS_MB,
32    MAX_CONCURRENT_CLI_INSTANCES,
33};
34use crate::errors::AppError;
35
36/// Job-type classification for `acquire_job_singleton`.
37///
38/// `Light` is intentionally NOT a variant here because lightweight
39/// commands (`recall`, `stats`, `read`, `list`) share the existing
40/// counting-semaphore in [`acquire_cli_slot`] and do not need a singleton.
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum JobType {
43    /// `enrich` command (LLM-driven entity/relation/body enrichment).
44    Enrich,
45    /// `ingest --mode claude-code` (LLM-curated ingestion).
46    IngestClaudeCode,
47    /// `ingest --mode codex` (OpenAI Codex CLI ingestion).
48    IngestCodex,
49}
50
51impl JobType {
52    /// Returns the kebab-case tag used inside the lock file name.
53    fn tag(self) -> &'static str {
54        match self {
55            JobType::Enrich => "enrich",
56            JobType::IngestClaudeCode => "ingest-claude-code",
57            JobType::IngestCodex => "ingest-codex",
58        }
59    }
60}
61
62/// Returns the lock file path for the given slot.
63///
64/// Precedence: CLI `--cache-dir` > XDG `cache.dir` (`config set`) > OS default
65/// cache directory via `directories::ProjectDirs`. No product env vars.
66/// The slot must be 1-based.
67fn slot_path(slot: usize) -> Result<PathBuf, AppError> {
68    let cache = cache_dir()?;
69    std::fs::create_dir_all(&cache)?;
70    Ok(cache.join(format!("cli-slot-{slot}.lock")))
71}
72
73/// Resolves the lock-file directory from XDG config or OS ProjectDirs.
74/// Cache root for CLI slot lock files.
75///
76/// GAP-SG-94: delegates to [`crate::paths::cache_dir`] so lock files, model
77/// files and LLM slots always share one directory. This function used to read
78/// its own key `paths.cache`, which meant `config set cache.dir X` moved the
79/// models but left the locks behind.
80fn cache_dir() -> Result<PathBuf, AppError> {
81    crate::paths::cache_dir()
82}
83
84/// Computes a short, filesystem-safe hash of the database path so two distinct
85/// databases (e.g. `/tmp/a.sqlite` and `/tmp/b.sqlite`) get distinct lock
86/// files in the shared cache directory. First 12 hex chars of BLAKE3 are
87/// sufficient for collision avoidance across the local filesystem.
88pub fn db_path_hash(db_path: &Path) -> String {
89    let canonical = db_path
90        .canonicalize()
91        .unwrap_or_else(|_| db_path.to_path_buf());
92    let hash = blake3::hash(canonical.to_string_lossy().as_bytes());
93    hash.to_hex().to_string()[..12].to_string()
94}
95
96/// Returns the singleton lock file path for a given (job_type, namespace, db_hash).
97///
98/// Layout: `job-singleton-{tag}-{namespace_slug}-{db_hash}.lock` in the same
99/// cache dir as the CLI slots. The namespace is sanitised to a filesystem-safe
100/// slug (lowercase, hyphens, alphanumeric) and defaults to `default` when
101/// empty. The `db_hash` is the BLAKE3 prefix returned by [`db_path_hash`].
102///
103/// G30 (v1.0.69): the previous implementation ignored the database path
104/// entirely, so two concurrent `enrich` invocations against different
105/// `graphrag.sqlite` files (production vs. test) collided on the same
106/// cache-dir lock. The db_hash scope makes the singleton per-database while
107/// still sharing the same cache dir.
108pub fn job_singleton_path(
109    job_type: JobType,
110    namespace: &str,
111    db_hash: &str,
112) -> Result<PathBuf, AppError> {
113    let cache = cache_dir()?;
114    std::fs::create_dir_all(&cache)?;
115    let slug = if namespace.is_empty() {
116        "default".to_string()
117    } else {
118        namespace
119            .chars()
120            .map(|c| {
121                if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
122                    c.to_ascii_lowercase()
123                } else {
124                    '-'
125                }
126            })
127            .collect::<String>()
128    };
129    let safe_hash: String = db_hash
130        .chars()
131        .filter(|c| c.is_ascii_alphanumeric())
132        .take(16)
133        .collect();
134    Ok(cache.join(format!(
135        "job-singleton-{}-{slug}-{safe_hash}.lock",
136        job_type.tag()
137    )))
138}
139
140/// Tries to open and exclusively lock the lock file for the given slot.
141///
142/// Returns `Ok(file)` if the slot is free, or `Err(io::Error)` if it is
143/// held by another instance (non-blocking).
144fn try_acquire_slot(slot: usize) -> Result<File, AppError> {
145    let path = slot_path(slot)?;
146    let file = OpenOptions::new()
147        .read(true)
148        .write(true)
149        .create(true)
150        .truncate(false)
151        .open(&path)?;
152    file.try_lock_exclusive().map_err(AppError::Io)?;
153    Ok(file)
154}
155
156/// Acquires a concurrency slot from the `max_concurrency`-position semaphore.
157///
158/// Iterates slots `1..=max_concurrency` attempting `try_lock_exclusive` on each
159/// `cli-slot-N.lock` file. When a free slot is found, returns `(File, slot_number)`.
160/// If all slots are occupied:
161///
162/// - If `wait_seconds` is `None` or `Some(0)`, returns immediately with
163///   `AppError::AllSlotsFull { max, waited_secs: 0 }`.
164/// - If `wait_seconds` is `Some(n) > 0`, enters a polling loop every
165///
166/// Returns the maximum number of parallel CLI instances the host can sustain
167/// without thrashing. The formula:
168///
169///   safe = min(cpus, available_mb / per_worker_mb) * 1.0
170///
171/// replaces the previous `... * 0.5` halving factor. The `* 0.5` was the
172/// root cause of G18: even on a 64 GB host the result was always
173/// clamped to 4 because of the division-by-2.
174///
175/// The per-worker cost is `LLM_WORKER_RSS_MB` (350): since v1.0.79 every
176/// build is LLM-only (the `embedding-legacy` feature and the ONNX path
177/// were removed), so the higher fastembed worker cost no longer applies.
178///
179/// Returns 1 as a defensive floor when system stats are unavailable.
180pub fn calculate_safe_concurrency() -> usize {
181    use sysinfo::System;
182    let mut sys = System::new();
183    sys.refresh_memory();
184    let available_mb = sys.available_memory() / 1_048_576;
185    let cpus = std::thread::available_parallelism()
186        .map(|n| n.get())
187        .unwrap_or(2);
188
189    let per_worker_mb = LLM_WORKER_RSS_MB;
190
191    let memory_bound = if available_mb == 0 {
192        cpus
193    } else {
194        (available_mb / per_worker_mb.max(1)) as usize
195    };
196    let raw = cpus.min(memory_bound).max(1);
197    raw.min(MAX_CONCURRENT_CLI_INSTANCES)
198}
199
200/// v1.0.75 — Returns the worker cost in MiB used by `calculate_safe_concurrency`.
201/// Exposed for telemetry and `--info` output.
202pub fn worker_cost_mb() -> u64 {
203    LLM_WORKER_RSS_MB
204}
205
206///   `AppError::AllSlotsFull { max, waited_secs: n }` if no slot opens.
207///
208/// The returned `File` MUST be kept alive until the process exits; dropping it
209/// releases the slot automatically via the implicit `flock` on close.
210pub fn acquire_cli_slot(
211    max_concurrency: usize,
212    wait_seconds: Option<u64>,
213) -> Result<(File, usize), AppError> {
214    // G18: use env override or 2*cpus as ceiling instead of hardcoded 4
215    let ncpus = std::thread::available_parallelism()
216        .map(|n| n.get())
217        .unwrap_or(4);
218    let ceiling = crate::config::get_setting("cli.max_instances")
219        .ok()
220        .flatten()
221        .and_then(|v| v.parse::<usize>().ok())
222        .unwrap_or_else(|| (2 * ncpus).max(MAX_CONCURRENT_CLI_INSTANCES));
223    let max = max_concurrency.clamp(1, ceiling);
224    let wait_secs = wait_seconds.unwrap_or(0);
225
226    // Tentativa inicial sem espera.
227    if let Some((file, slot)) = try_any_slot(max)? {
228        return Ok((file, slot));
229    }
230
231    if wait_secs == 0 {
232        return Err(AppError::AllSlotsFull {
233            max,
234            waited_secs: 0,
235        });
236    }
237
238    // Polling loop with progressive backoff until the deadline.
239    let deadline = Instant::now() + Duration::from_secs(wait_secs);
240    let mut polls: u64 = 0;
241    loop {
242        let poll_delay = CLI_LOCK_POLL_INTERVAL_MS
243            .saturating_mul(1 + polls / 4)
244            .min(CLI_LOCK_POLL_INTERVAL_MS * 4);
245        thread::sleep(Duration::from_millis(poll_delay));
246        polls += 1;
247        if let Some((file, slot)) = try_any_slot(max)? {
248            return Ok((file, slot));
249        }
250        if Instant::now() >= deadline {
251            return Err(AppError::AllSlotsFull {
252                max,
253                waited_secs: wait_secs,
254            });
255        }
256    }
257}
258
259/// Acquires a process-wide singleton lock for a heavy job type and namespace.
260///
261/// G28-B (v1.0.68): ensures at most one `enrich`, `ingest --mode
262/// claude-code`, or `ingest --mode codex` runs at a time per namespace.
263/// A second invocation in the same namespace either:
264///
265/// - Returns immediately with `AppError::JobSingletonLocked { job_type,
266///   namespace }` when `wait_seconds` is `None` or `Some(0)`.
267/// - Polls every [`JOB_SINGLETON_POLL_INTERVAL_MS`] ms until the lock
268///   drops or the deadline expires, returning the same error on timeout.
269///
270/// The returned `File` MUST be kept alive until the process exits;
271/// dropping it releases the singleton for the next invocation.
272pub fn acquire_job_singleton(
273    job_type: JobType,
274    namespace: &str,
275    db_path: &Path,
276    wait_seconds: Option<u64>,
277    force: bool,
278) -> Result<File, AppError> {
279    let db_hash = db_path_hash(db_path);
280    let path = job_singleton_path(job_type, namespace, &db_hash)?;
281
282    // G30+G09: when --force is set, attempt to break a stale lock by
283    // detecting and removing a pre-existing lock file. This is a last
284    // resort: only enabled by an explicit operator flag. A real orphan
285    // lock from a previous crash leaves a 0-byte file behind, which the
286    // next non-forced caller would still try to lock.
287    if force && path.exists() {
288        tracing::warn!(target: "lock",
289            path = %path.display(),
290            "force=true; removing pre-existing singleton lock file"
291        );
292        let _ = std::fs::remove_file(&path);
293    }
294
295    let file = OpenOptions::new()
296        .read(true)
297        .write(true)
298        .create(true)
299        .truncate(false)
300        .open(&path)?;
301    if let Err(e) = file.try_lock_exclusive() {
302        if !is_lock_contended(&e) {
303            return Err(AppError::Io(e));
304        }
305        // Already held by another instance.
306        let wait_secs = wait_seconds.unwrap_or(0);
307        if wait_secs == 0 {
308            return Err(AppError::JobSingletonLocked {
309                job_type: job_type.tag().to_string(),
310                namespace: namespace.to_string(),
311            });
312        }
313        let deadline = Instant::now() + Duration::from_secs(wait_secs);
314        // Drop the failed handle before polling; flock is per-process so we
315        // re-open each attempt to refresh contention state.
316        drop(file);
317        loop {
318            thread::sleep(Duration::from_millis(JOB_SINGLETON_POLL_INTERVAL_MS));
319            let file = OpenOptions::new()
320                .read(true)
321                .write(true)
322                .create(true)
323                .truncate(false)
324                .open(&path)?;
325            if file.try_lock_exclusive().is_ok() {
326                return Ok(file);
327            }
328            if Instant::now() >= deadline {
329                return Err(AppError::JobSingletonLocked {
330                    job_type: job_type.tag().to_string(),
331                    namespace: namespace.to_string(),
332                });
333            }
334        }
335    }
336    Ok(file)
337}
338
339/// G45: returns the lock file path for the embedding singleton
340/// of a `(namespace, db_hash)` pair. Layout:
341/// `embed-singleton-{namespace_slug}-{db_hash}.lock` in the same
342/// cache directory as the other singletons. The namespace is sanitised
343/// to a filesystem-safe slug the same way as [`job_singleton_path`].
344fn embedding_singleton_path(namespace: &str, db_hash: &str) -> Result<PathBuf, AppError> {
345    let cache = cache_dir()?;
346    std::fs::create_dir_all(&cache)?;
347    let slug = if namespace.is_empty() {
348        "default".to_string()
349    } else {
350        namespace
351            .chars()
352            .map(|c| {
353                if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
354                    c.to_ascii_lowercase()
355                } else {
356                    '-'
357                }
358            })
359            .collect::<String>()
360    };
361    let safe_hash: String = db_hash
362        .chars()
363        .filter(|c| c.is_ascii_alphanumeric())
364        .take(16)
365        .collect();
366    Ok(cache.join(format!("embed-singleton-{slug}-{safe_hash}.lock")))
367}
368
369/// G45: acquires a cross-process singleton lock for LLM embedding
370/// operations against a given `(namespace, db)` pair.
371///
372/// The lock is opened and held with `flock` (same mechanism as
373/// [`acquire_job_singleton`]). Two CLI invocations writing to the same
374/// database while both are calling the LLM on entity names will now
375/// serialise: the second one receives [`AppError::EmbeddingSingletonLocked`]
376/// (exit 75) instead of double-spawning `claude -p` / `codex exec`
377/// subprocesses.
378///
379/// Behaviour:
380/// - `wait_seconds = Some(0)` or `None` → fail immediately if held.
381/// - `wait_seconds = Some(n) > 0` → poll every
382///   [`JOB_SINGLETON_POLL_INTERVAL_MS`] ms until the lock drops or the
383///   deadline expires.
384/// - `force = true` → remove a stale lock file before acquiring
385///   (operator escape hatch, same contract as `acquire_job_singleton`).
386///
387/// The returned [`File`] MUST be kept alive for the duration of the
388/// embedding work; dropping it releases the singleton for the next
389/// process.
390pub fn acquire_embedding_singleton(
391    namespace: &str,
392    db_path: &Path,
393    wait_seconds: Option<u64>,
394    force: bool,
395) -> Result<File, AppError> {
396    let db_hash = db_path_hash(db_path);
397    let path = embedding_singleton_path(namespace, &db_hash)?;
398
399    if force && path.exists() {
400        tracing::warn!(target: "lock.g45",
401            path = %path.display(),
402            "force=true; removing pre-existing embedding singleton lock file"
403        );
404        let _ = std::fs::remove_file(&path);
405    }
406
407    let file = OpenOptions::new()
408        .read(true)
409        .write(true)
410        .create(true)
411        .truncate(false)
412        .open(&path)?;
413    if let Err(e) = file.try_lock_exclusive() {
414        if !is_lock_contended(&e) {
415            return Err(AppError::Io(e));
416        }
417        let wait_secs = wait_seconds.unwrap_or(0);
418        if wait_secs == 0 {
419            return Err(AppError::EmbeddingSingletonLocked {
420                namespace: namespace.to_string(),
421            });
422        }
423        let deadline = Instant::now() + Duration::from_secs(wait_secs);
424        drop(file);
425        loop {
426            thread::sleep(Duration::from_millis(JOB_SINGLETON_POLL_INTERVAL_MS));
427            let file = OpenOptions::new()
428                .read(true)
429                .write(true)
430                .create(true)
431                .truncate(false)
432                .open(&path)?;
433            if file.try_lock_exclusive().is_ok() {
434                return Ok(file);
435            }
436            if Instant::now() >= deadline {
437                return Err(AppError::EmbeddingSingletonLocked {
438                    namespace: namespace.to_string(),
439                });
440            }
441        }
442    }
443    Ok(file)
444}
445
446/// Tries to acquire any free slot in `1..=max`, returning the first available one.
447///
448/// Returns `Ok(Some((file, slot)))` if a slot was obtained, `Ok(None)` if all are
449/// occupied (`EWOULDBLOCK`). Propagates I/O errors other than "lock contended".
450fn try_any_slot(max: usize) -> Result<Option<(File, usize)>, AppError> {
451    for slot in 1..=max {
452        match try_acquire_slot(slot) {
453            Ok(file) => return Ok(Some((file, slot))),
454            Err(AppError::Io(e)) if is_lock_contended(&e) => continue,
455            Err(e) => return Err(e),
456        }
457    }
458    Ok(None)
459}
460
461fn is_lock_contended(error: &std::io::Error) -> bool {
462    if error.kind() == std::io::ErrorKind::WouldBlock {
463        return true;
464    }
465
466    #[cfg(windows)]
467    {
468        matches!(error.raw_os_error(), Some(32 | 33))
469    }
470
471    #[cfg(not(windows))]
472    {
473        false
474    }
475}
476
477#[cfg(test)]
478mod tests {
479    use super::*;
480    use std::sync::atomic::{AtomicUsize, Ordering};
481    static SEQ: AtomicUsize = AtomicUsize::new(0);
482
483    fn unique_ns() -> String {
484        let n = SEQ.fetch_add(1, Ordering::SeqCst);
485        let pid = std::process::id();
486        format!("test-{pid}-{n}")
487    }
488
489    #[test]
490    fn job_singleton_path_sanitises_namespace() {
491        let p = job_singleton_path(JobType::Enrich, "Foo Bar/Baz", "abc123def456")
492            .expect("path should resolve");
493        let name = p.file_name().unwrap().to_string_lossy().to_string();
494        assert!(name.contains("enrich"), "got {name}");
495        assert!(name.contains("foo-bar-baz"), "got {name}");
496        assert!(
497            name.contains("abc123def456"),
498            "must embed db_hash: got {name}"
499        );
500    }
501
502    #[test]
503    fn job_singleton_blocks_second_invocation_same_namespace() {
504        let ns = unique_ns();
505        let db = std::env::temp_dir().join(format!("test-{}.sqlite", unique_ns()));
506        let first = acquire_job_singleton(JobType::Enrich, &ns, &db, Some(0), false)
507            .expect("first acquire should succeed");
508        let second = acquire_job_singleton(JobType::Enrich, &ns, &db, Some(0), false);
509        assert!(
510            matches!(second, Err(AppError::JobSingletonLocked { .. })),
511            "expected JobSingletonLocked, got {second:?}"
512        );
513        drop(first);
514    }
515
516    #[test]
517    fn job_singleton_allows_different_namespaces() {
518        let ns_a = unique_ns();
519        let ns_b = unique_ns();
520        let db_a = std::env::temp_dir().join(format!("test-a-{}.sqlite", unique_ns()));
521        let db_b = std::env::temp_dir().join(format!("test-b-{}.sqlite", unique_ns()));
522        let first = acquire_job_singleton(JobType::IngestClaudeCode, &ns_a, &db_a, Some(0), false)
523            .expect("ns_a should acquire");
524        let second = acquire_job_singleton(JobType::IngestClaudeCode, &ns_b, &db_b, Some(0), false)
525            .expect("ns_b should acquire in parallel");
526        drop(first);
527        drop(second);
528    }
529
530    #[test]
531    fn job_singleton_scoped_by_db_hash() {
532        // G30: two databases, same namespace, different content. Both locks
533        // should succeed because the db_hash differs.
534        let ns = unique_ns();
535        let db_a = std::env::temp_dir().join(format!("test-x-{}.sqlite", unique_ns()));
536        let db_b = std::env::temp_dir().join(format!("test-y-{}.sqlite", unique_ns()));
537        let first = acquire_job_singleton(JobType::Enrich, &ns, &db_a, Some(0), false)
538            .expect("db_a should acquire");
539        let second = acquire_job_singleton(JobType::Enrich, &ns, &db_b, Some(0), false)
540            .expect("db_b should acquire independently (G30 fix)");
541        drop(first);
542        drop(second);
543    }
544
545    #[test]
546    fn db_path_hash_is_stable_for_same_path() {
547        let p = std::env::temp_dir().join("hashing-test.sqlite");
548        let h1 = db_path_hash(&p);
549        let h2 = db_path_hash(&p);
550        assert_eq!(h1, h2, "same path must produce same hash");
551        assert_eq!(h1.len(), 12, "BLAKE3 prefix must be 12 hex chars");
552    }
553
554    #[test]
555    fn db_path_hash_differs_for_different_paths() {
556        let a = std::env::temp_dir().join("hash-a.sqlite");
557        let b = std::env::temp_dir().join("hash-b.sqlite");
558        assert_ne!(db_path_hash(&a), db_path_hash(&b));
559    }
560
561    // G45: embedding singleton — cross-process coordination
562    #[test]
563    fn g45_embedding_singleton_blocks_second_invocation_same_db() {
564        let ns = unique_ns();
565        let db = std::env::temp_dir().join(format!("g45-{}.sqlite", unique_ns()));
566        let first = acquire_embedding_singleton(&ns, &db, Some(0), false)
567            .expect("first acquire should succeed");
568        let second = acquire_embedding_singleton(&ns, &db, Some(0), false);
569        assert!(
570            matches!(second, Err(AppError::EmbeddingSingletonLocked { .. })),
571            "expected EmbeddingSingletonLocked, got {second:?}"
572        );
573        drop(first);
574    }
575
576    #[test]
577    fn g45_embedding_singleton_allows_different_namespaces() {
578        let ns_a = unique_ns();
579        let ns_b = unique_ns();
580        let db = std::env::temp_dir().join(format!("g45-multi-{}.sqlite", unique_ns()));
581        let first =
582            acquire_embedding_singleton(&ns_a, &db, Some(0), false).expect("ns_a should acquire");
583        let second = acquire_embedding_singleton(&ns_b, &db, Some(0), false)
584            .expect("ns_b should acquire in parallel (different namespace)");
585        drop(first);
586        drop(second);
587    }
588
589    #[test]
590    fn g45_embedding_singleton_scoped_by_db_hash() {
591        // Same namespace, different databases → independent locks.
592        let ns = unique_ns();
593        let db_a = std::env::temp_dir().join(format!("g45-x-{}.sqlite", unique_ns()));
594        let db_b = std::env::temp_dir().join(format!("g45-y-{}.sqlite", unique_ns()));
595        let first =
596            acquire_embedding_singleton(&ns, &db_a, Some(0), false).expect("db_a should acquire");
597        let second = acquire_embedding_singleton(&ns, &db_b, Some(0), false)
598            .expect("db_b should acquire independently (G45 db_hash scope)");
599        drop(first);
600        drop(second);
601    }
602}