From 2064e11143955ada9ad6dfbb705d4a06c9e46e05 Mon Sep 17 00:00:00 2001 From: Orual Date: Sat, 13 Dec 2025 12:40:56 -0500 Subject: [PATCH] parallel tap --- crates/weaver-index/src/bin/weaver_indexer.rs | 23 +- .../src/clickhouse/resilient_inserter.rs | 1 - crates/weaver-index/src/config.rs | 15 +- crates/weaver-index/src/indexer.rs | 222 +-------------- crates/weaver-index/src/lib.rs | 4 +- crates/weaver-index/src/parallel_tap.rs | 261 ++++++++++++++++++ 6 files changed, 299 insertions(+), 227 deletions(-) create mode 100644 crates/weaver-index/src/parallel_tap.rs diff --git a/crates/weaver-index/src/bin/weaver_indexer.rs b/crates/weaver-index/src/bin/weaver_indexer.rs index 5a62249..fb50314 100644 --- a/crates/weaver-index/src/bin/weaver_indexer.rs +++ b/crates/weaver-index/src/bin/weaver_indexer.rs @@ -6,6 +6,7 @@ use weaver_index::config::{ }; use weaver_index::firehose::FirehoseConsumer; use weaver_index::server::{AppState, ServerConfig, TelemetryConfig, telemetry}; +use weaver_index::clickhouse::InserterConfig; use weaver_index::{FirehoseIndexer, TapIndexer, load_cursor}; #[derive(Parser)] @@ -166,8 +167,15 @@ async fn run_full() -> miette::Result<()> { } SourceMode::Tap => { let tap_config = TapConfig::from_env()?; - let indexer = TapIndexer::new(indexer_client, tap_config, indexer_config); - info!("Starting tap indexer"); + let num_workers = tap_config.num_workers; + let indexer = TapIndexer::new( + indexer_client, + tap_config, + InserterConfig::default(), + indexer_config, + num_workers, + ); + info!("Starting tap indexer with {} workers", num_workers); tokio::spawn(async move { indexer.run().await }) } }; @@ -239,9 +247,16 @@ async fn run_tap_indexer( tap_config: TapConfig, indexer_config: IndexerConfig, ) -> miette::Result<()> { - let indexer = TapIndexer::new(client, tap_config, indexer_config); + let num_workers = tap_config.num_workers; + let indexer = TapIndexer::new( + client, + tap_config, + InserterConfig::default(), + indexer_config, + num_workers, + ); - info!("Starting tap indexer"); + info!("Starting tap indexer with {} workers", num_workers); indexer.run().await?; Ok(()) diff --git a/crates/weaver-index/src/clickhouse/resilient_inserter.rs b/crates/weaver-index/src/clickhouse/resilient_inserter.rs index 3e5b895..3bda4cd 100644 --- a/crates/weaver-index/src/clickhouse/resilient_inserter.rs +++ b/crates/weaver-index/src/clickhouse/resilient_inserter.rs @@ -513,7 +513,6 @@ fn extract_failing_row(error: &clickhouse::error::Error) -> Option { #[cfg(test)] mod tests { - use super::*; #[test] fn test_extract_failing_row() { diff --git a/crates/weaver-index/src/config.rs b/crates/weaver-index/src/config.rs index 3b97ccb..7aa1cea 100644 --- a/crates/weaver-index/src/config.rs +++ b/crates/weaver-index/src/config.rs @@ -207,17 +207,21 @@ impl IndexerConfig { pub struct TapConfig { pub url: Url, pub send_acks: bool, + pub num_workers: usize, } impl TapConfig { /// Default tap URL (local) pub const DEFAULT_URL: &'static str = "ws://localhost:2480/channel"; + /// Default number of parallel workers + pub const DEFAULT_WORKERS: usize = 4; /// Load configuration from environment variables. /// /// Optional env vars: /// - `TAP_URL`: Tap WebSocket URL (default: ws://localhost:2480/channel) /// - `TAP_SEND_ACKS`: Whether to send acks (default: true) + /// - `TAP_WORKERS`: Number of parallel workers (default: 4) pub fn from_env() -> Result { let url_str = std::env::var("TAP_URL").unwrap_or_else(|_| Self::DEFAULT_URL.to_string()); @@ -230,7 +234,16 @@ impl TapConfig { .map(|s| s.to_lowercase() != "false") .unwrap_or(true); - Ok(Self { url, send_acks }) + let num_workers = std::env::var("TAP_WORKERS") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(Self::DEFAULT_WORKERS); + + Ok(Self { + url, + send_acks, + num_workers, + }) } } diff --git a/crates/weaver-index/src/indexer.rs b/crates/weaver-index/src/indexer.rs index d2dc480..a028480 100644 --- a/crates/weaver-index/src/indexer.rs +++ b/crates/weaver-index/src/indexer.rs @@ -5,7 +5,7 @@ use chrono::Utc; use dashmap::DashMap; use n0_future::StreamExt; use smol_str::{SmolStr, ToSmolStr}; -use tracing::{debug, info, trace, warn}; +use tracing::{debug, info, warn}; use chrono::DateTime; @@ -14,13 +14,11 @@ use crate::clickhouse::{ RawRecordInsert, ResilientRecordInserter, }; use crate::config::IndexerConfig; -use crate::config::TapConfig; -use crate::error::{ClickHouseError, IndexError, Result}; +use crate::error::{IndexError, Result}; use crate::firehose::{ Account, ExtractedRecord, FirehoseConsumer, Identity, MessageStream, SubscribeReposMessage, extract_records, }; -use crate::tap::{TapConfig as TapConsumerConfig, TapConsumer, TapEvent}; /// Default consumer ID for cursor tracking const CONSUMER_ID: &str = "main"; @@ -523,219 +521,3 @@ async fn handle_delete(client: &Client, record: ExtractedRecord) -> Result<()> { tokio::time::sleep(Duration::from_secs(1)).await; } } - -// ============================================================================ -// TapIndexer - consumes from tap websocket -// ============================================================================ - -/// Consumer ID for tap cursor tracking -const TAP_CONSUMER_ID: &str = "tap"; - -/// Tap indexer that consumes from tap websocket and writes to ClickHouse -pub struct TapIndexer { - client: Arc, - tap_config: TapConfig, - config: IndexerConfig, -} - -impl TapIndexer { - /// Create a new tap indexer - pub fn new(client: Client, tap_config: TapConfig, config: IndexerConfig) -> Self { - Self { - client: Arc::new(client), - tap_config, - config, - } - } - - /// Save tap cursor to ClickHouse for visibility - async fn save_cursor(&self, seq: u64) -> Result<()> { - let query = format!( - "INSERT INTO firehose_cursor (consumer_id, seq, event_time) VALUES ('{}', {}, now64(3))", - TAP_CONSUMER_ID, seq - ); - - self.client.execute(&query).await?; - debug!(seq, "saved tap cursor"); - Ok(()) - } - - /// Run the tap indexer loop - pub async fn run(&self) -> Result<()> { - info!(url = %self.tap_config.url, "connecting to tap..."); - - let consumer_config = TapConsumerConfig::new(self.tap_config.url.clone()) - .with_acks(self.tap_config.send_acks); - let consumer = TapConsumer::new(consumer_config); - - let (mut events, ack_tx) = consumer.connect().await?; - - // Use resilient inserter for records since that's where untrusted JSON enters - let mut records = - ResilientRecordInserter::new(self.client.inner().clone(), InserterConfig::default()); - let mut identities = self - .client - .inserter::("raw_identity_events"); - - let mut processed: u64 = 0; - let mut last_seq: u64 = 0; - let mut last_stats = Instant::now(); - let mut last_cursor_save = Instant::now(); - - info!("starting tap indexer loop"); - - loop { - // Get time until next required flush - must commit before socket timeout (30s) - let records_time = records.time_left().unwrap_or(Duration::from_secs(10)); - let identities_time = identities.time_left().unwrap_or(Duration::from_secs(10)); - let time_left = records_time.min(identities_time); - - let event = match tokio::time::timeout(time_left, events.recv()).await { - Ok(Some(event)) => event, - Ok(None) => { - // Channel closed, exit loop - break; - } - Err(_) => { - // Timeout - flush inserters to keep INSERT alive - trace!("flush timeout, committing inserters"); - records.commit().await?; - identities - .commit() - .await - .map_err(|e| ClickHouseError::Query { - message: "periodic identities commit failed".into(), - source: e, - })?; - continue; - } - }; - - let event_id = event.id(); - last_seq = event_id; - - match event { - TapEvent::Record(envelope) => { - let record = &envelope.record; - - // Collection filter - if !self.config.collections.matches(&record.collection) { - // Still ack even if filtered - let _ = ack_tx.send(event_id).await; - continue; - } - - let json = match &record.record { - Some(v) => match serde_json::to_string(v) { - Ok(s) => s, - Err(e) => { - warn!( - did = %record.did, - collection = %record.collection, - rkey = %record.rkey, - error = ?e, - "failed to serialize record, sending to DLQ" - ); - let raw_data = format!( - r#"{{"did":"{}","collection":"{}","rkey":"{}","cid":"{}","error":"serialization_failed"}}"#, - record.did, record.collection, record.rkey, record.cid - ); - records - .write_raw_to_dlq( - record.action.as_str().to_smolstr(), - raw_data, - e.to_string(), - event_id, - ) - .await?; - let _ = ack_tx.send(event_id).await; - continue; - } - }, - None => "{}".to_string(), - }; - - debug!( - op = record.action.as_str(), - id = event_id, - len = json.len(), - "writing record" - ); - - records - .write(RawRecordInsert { - did: record.did.clone(), - collection: record.collection.clone(), - rkey: record.rkey.clone(), - cid: record.cid.clone(), - rev: record.rev.clone(), - record: json.to_smolstr(), - operation: record.action.as_str().to_smolstr(), - seq: event_id, - event_time: Utc::now(), - is_live: record.live, - }) - .await?; - records.commit().await?; - - processed += 1; - } - TapEvent::Identity(envelope) => { - let identity = &envelope.identity; - - identities - .write(&RawIdentityEvent { - did: identity.did.clone(), - handle: identity.handle.clone(), - seq: event_id, - event_time: Utc::now(), - }) - .await - .map_err(|e| ClickHouseError::Query { - message: "identity write failed".into(), - source: e, - })?; - identities - .commit() - .await - .map_err(|e| ClickHouseError::Query { - message: "identity commit failed".into(), - source: e, - })?; - } - } - - // Send ack after successful write+commit - let _ = ack_tx.send(event_id).await; - - // Periodic stats - if last_stats.elapsed() >= Duration::from_secs(10) { - info!(processed, last_seq, "tap indexer stats"); - last_stats = Instant::now(); - } - - // Save cursor every 30s for visibility - if last_cursor_save.elapsed() >= Duration::from_secs(30) && last_seq > 0 { - if let Err(e) = self.save_cursor(last_seq).await { - warn!(error = ?e, "failed to save tap cursor"); - } - last_cursor_save = Instant::now(); - } - } - - // Final flush - records.end().await?; - identities.end().await.map_err(|e| ClickHouseError::Query { - message: "final identities flush failed".into(), - source: e, - })?; - - // Final cursor save - if last_seq > 0 { - self.save_cursor(last_seq).await?; - } - - info!(last_seq, "tap stream ended"); - Ok(()) - } -} diff --git a/crates/weaver-index/src/lib.rs b/crates/weaver-index/src/lib.rs index 93a4e7c..6bb66ee 100644 --- a/crates/weaver-index/src/lib.rs +++ b/crates/weaver-index/src/lib.rs @@ -4,12 +4,14 @@ pub mod endpoints; pub mod error; pub mod firehose; pub mod indexer; +pub mod parallel_tap; pub mod server; pub mod sqlite; pub mod tap; pub use config::Config; pub use error::{IndexError, Result}; -pub use indexer::{FirehoseIndexer, TapIndexer, load_cursor}; +pub use indexer::{FirehoseIndexer, load_cursor}; +pub use parallel_tap::TapIndexer; pub use server::{AppState, ServerConfig}; pub use sqlite::{ShardKey, ShardRouter, SqliteShard}; diff --git a/crates/weaver-index/src/parallel_tap.rs b/crates/weaver-index/src/parallel_tap.rs new file mode 100644 index 0000000..974b57c --- /dev/null +++ b/crates/weaver-index/src/parallel_tap.rs @@ -0,0 +1,261 @@ +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use chrono::Utc; +use smol_str::ToSmolStr; +use tokio::task::JoinHandle; +use tracing::{debug, error, info, trace, warn}; + +use crate::clickhouse::{ + Client, InserterConfig, RawIdentityEvent, RawRecordInsert, ResilientRecordInserter, +}; +use crate::config::{IndexerConfig, TapConfig}; +use crate::error::{ClickHouseError, Result}; +use crate::tap::{TapConfig as TapConsumerConfig, TapConsumer, TapEvent}; + +/// TAP indexer with multiple parallel websocket connections +/// +/// Each worker maintains its own websocket connection to TAP and its own +/// ClickHouse inserter. TAP distributes events across connected clients, +/// and its ack-gating mechanism ensures per-DID ordering is preserved +/// regardless of which worker handles which events. +pub struct TapIndexer { + client: Arc, + tap_config: TapConfig, + inserter_config: InserterConfig, + config: Arc, + num_workers: usize, +} + +impl TapIndexer { + pub fn new( + client: Client, + tap_config: TapConfig, + inserter_config: InserterConfig, + config: IndexerConfig, + num_workers: usize, + ) -> Self { + Self { + client: Arc::new(client), + tap_config, + inserter_config, + config: Arc::new(config), + num_workers, + } + } + + pub async fn run(&self) -> Result<()> { + info!( + num_workers = self.num_workers, + url = %self.tap_config.url, + "starting parallel tap indexer" + ); + + let mut handles: Vec>> = Vec::with_capacity(self.num_workers); + + for worker_id in 0..self.num_workers { + let client = self.client.clone(); + let tap_config = self.tap_config.clone(); + let inserter_config = self.inserter_config.clone(); + let config = self.config.clone(); + + let handle = tokio::spawn(async move { + run_tap_worker(worker_id, client, tap_config, inserter_config, config).await + }); + + handles.push(handle); + } + + // Wait for all workers + // TODO: Implement proper supervision - restart failed workers instead of propagating + for (i, handle) in handles.into_iter().enumerate() { + match handle.await { + Ok(Ok(())) => { + info!(worker_id = i, "tap worker finished cleanly"); + } + Ok(Err(e)) => { + error!(worker_id = i, error = ?e, "tap worker failed"); + return Err(e); + } + Err(e) => { + error!(worker_id = i, error = ?e, "tap worker panicked"); + return Err(crate::error::FirehoseError::Stream { + message: format!("worker {} panicked: {}", i, e), + } + .into()); + } + } + } + + Ok(()) + } +} + +async fn run_tap_worker( + worker_id: usize, + client: Arc, + tap_config: TapConfig, + inserter_config: InserterConfig, + config: Arc, +) -> Result<()> { + info!(worker_id, url = %tap_config.url, "tap worker starting"); + + let consumer_config = + TapConsumerConfig::new(tap_config.url.clone()).with_acks(tap_config.send_acks); + let consumer = TapConsumer::new(consumer_config); + + let (mut events, ack_tx) = consumer.connect().await?; + + // Each worker has its own resilient inserter + let mut records = ResilientRecordInserter::new(client.inner().clone(), inserter_config); + let mut identities = client.inserter::("raw_identity_events"); + + let mut processed: u64 = 0; + let mut last_stats = Instant::now(); + + info!(worker_id, "tap worker connected, starting event loop"); + + loop { + // Get time until next required flush + let records_time = records.time_left().unwrap_or(Duration::from_secs(10)); + let identities_time = identities.time_left().unwrap_or(Duration::from_secs(10)); + let time_left = records_time.min(identities_time); + + let event = match tokio::time::timeout(time_left, events.recv()).await { + Ok(Some(event)) => event, + Ok(None) => { + info!(worker_id, "tap channel closed, exiting"); + break; + } + Err(_) => { + // Timeout - flush inserters + trace!(worker_id, "flush timeout, committing inserters"); + records.commit().await?; + identities + .commit() + .await + .map_err(|e| ClickHouseError::Query { + message: "periodic identities commit failed".into(), + source: e, + })?; + continue; + } + }; + + let event_id = event.id(); + + match event { + TapEvent::Record(envelope) => { + let record = &envelope.record; + + // Collection filter + if !config.collections.matches(&record.collection) { + let _ = ack_tx.send(event_id).await; + continue; + } + + // Serialize record + let json = match &record.record { + Some(v) => match serde_json::to_string(v) { + Ok(s) => s, + Err(e) => { + warn!( + worker_id, + did = %record.did, + collection = %record.collection, + rkey = %record.rkey, + error = ?e, + "failed to serialize record, sending to DLQ" + ); + let raw_data = format!( + r#"{{"did":"{}","collection":"{}","rkey":"{}","cid":"{}","error":"serialization_failed"}}"#, + record.did, record.collection, record.rkey, record.cid + ); + records + .write_raw_to_dlq( + record.action.as_str().to_smolstr(), + raw_data, + e.to_string(), + event_id, + ) + .await?; + let _ = ack_tx.send(event_id).await; + continue; + } + }, + None => "{}".to_string(), + }; + + debug!( + worker_id, + op = record.action.as_str(), + id = event_id, + len = json.len(), + "writing record" + ); + + records + .write(RawRecordInsert { + did: record.did.clone(), + collection: record.collection.clone(), + rkey: record.rkey.clone(), + cid: record.cid.clone(), + rev: record.rev.clone(), + record: json.to_smolstr(), + operation: record.action.as_str().to_smolstr(), + seq: event_id, + event_time: Utc::now(), + is_live: record.live, + }) + .await?; + records.commit().await?; + + // Ack after successful processing + let _ = ack_tx.send(event_id).await; + + processed += 1; + } + TapEvent::Identity(envelope) => { + let identity = &envelope.identity; + + identities + .write(&RawIdentityEvent { + did: identity.did.clone(), + handle: identity.handle.clone(), + seq: event_id, + event_time: Utc::now(), + }) + .await + .map_err(|e| ClickHouseError::Query { + message: "identity write failed".into(), + source: e, + })?; + identities + .commit() + .await + .map_err(|e| ClickHouseError::Query { + message: "identity commit failed".into(), + source: e, + })?; + + let _ = ack_tx.send(event_id).await; + } + } + + // Periodic stats + if last_stats.elapsed() > Duration::from_secs(30) { + info!(worker_id, processed, "tap worker stats"); + last_stats = Instant::now(); + } + } + + // Clean shutdown + records.end().await?; + identities.end().await.map_err(|e| ClickHouseError::Query { + message: "identities end failed".into(), + source: e, + })?; + + info!(worker_id, processed, "tap worker finished"); + Ok(()) +} -- knotmirror 0.0.1