use crate::error::Result;
use crate::metrics;
use crate::models::stream_event::Event;
use crate::source::adapter::SourceAdapter;
use futures::StreamExt;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
pub struct CdcCoordinator {
source_adapter: Box<dyn SourceAdapter>,
task_channels: HashMap<String, mpsc::Sender<Result<Event>>>,
shutdown_rx: watch::Receiver<bool>,
is_paused: Arc<RwLock<bool>>,
paused_tables: Arc<RwLock<std::collections::HashSet<String>>>,
}
impl CdcCoordinator {
pub fn new(source_adapter: Box<dyn SourceAdapter>, shutdown_rx: watch::Receiver<bool>) -> Self {
Self {
source_adapter,
task_channels: HashMap::new(),
shutdown_rx,
is_paused: Arc::new(RwLock::new(false)),
paused_tables: Arc::new(RwLock::new(std::collections::HashSet::new())),
}
}
pub fn register_task(&mut self, table_name: String, tx: mpsc::Sender<Result<Event>>) {
info!(
"Coordinator: Registering sync task for table '{}'",
table_name
);
self.task_channels.insert(table_name.clone(), tx);
debug!(
"Coordinator: Registered tasks: {:?}",
self.task_channels.keys().collect::<Vec<_>>()
);
}
pub async fn run(&mut self) -> Result<()> {
info!("Starting CDC coordinator");
let mut event_stream = self.source_adapter.get_changes().await?;
loop {
tokio::select! {
Some(event_result) = event_stream.next() => {
match event_result {
Ok(event) => {
debug!("Coordinator: Received event from stream");
if *self.is_paused.read().await {
debug!("Coordinator is paused, skipping event");
continue;
}
if let Event::Cdc(ref cdc_event) = event {
let table = format!("{}.{}", cdc_event.schema, cdc_event.table);
let event_type = format!("{:?}", cdc_event.event_type).to_lowercase();
metrics::record_cdc_event_received(&table, &event_type);
}
self.distribute_event(event).await;
}
Err(e) => {
error!("Coordinator: Error receiving CDC event: {}", e);
}
}
}
_ = self.shutdown_rx.changed() => {
if *self.shutdown_rx.borrow() {
info!("CDC coordinator received shutdown signal");
break;
}
}
}
}
info!("CDC coordinator stopped");
Ok(())
}
async fn distribute_event(&self, event: Event) {
let table_name = match &event {
Event::Cdc(cdc_event) => {
if cdc_event.schema.is_empty() {
cdc_event.table.clone()
} else {
format!("{}.{}", cdc_event.schema, cdc_event.table)
}
}
_ => {
debug!("Skipping non-CDC event distribution");
return;
}
};
{
let paused_tables = self.paused_tables.read().await;
if paused_tables.contains(&table_name)
|| (table_name.contains('.')
&& paused_tables.contains(table_name.split('.').next_back().unwrap_or("")))
{
debug!("Table '{}' is paused, skipping event", table_name);
return;
}
}
debug!("Distributing event for table: {}", table_name);
if let Some(tx) = self.task_channels.get(&table_name) {
debug!("Sending event to task for table '{}'", table_name);
match tx.try_send(Ok(event.clone())) {
Ok(()) => {
}
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
warn!(
"Channel full for table '{}', activating backpressure",
table_name
);
self.pause_table(&table_name).await;
if tx.send(Ok(event.clone())).await.is_err() {
warn!("Failed to send event to task for table '{}'", table_name);
}
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
warn!("Channel closed for table '{}'", table_name);
}
}
} else {
debug!("No task registered for table '{}'", table_name);
}
if table_name.contains('.') {
let table_only = table_name.split('.').next_back().unwrap_or("");
if let Some(tx) = self.task_channels.get(table_only) {
debug!(
"Sending event to task for table '{}' (without schema)",
table_only
);
match tx.try_send(Ok(event)) {
Ok(()) => {
}
Err(tokio::sync::mpsc::error::TrySendError::Full(msg)) => {
warn!(
"Channel full for table '{}' (without schema), activating backpressure",
table_only
);
self.pause_table(table_only).await;
if tx.send(msg).await.is_err() {
warn!("Failed to send event to task for table '{}'", table_only);
}
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
warn!("Channel closed for table '{}'", table_only);
}
}
}
}
}
pub async fn pause_all(&self) {
*self.is_paused.write().await = true;
info!("CDC coordinator paused all event processing");
metrics::CDC_CONTROL_OPERATIONS
.with_label_values(&["pause", "global", "all"])
.inc();
}
pub async fn resume_all(&self) {
*self.is_paused.write().await = false;
info!("CDC coordinator resumed all event processing");
metrics::CDC_CONTROL_OPERATIONS
.with_label_values(&["resume", "global", "all"])
.inc();
}
pub async fn is_paused(&self) -> bool {
*self.is_paused.read().await
}
pub async fn pause_table(&self, table_name: &str) {
self.paused_tables
.write()
.await
.insert(table_name.to_string());
info!(
"CDC coordinator paused processing for table '{}'",
table_name
);
metrics::CDC_CONTROL_OPERATIONS
.with_label_values(&["pause", "table", table_name])
.inc();
}
pub async fn resume_table(&self, table_name: &str) {
self.paused_tables.write().await.remove(table_name);
info!(
"CDC coordinator resumed processing for table '{}'",
table_name
);
metrics::CDC_CONTROL_OPERATIONS
.with_label_values(&["resume", "table", table_name])
.inc();
}
pub async fn get_paused_tables(&self) -> Vec<String> {
self.paused_tables.read().await.iter().cloned().collect()
}
pub fn unregister_task(&mut self, table_name: &str) {
info!(
"Coordinator: Unregistering sync task for table '{}'",
table_name
);
self.task_channels.remove(table_name);
debug!(
"Coordinator: Remaining tasks: {:?}",
self.task_channels.keys().collect::<Vec<_>>()
);
}
}