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

//! Parallel Dual-Write Indexer (Refactored)
//!
//! pgvector (primary, must succeed) + OpenSearch (secondary, fire-and-forget)
//! Both execute concurrently via tokio::join!

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 serde::{Deserialize, Serialize};

#[derive(Clone)]
pub struct ParallelDualWriteIndexer {
    pool: PgPool,
    opensearch: Option<Arc<OpenSearchClient>>,
}

/// Chunk to index
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexableChunk {
    pub chunk_id: String,
    pub content: String,
    pub source: String,
    pub project: String,
    pub level: String,
    pub breadcrumb: Vec<String>,
}

/// Result of parallel dual-write
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DualWriteResult {
    pub chunk_id: String,
    pub pgvector_success: bool,
    pub opensearch_success: bool,
    pub error: Option<String>,
}

impl ParallelDualWriteIndexer {
    pub fn new(pool: PgPool, opensearch: Option<Arc<OpenSearchClient>>) -> Self {
        Self { pool, opensearch }
    }

    /// Index chunk to both pgvector AND OpenSearch in parallel
    pub async fn index_parallel(&self, chunk: &IndexableChunk, embedding: &[f32]) -> Result<DualWriteResult> {
        let chunk_id = chunk.chunk_id.clone();

        // PARALLEL: Execute both writes concurrently
        let (pgvector_result, opensearch_result) = tokio::join!(
            self.write_pgvector(chunk, embedding),
            self.write_opensearch(chunk, embedding)
        );

        let pgvector_success = pgvector_result.is_ok();
        let opensearch_success = opensearch_result.is_ok();

        let error = if !pgvector_success {
            pgvector_result.err().map(|e| e.to_string())
        } else if !opensearch_success {
            opensearch_result.err().map(|e| e.to_string())
        } else {
            None
        };

        // Primary (pgvector) success = operation success
        if !pgvector_success {
            return Err(anyhow!("pgvector write failed: {:?}", error));
        }

        Ok(DualWriteResult {
            chunk_id,
            pgvector_success,
            opensearch_success,
            error,
        })
    }

    /// Write to pgvector (PRIMARY - must succeed)
    async fn write_pgvector(&self, chunk: &IndexableChunk, embedding: &[f32]) -> Result<()> {
        let vector = Vector::from(embedding.to_vec());
        let chunk_hash = self.compute_hash(&chunk.content);

        sqlx::query(
            "INSERT INTO memory_vector (id, project, level, text, embedding, breadcrumb, source, chunk_hash, indexed_at)
             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, now())
             ON CONFLICT (id) DO UPDATE SET
               indexed_at = now(),
               embedding = $5"
        )
        .bind(&chunk.chunk_id)
        .bind(&chunk.project)
        .bind(&chunk.level)
        .bind(&chunk.content)
        .bind(&vector)
        .bind(chunk.breadcrumb.join(" > "))
        .bind(&chunk.source)
        .bind(&chunk_hash)
        .execute(&self.pool)
        .await?;

        tracing::debug!("pgvector indexed: {}", chunk.chunk_id);
        Ok(())
    }

    /// Write to OpenSearch (SECONDARY - fire-and-forget)
    async fn write_opensearch(&self, chunk: &IndexableChunk, _embedding: &[f32]) -> Result<()> {
        if self.opensearch.is_none() {
            return Ok(());
        }

        let opensearch = self.opensearch.clone().unwrap();
        let chunk_id = chunk.chunk_id.clone();
        let chunk = chunk.clone();

        // Spawn background task (non-blocking)
        tokio::spawn(async move {
            let result = opensearch.index_document(
                &chunk_id,
                &chunk.content,
                &chunk.source,
                &chunk.level,
                chunk.breadcrumb.clone(),
                "",  // jwt_token - not available in background task
            ).await;

            match result {
                Ok(_) => tracing::debug!("OpenSearch indexed (async): {}", chunk_id),
                Err(e) => tracing::warn!("OpenSearch index failed (async, non-blocking): {}: {}", chunk_id, e),
            }
        });

        Ok(())
    }

    /// Batch parallel index (multiple chunks)
    pub async fn index_batch_parallel(
        &self,
        chunks: Vec<(&IndexableChunk, Vec<f32>)>,
    ) -> Vec<DualWriteResult> {
        let futures = chunks.into_iter().map(|(chunk, embedding)| async move {
            self.index_parallel(chunk, &embedding).await
        });

        futures::future::join_all(futures)
            .await
            .into_iter()
            .filter_map(|r| r.ok())
            .collect()
    }

    /// Compute SHA256 hash
    fn compute_hash(&self, content: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(content.as_bytes());
        format!("{:x}", hasher.finalize())
    }
}