feat: M8.2 Gateway Queue Adapter for SQS via api.riotpiao.com
- Unified QueueAdapter trait for concurrent dual-write operations - GatewayQueueAdapter routes messages via api.riotpiao.com with X-Service: sqs header - TokenProvider abstraction: StaticTokenProvider + AuthentikTokenProvider - JWT bearer token support (from Authentik OAuth2) - InMemoryQueueAdapter for testing - Base64 encoding/decoding for SQS message bodies - HTTP/REST integration (no direct gRPC complexity) - 8 unit tests + comprehensive documentation - Supports long-polling (ReceiveMessage), visibility timeout, DLQ Uses standard SQS API patterns: - SendMessage: Queue chunk for dual-write processing - ReceiveMessage: Long-poll up to 10 messages, 20s wait - DeleteMessage: Acknowledge on success - ChangeMessageVisibility: Retry on failure - SendToDLQ: After max retries Files: - crates/mem-cli/src/queue_adapter.rs (310 LOC) - crates/mem-cli/src/gateway_queue_adapter.rs (530 LOC) - tests/it_gateway_queue_adapter.rs (110 LOC) - docs/M8.2-GATEWAY_QUEUE_ADAPTER.md (400 LOC)
This commit is contained in:
@@ -0,0 +1,543 @@
|
||||
//! 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<Arc<OpenSearchClient>>,
|
||||
/// Queue adapter for concurrent dual-write processing
|
||||
/// Can be: kmsvc (production), in-memory (testing), or SQS (future)
|
||||
queue: Arc<dyn QueueAdapter>,
|
||||
}
|
||||
|
||||
/// 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<String>,
|
||||
}
|
||||
|
||||
/// 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<String>,
|
||||
}
|
||||
|
||||
impl DualWriteIndexer {
|
||||
/// Create dual-write indexer with queue adapter
|
||||
pub fn new(
|
||||
pool: PgPool,
|
||||
opensearch: Option<Arc<OpenSearchClient>>,
|
||||
queue: Arc<dyn QueueAdapter>,
|
||||
) -> 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<String> {
|
||||
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<DualWriteResult> {
|
||||
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::<Uuid>()?;
|
||||
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<String> = 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<DualWriteResult> {
|
||||
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<bool> {
|
||||
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<OpenSearchClient>,
|
||||
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<usize> {
|
||||
if self.opensearch.is_none() {
|
||||
return Ok(0); // Skip if OpenSearch not configured
|
||||
}
|
||||
|
||||
let pending = sqlx::query_as::<_, (Uuid, String, String, String, Vec<String>)>(
|
||||
"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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user