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        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    }
79
80    fn return_conn(&self, conn: Connection) {
81        if let Ok(mut guard) = self.idle.lock() {
82            guard.push(conn);
83        }
84        // If lock fails (poisoned), drop the connection — it will be recreated on next take.
85    }
86}
87
88// ── AsyncPoolGraph ──────────────────────────────────
89
90/// Bounded async connection pool for the knowledge graph.
91///
92/// Uses a `Semaphore` to bound concurrency and a `Mutex<Vec<Connection>>`
93/// for idle connection reuse. Each method acquires a permit, takes (or creates)
94/// a connection, runs the operation via `spawn_blocking`, then returns the
95/// connection to the pool.
96#[derive(Clone)]
97pub struct AsyncPoolGraph {
98    inner: Arc<PoolInner>,
99    sem: Arc<Semaphore>,
100}
101
102impl AsyncPoolGraph {
103    /// Open (or create) a database and initialise the graph schema.
104    ///
105    /// `max_conns` bounds the number of concurrent operations.
106    pub async fn open(path: impl AsRef<Path>, max_conns: usize) -> Result<Self> {
107        let path_str = path
108            .as_ref()
109            .to_str()
110            .ok_or_else(|| KernelError::Store("invalid path".into()))?
111            .to_string();
112
113        // Create parent dirs + open first connection + apply schema
114        let path_for_open = path_str.clone();
115        let first_conn = task::spawn_blocking(move || -> Result<Connection> {
116            if let Some(parent) = Path::new(&path_for_open).parent() {
117                std::fs::create_dir_all(parent)?;
118            }
119            let conn =
120                Connection::open(&path_for_open).map_err(|e| KernelError::Store(e.to_string()))?;
121            crate::graph::schema::init_graph_schema(&conn)?;
122            Ok(conn)
123        })
124        .await
125        .map_err(|e| KernelError::Store(e.to_string()))??;
126
127        let inner = Arc::new(PoolInner {
128            idle: Mutex::new(vec![first_conn]),
129            path: path_str,
130            shared_mem: false,
131        });
132
133        Ok(Self {
134            inner,
135            sem: Arc::new(Semaphore::new(max_conns.max(1))),
136        })
137    }
138
139    /// Create an in-memory pool with schema applied. Useful for tests.
140    ///
141    /// Uses SQLite shared-cache mode so all connections in the pool see the
142    /// same data (plain `:memory:` creates an independent DB per connection).
143    pub async fn open_in_memory(max_conns: usize) -> Result<Self> {
144        let id = MEM_POOL_ID.fetch_add(1, Ordering::Relaxed);
145        let uri = format!("file:llm_kernel_pool_{id}?mode=memory&cache=shared");
146        let uri_clone = uri.clone();
147        let conn = task::spawn_blocking(move || -> Result<Connection> {
148            let conn = Connection::open_with_flags(
149                &uri_clone,
150                OpenFlags::SQLITE_OPEN_READ_WRITE
151                    | OpenFlags::SQLITE_OPEN_CREATE
152                    | OpenFlags::SQLITE_OPEN_URI
153                    | OpenFlags::SQLITE_OPEN_NO_MUTEX,
154            )
155            .map_err(|e| KernelError::Store(e.to_string()))?;
156            crate::graph::schema::init_graph_schema(&conn)?;
157            Ok(conn)
158        })
159        .await
160        .map_err(|e| KernelError::Store(e.to_string()))??;
161
162        let inner = Arc::new(PoolInner {
163            idle: Mutex::new(vec![conn]),
164            path: uri,
165            shared_mem: true,
166        });
167
168        Ok(Self {
169            inner,
170            sem: Arc::new(Semaphore::new(max_conns.max(1))),
171        })
172    }
173
174    /// Execute a closure with a pooled connection.
175    async fn with_conn<F, T>(&self, f: F) -> Result<T>
176    where
177        F: FnOnce(&Connection) -> Result<T> + Send + 'static,
178        T: Send + 'static,
179    {
180        let _permit = self
181            .sem
182            .acquire()
183            .await
184            .map_err(|_| KernelError::Store("semaphore closed".into()))?;
185
186        let inner = Arc::clone(&self.inner);
187        task::spawn_blocking(move || {
188            let conn = inner.take()?;
189            let result = f(&conn);
190            inner.return_conn(conn);
191            result
192        })
193        .await
194        .map_err(|e| KernelError::Store(e.to_string()))?
195    }
196
197    // ── Node CRUD ───────────────────────────────────
198
199    /// Insert or replace a node.
200    pub async fn upsert_node(&self, node: GraphNode) -> Result<()> {
201        self.with_conn(move |c| crate::graph::store::upsert_node(c, &node))
202            .await
203    }
204
205    /// Read a node by ID. Returns `None` if not found.
206    pub async fn read_node(&self, id: impl Into<String>) -> Result<Option<GraphNode>> {
207        let id = id.into();
208        self.with_conn(move |c| crate::graph::store::read_node(c, &id))
209            .await
210    }
211
212    /// Read all nodes (limited to 10 000).
213    pub async fn read_nodes(&self) -> Result<Vec<GraphNode>> {
214        self.with_conn(|c| crate::graph::store::read_nodes_limited(c, 10_000))
215            .await
216    }
217
218    /// Delete a node by ID. Returns `true` if a row was deleted.
219    pub async fn delete_node(&self, id: impl Into<String>) -> Result<bool> {
220        let id = id.into();
221        self.with_conn(move |c| crate::graph::store::delete_node(c, &id))
222            .await
223    }
224
225    // ── Edge CRUD ───────────────────────────────────
226
227    /// Append an edge (duplicates by ID are ignored).
228    pub async fn append_edge(&self, edge: GraphEdge) -> Result<()> {
229        self.with_conn(move |c| crate::graph::store::append_edge(c, &edge))
230            .await
231    }
232
233    /// Read all edges (limited to 10 000).
234    pub async fn read_edges(&self) -> Result<Vec<GraphEdge>> {
235        self.with_conn(|c| crate::graph::store::read_edges(c, 10_000))
236            .await
237    }
238
239    /// Delete an edge by ID. Returns `true` if a row was deleted.
240    pub async fn delete_edge(&self, id: impl Into<String>) -> Result<bool> {
241        let id = id.into();
242        self.with_conn(move |c| crate::graph::store::delete_edge(c, &id))
243            .await
244    }
245
246    // ── Search & Recall ─────────────────────────────
247
248    /// Full-text search over node titles and bodies.
249    pub async fn search_nodes(
250        &self,
251        query: impl Into<String>,
252        limit: usize,
253    ) -> Result<Vec<GraphNode>> {
254        let query = query.into();
255        self.with_conn(move |c| crate::graph::search::search_nodes(c, &query, limit))
256            .await
257    }
258
259    /// Smart recall with composite scoring.
260    pub async fn smart_recall(
261        &self,
262        project: Option<String>,
263        hint: Option<String>,
264        limit: usize,
265    ) -> Result<Vec<ScoredNode>> {
266        self.with_conn(move |c| {
267            crate::graph::recall::smart_recall(c, project.as_deref(), hint.as_deref(), limit)
268        })
269        .await
270    }
271
272    // ── Stats ───────────────────────────────────────
273
274    /// Compute graph statistics (node/edge counts, avg importance).
275    pub async fn stats(&self) -> Result<GraphStats> {
276        self.with_conn(crate::graph::lifecycle::compute_stats).await
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use super::*;
283
284    fn node(id: &str) -> GraphNode {
285        GraphNode {
286            id: id.into(),
287            node_type: "concept".into(),
288            title: format!("Node {id}"),
289            body: "body".into(),
290            tags: vec![],
291            projects: vec![],
292            agents: vec![],
293            created: "2026-01-01T00:00:00Z".into(),
294            updated: "2026-01-01T00:00:00Z".into(),
295            importance: 0.5,
296            access_count: 0,
297            accessed_at: String::new(),
298        }
299    }
300
301    async fn mem() -> AsyncPoolGraph {
302        AsyncPoolGraph::open_in_memory(2).await.unwrap()
303    }
304
305    #[tokio::test]
306    async fn upsert_and_read_node() {
307        let pool = mem().await;
308        pool.upsert_node(node("n1")).await.unwrap();
309        let n = pool.read_node("n1").await.unwrap().unwrap();
310        assert_eq!(n.id, "n1");
311    }
312
313    #[tokio::test]
314    async fn read_missing_returns_none() {
315        let pool = mem().await;
316        assert!(pool.read_node("ghost").await.unwrap().is_none());
317    }
318
319    #[tokio::test]
320    async fn delete_node() {
321        let pool = mem().await;
322        pool.upsert_node(node("n1")).await.unwrap();
323        assert!(pool.delete_node("n1").await.unwrap());
324        assert!(pool.read_node("n1").await.unwrap().is_none());
325    }
326
327    #[tokio::test]
328    async fn append_and_read_edges() {
329        let pool = mem().await;
330        pool.upsert_node(node("a")).await.unwrap();
331        pool.upsert_node(node("b")).await.unwrap();
332        pool.append_edge(GraphEdge {
333            id: "e1".into(),
334            source: "a".into(),
335            target: "b".into(),
336            relation: "related".into(),
337            weight: 1.0,
338            ts: "2026-01-01T00:00:00Z".into(),
339        })
340        .await
341        .unwrap();
342        let edges = pool.read_edges().await.unwrap();
343        assert_eq!(edges.len(), 1);
344    }
345
346    #[tokio::test]
347    async fn delete_edge() {
348        let pool = mem().await;
349        pool.append_edge(GraphEdge {
350            id: "e1".into(),
351            source: "a".into(),
352            target: "b".into(),
353            relation: "related".into(),
354            weight: 1.0,
355            ts: "2026-01-01T00:00:00Z".into(),
356        })
357        .await
358        .unwrap();
359        assert!(pool.delete_edge("e1").await.unwrap());
360        assert!(pool.read_edges().await.unwrap().is_empty());
361    }
362
363    #[tokio::test]
364    async fn search_finds_nodes() {
365        let pool = mem().await;
366        let mut n = node("n1");
367        n.title = "Rust ownership".to_string();
368        pool.upsert_node(n).await.unwrap();
369        let results = pool.search_nodes("Rust", 10).await.unwrap();
370        assert_eq!(results.len(), 1);
371    }
372
373    #[tokio::test]
374    async fn stats_returns_counts() {
375        let pool = mem().await;
376        pool.upsert_node(node("a")).await.unwrap();
377        pool.upsert_node(node("b")).await.unwrap();
378        let s = pool.stats().await.unwrap();
379        assert_eq!(s.total_nodes, 2);
380        assert_eq!(s.total_edges, 0);
381    }
382
383    #[tokio::test]
384    async fn clone_shares_pool() {
385        let pool = mem().await;
386        let pool2 = pool.clone();
387        pool.upsert_node(node("n1")).await.unwrap();
388        assert!(pool2.read_node("n1").await.unwrap().is_some());
389    }
390
391    #[tokio::test]
392    async fn concurrent_reads() {
393        let pool = mem().await;
394        pool.upsert_node(node("n1")).await.unwrap();
395
396        let mut handles = vec![];
397        for _ in 0..4 {
398            let p = pool.clone();
399            handles.push(tokio::spawn(async move {
400                p.read_node("n1").await.unwrap().is_some()
401            }));
402        }
403        for h in handles {
404            assert!(h.await.unwrap());
405        }
406    }
407
408    #[tokio::test]
409    async fn open_creates_file() {
410        let dir = tempfile::tempdir().unwrap();
411        let path = dir.path().join("sub").join("test.db");
412        let pool = AsyncPoolGraph::open(&path, 2).await.unwrap();
413        pool.upsert_node(node("n1")).await.unwrap();
414        assert!(path.exists());
415        drop(pool);
416    }
417}
llm_kernel/graph/async_pool.rs

llm_kernel/graph/
async_pool.rs
