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::{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    /// Read all edges (limited to 10 000).
265    pub async fn read_edges(&self) -> Result<Vec<GraphEdge>> {
266        self.with_conn(|c| crate::graph::store::read_edges(c, 10_000))
267            .await
268    }
269
270    /// Delete an edge by ID. Returns `true` if a row was deleted.
271    pub async fn delete_edge(&self, id: impl Into<String>) -> Result<bool> {
272        let id = id.into();
273        self.with_conn(move |c| crate::graph::store::delete_edge(c, &id))
274            .await
275    }
276
277    // ── Search & Recall ─────────────────────────────
278
279    /// Full-text search over node titles and bodies.
280    pub async fn search_nodes(
281        &self,
282        query: impl Into<String>,
283        limit: usize,
284    ) -> Result<Vec<GraphNode>> {
285        let query = query.into();
286        self.with_conn(move |c| crate::graph::search::search_nodes(c, &query, limit))
287            .await
288    }
289
290    /// Smart recall with composite scoring.
291    pub async fn smart_recall(
292        &self,
293        project: Option<String>,
294        hint: Option<String>,
295        limit: usize,
296    ) -> Result<Vec<ScoredNode>> {
297        self.with_conn(move |c| {
298            crate::graph::recall::smart_recall(c, project.as_deref(), hint.as_deref(), limit)
299        })
300        .await
301    }
302
303    // ── Stats ───────────────────────────────────────
304
305    /// Compute graph statistics (node/edge counts, avg importance).
306    pub async fn stats(&self) -> Result<GraphStats> {
307        self.with_conn(crate::graph::lifecycle::compute_stats).await
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314
315    fn node(id: &str) -> GraphNode {
316        GraphNode {
317            id: id.into(),
318            node_type: "concept".into(),
319            title: format!("Node {id}"),
320            body: "body".into(),
321            tags: vec![],
322            projects: vec![],
323            agents: vec![],
324            created: "2026-01-01T00:00:00Z".into(),
325            updated: "2026-01-01T00:00:00Z".into(),
326            importance: 0.5,
327            access_count: 0,
328            accessed_at: String::new(),
329        }
330    }
331
332    async fn mem() -> AsyncPoolGraph {
333        AsyncPoolGraph::open_in_memory(2).await.unwrap()
334    }
335
336    #[tokio::test]
337    async fn upsert_and_read_node() {
338        let pool = mem().await;
339        pool.upsert_node(node("n1")).await.unwrap();
340        let n = pool.read_node("n1").await.unwrap().unwrap();
341        assert_eq!(n.id, "n1");
342    }
343
344    #[tokio::test]
345    async fn read_missing_returns_none() {
346        let pool = mem().await;
347        assert!(pool.read_node("ghost").await.unwrap().is_none());
348    }
349
350    #[tokio::test]
351    async fn delete_node() {
352        let pool = mem().await;
353        pool.upsert_node(node("n1")).await.unwrap();
354        assert!(pool.delete_node("n1").await.unwrap());
355        assert!(pool.read_node("n1").await.unwrap().is_none());
356    }
357
358    #[tokio::test]
359    async fn append_and_read_edges() {
360        let pool = mem().await;
361        pool.upsert_node(node("a")).await.unwrap();
362        pool.upsert_node(node("b")).await.unwrap();
363        pool.append_edge(GraphEdge {
364            id: "e1".into(),
365            source: "a".into(),
366            target: "b".into(),
367            relation: "related".into(),
368            weight: 1.0,
369            ts: "2026-01-01T00:00:00Z".into(),
370        })
371        .await
372        .unwrap();
373        let edges = pool.read_edges().await.unwrap();
374        assert_eq!(edges.len(), 1);
375    }
376
377    #[tokio::test]
378    async fn delete_edge() {
379        let pool = mem().await;
380        pool.append_edge(GraphEdge {
381            id: "e1".into(),
382            source: "a".into(),
383            target: "b".into(),
384            relation: "related".into(),
385            weight: 1.0,
386            ts: "2026-01-01T00:00:00Z".into(),
387        })
388        .await
389        .unwrap();
390        assert!(pool.delete_edge("e1").await.unwrap());
391        assert!(pool.read_edges().await.unwrap().is_empty());
392    }
393
394    #[tokio::test]
395    async fn search_finds_nodes() {
396        let pool = mem().await;
397        let mut n = node("n1");
398        n.title = "Rust ownership".to_string();
399        pool.upsert_node(n).await.unwrap();
400        let results = pool.search_nodes("Rust", 10).await.unwrap();
401        assert_eq!(results.len(), 1);
402    }
403
404    #[tokio::test]
405    async fn stats_returns_counts() {
406        let pool = mem().await;
407        pool.upsert_node(node("a")).await.unwrap();
408        pool.upsert_node(node("b")).await.unwrap();
409        let s = pool.stats().await.unwrap();
410        assert_eq!(s.total_nodes, 2);
411        assert_eq!(s.total_edges, 0);
412    }
413
414    #[tokio::test]
415    async fn clone_shares_pool() {
416        let pool = mem().await;
417        let pool2 = pool.clone();
418        pool.upsert_node(node("n1")).await.unwrap();
419        assert!(pool2.read_node("n1").await.unwrap().is_some());
420    }
421
422    #[tokio::test]
423    async fn concurrent_reads() {
424        let pool = mem().await;
425        pool.upsert_node(node("n1")).await.unwrap();
426
427        let mut handles = vec![];
428        for _ in 0..4 {
429            let p = pool.clone();
430            handles.push(tokio::spawn(async move {
431                p.read_node("n1").await.unwrap().is_some()
432            }));
433        }
434        for h in handles {
435            assert!(h.await.unwrap());
436        }
437    }
438
439    #[tokio::test]
440    async fn open_creates_file() {
441        let dir = tempfile::tempdir().unwrap();
442        let path = dir.path().join("sub").join("test.db");
443        let pool = AsyncPoolGraph::open(&path, 2).await.unwrap();
444        pool.upsert_node(node("n1")).await.unwrap();
445        assert!(path.exists());
446        drop(pool);
447    }
448}
llm_kernel/graph/async_pool.rs

llm_kernel/graph/
async_pool.rs
