use crate::destination::{Destination, DestinationError};
use crate::event::ChangeEvent;
use crate::metrics;
use crate::state::StateStore;
use crate::stream::{ChangeStreamConfig, ChangeStreamListener};
use crate::watch_level::WatchLevel;
use futures::StreamExt;
use mongodb::bson::Document;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{broadcast, Mutex, RwLock};
use tokio::task::JoinHandle;
use tokio::time::{interval, Instant};
use tracing::{debug, error, info, instrument, warn};
#[derive(Debug, Clone)]
pub struct DistributedLockConfig {
pub enabled: bool,
pub ttl: Duration,
pub refresh_interval: Duration,
pub retry_interval: Duration,
}
impl Default for DistributedLockConfig {
fn default() -> Self {
Self {
enabled: true, ttl: Duration::from_secs(30),
refresh_interval: Duration::from_secs(10),
retry_interval: Duration::from_secs(5),
}
}
}
impl DistributedLockConfig {
#[must_use]
pub fn builder() -> DistributedLockConfigBuilder {
DistributedLockConfigBuilder::default()
}
#[must_use]
pub fn disabled() -> Self {
Self {
enabled: false,
..Default::default()
}
}
}
#[derive(Debug, Default)]
pub struct DistributedLockConfigBuilder {
enabled: Option<bool>,
ttl: Option<Duration>,
refresh_interval: Option<Duration>,
retry_interval: Option<Duration>,
}
impl DistributedLockConfigBuilder {
#[must_use]
pub fn enabled(mut self, enabled: bool) -> Self {
self.enabled = Some(enabled);
self
}
#[must_use]
pub fn ttl(mut self, ttl: Duration) -> Self {
self.ttl = Some(ttl);
self
}
#[must_use]
pub fn refresh_interval(mut self, interval: Duration) -> Self {
self.refresh_interval = Some(interval);
self
}
#[must_use]
pub fn retry_interval(mut self, interval: Duration) -> Self {
self.retry_interval = Some(interval);
self
}
pub fn build(self) -> Result<DistributedLockConfig, ConfigError> {
let ttl = self.ttl.unwrap_or(Duration::from_secs(30));
let refresh_interval = self.refresh_interval.unwrap_or(Duration::from_secs(10));
if refresh_interval >= ttl {
return Err(ConfigError::InvalidLockConfig {
reason: format!(
"refresh_interval ({:?}) must be less than ttl ({:?})",
refresh_interval, ttl
),
});
}
Ok(DistributedLockConfig {
enabled: self.enabled.unwrap_or(true),
ttl,
refresh_interval,
retry_interval: self.retry_interval.unwrap_or(Duration::from_secs(5)),
})
}
}
#[derive(Debug, Clone)]
pub struct PipelineConfig {
pub mongodb_uri: String,
pub database: String,
pub watch_level: WatchLevel,
pub batch_size: usize,
pub batch_timeout: Duration,
pub max_retries: usize,
pub retry_delay: Duration,
pub max_retry_delay: Duration,
pub channel_buffer_size: usize,
pub stream_config: ChangeStreamConfig,
pub distributed_lock: DistributedLockConfig,
}
impl PipelineConfig {
#[must_use]
pub fn builder() -> PipelineConfigBuilder {
PipelineConfigBuilder::default()
}
}
#[derive(Debug, Default)]
pub struct PipelineConfigBuilder {
mongodb_uri: Option<String>,
database: Option<String>,
watch_level: Option<WatchLevel>,
batch_size: usize,
batch_timeout: Duration,
max_retries: usize,
retry_delay: Duration,
max_retry_delay: Duration,
channel_buffer_size: usize,
stream_config: Option<ChangeStreamConfig>,
distributed_lock: Option<DistributedLockConfig>,
}
impl PipelineConfigBuilder {
#[must_use]
pub fn mongodb_uri(mut self, uri: impl Into<String>) -> Self {
self.mongodb_uri = Some(uri.into());
self
}
#[must_use]
pub fn database(mut self, database: impl Into<String>) -> Self {
self.database = Some(database.into());
self
}
#[must_use]
#[deprecated(since = "0.2.0", note = "Use watch_collections() instead")]
pub fn collections(mut self, collections: Vec<String>) -> Self {
self.watch_level = Some(WatchLevel::Collection(collections));
self
}
#[must_use]
pub fn watch_collections(mut self, collections: Vec<String>) -> Self {
self.watch_level = Some(WatchLevel::Collection(collections));
self
}
#[must_use]
pub fn watch_database(mut self) -> Self {
self.watch_level = Some(WatchLevel::Database);
self
}
#[must_use]
pub fn watch_deployment(mut self) -> Self {
self.watch_level = Some(WatchLevel::Deployment);
self
}
#[must_use]
pub fn batch_size(mut self, size: usize) -> Self {
self.batch_size = size;
self
}
#[must_use]
pub fn batch_timeout(mut self, timeout: Duration) -> Self {
self.batch_timeout = timeout;
self
}
#[must_use]
pub fn max_retries(mut self, retries: usize) -> Self {
self.max_retries = retries;
self
}
#[must_use]
pub fn retry_delay(mut self, delay: Duration) -> Self {
self.retry_delay = delay;
self
}
#[must_use]
pub fn max_retry_delay(mut self, delay: Duration) -> Self {
self.max_retry_delay = delay;
self
}
#[must_use]
pub fn channel_buffer_size(mut self, size: usize) -> Self {
self.channel_buffer_size = size;
self
}
#[must_use]
pub fn stream_config(mut self, config: ChangeStreamConfig) -> Self {
self.stream_config = Some(config);
self
}
#[must_use]
pub fn distributed_lock(mut self, config: DistributedLockConfig) -> Self {
self.distributed_lock = Some(config);
self
}
#[must_use]
pub fn disable_distributed_lock(mut self) -> Self {
self.distributed_lock = Some(DistributedLockConfig::disabled());
self
}
pub fn build(self) -> Result<PipelineConfig, ConfigError> {
let mongodb_uri = self.mongodb_uri.ok_or(ConfigError::MissingMongoUri)?;
let database = self.database.ok_or(ConfigError::MissingDatabase)?;
let watch_level = self.watch_level.unwrap_or_default();
let batch_size = match self.batch_size {
0 => 100, size if size > 10_000 => {
return Err(ConfigError::InvalidBatchSize {
value: size,
reason: "batch_size exceeds maximum (10,000)",
})
}
size => size,
};
let batch_timeout = if self.batch_timeout.is_zero() {
Duration::from_secs(5) } else {
self.batch_timeout
};
let retry_delay = if self.retry_delay.is_zero() {
Duration::from_millis(100)
} else {
self.retry_delay
};
let max_retry_delay = if self.max_retry_delay.is_zero() {
Duration::from_secs(30)
} else {
self.max_retry_delay
};
if retry_delay > max_retry_delay {
return Err(ConfigError::RetryDelayExceedsMax {
retry_delay,
max_retry_delay,
});
}
let channel_buffer_size = match self.channel_buffer_size {
0 => 1000, size if size < 10 => {
return Err(ConfigError::InvalidChannelBufferSize {
value: size,
reason: "channel_buffer_size must be at least 10",
})
}
size => size,
};
let stream_config = self.stream_config.unwrap_or_else(|| {
ChangeStreamConfig::builder()
.build()
.expect("Default stream config should build")
});
let distributed_lock = self.distributed_lock.unwrap_or_default();
Ok(PipelineConfig {
mongodb_uri,
database,
watch_level,
batch_size,
batch_timeout,
max_retries: self.max_retries,
retry_delay,
max_retry_delay,
channel_buffer_size,
stream_config,
distributed_lock,
})
}
}
#[derive(Debug, Clone, Default)]
pub struct PipelineStats {
pub events_processed: u64,
pub batches_written: u64,
pub write_errors: u64,
pub retries: u64,
}
type WorkerHandle = JoinHandle<Result<(), PipelineError>>;
type LockRefreshHandle = JoinHandle<()>;
pub struct Pipeline<S: StateStore, D: Destination> {
config: PipelineConfig,
store: Arc<S>,
destination: Arc<Mutex<D>>,
shutdown_tx: Option<broadcast::Sender<()>>,
workers: Arc<RwLock<Vec<WorkerHandle>>>,
lock_refresh_handles: Arc<RwLock<Vec<LockRefreshHandle>>>,
stats: Arc<RwLock<PipelineStats>>,
running: Arc<RwLock<bool>>,
owner_id: String,
locked_collections: Arc<RwLock<Vec<String>>>,
}
impl<S: StateStore + Send + Sync + 'static, D: Destination + Send + Sync + 'static> Pipeline<S, D> {
pub async fn new(
config: PipelineConfig,
store: S,
destination: D,
) -> Result<Self, PipelineError> {
let hostname = hostname::get()
.map(|h| h.to_string_lossy().to_string())
.unwrap_or_else(|_| "unknown".to_string());
let owner_id = format!("{}-{}", hostname, uuid::Uuid::new_v4());
info!(
database = %config.database,
watch_level = %config.watch_level,
batch_size = config.batch_size,
batch_timeout = ?config.batch_timeout,
owner_id = %owner_id,
distributed_lock_enabled = config.distributed_lock.enabled,
"Creating pipeline"
);
Ok(Self {
config,
store: Arc::new(store),
destination: Arc::new(Mutex::new(destination)),
shutdown_tx: None,
workers: Arc::new(RwLock::new(Vec::new())),
lock_refresh_handles: Arc::new(RwLock::new(Vec::new())),
stats: Arc::new(RwLock::new(PipelineStats::default())),
running: Arc::new(RwLock::new(false)),
owner_id,
locked_collections: Arc::new(RwLock::new(Vec::new())),
})
}
#[must_use]
pub fn owner_id(&self) -> &str {
&self.owner_id
}
fn lock_key(&self, collection: &str) -> String {
format!("rigatoni:lock:{}:{}", self.config.database, collection)
}
fn database_lock_key(&self) -> String {
format!("rigatoni:lock:{}:__database__", self.config.database)
}
fn deployment_lock_key(&self) -> String {
"rigatoni:lock:__deployment__".to_string()
}
#[instrument(skip(self), fields(database = %self.config.database, owner_id = %self.owner_id))]
pub async fn start(&mut self) -> Result<(), PipelineError> {
let mut running = self.running.write().await;
if *running {
return Err(PipelineError::AlreadyRunning);
}
info!(
watch_level = %self.config.watch_level,
distributed_lock_enabled = self.config.distributed_lock.enabled,
"Starting pipeline"
);
let (shutdown_tx, _) = broadcast::channel(1);
self.shutdown_tx = Some(shutdown_tx.clone());
let mut workers = self.workers.write().await;
let mut lock_refresh_handles = self.lock_refresh_handles.write().await;
let mut locked_collections = self.locked_collections.write().await;
let mut num_workers = 0;
match &self.config.watch_level {
WatchLevel::Collection(collections) => {
if collections.is_empty() {
return Err(PipelineError::Configuration(
"No collections specified. Either provide collection names with \
watch_collections() or use watch_database() to watch all collections."
.to_string(),
));
}
info!(
collections = ?collections,
"Starting collection-level watching"
);
for collection in collections {
let lock_key = self.lock_key(collection);
if self.config.distributed_lock.enabled {
let acquired = self
.store
.try_acquire_lock(
&lock_key,
&self.owner_id,
self.config.distributed_lock.ttl,
)
.await
.map_err(|e| PipelineError::StateStore(e.to_string()))?;
if !acquired {
info!(
collection = %collection,
"Collection is locked by another instance, skipping"
);
metrics::increment_lock_acquisition_failures(
metrics::LockFailureReason::AlreadyHeld,
);
continue;
}
info!(collection = %collection, "Acquired lock for collection");
metrics::increment_lock_acquisitions();
locked_collections.push(collection.clone());
let refresh_handle = self.start_lock_refresh_task(
lock_key,
collection.clone(),
shutdown_tx.subscribe(),
);
lock_refresh_handles.push(refresh_handle);
}
let shutdown_rx = shutdown_tx.subscribe();
let worker = self
.spawn_collection_worker(collection.clone(), shutdown_rx)
.await?;
workers.push(worker);
num_workers += 1;
}
metrics::set_active_collections(num_workers);
}
WatchLevel::Database => {
info!(
database = %self.config.database,
"Starting database-level watching"
);
let lock_key = self.database_lock_key();
if self.config.distributed_lock.enabled {
let acquired = self
.store
.try_acquire_lock(
&lock_key,
&self.owner_id,
self.config.distributed_lock.ttl,
)
.await
.map_err(|e| PipelineError::StateStore(e.to_string()))?;
if !acquired {
info!("Database is locked by another instance, cannot start");
metrics::increment_lock_acquisition_failures(
metrics::LockFailureReason::AlreadyHeld,
);
return Err(PipelineError::Configuration(
"Database is locked by another instance. \
Wait for the lock to expire or use collection-level watching."
.to_string(),
));
}
info!("Acquired lock for database");
metrics::increment_lock_acquisitions();
locked_collections.push("__database__".to_string());
let refresh_handle = self.start_lock_refresh_task(
lock_key,
"__database__".to_string(),
shutdown_tx.subscribe(),
);
lock_refresh_handles.push(refresh_handle);
}
let shutdown_rx = shutdown_tx.subscribe();
let worker = self.spawn_database_worker(shutdown_rx).await?;
workers.push(worker);
num_workers = 1;
metrics::set_active_collections(1);
}
WatchLevel::Deployment => {
info!("Starting deployment-level watching (cluster-wide)");
let lock_key = self.deployment_lock_key();
if self.config.distributed_lock.enabled {
let acquired = self
.store
.try_acquire_lock(
&lock_key,
&self.owner_id,
self.config.distributed_lock.ttl,
)
.await
.map_err(|e| PipelineError::StateStore(e.to_string()))?;
if !acquired {
info!("Deployment is locked by another instance, cannot start");
metrics::increment_lock_acquisition_failures(
metrics::LockFailureReason::AlreadyHeld,
);
return Err(PipelineError::Configuration(
"Deployment is locked by another instance. \
Wait for the lock to expire or use collection/database-level watching."
.to_string(),
));
}
info!("Acquired lock for deployment");
metrics::increment_lock_acquisitions();
locked_collections.push("__deployment__".to_string());
let refresh_handle = self.start_lock_refresh_task(
lock_key,
"__deployment__".to_string(),
shutdown_tx.subscribe(),
);
lock_refresh_handles.push(refresh_handle);
}
let shutdown_rx = shutdown_tx.subscribe();
let worker = self.spawn_deployment_worker(shutdown_rx).await?;
workers.push(worker);
num_workers = 1;
metrics::set_active_collections(1);
}
}
*running = true;
info!(
workers = num_workers,
locks_held = locked_collections.len(),
"Pipeline started"
);
metrics::set_pipeline_status(metrics::PipelineStatus::Running);
metrics::set_locks_held(locked_collections.len());
Ok(())
}
fn start_lock_refresh_task(
&self,
lock_key: String,
collection_name: String,
mut shutdown_rx: broadcast::Receiver<()>,
) -> LockRefreshHandle {
let store = Arc::clone(&self.store);
let owner_id = self.owner_id.clone();
let refresh_interval = self.config.distributed_lock.refresh_interval;
let ttl = self.config.distributed_lock.ttl;
tokio::spawn(async move {
let mut interval_timer = interval(refresh_interval);
interval_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tokio::select! {
_ = shutdown_rx.recv() => {
debug!(
collection = %collection_name,
"Lock refresh task received shutdown signal"
);
break;
}
_ = interval_timer.tick() => {
match store.refresh_lock(&lock_key, &owner_id, ttl).await {
Ok(true) => {
debug!(
collection = %collection_name,
"Lock refreshed successfully"
);
metrics::increment_lock_refreshes();
}
Ok(false) => {
error!(
collection = %collection_name,
"Lost lock (acquired by another instance or expired)"
);
metrics::increment_locks_lost();
break;
}
Err(e) => {
warn!(
collection = %collection_name,
error = %e,
"Failed to refresh lock (transient error, will retry)"
);
}
}
}
}
}
})
}
async fn spawn_collection_worker(
&self,
collection: String,
shutdown_rx: broadcast::Receiver<()>,
) -> Result<WorkerHandle, PipelineError> {
let config = self.config.clone();
let store = Arc::clone(&self.store);
let destination = Arc::clone(&self.destination);
let stats = Arc::clone(&self.stats);
let handle = tokio::spawn(async move {
Self::collection_worker(collection, config, store, destination, stats, shutdown_rx)
.await
});
Ok(handle)
}
async fn spawn_database_worker(
&self,
shutdown_rx: broadcast::Receiver<()>,
) -> Result<WorkerHandle, PipelineError> {
let config = self.config.clone();
let store = Arc::clone(&self.store);
let destination = Arc::clone(&self.destination);
let stats = Arc::clone(&self.stats);
let handle = tokio::spawn(async move {
Self::database_worker(config, store, destination, stats, shutdown_rx).await
});
Ok(handle)
}
async fn spawn_deployment_worker(
&self,
shutdown_rx: broadcast::Receiver<()>,
) -> Result<WorkerHandle, PipelineError> {
let config = self.config.clone();
let store = Arc::clone(&self.store);
let destination = Arc::clone(&self.destination);
let stats = Arc::clone(&self.stats);
let handle = tokio::spawn(async move {
Self::deployment_worker(config, store, destination, stats, shutdown_rx).await
});
Ok(handle)
}
#[allow(clippy::too_many_lines)]
#[instrument(skip(config, store, destination, stats, shutdown_rx), fields(collection = %collection))]
async fn collection_worker(
collection: String,
config: PipelineConfig,
store: Arc<S>,
destination: Arc<Mutex<D>>,
stats: Arc<RwLock<PipelineStats>>,
mut shutdown_rx: broadcast::Receiver<()>,
) -> Result<(), PipelineError> {
info!("Starting collection worker");
let resume_token_key = config
.watch_level
.resume_token_key(&config.database, Some(&collection));
let resume_token = store
.get_resume_token(&resume_token_key)
.await
.map_err(|e| PipelineError::StateStore(e.to_string()))?;
if let Some(ref token) = resume_token {
info!(?token, "Resuming from saved token");
}
let client = mongodb::Client::with_uri_str(&config.mongodb_uri)
.await
.map_err(|e| PipelineError::MongoDB(e.to_string()))?;
let db = client.database(&config.database);
let mongo_collection = db.collection(&collection);
let store_clone = Arc::clone(&store);
let resume_key = resume_token_key.clone();
let resume_token_callback = move |token: Document| {
let store = Arc::clone(&store_clone);
let key = resume_key.clone();
Box::pin(async move {
store
.save_resume_token(&key, &token)
.await
.map_err(|e| e.to_string())
}) as Pin<Box<dyn Future<Output = Result<(), String>> + Send>>
};
let mut listener = ChangeStreamListener::new(
mongo_collection,
config.stream_config.clone(),
resume_token_callback,
)
.await
.map_err(|e| PipelineError::ChangeStream(e.to_string()))?;
let mut batch: Vec<ChangeEvent> = Vec::with_capacity(config.batch_size);
let mut last_resume_token: Option<Document> = None;
let mut batch_timer = interval(config.batch_timeout);
batch_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
info!(
batch_size = config.batch_size,
batch_timeout = ?config.batch_timeout,
"Worker event loop started"
);
loop {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("Received shutdown signal");
if !batch.is_empty() {
info!(batch_size = batch.len(), "Flushing pending batch on shutdown");
if let Err(e) = Self::flush_batch(
&collection,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush batch on shutdown");
}
}
info!("Worker shutting down gracefully");
break;
}
_ = batch_timer.tick() => {
if !batch.is_empty() {
debug!(batch_size = batch.len(), "Batch timeout - flushing");
if let Err(e) = Self::flush_batch(
&collection,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush batch on timeout");
}
}
}
event_result = listener.next() => {
match event_result {
Some(Ok(ackable_event)) => {
let event = ackable_event.event.clone();
debug!(
operation = ?event.operation,
collection = %event.namespace.collection,
"Received event"
);
last_resume_token = Some(event.resume_token.clone());
ackable_event.ack();
batch.push(event.clone());
metrics::increment_batch_queue_size(&collection);
if batch.len() >= config.batch_size {
debug!(batch_size = batch.len(), "Batch full - flushing");
if let Err(e) = Self::flush_batch(
&collection,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush full batch");
}
}
}
Some(Err(e)) => {
error!(?e, "Error reading from change stream");
tokio::time::sleep(Duration::from_secs(1)).await;
}
None => {
warn!("Change stream ended unexpectedly");
break;
}
}
}
}
}
Ok(())
}
#[allow(clippy::too_many_lines)]
#[instrument(skip(config, store, destination, stats, shutdown_rx), fields(database = %config.database))]
async fn database_worker(
config: PipelineConfig,
store: Arc<S>,
destination: Arc<Mutex<D>>,
stats: Arc<RwLock<PipelineStats>>,
mut shutdown_rx: broadcast::Receiver<()>,
) -> Result<(), PipelineError> {
info!("Starting database worker");
let resume_token_key = config.watch_level.resume_token_key(&config.database, None);
let resume_token = store
.get_resume_token(&resume_token_key)
.await
.map_err(|e| PipelineError::StateStore(e.to_string()))?;
if let Some(ref token) = resume_token {
info!(?token, "Resuming from saved token");
}
let client = mongodb::Client::with_uri_str(&config.mongodb_uri)
.await
.map_err(|e| PipelineError::MongoDB(e.to_string()))?;
let db = client.database(&config.database);
let mut options = mongodb::options::ChangeStreamOptions::default();
if config.stream_config.full_document_on_update {
options.full_document = Some(mongodb::options::FullDocumentType::UpdateLookup);
}
if config.stream_config.full_document_before_change {
options.full_document_before_change =
Some(mongodb::options::FullDocumentBeforeChangeType::WhenAvailable);
}
options.batch_size = config.stream_config.batch_size;
if let Some(ref token_doc) = resume_token {
if let Ok(bytes) = bson::to_vec(token_doc) {
if let Ok(resume_token) =
bson::from_slice::<mongodb::change_stream::event::ResumeToken>(&bytes)
{
options.resume_after = Some(resume_token);
}
}
}
let mut stream = if config.stream_config.pipeline.is_empty() {
db.watch().with_options(options).await
} else {
db.watch()
.pipeline(config.stream_config.pipeline.clone())
.with_options(options)
.await
}
.map_err(|e| PipelineError::MongoDB(format!("Failed to create database watch: {}", e)))?;
info!("Database change stream created successfully");
let mut batch: Vec<ChangeEvent> = Vec::with_capacity(config.batch_size);
let mut last_resume_token: Option<Document> = None;
let mut batch_timer = interval(config.batch_timeout);
batch_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let metrics_label = format!("__db:{}", config.database);
info!(
batch_size = config.batch_size,
batch_timeout = ?config.batch_timeout,
"Database worker event loop started"
);
loop {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("Received shutdown signal");
if !batch.is_empty() {
info!(batch_size = batch.len(), "Flushing pending batch on shutdown");
if let Err(e) = Self::flush_batch(
&metrics_label,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush batch on shutdown");
}
}
info!("Database worker shutting down gracefully");
break;
}
_ = batch_timer.tick() => {
if !batch.is_empty() {
debug!(batch_size = batch.len(), "Batch timeout - flushing");
if let Err(e) = Self::flush_batch(
&metrics_label,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush batch on timeout");
}
}
}
event_result = stream.next() => {
match event_result {
Some(Ok(change_event)) => {
let resume_token = match bson::to_document(&change_event.id) {
Ok(token) => token,
Err(e) => {
error!(?e, "Failed to serialize resume token");
continue;
}
};
let event = match ChangeEvent::try_from(change_event) {
Ok(evt) => evt,
Err(e) => {
error!(?e, "Failed to convert change event");
continue;
}
};
debug!(
operation = ?event.operation,
database = %event.namespace.database,
collection = %event.namespace.collection,
"Received database event"
);
last_resume_token = Some(resume_token.clone());
if let Err(e) = store.save_resume_token(&resume_token_key, &resume_token).await {
warn!(?e, "Failed to save resume token");
}
batch.push(event);
metrics::increment_batch_queue_size(&metrics_label);
if batch.len() >= config.batch_size {
debug!(batch_size = batch.len(), "Batch full - flushing");
if let Err(e) = Self::flush_batch(
&metrics_label,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush full batch");
}
}
}
Some(Err(e)) => {
error!(?e, "Error reading from database change stream");
tokio::time::sleep(Duration::from_secs(1)).await;
}
None => {
warn!("Database change stream ended unexpectedly");
break;
}
}
}
}
}
Ok(())
}
#[allow(clippy::too_many_lines)]
#[instrument(skip(config, store, destination, stats, shutdown_rx))]
async fn deployment_worker(
config: PipelineConfig,
store: Arc<S>,
destination: Arc<Mutex<D>>,
stats: Arc<RwLock<PipelineStats>>,
mut shutdown_rx: broadcast::Receiver<()>,
) -> Result<(), PipelineError> {
info!("Starting deployment worker (cluster-wide)");
let resume_token_key = config.watch_level.resume_token_key(&config.database, None);
let resume_token = store
.get_resume_token(&resume_token_key)
.await
.map_err(|e| PipelineError::StateStore(e.to_string()))?;
if let Some(ref token) = resume_token {
info!(?token, "Resuming from saved token");
}
let client = mongodb::Client::with_uri_str(&config.mongodb_uri)
.await
.map_err(|e| PipelineError::MongoDB(e.to_string()))?;
let mut options = mongodb::options::ChangeStreamOptions::default();
if config.stream_config.full_document_on_update {
options.full_document = Some(mongodb::options::FullDocumentType::UpdateLookup);
}
if config.stream_config.full_document_before_change {
options.full_document_before_change =
Some(mongodb::options::FullDocumentBeforeChangeType::WhenAvailable);
}
options.batch_size = config.stream_config.batch_size;
if let Some(ref token_doc) = resume_token {
if let Ok(bytes) = bson::to_vec(token_doc) {
if let Ok(resume_token) =
bson::from_slice::<mongodb::change_stream::event::ResumeToken>(&bytes)
{
options.resume_after = Some(resume_token);
}
}
}
let mut stream = if config.stream_config.pipeline.is_empty() {
client.watch().with_options(options).await
} else {
client
.watch()
.pipeline(config.stream_config.pipeline.clone())
.with_options(options)
.await
}
.map_err(|e| PipelineError::MongoDB(format!("Failed to create deployment watch: {}", e)))?;
info!("Deployment change stream created successfully");
let mut batch: Vec<ChangeEvent> = Vec::with_capacity(config.batch_size);
let mut last_resume_token: Option<Document> = None;
let mut batch_timer = interval(config.batch_timeout);
batch_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let metrics_label = "__deployment__".to_string();
info!(
batch_size = config.batch_size,
batch_timeout = ?config.batch_timeout,
"Deployment worker event loop started"
);
loop {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("Received shutdown signal");
if !batch.is_empty() {
info!(batch_size = batch.len(), "Flushing pending batch on shutdown");
if let Err(e) = Self::flush_batch(
&metrics_label,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush batch on shutdown");
}
}
info!("Deployment worker shutting down gracefully");
break;
}
_ = batch_timer.tick() => {
if !batch.is_empty() {
debug!(batch_size = batch.len(), "Batch timeout - flushing");
if let Err(e) = Self::flush_batch(
&metrics_label,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush batch on timeout");
}
}
}
event_result = stream.next() => {
match event_result {
Some(Ok(change_event)) => {
let resume_token = match bson::to_document(&change_event.id) {
Ok(token) => token,
Err(e) => {
error!(?e, "Failed to serialize resume token");
continue;
}
};
let event = match ChangeEvent::try_from(change_event) {
Ok(evt) => evt,
Err(e) => {
error!(?e, "Failed to convert change event");
continue;
}
};
debug!(
operation = ?event.operation,
database = %event.namespace.database,
collection = %event.namespace.collection,
"Received deployment event"
);
last_resume_token = Some(resume_token.clone());
if let Err(e) = store.save_resume_token(&resume_token_key, &resume_token).await {
warn!(?e, "Failed to save resume token");
}
batch.push(event);
metrics::increment_batch_queue_size(&metrics_label);
if batch.len() >= config.batch_size {
debug!(batch_size = batch.len(), "Batch full - flushing");
if let Err(e) = Self::flush_batch(
&metrics_label,
&mut batch,
last_resume_token.as_ref(),
&destination,
&store,
&stats,
&config,
)
.await
{
error!(?e, "Failed to flush full batch");
}
}
}
Some(Err(e)) => {
error!(?e, "Error reading from deployment change stream");
tokio::time::sleep(Duration::from_secs(1)).await;
}
None => {
warn!("Deployment change stream ended unexpectedly");
break;
}
}
}
}
}
Ok(())
}
#[instrument(skip(batch, last_resume_token, destination, store, stats, config), fields(collection = %collection, batch_size = batch.len()))]
async fn flush_batch(
collection: &str,
batch: &mut Vec<ChangeEvent>,
last_resume_token: Option<&Document>,
destination: &Arc<Mutex<D>>,
store: &Arc<S>,
stats: &Arc<RwLock<PipelineStats>>,
config: &PipelineConfig,
) -> Result<(), PipelineError> {
if batch.is_empty() {
return Ok(());
}
let batch_size = batch.len();
let start_time = Instant::now();
debug!("Flushing batch to destination");
metrics::record_batch_size(batch_size, collection);
Self::write_with_retry(batch, destination, config, stats).await?;
let elapsed = start_time.elapsed();
info!(
batch_size,
elapsed_ms = elapsed.as_millis(),
"Batch written successfully"
);
metrics::record_batch_duration(elapsed.as_secs_f64(), collection);
if let Some(token) = last_resume_token {
store
.save_resume_token(collection, token)
.await
.map_err(|e| PipelineError::StateStore(e.to_string()))?;
debug!("Resume token saved");
}
let mut operation_counts = std::collections::HashMap::new();
for event in batch.iter() {
*operation_counts.entry(&event.operation).or_insert(0u64) += 1;
}
for (operation, count) in operation_counts {
metrics::increment_events_processed_by(count, collection, operation.as_str());
}
let mut s = stats.write().await;
s.events_processed += batch_size as u64;
s.batches_written += 1;
metrics::decrement_batch_queue_size(batch_size, collection);
batch.clear();
Ok(())
}
#[instrument(skip(batch, destination, config, stats), fields(batch_size = batch.len()))]
async fn write_with_retry(
batch: &[ChangeEvent],
destination: &Arc<Mutex<D>>,
config: &PipelineConfig,
stats: &Arc<RwLock<PipelineStats>>,
) -> Result<(), PipelineError> {
let mut retry_delay = config.retry_delay;
let mut attempt = 0;
loop {
let result = {
let mut dest = destination.lock().await;
match dest.write_batch(batch).await {
Ok(()) => dest.flush().await,
Err(e) => Err(e),
}
};
match result {
Ok(()) => {
if attempt > 0 {
info!(attempts = attempt + 1, "Write succeeded after retries");
}
return Ok(());
}
Err(e) => {
{
let mut s = stats.write().await;
s.write_errors += 1;
}
let error_category = Self::categorize_error(&e);
let destination_type = {
let dest = destination.lock().await;
dest.metadata().destination_type.clone()
};
metrics::increment_destination_errors(&destination_type, error_category);
attempt += 1;
if attempt > config.max_retries {
error!(attempts = attempt, ?e, "Write failed after max retries");
return Err(PipelineError::Destination(e.to_string()));
}
if !Self::is_retryable_error(&e) {
error!(?e, "Non-retryable error encountered");
return Err(PipelineError::Destination(e.to_string()));
}
{
let mut s = stats.write().await;
s.retries += 1;
}
metrics::increment_retries(error_category);
warn!(
attempt,
max_retries = config.max_retries,
retry_delay_ms = retry_delay.as_millis(),
?e,
"Write failed, retrying"
);
tokio::time::sleep(retry_delay).await;
retry_delay = std::cmp::min(retry_delay * 2, config.max_retry_delay);
}
}
}
}
fn is_retryable_error(error: &DestinationError) -> bool {
error.to_string().contains("retryable") || error.to_string().contains("timeout")
}
fn categorize_error(error: &DestinationError) -> metrics::ErrorCategory {
let error_str = error.to_string().to_lowercase();
if error_str.contains("timeout") {
metrics::ErrorCategory::Timeout
} else if error_str.contains("connection") || error_str.contains("network") {
metrics::ErrorCategory::Connection
} else if error_str.contains("serialization") || error_str.contains("encode") {
metrics::ErrorCategory::Serialization
} else if error_str.contains("permission") || error_str.contains("auth") {
metrics::ErrorCategory::Permission
} else if error_str.contains("validation") {
metrics::ErrorCategory::Validation
} else if error_str.contains("not found") || error_str.contains("404") {
metrics::ErrorCategory::NotFound
} else if error_str.contains("rate limit") || error_str.contains("throttle") {
metrics::ErrorCategory::RateLimit
} else {
metrics::ErrorCategory::Unknown
}
}
#[instrument(skip(self), fields(owner_id = %self.owner_id))]
pub async fn stop(&mut self) -> Result<(), PipelineError> {
info!("Stopping pipeline");
let mut running = self.running.write().await;
if !*running {
warn!("Pipeline is not running");
return Ok(());
}
if let Some(tx) = self.shutdown_tx.take() {
let _ = tx.send(());
}
let mut workers = self.workers.write().await;
for worker in workers.drain(..) {
match worker.await {
Ok(Ok(())) => {
debug!("Worker stopped successfully");
}
Ok(Err(e)) => {
error!(?e, "Worker stopped with error");
}
Err(e) => {
error!(?e, "Worker panicked");
}
}
}
let mut lock_refresh_handles = self.lock_refresh_handles.write().await;
for handle in lock_refresh_handles.drain(..) {
if let Err(e) = handle.await {
warn!(?e, "Lock refresh task panicked");
}
}
if self.config.distributed_lock.enabled {
let locked_collections = self.locked_collections.read().await;
for collection in locked_collections.iter() {
let lock_key = match collection.as_str() {
"__database__" => self.database_lock_key(),
"__deployment__" => self.deployment_lock_key(),
_ => self.lock_key(collection),
};
match self.store.release_lock(&lock_key, &self.owner_id).await {
Ok(true) => {
info!(collection = %collection, "Released lock");
metrics::increment_locks_released();
}
Ok(false) => {
warn!(
collection = %collection,
"Lock was not held by this instance (already released or stolen)"
);
}
Err(e) => {
warn!(
collection = %collection,
error = %e,
"Failed to release lock (will expire via TTL)"
);
}
}
}
}
let mut dest = self.destination.lock().await;
dest.flush()
.await
.map_err(|e| PipelineError::Destination(e.to_string()))?;
dest.close()
.await
.map_err(|e| PipelineError::Destination(e.to_string()))?;
*running = false;
let mut locked_collections = self.locked_collections.write().await;
locked_collections.clear();
metrics::set_pipeline_status(metrics::PipelineStatus::Stopped);
metrics::set_active_collections(0);
metrics::set_locks_held(0);
let stats = self.stats.read().await;
info!(
events_processed = stats.events_processed,
batches_written = stats.batches_written,
write_errors = stats.write_errors,
retries = stats.retries,
"Pipeline stopped"
);
Ok(())
}
#[must_use]
pub async fn stats(&self) -> PipelineStats {
self.stats.read().await.clone()
}
#[must_use]
pub async fn is_running(&self) -> bool {
*self.running.read().await
}
}
#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
#[error("mongodb_uri is required")]
MissingMongoUri,
#[error("database is required")]
MissingDatabase,
#[error("Invalid batch_size: {value} ({reason})")]
InvalidBatchSize { value: usize, reason: &'static str },
#[error("Invalid batch_timeout: {reason}")]
InvalidBatchTimeout { reason: &'static str },
#[error("retry_delay ({retry_delay:?}) exceeds max_retry_delay ({max_retry_delay:?})")]
RetryDelayExceedsMax {
retry_delay: Duration,
max_retry_delay: Duration,
},
#[error("Invalid channel_buffer_size: {value} ({reason})")]
InvalidChannelBufferSize { value: usize, reason: &'static str },
#[error("Invalid distributed lock config: {reason}")]
InvalidLockConfig { reason: String },
}
#[derive(Debug, thiserror::Error)]
pub enum PipelineError {
#[error("Pipeline is already running")]
AlreadyRunning,
#[error("MongoDB error: {0}")]
MongoDB(String),
#[error("Change stream error: {0}")]
ChangeStream(String),
#[error("Destination error: {0}")]
Destination(String),
#[error("State store error: {0}")]
StateStore(String),
#[error("Configuration error: {0}")]
Configuration(String),
#[error("Pipeline error: {0}")]
Other(String),
#[error("Lock lost for collection '{collection}': {reason}")]
LockLost { collection: String, reason: String },
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_watch_collections_builds_collection_level() {
let config = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.watch_collections(vec!["users".to_string(), "orders".to_string()])
.build()
.unwrap();
assert!(matches!(config.watch_level, WatchLevel::Collection(_)));
if let WatchLevel::Collection(collections) = config.watch_level {
assert_eq!(collections.len(), 2);
assert!(collections.contains(&"users".to_string()));
assert!(collections.contains(&"orders".to_string()));
}
}
#[test]
fn test_watch_database_builds_database_level() {
let config = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.watch_database()
.build()
.unwrap();
assert!(matches!(config.watch_level, WatchLevel::Database));
}
#[test]
fn test_watch_deployment_builds_deployment_level() {
let config = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.watch_deployment()
.build()
.unwrap();
assert!(matches!(config.watch_level, WatchLevel::Deployment));
}
#[test]
fn test_default_watch_level_is_database() {
let config = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.build()
.unwrap();
assert!(matches!(config.watch_level, WatchLevel::Database));
}
#[test]
#[allow(deprecated)]
fn test_deprecated_collections_method_still_works() {
let config = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.collections(vec!["users".to_string()])
.build()
.unwrap();
assert!(matches!(config.watch_level, WatchLevel::Collection(_)));
if let WatchLevel::Collection(collections) = config.watch_level {
assert_eq!(collections.len(), 1);
assert_eq!(collections[0], "users");
}
}
#[test]
fn test_watch_level_can_be_overridden() {
let config = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.watch_collections(vec!["users".to_string()])
.watch_database() .build()
.unwrap();
assert!(matches!(config.watch_level, WatchLevel::Database));
}
#[test]
fn test_config_builder_defaults() {
let config = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.build()
.unwrap();
assert_eq!(config.batch_size, 100);
assert_eq!(config.batch_timeout, Duration::from_secs(5));
assert_eq!(config.max_retries, 0);
assert_eq!(config.retry_delay, Duration::from_millis(100));
assert_eq!(config.max_retry_delay, Duration::from_secs(30));
assert_eq!(config.channel_buffer_size, 1000);
assert!(matches!(config.watch_level, WatchLevel::Database));
}
#[test]
fn test_config_builder_validates_batch_size() {
let result = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.batch_size(20_000) .build();
assert!(result.is_err());
if let Err(e) = result {
assert!(matches!(e, ConfigError::InvalidBatchSize { .. }));
}
}
#[test]
fn test_config_builder_validates_channel_buffer_size() {
let result = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.channel_buffer_size(5) .build();
assert!(result.is_err());
if let Err(e) = result {
assert!(matches!(e, ConfigError::InvalidChannelBufferSize { .. }));
}
}
#[test]
fn test_config_builder_requires_mongodb_uri() {
let result = PipelineConfig::builder().database("testdb").build();
assert!(result.is_err());
if let Err(e) = result {
assert!(matches!(e, ConfigError::MissingMongoUri));
}
}
#[test]
fn test_config_builder_requires_database() {
let result = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.build();
assert!(result.is_err());
if let Err(e) = result {
assert!(matches!(e, ConfigError::MissingDatabase));
}
}
#[test]
fn test_retry_delay_validation() {
let result = PipelineConfig::builder()
.mongodb_uri("mongodb://localhost:27017")
.database("testdb")
.retry_delay(Duration::from_secs(60))
.max_retry_delay(Duration::from_secs(30)) .build();
assert!(result.is_err());
if let Err(e) = result {
assert!(matches!(e, ConfigError::RetryDelayExceedsMax { .. }));
}
}
}