rock
cd3d00048a
- 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)
337 lines
9.8 KiB
Rust
337 lines
9.8 KiB
Rust
//! M8.2 — Unified Queue Adapter (SQS-compatible interface)
|
|
//!
|
|
//! Abstraction over external queue services (SQS, kmsvc, RabbitMQ, etc.)
|
|
//! Enables concurrent dual-write processing without database overhead.
|
|
//!
|
|
//! # Design
|
|
//!
|
|
//! Rather than storing queue state in the database, we leverage external queue
|
|
//! services via a unified API. This enables true horizontal scalability:
|
|
//!
|
|
//! ```text
|
|
//! Ingest Worker Queue Service (SQS/kmsvc) Dual-Write Workers
|
|
//! │ │ │
|
|
//! │─── send_chunk() ────────────>│ │
|
|
//! │ │ │
|
|
//! └──────────────────────────────┤<─── receive_chunks(10) ────────┤
|
|
//! │ │
|
|
//! │<─── delete_chunk() ────────────┤
|
|
//! │ (on success) │
|
|
//! │ │
|
|
//! │<─── change_visibility() ───────┤
|
|
//! │ (on retry) │
|
|
//! ```
|
|
//!
|
|
//! # Implementations
|
|
//! - `SqsQueueAdapter`: AWS SQS backend
|
|
//! - `KmsvcQueueAdapter`: Kubernetes native messaging service
|
|
//! - In-memory for testing
|
|
|
|
use async_trait::async_trait;
|
|
use serde::{Deserialize, Serialize};
|
|
use uuid::Uuid;
|
|
use anyhow::Result;
|
|
|
|
/// SQS-compatible message envelope
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct QueueMessage {
|
|
/// Unique message ID (from queue service)
|
|
pub message_id: String,
|
|
|
|
/// Original chunk UUID
|
|
pub chunk_id: Uuid,
|
|
|
|
/// Message body (serialized JSON)
|
|
pub body: String,
|
|
|
|
/// Receive count (number of times retrieved)
|
|
pub receive_count: i32,
|
|
|
|
/// Receipt handle (for delete/change_visibility)
|
|
pub receipt_handle: String,
|
|
|
|
/// Project context
|
|
pub project: String,
|
|
|
|
/// Metadata
|
|
pub attributes: std::collections::HashMap<String, String>,
|
|
}
|
|
|
|
/// Queue statistics
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct QueueStats {
|
|
pub available_messages: i64,
|
|
pub in_flight_messages: i64,
|
|
pub dead_letter_messages: i64,
|
|
pub total_processed: i64,
|
|
pub average_delay_secs: i64,
|
|
}
|
|
|
|
/// Unified queue adapter trait (SQS-like interface)
|
|
#[async_trait]
|
|
pub trait QueueAdapter: Send + Sync {
|
|
/// Send chunk message to queue
|
|
///
|
|
/// # Arguments
|
|
/// * `chunk_id` — Unique chunk identifier
|
|
/// * `body` — Serialized message body (JSON)
|
|
/// * `project` — Project context
|
|
/// * `attributes` — Optional metadata (e.g., source, level, breadcrumb)
|
|
///
|
|
/// # Returns
|
|
/// Message ID from queue service
|
|
async fn send_chunk(
|
|
&self,
|
|
chunk_id: Uuid,
|
|
body: String,
|
|
project: String,
|
|
attributes: std::collections::HashMap<String, String>,
|
|
) -> Result<String>;
|
|
|
|
/// Receive chunk messages from queue
|
|
///
|
|
/// # Arguments
|
|
/// * `max_messages` — Max number of messages (1-10)
|
|
/// * `visibility_timeout_secs` — Visibility timeout duration
|
|
/// * `project` — Project filter (optional)
|
|
///
|
|
/// # Returns
|
|
/// List of available messages
|
|
async fn receive_chunks(
|
|
&self,
|
|
max_messages: i32,
|
|
visibility_timeout_secs: i32,
|
|
project: Option<&str>,
|
|
) -> Result<Vec<QueueMessage>>;
|
|
|
|
/// Delete message from queue (after successful processing)
|
|
///
|
|
/// # Arguments
|
|
/// * `message_id` — Message to delete
|
|
/// * `receipt_handle` — Receipt handle (for idempotency)
|
|
async fn delete_chunk(&self, message_id: &str, receipt_handle: &str) -> Result<()>;
|
|
|
|
/// Change message visibility timeout
|
|
///
|
|
/// Called when processing takes longer than expected.
|
|
async fn change_visibility(
|
|
&self,
|
|
message_id: &str,
|
|
receipt_handle: &str,
|
|
visibility_timeout_secs: i32,
|
|
) -> Result<()>;
|
|
|
|
/// Send message to dead-letter queue
|
|
///
|
|
/// Called when message exceeds max receive count.
|
|
async fn send_to_dlq(&self, message_id: &str, receipt_handle: &str, reason: &str) -> Result<()>;
|
|
|
|
/// Get queue statistics
|
|
async fn get_stats(&self, project: Option<&str>) -> Result<QueueStats>;
|
|
|
|
/// Purge queue (test/admin only)
|
|
async fn purge(&self, project: Option<&str>) -> Result<usize>;
|
|
|
|
/// Health check
|
|
async fn health_check(&self) -> Result<()>;
|
|
}
|
|
|
|
/// In-memory queue adapter (for testing and local development)
|
|
pub struct InMemoryQueueAdapter {
|
|
messages: std::sync::Arc<tokio::sync::Mutex<Vec<QueueMessage>>>,
|
|
}
|
|
|
|
impl InMemoryQueueAdapter {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
messages: std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new())),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for InMemoryQueueAdapter {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl QueueAdapter for InMemoryQueueAdapter {
|
|
async fn send_chunk(
|
|
&self,
|
|
chunk_id: Uuid,
|
|
body: String,
|
|
project: String,
|
|
attributes: std::collections::HashMap<String, String>,
|
|
) -> Result<String> {
|
|
let message_id = format!("msg-{}", Uuid::new_v4());
|
|
let receipt_handle = format!("handle-{}", Uuid::new_v4());
|
|
|
|
let msg = QueueMessage {
|
|
message_id: message_id.clone(),
|
|
chunk_id,
|
|
body,
|
|
receive_count: 0,
|
|
receipt_handle,
|
|
project,
|
|
attributes,
|
|
};
|
|
|
|
let mut msgs = self.messages.lock().await;
|
|
msgs.push(msg);
|
|
|
|
Ok(message_id)
|
|
}
|
|
|
|
async fn receive_chunks(
|
|
&self,
|
|
max_messages: i32,
|
|
_visibility_timeout_secs: i32,
|
|
project: Option<&str>,
|
|
) -> Result<Vec<QueueMessage>> {
|
|
let mut msgs = self.messages.lock().await;
|
|
let max = max_messages.min(10).max(1) as usize;
|
|
let drain_count = msgs.len().min(max);
|
|
|
|
let result: Vec<_> = msgs
|
|
.drain(..drain_count)
|
|
.filter(|m| project.is_none() || m.project.as_str() == project.unwrap())
|
|
.collect();
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
async fn delete_chunk(&self, message_id: &str, _receipt_handle: &str) -> Result<()> {
|
|
let mut msgs = self.messages.lock().await;
|
|
msgs.retain(|m| m.message_id != message_id);
|
|
Ok(())
|
|
}
|
|
|
|
async fn change_visibility(
|
|
&self,
|
|
_message_id: &str,
|
|
_receipt_handle: &str,
|
|
_visibility_timeout_secs: i32,
|
|
) -> Result<()> {
|
|
// No-op for in-memory
|
|
Ok(())
|
|
}
|
|
|
|
async fn send_to_dlq(&self, message_id: &str, _receipt_handle: &str, _reason: &str) -> Result<()> {
|
|
let mut msgs = self.messages.lock().await;
|
|
msgs.retain(|m| m.message_id != message_id);
|
|
Ok(())
|
|
}
|
|
|
|
async fn get_stats(&self, _project: Option<&str>) -> Result<QueueStats> {
|
|
let msgs = self.messages.lock().await;
|
|
Ok(QueueStats {
|
|
available_messages: msgs.len() as i64,
|
|
in_flight_messages: 0,
|
|
dead_letter_messages: 0,
|
|
total_processed: 0,
|
|
average_delay_secs: 0,
|
|
})
|
|
}
|
|
|
|
async fn purge(&self, _project: Option<&str>) -> Result<usize> {
|
|
let mut msgs = self.messages.lock().await;
|
|
let count = msgs.len();
|
|
msgs.clear();
|
|
Ok(count)
|
|
}
|
|
|
|
async fn health_check(&self) -> Result<()> {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn test_in_memory_send_chunk() {
|
|
let queue = InMemoryQueueAdapter::new();
|
|
let msg_id = queue
|
|
.send_chunk(
|
|
Uuid::new_v4(),
|
|
r#"{"content": "test"}"#.to_string(),
|
|
"test-project".to_string(),
|
|
std::collections::HashMap::new(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(msg_id.starts_with("msg-"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_in_memory_receive_chunks() {
|
|
let queue = InMemoryQueueAdapter::new();
|
|
|
|
for i in 0..5 {
|
|
queue
|
|
.send_chunk(
|
|
Uuid::new_v4(),
|
|
format!(r#"{{"content": "test{}"}}"#, i),
|
|
"test-project".to_string(),
|
|
std::collections::HashMap::new(),
|
|
)
|
|
.await
|
|
.ok();
|
|
}
|
|
|
|
let messages = queue
|
|
.receive_chunks(3, 30, Some("test-project"))
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(messages.len(), 3);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_in_memory_delete_chunk() {
|
|
let queue = InMemoryQueueAdapter::new();
|
|
|
|
let msg_id = queue
|
|
.send_chunk(
|
|
Uuid::new_v4(),
|
|
"body".to_string(),
|
|
"test".to_string(),
|
|
std::collections::HashMap::new(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
queue.delete_chunk(&msg_id, "handle").await.unwrap();
|
|
|
|
let msgs = queue.receive_chunks(10, 30, None).await.unwrap();
|
|
assert_eq!(msgs.len(), 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_queue_stats() {
|
|
let queue = InMemoryQueueAdapter::new();
|
|
|
|
queue
|
|
.send_chunk(
|
|
Uuid::new_v4(),
|
|
"body".to_string(),
|
|
"test".to_string(),
|
|
std::collections::HashMap::new(),
|
|
)
|
|
.await
|
|
.ok();
|
|
|
|
let stats = queue.get_stats(None).await.unwrap();
|
|
assert_eq!(stats.available_messages, 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_health_check() {
|
|
let queue = InMemoryQueueAdapter::new();
|
|
assert!(queue.health_check().await.is_ok());
|
|
}
|
|
}
|