//! M8.2 — Dual-write indexing pipeline //! //! Coordinates atomic writes to both pgvector (embedding search) and OpenSearch (lexical search). //! Same chunk_id in both stores. If OpenSearch fails, marks `opensearch_pending=true` for eventual //! consistency retry loop. use anyhow::{anyhow, Result}; use sha2::{Digest, Sha256}; use sqlx::PgPool; use uuid::Uuid; use pgvector::Vector; use std::sync::Arc; use crate::opensearch_client::OpenSearchClient; use crate::queue_adapter::QueueAdapter; #[derive(Clone)] pub struct DualWriteIndexer { pool: PgPool, opensearch: Option>, /// Queue adapter for concurrent dual-write processing /// Can be: kmsvc (production), in-memory (testing), or SQS (future) queue: Arc, } /// Input chunk for dual-write #[derive(Debug, Clone)] pub struct ChunkInput { pub content: String, pub source: String, pub project: String, pub level: String, // "L0", "L1", "L2", "R" pub breadcrumb: Vec, } /// Result of dual-write operation #[derive(Debug, Clone)] pub struct DualWriteResult { pub chunk_id: Uuid, pub chunk_hash: String, pub pgvector_success: bool, pub opensearch_success: bool, pub opensearch_pending: bool, // true if OpenSearch failed pub error: Option, } impl DualWriteIndexer { /// Create dual-write indexer with queue adapter pub fn new( pool: PgPool, opensearch: Option>, queue: Arc, ) -> Self { Self { pool, opensearch, queue, } } /// Queue chunk for dual-write processing /// /// Sequence: /// 1. Check dedup (chunk_hash exists AND indexed_in_pgvector AND indexed_in_opensearch) /// 2. Queue message to external queue service (kmsvc/SQS/etc) /// 3. Concurrent workers receive from queue and perform dual-write /// /// Returns message_id for tracking progress pub async fn queue_chunk( &self, chunk: &ChunkInput, embedding: &[f32], ) -> Result { let chunk_id = Uuid::new_v4(); let chunk_hash = self.compute_hash(&chunk.content); // Check deduplication if self.is_already_indexed(&chunk_hash, &chunk.project).await? { tracing::debug!("Chunk already indexed (dedup): {}", chunk_hash); return Ok(Uuid::nil().to_string()); } // Build message attributes let mut attributes = std::collections::HashMap::new(); attributes.insert("source".to_string(), chunk.source.clone()); attributes.insert("level".to_string(), chunk.level.clone()); attributes.insert("breadcrumb".to_string(), serde_json::to_string(&chunk.breadcrumb)?); attributes.insert("embedding_size".to_string(), embedding.len().to_string()); // Build message body let body = serde_json::json!({ "chunk_id": chunk_id, "content": chunk.content, "source": chunk.source, "level": chunk.level, "breadcrumb": chunk.breadcrumb, "embedding": embedding, }).to_string(); // Queue message let message_id = self.queue.send_chunk( chunk_id, body, chunk.project.clone(), attributes, ).await?; tracing::info!("Chunk queued for dual-write: message_id={}, chunk_hash={}", message_id, chunk_hash); Ok(message_id) } /// Worker: Process queued chunk for dual-write /// /// Called by concurrent workers receiving from queue. /// Sequence: /// 1. Receive message from queue /// 2. Write to pgvector with embedding /// 3. Write to OpenSearch (fail-soft) /// 4. Delete from queue on success, or extend visibility on retry pub async fn process_queued_chunk( &self, message: &crate::queue_adapter::QueueMessage, embedding: &[f32], ) -> Result { let body: serde_json::Value = serde_json::from_str(&message.body)?; let chunk_id = body["chunk_id"].as_str().ok_or_else(|| anyhow!("Missing chunk_id"))? .parse::()?; let content = body["content"].as_str().ok_or_else(|| anyhow!("Missing content"))?.to_string(); let source = body["source"].as_str().ok_or_else(|| anyhow!("Missing source"))?.to_string(); let project = message.project.clone(); let level = body["level"].as_str().ok_or_else(|| anyhow!("Missing level"))?.to_string(); let breadcrumb: Vec = serde_json::from_value(body["breadcrumb"].clone())?; let chunk_hash = self.compute_hash(&content); // Write to pgvector let pgvector_success = self .write_pgvector( &chunk_id, &chunk_hash, &content, &source, &project, &level, &breadcrumb, embedding, ) .await; if !pgvector_success.is_ok() { tracing::error!("pgvector write failed: {}", pgvector_success.as_ref().err().unwrap()); // Extend visibility timeout for retry self.queue.change_visibility(&message.message_id, &message.receipt_handle, 300).await.ok(); return Ok(DualWriteResult { chunk_id, chunk_hash, pgvector_success: false, opensearch_success: false, opensearch_pending: false, error: Some(format!("{:?}", pgvector_success.err())), }); } // Write to OpenSearch (fail-soft) let opensearch_success = if let Some(os_client) = &self.opensearch { self.write_opensearch( os_client, &chunk_id, &content, &source, &project, &level, &breadcrumb, ) .await } else { Ok(()) }; let opensearch_pending = opensearch_success.is_err(); if opensearch_pending { tracing::warn!( "OpenSearch write failed, marking for retry: {}", opensearch_success.as_ref().err().unwrap() ); self.queue.change_visibility(&message.message_id, &message.receipt_handle, 300).await.ok(); } else { // Success: delete from queue self.queue.delete_chunk(&message.message_id, &message.receipt_handle).await.ok(); } Ok(DualWriteResult { chunk_id, chunk_hash, pgvector_success: pgvector_success.is_ok(), opensearch_success: opensearch_success.is_ok(), opensearch_pending, error: if opensearch_pending { Some(format!("{:?}", opensearch_success.err())) } else { None }, }) } /// Legacy: Direct dual-write (for backward compatibility) /// /// If queue adapter is not available, use this for synchronous processing. pub async fn dual_write( &self, chunk: &ChunkInput, embedding: &[f32], ) -> Result { let chunk_id = Uuid::new_v4(); let chunk_hash = self.compute_hash(&chunk.content); // Step 1: Check deduplication if self.is_already_indexed(&chunk_hash, &chunk.project).await? { tracing::debug!("Chunk already indexed (dedup): {}", chunk_hash); return Ok(DualWriteResult { chunk_id: Uuid::nil(), // Placeholder chunk_hash, pgvector_success: true, opensearch_success: true, opensearch_pending: false, error: Some("already_indexed".to_string()), }); } // Step 2: Write to pgvector let pgvector_success = self.write_pgvector( &chunk_id, &chunk_hash, &chunk.content, &chunk.source, &chunk.project, &chunk.level, &chunk.breadcrumb, embedding, ) .await; if !pgvector_success.is_ok() { tracing::error!("pgvector write failed: {}", pgvector_success.as_ref().err().unwrap()); return Ok(DualWriteResult { chunk_id, chunk_hash, pgvector_success: false, opensearch_success: false, opensearch_pending: false, error: Some(format!("{:?}", pgvector_success.err())), }); } // Step 3: Write to OpenSearch (fail-soft) let opensearch_success = if let Some(os_client) = &self.opensearch { self.write_opensearch( os_client, &chunk_id, &chunk.content, &chunk.source, &chunk.project, &chunk.level, &chunk.breadcrumb, ) .await } else { // OpenSearch not configured, skip Ok(()) }; let opensearch_pending = opensearch_success.is_err(); if opensearch_pending { tracing::warn!( "OpenSearch write failed for chunk {}, marked for retry: {}", chunk_id, opensearch_success.as_ref().err().unwrap() ); // Mark as pending in pgvector self.mark_opensearch_pending(&chunk_id).await.ok(); } // Step 4: Update indexed flags let pgvector_ok = pgvector_success.is_ok(); let opensearch_ok = opensearch_success.is_ok(); if pgvector_ok { self.update_pgvector_indexed(&chunk_id).await.ok(); } if opensearch_ok { self.update_opensearch_indexed(&chunk_id).await.ok(); } Ok(DualWriteResult { chunk_id, chunk_hash, pgvector_success: pgvector_ok, opensearch_success: opensearch_ok, opensearch_pending, error: if opensearch_pending { Some(format!("{:?}", opensearch_success.err())) } else { None }, }) } /// Compute SHA256 hash of content for deduplication fn compute_hash(&self, content: &str) -> String { let mut hasher = Sha256::new(); hasher.update(content.as_bytes()); format!("{:x}", hasher.finalize()) } /// Check if chunk is already fully indexed async fn is_already_indexed(&self, chunk_hash: &str, project: &str) -> Result { let row = sqlx::query_scalar::<_, bool>( "SELECT (indexed_in_pgvector AND indexed_in_opensearch) FROM chunks WHERE chunk_hash = $1 AND project = $2 LIMIT 1" ) .bind(chunk_hash) .bind(project) .fetch_optional(&self.pool) .await?; Ok(row.unwrap_or(false)) } /// Write chunk to pgvector async fn write_pgvector( &self, chunk_id: &Uuid, chunk_hash: &str, content: &str, source: &str, project: &str, level: &str, breadcrumb: &[String], embedding: &[f32], ) -> Result<()> { let embedding_vec = Vector::from(embedding.to_vec()); sqlx::query( "INSERT INTO chunks (id, chunk_hash, content, source, project, level, breadcrumb, embedding, indexed_in_pgvector, pgvector_indexed_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, true, now()) ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding, indexed_in_pgvector = true, pgvector_indexed_at = now()" ) .bind(chunk_id) .bind(chunk_hash) .bind(content) .bind(source) .bind(project) .bind(level) .bind(breadcrumb) .bind(embedding_vec) .execute(&self.pool) .await?; Ok(()) } /// Write chunk to OpenSearch async fn write_opensearch( &self, os_client: &Arc, chunk_id: &Uuid, content: &str, source: &str, project: &str, level: &str, breadcrumb: &[String], ) -> Result<()> { // Note: JWT token handling would come from AppState in http_server // For now, we'll pass empty token—production code should inject from context os_client .index_document( &chunk_id.to_string(), content, source, level, breadcrumb.to_vec(), "", // TODO: inject JWT from AppState ) .await?; Ok(()) } /// Mark chunk as pending OpenSearch retry async fn mark_opensearch_pending(&self, chunk_id: &Uuid) -> Result<()> { sqlx::query( "UPDATE chunks SET opensearch_pending = true, opensearch_retry_count = opensearch_retry_count + 1, opensearch_last_retry_at = now() WHERE id = $1" ) .bind(chunk_id) .execute(&self.pool) .await?; Ok(()) } /// Mark chunk as pgvector indexed async fn update_pgvector_indexed(&self, chunk_id: &Uuid) -> Result<()> { sqlx::query( "UPDATE chunks SET indexed_in_pgvector = true, pgvector_indexed_at = now() WHERE id = $1" ) .bind(chunk_id) .execute(&self.pool) .await?; Ok(()) } /// Mark chunk as OpenSearch indexed async fn update_opensearch_indexed(&self, chunk_id: &Uuid) -> Result<()> { sqlx::query( "UPDATE chunks SET indexed_in_opensearch = true, opensearch_pending = false, opensearch_indexed_at = now() WHERE id = $1" ) .bind(chunk_id) .execute(&self.pool) .await?; Ok(()) } /// Retry failed OpenSearch writes (background task) /// /// Polls for chunks where opensearch_pending=true and retries up to 3 times. /// Runs every 5 minutes. pub async fn retry_pending_chunks(&self, project: &str, max_retries: i32) -> Result { if self.opensearch.is_none() { return Ok(0); // Skip if OpenSearch not configured } let pending = sqlx::query_as::<_, (Uuid, String, String, String, Vec)>( "SELECT id, content, source, level, breadcrumb FROM chunks WHERE project = $1 AND opensearch_pending = true AND opensearch_retry_count < $2 ORDER BY opensearch_last_retry_at ASC LIMIT 100" ) .bind(project) .bind(max_retries) .fetch_all(&self.pool) .await?; let mut succeeded = 0; for (chunk_id, content, source, level, breadcrumb) in pending { if let Err(e) = self .write_opensearch( self.opensearch.as_ref().unwrap(), &chunk_id, &content, &source, project, &level, &breadcrumb, ) .await { tracing::warn!("Retry failed for chunk {}: {}", chunk_id, e); // Increment retry count sqlx::query( "UPDATE chunks SET opensearch_retry_count = opensearch_retry_count + 1, opensearch_last_retry_at = now() WHERE id = $1" ) .bind(&chunk_id) .execute(&self.pool) .await .ok(); } else { tracing::info!("Retry succeeded for chunk {}", chunk_id); self.update_opensearch_indexed(&chunk_id).await.ok(); succeeded += 1; } } Ok(succeeded) } /// Get retry statistics pub async fn retry_stats(&self, project: &str) -> Result<(usize, usize)> { let pending: (i64,) = sqlx::query_as( "SELECT COUNT(*) FROM chunks WHERE project = $1 AND opensearch_pending = true" ) .bind(project) .fetch_one(&self.pool) .await?; let failed: (i64,) = sqlx::query_as( "SELECT COUNT(*) FROM chunks WHERE project = $1 AND opensearch_retry_count >= 3" ) .bind(project) .fetch_one(&self.pool) .await?; Ok((pending.0 as usize, failed.0 as usize)) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_compute_hash() { let indexer = DualWriteIndexer::new( sqlx::pool::PoolOptions::new().max_connections(1).connect_lazy("postgresql://localhost").unwrap(), None, ); let hash1 = indexer.compute_hash("same content"); let hash2 = indexer.compute_hash("same content"); assert_eq!(hash1, hash2, "Same content must produce same hash"); let hash3 = indexer.compute_hash("different"); assert_ne!(hash1, hash3, "Different content must produce different hash"); } #[test] fn test_hash_deterministic() { let indexer = DualWriteIndexer::new( sqlx::pool::PoolOptions::new().max_connections(1).connect_lazy("postgresql://localhost").unwrap(), None, ); let content = "ERROR: permission denied\nStack trace..."; let hash1 = indexer.compute_hash(content); let hash2 = indexer.compute_hash(content); assert_eq!(hash1, hash2); assert_eq!(hash1.len(), 64); // SHA256 hex is 64 chars } }