//! M8.2 — Queue Worker for Concurrent Dual-Write Processing //! //! Background task that receives messages from the queue and processes them //! via DualWriteIndexer. Runs concurrently with ingest, improving throughput. //! //! # Architecture //! //! ``` //! IngestWorker (fast path) QueueWorker (background) //! │ │ //! ├─ chunk_input │ //! │ (embedding) │ //! │ │ //! ├─ queue.send_chunk()────┐ │ //! │ (returns immediately) │ │ //! │ │ │ //! └─ continues... │ │ //! │ │ //! ├─ queue.receive_chunks(10, 30) //! │ (long-poll, up to 30s) //! │ //! ├─ for each message: //! │ - process_queued_chunk() //! │ - embed_one() [happens here] //! │ - write_pgvector() //! │ - write_opensearch() //! │ - delete_chunk() on success //! │ - change_visibility() on retry //! │ //! └─ loop back to receive //! ``` //! //! Benefits: //! - Ingest path is decoupled from embedding/pgvector/OpenSearch writes //! - Multiple workers can process messages concurrently //! - Non-blocking: queue.send_chunk() returns immediately //! - Fault-tolerant: failed messages auto-retry with exponential backoff use anyhow::{anyhow, Result}; use std::sync::Arc; use std::time::Duration; use tokio::time::sleep; use tracing::{debug, error, info, warn}; use crate::dual_write_indexer::DualWriteIndexer; use crate::queue_adapter::QueueAdapter; use mem_llm::EmbeddingsClient; /// Configuration for queue worker #[derive(Debug, Clone)] pub struct QueueWorkerConfig { /// Max messages per receive (1-10) pub max_messages_per_batch: i32, /// Visibility timeout for processing (seconds) pub visibility_timeout_secs: i32, /// Time to wait for messages (0-20 seconds) pub wait_time_secs: i32, /// Project to process (None = all projects) pub project: Option, /// Max retries before DLQ pub max_retries: i32, /// Retry backoff: exponential starting from this value (seconds) pub retry_backoff_initial_secs: i32, /// Poll interval when queue is empty (seconds) pub empty_poll_interval_secs: u64, /// Enable metrics collection pub enable_metrics: bool, } impl Default for QueueWorkerConfig { fn default() -> Self { Self { max_messages_per_batch: 10, visibility_timeout_secs: 300, // 5 minutes wait_time_secs: 20, // Long-poll timeout project: None, max_retries: 3, retry_backoff_initial_secs: 60, empty_poll_interval_secs: 5, enable_metrics: true, } } } /// Metrics for worker execution #[derive(Debug, Clone, Default)] pub struct WorkerMetrics { pub messages_received: u64, pub messages_processed: u64, pub messages_failed: u64, pub messages_dlq: u64, pub total_processing_time_ms: u64, } /// Queue worker for processing dual-write messages pub struct QueueWorker { indexer: Arc, embeddings: Arc, config: QueueWorkerConfig, metrics: Arc>, } impl QueueWorker { /// Create new queue worker pub fn new( indexer: Arc, embeddings: Arc, config: QueueWorkerConfig, ) -> Self { Self { indexer, embeddings, config, metrics: Arc::new(tokio::sync::RwLock::new(WorkerMetrics::default())), } } /// Start worker (blocking loop) pub async fn start(&self) -> Result<()> { info!("Queue worker starting: config={:?}", self.config); loop { match self.process_batch().await { Ok(count) => { if count == 0 { // Empty batch: sleep before retrying debug!( "Queue empty, waiting {}s before retry", self.config.empty_poll_interval_secs ); sleep(Duration::from_secs(self.config.empty_poll_interval_secs)).await; } } Err(e) => { error!("Worker error (will retry): {}", e); sleep(Duration::from_secs(5)).await; } } } } /// Process one batch of messages from queue async fn process_batch(&self) -> Result { let queue = &self.indexer.queue; // Receive messages let messages = queue .receive_chunks( self.config.max_messages_per_batch, self.config.visibility_timeout_secs, self.config.project.as_deref(), ) .await?; let batch_size = messages.len(); if batch_size == 0 { return Ok(0); } let mut metrics = self.metrics.write().await; metrics.messages_received += batch_size as u64; drop(metrics); // Process each message concurrently let handles: Vec<_> = messages .into_iter() .map(|msg| { let indexer = self.indexer.clone(); let embeddings = self.embeddings.clone(); let config = self.config.clone(); let metrics = self.metrics.clone(); tokio::spawn(async move { Self::process_message(indexer, embeddings, config, metrics, msg).await }) }) .collect(); // Wait for all to complete for handle in handles { if let Err(e) = handle.await { error!("Worker task panicked: {}", e); } } Ok(batch_size) } /// Process a single message async fn process_message( indexer: Arc, embeddings: Arc, config: QueueWorkerConfig, metrics: Arc>, message: crate::queue_adapter::QueueMessage, ) -> Result<()> { let start = std::time::Instant::now(); let message_id = message.message_id.clone(); let receipt_handle = message.receipt_handle.clone(); debug!("Processing message: {}", message_id); // Parse message body let body: serde_json::Value = match serde_json::from_str(&message.body) { Ok(b) => b, Err(e) => { error!("Failed to parse message body: {}", e); indexer .queue .send_to_dlq(&message_id, &receipt_handle, "invalid_json") .await .ok(); let mut m = metrics.write().await; m.messages_dlq += 1; return Err(e.into()); } }; // Extract chunk_id let chunk_id = match body["chunk_id"].as_str() { Some(id) => match uuid::Uuid::parse_str(id) { Ok(u) => u, Err(e) => { error!("Invalid chunk_id: {}", e); indexer .queue .send_to_dlq(&message_id, &receipt_handle, "invalid_uuid") .await .ok(); let mut m = metrics.write().await; m.messages_dlq += 1; return Err(e.into()); } }, None => { error!("Missing chunk_id in message"); indexer .queue .send_to_dlq(&message_id, &receipt_handle, "missing_chunk_id") .await .ok(); let mut m = metrics.write().await; m.messages_dlq += 1; return Err(anyhow!("Missing chunk_id")); } }; // Extract content let content = match body["content"].as_str() { Some(c) => c.to_string(), None => { error!("Missing content in message"); indexer .queue .send_to_dlq(&message_id, &receipt_handle, "missing_content") .await .ok(); let mut m = metrics.write().await; m.messages_dlq += 1; return Err(anyhow!("Missing content")); } }; // Compute embedding let embedding_vec = match embeddings.embed_one(&content).await { Ok(vec) => vec, Err(e) => { warn!("Embedding failed, extending visibility for retry: {}", e); indexer .queue .change_visibility(&message_id, &receipt_handle, 300) .await .ok(); let mut m = metrics.write().await; m.messages_failed += 1; return Err(e); } }; // Convert pgvector::Vector to Vec let embedding: Vec = embedding_vec.to_vec(); // Process dual-write match indexer.process_queued_chunk(&message, &embedding).await { Ok(result) => { if result.pgvector_success && !result.opensearch_pending { // Success: already deleted by process_queued_chunk debug!("Message processed successfully: {}", message_id); let elapsed = start.elapsed().as_millis() as u64; let mut m = metrics.write().await; m.messages_processed += 1; m.total_processing_time_ms += elapsed; } else if result.pgvector_success && result.opensearch_pending { // pgvector OK, OpenSearch pending: visibility already extended warn!("Message will retry: {}", message_id); let mut m = metrics.write().await; m.messages_failed += 1; } else { // pgvector failed: visibility already extended warn!("pgvector write failed, will retry: {}", message_id); let mut m = metrics.write().await; m.messages_failed += 1; } Ok(()) } Err(e) => { // Check receive count if message.receive_count >= config.max_retries { error!( "Message max retries exceeded ({}), sending to DLQ: {}", message.receive_count, message_id ); indexer .queue .send_to_dlq(&message_id, &receipt_handle, "max_retries") .await .ok(); let mut m = metrics.write().await; m.messages_dlq += 1; } else { // Extend visibility for retry warn!( "Message processing failed (retry {}), extending visibility: {}", message.receive_count, message_id ); indexer .queue .change_visibility(&message_id, &receipt_handle, 300) .await .ok(); let mut m = metrics.write().await; m.messages_failed += 1; } Err(e) } } } /// Get current metrics pub async fn metrics(&self) -> WorkerMetrics { self.metrics.read().await.clone() } /// Reset metrics pub async fn reset_metrics(&self) { let mut m = self.metrics.write().await; *m = WorkerMetrics::default(); } } #[cfg(test)] mod tests { use super::*; #[test] fn test_queue_worker_config_default() { let config = QueueWorkerConfig::default(); assert_eq!(config.max_messages_per_batch, 10); assert_eq!(config.visibility_timeout_secs, 300); assert_eq!(config.wait_time_secs, 20); assert_eq!(config.max_retries, 3); } #[test] fn test_worker_metrics_default() { let metrics = WorkerMetrics::default(); assert_eq!(metrics.messages_received, 0); assert_eq!(metrics.messages_processed, 0); } #[test] fn test_queue_worker_config_custom() { let config = QueueWorkerConfig { max_messages_per_batch: 5, visibility_timeout_secs: 600, project: Some("test-proj".to_string()), ..Default::default() }; assert_eq!(config.max_messages_per_batch, 5); assert_eq!(config.project, Some("test-proj".to_string())); } }