Skip to main content

store_sqlite/
claim_store.rs

1//! `SqliteClaimStore`: atomic work-queue with fuzzy near-duplicate detection.
2
3use std::collections::HashSet;
4use std::sync::Mutex;
5
6use chrono::Utc;
7use rusqlite::{params, Connection, OptionalExtension, TransactionBehavior};
8use strsim::levenshtein;
9use uuid::Uuid;
10
11use substrate_core::claim_port::{ClaimPort, WorkItem, WorkItemState};
12
13use crate::error::StoreError;
14use crate::schema;
15
16/// Jaccard similarity threshold for token sets.
17const JACCARD_THRESHOLD: f64 = 0.75;
18/// Normalized Levenshtein similarity threshold.
19const LEVENSHTEIN_THRESHOLD: f64 = 0.85;
20
21/// SQLite-backed [`ClaimPort`] with `BEGIN IMMEDIATE` CAS claiming.
22pub struct SqliteClaimStore {
23    conn: Mutex<Connection>,
24}
25
26impl SqliteClaimStore {
27    /// Open (or create) a claim store at the given file path.
28    pub fn open(path: &str) -> Result<Self, StoreError> {
29        let conn = Connection::open(path)?;
30        schema::init(&conn)?;
31        Ok(Self {
32            conn: Mutex::new(conn),
33        })
34    }
35
36    /// Open a transient in-memory store (useful in tests).
37    pub fn open_in_memory() -> Result<Self, StoreError> {
38        let conn = Connection::open_in_memory()?;
39        schema::init(&conn)?;
40        Ok(Self {
41            conn: Mutex::new(conn),
42        })
43    }
44}
45
46fn token_jaccard(a: &str, b: &str) -> f64 {
47    let a_lower = a.to_lowercase();
48    let b_lower = b.to_lowercase();
49    let ta: HashSet<&str> = a_lower.split_whitespace().collect();
50    let tb: HashSet<&str> = b_lower.split_whitespace().collect();
51    if ta.is_empty() && tb.is_empty() {
52        return 1.0;
53    }
54    let inter = ta.intersection(&tb).count();
55    let union = ta.union(&tb).count();
56    if union == 0 {
57        return 0.0;
58    }
59    inter as f64 / union as f64
60}
61
62fn levenshtein_similarity(a: &str, b: &str) -> f64 {
63    let dist = levenshtein(a, b);
64    let max_len = a.chars().count().max(b.chars().count()).max(1);
65    1.0 - (dist as f64 / max_len as f64)
66}
67
68/// Returns true when bodies are near-duplicates by token Jaccard or Levenshtein.
69pub fn bodies_are_near_duplicate(a: &str, b: &str) -> bool {
70    token_jaccard(a, b) >= JACCARD_THRESHOLD
71        || levenshtein_similarity(a, b) >= LEVENSHTEIN_THRESHOLD
72}
73
74impl ClaimPort for SqliteClaimStore {
75    type Error = StoreError;
76
77    fn enqueue(&self, queue: &str, body: &str) -> Result<Uuid, StoreError> {
78        if self.is_near_duplicate(queue, body)? {
79            return Err(StoreError::Duplicate(format!(
80                "near-duplicate in queue {queue}"
81            )));
82        }
83        let id = Uuid::new_v4();
84        let conn = self.conn.lock().unwrap();
85        conn.execute(
86            "INSERT INTO work_queue (id, queue, body, state, claimed_by, created_at) \
87             VALUES (?1, ?2, ?3, 'pending', NULL, ?4)",
88            params![id.to_string(), queue, body, Utc::now().to_rfc3339()],
89        )?;
90        Ok(id)
91    }
92
93    fn claim_next(&self, queue: &str, worker_id: &str) -> Result<Option<WorkItem>, StoreError> {
94        let mut conn = self.conn.lock().unwrap();
95        let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
96
97        let row: Option<(String, String, String)> = tx
98            .query_row(
99                "SELECT id, queue, body FROM work_queue \
100                 WHERE queue=?1 AND state='pending' \
101                 ORDER BY created_at ASC LIMIT 1",
102                params![queue],
103                |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
104            )
105            .optional()?;
106
107        let Some((id, q, body)) = row else {
108            tx.rollback()?;
109            return Ok(None);
110        };
111
112        let updated = tx.execute(
113            "UPDATE work_queue SET state='claimed', claimed_by=?2 \
114             WHERE id=?1 AND state='pending'",
115            params![id, worker_id],
116        )?;
117
118        if updated != 1 {
119            tx.rollback()?;
120            return Ok(None);
121        }
122
123        tx.commit()?;
124
125        Ok(Some(WorkItem {
126            id: Uuid::parse_str(&id).unwrap_or_else(|_| Uuid::new_v4()),
127            queue: q,
128            body,
129            state: WorkItemState::Claimed,
130            claimed_by: Some(worker_id.to_string()),
131        }))
132    }
133
134    fn is_near_duplicate(&self, queue: &str, body: &str) -> Result<bool, StoreError> {
135        let conn = self.conn.lock().unwrap();
136        let mut stmt = conn.prepare(
137            "SELECT body FROM work_queue \
138             WHERE queue=?1 AND state IN ('pending', 'claimed')",
139        )?;
140        let rows = stmt.query_map(params![queue], |row| row.get::<_, String>(0))?;
141        for row in rows {
142            let existing = row?;
143            if bodies_are_near_duplicate(body, &existing) {
144                return Ok(true);
145            }
146        }
147        Ok(false)
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn near_duplicate_detection() {
157        assert!(bodies_are_near_duplicate(
158            "refactor auth module login flow",
159            "refactor auth module login flo"
160        ));
161        assert!(bodies_are_near_duplicate(
162            "fix the login bug in auth module",
163            "fix the login bug in auth modul"
164        ));
165        assert!(!bodies_are_near_duplicate(
166            "implement payment gateway",
167            "write unit tests for scheduler"
168        ));
169    }
170}