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

llm_kernel/graph/
async_pool.rs
