mentedb_server/
extraction_queue.rs

1//! Bounded extraction queue with backpressure and graceful shutdown.
2//!
3//! Instead of fire-and-forget `tokio::spawn` per extraction request,
4//! this module provides a bounded channel + worker pool. Benefits:
5//! - Backpressure: at most `QUEUE_CAPACITY` pending extractions
6//! - Bounded concurrency: `MAX_CONCURRENT` parallel LLM calls
7//! - Graceful drain: on shutdown, the worker finishes in-flight work
8
9use std::sync::Arc;
10
11use mentedb_core::types::{AgentId, SpaceId};
12use mentedb_extraction::ExtractionConfig;
13use tokio::sync::mpsc;
14use tokio::task::JoinHandle;
15
16const QUEUE_CAPACITY: usize = 64;
17const MAX_CONCURRENT: usize = 4;
18
19/// What we send into the queue.
20pub struct ExtractionRequest {
21    pub config: ExtractionConfig,
22    pub content: String,
23    pub agent_id: AgentId,
24    pub space_id: SpaceId,
25    pub db: Arc<mentedb::MenteDb>,
26}
27
28pub type ExtractionSender = mpsc::Sender<ExtractionRequest>;
29type ExtractionReceiver = mpsc::Receiver<ExtractionRequest>;
30
31/// Spawn the extraction worker. Returns the sender half and a join handle
32/// that resolves when all in-flight work is drained.
33pub fn spawn_worker() -> (ExtractionSender, JoinHandle<()>) {
34    let (tx, rx) = mpsc::channel(QUEUE_CAPACITY);
35    let handle = tokio::spawn(worker_loop(rx));
36    (tx, handle)
37}
38
39async fn worker_loop(mut rx: ExtractionReceiver) {
40    let semaphore = Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT));
41
42    while let Some(req) = rx.recv().await {
43        let permit = semaphore.clone().acquire_owned().await;
44        if permit.is_err() {
45            break; // semaphore closed
46        }
47        let permit = permit.unwrap();
48        tokio::spawn(async move {
49            if let Err(e) = process_extraction(req).await {
50                tracing::warn!(error = %e, "background extraction failed");
51            }
52            drop(permit);
53        });
54    }
55
56    // Channel closed — wait for all in-flight extractions to finish
57    let _ = semaphore.acquire_many(MAX_CONCURRENT as u32).await;
58    tracing::info!("extraction worker drained");
59}
60
61async fn process_extraction(req: ExtractionRequest) -> Result<(), String> {
62    use mentedb_core::MemoryNode;
63    use mentedb_core::types::MemoryId;
64    use mentedb_extraction::{ExtractionPipeline, HttpExtractionProvider};
65
66    let provider = HttpExtractionProvider::new(req.config.clone())
67        .map_err(|e| format!("extraction provider init: {e}"))?;
68    let pipeline = ExtractionPipeline::new(provider, req.config);
69
70    let all_memories = pipeline
71        .extract_from_conversation(&req.content)
72        .await
73        .map_err(|e| format!("extraction failed: {e}"))?;
74
75    let quality_passed = pipeline.filter_quality(&all_memories);
76
77    let now = std::time::SystemTime::now()
78        .duration_since(std::time::UNIX_EPOCH)
79        .unwrap_or_default()
80        .as_micros() as u64;
81
82    let mut stored = 0usize;
83    for memory in &quality_passed {
84        let memory_type =
85            mentedb_extraction::map_extraction_type_to_memory_type(&memory.memory_type);
86        let id = MemoryId::new();
87        let node = MemoryNode {
88            id,
89            agent_id: req.agent_id,
90            memory_type,
91            embedding: vec![],
92            content: memory.content.clone(),
93            created_at: now,
94            accessed_at: now,
95            access_count: 0,
96            salience: memory.confidence,
97            confidence: memory.confidence,
98            space_id: req.space_id,
99            attributes: std::collections::HashMap::new(),
100            tags: memory.tags.clone(),
101            valid_from: None,
102            valid_until: None,
103        };
104        match req.db.store(node) {
105            Ok(()) => stored += 1,
106            Err(e) => {
107                tracing::warn!(error = %e, "failed to store extracted memory, skipping");
108            }
109        }
110    }
111
112    tracing::debug!(
113        total = all_memories.len(),
114        quality = quality_passed.len(),
115        stored,
116        "extraction complete"
117    );
118    Ok(())
119}
mentedb_server/extraction_queue.rs

mentedb_server/
extraction_queue.rs
