llm_kernel/graph/
async_pool.rs

1//! Multi-connection async pool for the knowledge graph.
2//!
3//! Unlike `AsyncGraph` (single `Arc<Mutex<Connection>>`),
4//! this module maintains a bounded pool of rusqlite connections gated by a
5//! tokio `Semaphore`. Multiple read queries can execute concurrently in WAL
6//! mode, while the semaphore bounds total concurrency.
7//!
8//! ```no_run
9//! use llm_kernel::graph::AsyncPoolGraph;
10//!
11//! # #[tokio::main]
12//! # async fn main() -> llm_kernel::error::Result<()> {
13//! let pool = AsyncPoolGraph::open("my.db", 4).await?;
14//! pool.upsert_node(llm_kernel::graph::GraphNode {
15//!     id: "n1".into(),
16//!     node_type: "concept".into(),
17//!     title: "Example".into(),
18//!     body: String::new(),
19//!     tags: vec![],
20//!     projects: vec![],
21//!     agents: vec![],
22//!     created: "2026-01-01T00:00:00Z".into(),
23//!     updated: "2026-01-01T00:00:00Z".into(),
24//!     importance: 0.5,
25//!     access_count: 0,
26//!     accessed_at: String::new(),
27//!     ..Default::default()
28//! }).await?;
29//! # Ok(())
30//! # }
31//! ```
32
33use std::path::Path;
34use std::sync::LazyLock;
35use std::sync::{
36    Arc, Mutex,
37    atomic::{AtomicU64, Ordering},
38};
39
40use rusqlite::{Connection, OpenFlags};
41use tokio::sync::Semaphore;
42use tokio::task;
43
44/// Monotonic counter for unique shared-memory database names.
45static MEM_POOL_ID: LazyLock<AtomicU64> = LazyLock::new(|| AtomicU64::new(0));
46
47use crate::error::{KernelError, Result};
48use crate::graph::types::{EdgeDirection, GraphEdge, GraphNode, GraphStats, ScoredNode};
49
50// ── Pool inner state ────────────────────────────────
51
52struct PoolInner {
53    idle: Mutex<Vec<Connection>>,
54    path: String,
55    /// true for in-memory pools — uses shared-cache URI so all connections
56    /// see the same data.
57    shared_mem: bool,
58}
59
60impl PoolInner {
61    fn take(&self) -> Result<Connection> {
62        if let Ok(mut guard) = self.idle.lock()
63            && let Some(conn) = guard.pop()
64        {
65            return Ok(conn);
66        }
67        let mut conn = if self.shared_mem {
68            Connection::open_with_flags(
69                &self.path,
70                OpenFlags::SQLITE_OPEN_READ_WRITE
71                    | OpenFlags::SQLITE_OPEN_CREATE
72                    | OpenFlags::SQLITE_OPEN_URI
73                    | OpenFlags::SQLITE_OPEN_NO_MUTEX,
74            )
75            .map_err(|e| KernelError::Store(e.to_string()))?
76        } else {
77            Connection::open(&self.path).map_err(|e| KernelError::Store(e.to_string()))?
78        };
79        // busy_timeout is per-connection and does NOT persist, so every newly
80        // opened connection must set it — otherwise concurrent writers get an
81        // immediate SQLITE_BUSY instead of waiting (see open() for WAL setup).
82        apply_concurrency_pragmas(&mut conn)?;
83        Ok(conn)
84    }
85
86    fn return_conn(&self, conn: Connection) {
87        if let Ok(mut guard) = self.idle.lock() {
88            guard.push(conn);
89        }
90        // If lock fails (poisoned), drop the connection — it will be recreated on next take.
91    }
92}
93
94/// Per-connection concurrency PRAGMAs. `journal_mode = WAL` is set once on the
95/// first file connection in [`AsyncPoolGraph::open`] (it persists to the DB
96/// file, so every later connection inherits it); `busy_timeout` and
97/// `synchronous` do **not** persist and must be applied to each connection.
98/// Without these the pool runs under the default DELETE journal with no busy
99/// timeout — writes block readers and concurrent writers fail immediately with
100/// SQLITE_BUSY, defeating the pool's reason to exist.
101fn apply_concurrency_pragmas(conn: &mut Connection) -> Result<()> {
102    conn.execute_batch(
103        "PRAGMA busy_timeout = 5000;\n\
104         PRAGMA synchronous = NORMAL;",
105    )
106    .map_err(|e| KernelError::Store(format!("PRAGMA failed: {e}")))?;
107    Ok(())
108}
109
110// ── AsyncPoolGraph ──────────────────────────────────
111
112/// Bounded async connection pool for the knowledge graph.
113///
114/// Uses a `Semaphore` to bound concurrency and a `Mutex<Vec<Connection>>`
115/// for idle connection reuse. Each method acquires a permit, takes (or creates)
116/// a connection, runs the operation via `spawn_blocking`, then returns the
117/// connection to the pool.
118#[derive(Clone)]
119pub struct AsyncPoolGraph {
120    inner: Arc<PoolInner>,
121    sem: Arc<Semaphore>,
122}
123
124impl AsyncPoolGraph {
125    /// Open (or create) a database and initialise the graph schema.
126    ///
127    /// `max_conns` bounds the number of concurrent operations.
128    pub async fn open(path: impl AsRef<Path>, max_conns: usize) -> Result<Self> {
129        let path_str = path
130            .as_ref()
131            .to_str()
132            .ok_or_else(|| KernelError::Store("invalid path".into()))?
133            .to_string();
134
135        // Create parent dirs + open first connection + apply schema
136        let path_for_open = path_str.clone();
137        let first_conn = task::spawn_blocking(move || -> Result<Connection> {
138            if let Some(parent) = Path::new(&path_for_open).parent() {
139                std::fs::create_dir_all(parent)?;
140            }
141            let mut conn =
142                Connection::open(&path_for_open).map_err(|e| KernelError::Store(e.to_string()))?;
143            crate::graph::schema::init_graph_schema(&conn)?;
144            // WAL persists to the DB file, so all later pool connections inherit
145            // it; busy_timeout + synchronous are per-connection (set on each via
146            // apply_concurrency_pragmas). Without WAL the module's "concurrent
147            // reads during writes" claim does not hold.
148            conn.execute_batch("PRAGMA journal_mode = WAL;")
149                .map_err(|e| KernelError::Store(format!("PRAGMA failed: {e}")))?;
150            apply_concurrency_pragmas(&mut conn)?;
151            Ok(conn)
152        })
153        .await
154        .map_err(|e| KernelError::Store(e.to_string()))??;
155
156        let inner = Arc::new(PoolInner {
157            idle: Mutex::new(vec![first_conn]),
158            path: path_str,
159            shared_mem: false,
160        });
161
162        Ok(Self {
163            inner,
164            sem: Arc::new(Semaphore::new(max_conns.max(1))),
165        })
166    }
167
168    /// Create an in-memory pool with schema applied. Useful for tests.
169    ///
170    /// Uses SQLite shared-cache mode so all connections in the pool see the
171    /// same data (plain `:memory:` creates an independent DB per connection).
172    pub async fn open_in_memory(max_conns: usize) -> Result<Self> {
173        let id = MEM_POOL_ID.fetch_add(1, Ordering::Relaxed);
174        let uri = format!("file:llm_kernel_pool_{id}?mode=memory&cache=shared");
175        let uri_clone = uri.clone();
176        let conn = task::spawn_blocking(move || -> Result<Connection> {
177            let mut conn = Connection::open_with_flags(
178                &uri_clone,
179                OpenFlags::SQLITE_OPEN_READ_WRITE
180                    | OpenFlags::SQLITE_OPEN_CREATE
181                    | OpenFlags::SQLITE_OPEN_URI
182                    | OpenFlags::SQLITE_OPEN_NO_MUTEX,
183            )
184            .map_err(|e| KernelError::Store(e.to_string()))?;
185            crate::graph::schema::init_graph_schema(&conn)?;
186            // In-memory DBs ignore journal_mode, but busy_timeout still matters
187            // for the shared-cache connections the pool spawns on demand.
188            apply_concurrency_pragmas(&mut conn)?;
189            Ok(conn)
190        })
191        .await
192        .map_err(|e| KernelError::Store(e.to_string()))??;
193
194        let inner = Arc::new(PoolInner {
195            idle: Mutex::new(vec![conn]),
196            path: uri,
197            shared_mem: true,
198        });
199
200        Ok(Self {
201            inner,
202            sem: Arc::new(Semaphore::new(max_conns.max(1))),
203        })
204    }
205
206    /// Execute a closure with a pooled connection.
207    async fn with_conn<F, T>(&self, f: F) -> Result<T>
208    where
209        F: FnOnce(&Connection) -> Result<T> + Send + 'static,
210        T: Send + 'static,
211    {
212        let _permit = self
213            .sem
214            .acquire()
215            .await
216            .map_err(|_| KernelError::Store("semaphore closed".into()))?;
217
218        let inner = Arc::clone(&self.inner);
219        task::spawn_blocking(move || {
220            let conn = inner.take()?;
221            let result = f(&conn);
222            inner.return_conn(conn);
223            result
224        })
225        .await
226        .map_err(|e| KernelError::Store(e.to_string()))?
227    }
228
229    // ── Node CRUD ───────────────────────────────────
230
231    /// Insert or replace a node.
232    pub async fn upsert_node(&self, node: GraphNode) -> Result<()> {
233        self.with_conn(move |c| crate::graph::store::upsert_node(c, &node))
234            .await
235    }
236
237    /// Read a node by ID. Returns `None` if not found.
238    pub async fn read_node(&self, id: impl Into<String>) -> Result<Option<GraphNode>> {
239        let id = id.into();
240        self.with_conn(move |c| crate::graph::store::read_node(c, &id))
241            .await
242    }
243
244    /// Read all nodes (limited to 10 000).
245    pub async fn read_nodes(&self) -> Result<Vec<GraphNode>> {
246        self.with_conn(|c| crate::graph::store::read_nodes_limited(c, 10_000))
247            .await
248    }
249
250    /// Structured node query — paging, time range, ordering — at the SQL level.
251    /// Prefer this over `read_nodes` for anything that filters or paginates.
252    pub async fn query_nodes_ex(&self, q: crate::graph::NodeQuery) -> Result<Vec<GraphNode>> {
253        self.with_conn(move |c| crate::graph::search::query_nodes_ex(c, &q))
254            .await
255    }
256
257    /// Delete a node by ID. Returns `true` if a row was deleted.
258    pub async fn delete_node(&self, id: impl Into<String>) -> Result<bool> {
259        let id = id.into();
260        self.with_conn(move |c| crate::graph::store::delete_node(c, &id))
261            .await
262    }
263
264    // ── Edge CRUD ───────────────────────────────────
265
266    /// Append an edge (duplicates by ID are ignored).
267    pub async fn append_edge(&self, edge: GraphEdge) -> Result<()> {
268        self.with_conn(move |c| crate::graph::store::append_edge(c, &edge))
269            .await
270    }
271
272    /// Append many edges in one transaction (duplicates by ID *or* by the
273    /// `(source, target, relation)` unique index are ignored).
274    pub async fn append_edges(&self, edges: Vec<GraphEdge>) -> Result<()> {
275        self.with_conn(move |c| crate::graph::store::append_edges(c, &edges))
276            .await
277    }
278
279    /// Read edges touching `node_id`, filtered by direction and optional relation.
280    pub async fn edges_for_node_dir(
281        &self,
282        node_id: impl Into<String>,
283        dir: EdgeDirection,
284        relation: Option<String>,
285    ) -> Result<Vec<GraphEdge>> {
286        let node_id = node_id.into();
287        self.with_conn(move |c| {
288            crate::graph::store::edges_for_node_dir(c, &node_id, dir, relation.as_deref())
289        })
290        .await
291    }
292
293    /// 1-hop neighbors of `seed_ids` (weighted sum), filtered by direction and
294    /// an optional relation. Seed nodes are excluded.
295    pub async fn neighbors_weighted(
296        &self,
297        seed_ids: Vec<String>,
298        dir: EdgeDirection,
299        relation: Option<String>,
300    ) -> Result<Vec<(String, f64)>> {
301        self.with_conn(move |c| {
302            Ok(crate::graph::traversal::neighbors_weighted(
303                c,
304                &seed_ids,
305                dir,
306                relation.as_deref(),
307            ))
308        })
309        .await
310    }
311
312    /// Read all edges (limited to 10 000).
313    pub async fn read_edges(&self) -> Result<Vec<GraphEdge>> {
314        self.with_conn(|c| crate::graph::store::read_edges(c, 10_000))
315            .await
316    }
317
318    /// Delete an edge by ID. Returns `true` if a row was deleted.
319    pub async fn delete_edge(&self, id: impl Into<String>) -> Result<bool> {
320        let id = id.into();
321        self.with_conn(move |c| crate::graph::store::delete_edge(c, &id))
322            .await
323    }
324
325    // ── Search & Recall ─────────────────────────────
326
327    /// Full-text search over node titles and bodies.
328    pub async fn search_nodes(
329        &self,
330        query: impl Into<String>,
331        limit: usize,
332    ) -> Result<Vec<GraphNode>> {
333        let query = query.into();
334        self.with_conn(move |c| crate::graph::search::search_nodes(c, &query, limit))
335            .await
336    }
337
338    /// Smart recall with composite scoring.
339    pub async fn smart_recall(
340        &self,
341        project: Option<String>,
342        hint: Option<String>,
343        limit: usize,
344    ) -> Result<Vec<ScoredNode>> {
345        self.with_conn(move |c| {
346            crate::graph::recall::smart_recall(c, project.as_deref(), hint.as_deref(), limit)
347        })
348        .await
349    }
350
351    /// Structured recall — node-type / tag / time filters + touch gating.
352    pub async fn smart_recall_with(
353        &self,
354        opts: crate::graph::RecallOptions,
355    ) -> Result<Vec<ScoredNode>> {
356        self.with_conn(move |c| crate::graph::recall::smart_recall_with(c, &opts))
357            .await
358    }
359
360    // ── Stats ───────────────────────────────────────
361
362    /// Compute graph statistics (node/edge counts, avg importance).
363    pub async fn stats(&self) -> Result<GraphStats> {
364        self.with_conn(crate::graph::lifecycle::compute_stats).await
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    fn node(id: &str) -> GraphNode {
373        GraphNode {
374            id: id.into(),
375            node_type: "concept".into(),
376            title: format!("Node {id}"),
377            body: "body".into(),
378            tags: vec![],
379            projects: vec![],
380            agents: vec![],
381            created: "2026-01-01T00:00:00Z".into(),
382            updated: "2026-01-01T00:00:00Z".into(),
383            importance: 0.5,
384            access_count: 0,
385            accessed_at: String::new(),
386            ..Default::default()
387        }
388    }
389
390    async fn mem() -> AsyncPoolGraph {
391        AsyncPoolGraph::open_in_memory(2).await.unwrap()
392    }
393
394    #[tokio::test]
395    async fn upsert_and_read_node() {
396        let pool = mem().await;
397        pool.upsert_node(node("n1")).await.unwrap();
398        let n = pool.read_node("n1").await.unwrap().unwrap();
399        assert_eq!(n.id, "n1");
400    }
401
402    #[tokio::test]
403    async fn read_missing_returns_none() {
404        let pool = mem().await;
405        assert!(pool.read_node("ghost").await.unwrap().is_none());
406    }
407
408    #[tokio::test]
409    async fn delete_node() {
410        let pool = mem().await;
411        pool.upsert_node(node("n1")).await.unwrap();
412        assert!(pool.delete_node("n1").await.unwrap());
413        assert!(pool.read_node("n1").await.unwrap().is_none());
414    }
415
416    #[tokio::test]
417    async fn append_and_read_edges() {
418        let pool = mem().await;
419        pool.upsert_node(node("a")).await.unwrap();
420        pool.upsert_node(node("b")).await.unwrap();
421        pool.append_edge(GraphEdge {
422            id: "e1".into(),
423            source: "a".into(),
424            target: "b".into(),
425            relation: "related".into(),
426            weight: 1.0,
427            ts: "2026-01-01T00:00:00Z".into(),
428        })
429        .await
430        .unwrap();
431        let edges = pool.read_edges().await.unwrap();
432        assert_eq!(edges.len(), 1);
433    }
434
435    #[tokio::test]
436    async fn delete_edge() {
437        let pool = mem().await;
438        pool.append_edge(GraphEdge {
439            id: "e1".into(),
440            source: "a".into(),
441            target: "b".into(),
442            relation: "related".into(),
443            weight: 1.0,
444            ts: "2026-01-01T00:00:00Z".into(),
445        })
446        .await
447        .unwrap();
448        assert!(pool.delete_edge("e1").await.unwrap());
449        assert!(pool.read_edges().await.unwrap().is_empty());
450    }
451
452    #[tokio::test]
453    async fn search_finds_nodes() {
454        let pool = mem().await;
455        let mut n = node("n1");
456        n.title = "Rust ownership".to_string();
457        pool.upsert_node(n).await.unwrap();
458        let results = pool.search_nodes("Rust", 10).await.unwrap();
459        assert_eq!(results.len(), 1);
460    }
461
462    #[tokio::test]
463    async fn stats_returns_counts() {
464        let pool = mem().await;
465        pool.upsert_node(node("a")).await.unwrap();
466        pool.upsert_node(node("b")).await.unwrap();
467        let s = pool.stats().await.unwrap();
468        assert_eq!(s.total_nodes, 2);
469        assert_eq!(s.total_edges, 0);
470    }
471
472    #[tokio::test]
473    async fn clone_shares_pool() {
474        let pool = mem().await;
475        let pool2 = pool.clone();
476        pool.upsert_node(node("n1")).await.unwrap();
477        assert!(pool2.read_node("n1").await.unwrap().is_some());
478    }
479
480    #[tokio::test]
481    async fn concurrent_reads() {
482        let pool = mem().await;
483        pool.upsert_node(node("n1")).await.unwrap();
484
485        let mut handles = vec![];
486        for _ in 0..4 {
487            let p = pool.clone();
488            handles.push(tokio::spawn(async move {
489                p.read_node("n1").await.unwrap().is_some()
490            }));
491        }
492        for h in handles {
493            assert!(h.await.unwrap());
494        }
495    }
496
497    #[tokio::test]
498    async fn open_creates_file() {
499        let dir = tempfile::tempdir().unwrap();
500        let path = dir.path().join("sub").join("test.db");
501        let pool = AsyncPoolGraph::open(&path, 2).await.unwrap();
502        pool.upsert_node(node("n1")).await.unwrap();
503        assert!(path.exists());
504        drop(pool);
505    }
506}
llm_kernel/graph/async_pool.rs

llm_kernel/graph/
async_pool.rs
