poimen-memory/crates/mem-cli/src/queue_worker.rs

//! 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<String>,

    /// 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<DualWriteIndexer>,
    embeddings: Arc<EmbeddingsClient>,
    config: QueueWorkerConfig,
    metrics: Arc<tokio::sync::RwLock<WorkerMetrics>>,
}

impl QueueWorker {
    /// Create new queue worker
    pub fn new(
        indexer: Arc<DualWriteIndexer>,
        embeddings: Arc<EmbeddingsClient>,
        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<usize> {
        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<DualWriteIndexer>,
        embeddings: Arc<EmbeddingsClient>,
        config: QueueWorkerConfig,
        metrics: Arc<tokio::sync::RwLock<WorkerMetrics>>,
        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<f32>
        let embedding: Vec<f32> = 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()));
    }
}