#![forbid(unsafe_code)]
use std::{
collections::{BTreeMap, BTreeSet},
path::PathBuf,
sync::{Arc, Mutex, mpsc},
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};
use datum::{Keep, Sink};
use datum_cdc::{
CdcCheckpointHandle, CdcSource, ChangeEvent, ChangeOperation, ColumnValue, PostgresCdcConfig,
SlotLifecycle,
};
use serde::{Deserialize, Serialize};
#[derive(Debug)]
struct Args {
database_url: String,
slot: String,
publication: String,
run_id: String,
expected: PathBuf,
done_file: PathBuf,
out: PathBuf,
timeout_after_done: Duration,
}
#[derive(Debug, Clone, Deserialize)]
struct ExpectedEvent {
key: i64,
op: String,
commit_ns: i64,
}
#[derive(Debug, Clone)]
struct ObservedEvent {
observed_ns: i64,
key: i64,
op: String,
}
#[derive(Debug, Clone)]
struct MetricSample {
time_ns: i64,
cpu_seconds: f64,
rss_kib: u64,
threads: u64,
retained_wal_bytes: Option<i64>,
confirmed_lag_bytes: Option<i64>,
slot_active: Option<bool>,
}
#[derive(Debug, Serialize)]
struct ObserverSummary {
run_id: String,
expected_events: usize,
observed_events: usize,
throughput_events_per_sec: f64,
latency_ms: LatencySummary,
counts: CountsSummary,
correctness: CorrectnessSummary,
datum: ProcessSummary,
postgres: PostgresSummary,
timed_out_after_done: bool,
}
#[derive(Debug, Serialize)]
struct LatencySummary {
count: usize,
min: Option<f64>,
p50: Option<f64>,
p95: Option<f64>,
p99: Option<f64>,
max: Option<f64>,
mean: Option<f64>,
}
#[derive(Debug, Serialize)]
struct CountsSummary {
expected: BTreeMap<String, usize>,
observed: BTreeMap<String, usize>,
}
#[derive(Debug, Serialize)]
struct CorrectnessSummary {
events_observed_eq_committed: bool,
ordered_per_key: bool,
duplicates: usize,
missing: usize,
mismatched: usize,
}
#[derive(Debug, Serialize)]
struct ProcessSummary {
sample_count: usize,
cpu_core_seconds: f64,
avg_cpu_cores: f64,
peak_rss_mib: f64,
peak_threads: u64,
}
#[derive(Debug, Serialize)]
struct PostgresSummary {
sample_count: usize,
max_slot_retained_wal_bytes: Option<i64>,
max_confirmed_lag_bytes: Option<i64>,
slot_active_seen: bool,
}
#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
async fn main() -> datum_cdc::CdcResult<()> {
let args = parse_args()?;
if let Some(parent) = args.out.parent() {
std::fs::create_dir_all(parent)?;
}
let (events_tx, events_rx) = mpsc::channel::<ChangeEvent>();
let checkpoint_slot = Arc::new(Mutex::new(None::<CdcCheckpointHandle>));
let checkpoint_for_sink = Arc::clone(&checkpoint_slot);
let source = CdcSource::postgres()
.connect(PostgresCdcConfig::from_url(&args.database_url)?)
.slot(&args.slot)
.publication(&args.publication)
.slot_lifecycle(SlotLifecycle::CreateOwned)
.status_interval(Duration::from_millis(50))
.idle_wakeup_interval(Duration::from_millis(50))
.build()?;
let graph = source.to_mat(
Sink::foreach(move |event: ChangeEvent| {
if let Some(checkpoint) = checkpoint_for_sink
.lock()
.expect("checkpoint slot poisoned")
.clone()
{
let _ = checkpoint.checkpoint(event.lsn.clone());
}
let _ = events_tx.send(event);
}),
Keep::both,
);
let (handle, completion) = graph.run()?;
*checkpoint_slot
.lock()
.expect("checkpoint slot poisoned after materialization") =
Some(handle.checkpoint_handle());
let mut samples = Vec::new();
let mut observed = Vec::new();
let mut last_sample = Instant::now() - Duration::from_secs(1);
let started = Instant::now();
let mut done_seen = None::<Instant>;
let mut timed_out = false;
loop {
match events_rx.recv_timeout(Duration::from_millis(100)) {
Ok(event) => {
if let Some(observed_event) = extract_event(&event, &args.run_id) {
observed.push(observed_event);
}
}
Err(mpsc::RecvTimeoutError::Timeout) => {}
Err(mpsc::RecvTimeoutError::Disconnected) => break,
}
if last_sample.elapsed() >= Duration::from_secs(1) {
samples.push(sample_metrics(&handle).await);
last_sample = Instant::now();
}
if args.done_file.exists() && done_seen.is_none() {
done_seen = Some(Instant::now());
}
let expected = load_expected(&args.expected);
if done_seen.is_some() && !expected.is_empty() && observed.len() >= expected.len() {
break;
}
if let Some(done_seen) = done_seen
&& done_seen.elapsed() >= args.timeout_after_done
{
timed_out = true;
break;
}
if started.elapsed() >= args.timeout_after_done + Duration::from_secs(600) {
timed_out = true;
break;
}
}
samples.push(sample_metrics(&handle).await);
let expected = load_expected(&args.expected);
let summary = build_summary(&args.run_id, &expected, &observed, &samples, timed_out);
std::fs::write(&args.out, serde_json::to_vec_pretty(&summary)?)?;
let _ = handle.stop();
let _ = completion.wait();
if timed_out {
std::process::exit(1);
}
Ok(())
}
fn parse_args() -> datum_cdc::CdcResult<Args> {
let mut database_url = None;
let mut slot = None;
let mut publication = None;
let mut run_id = None;
let mut expected = None;
let mut done_file = None;
let mut out = None;
let mut timeout_after_done = Duration::from_secs(120);
let mut args = std::env::args().skip(1);
while let Some(arg) = args.next() {
let value = match arg.as_str() {
"--database-url" => &mut database_url,
"--slot" => &mut slot,
"--publication" => &mut publication,
"--run-id" => &mut run_id,
"--expected" => &mut expected,
"--done-file" => &mut done_file,
"--out" => &mut out,
"--timeout-after-done" => {
let value = args.next().ok_or_else(|| {
datum_cdc::CdcError::Config("--timeout-after-done needs a value".into())
})?;
timeout_after_done = Duration::from_secs(value.parse().map_err(|err| {
datum_cdc::CdcError::Config(format!("invalid timeout: {err}"))
})?);
continue;
}
"-h" | "--help" => {
print_usage();
std::process::exit(0);
}
_ => {
return Err(datum_cdc::CdcError::Config(format!(
"unknown argument {arg:?}"
)));
}
};
*value = Some(args.next().ok_or_else(|| {
datum_cdc::CdcError::Config(format!("argument {arg:?} needs a value"))
})?);
}
Ok(Args {
database_url: required(database_url, "--database-url")?,
slot: required(slot, "--slot")?,
publication: required(publication, "--publication")?,
run_id: required(run_id, "--run-id")?,
expected: PathBuf::from(required(expected, "--expected")?),
done_file: PathBuf::from(required(done_file, "--done-file")?),
out: PathBuf::from(required(out, "--out")?),
timeout_after_done,
})
}
fn required(value: Option<String>, name: &str) -> datum_cdc::CdcResult<String> {
value.ok_or_else(|| datum_cdc::CdcError::Config(format!("missing {name}")))
}
fn print_usage() {
eprintln!(
"Usage: cdc_observer --database-url URL --slot SLOT --publication PUB --run-id ID --expected PATH --done-file PATH --out PATH [--timeout-after-done SECS]"
);
}
fn extract_event(event: &ChangeEvent, run_id: &str) -> Option<ObservedEvent> {
let row = match event.op {
ChangeOperation::Insert | ChangeOperation::Update => event.after.as_ref()?,
ChangeOperation::Delete => event.before.as_ref()?,
ChangeOperation::Truncate => return None,
};
if row.get_text(&event.relation, "run_id")? != run_id {
return None;
}
let key = row.get_text(&event.relation, "id")?.parse().ok()?;
Some(ObservedEvent {
observed_ns: now_ns(),
key,
op: event.op.debezium_letter().to_owned(),
})
}
fn load_expected(path: &PathBuf) -> Vec<ExpectedEvent> {
let Ok(bytes) = std::fs::read_to_string(path) else {
return Vec::new();
};
bytes
.lines()
.filter_map(|line| serde_json::from_str::<ExpectedEvent>(line).ok())
.collect()
}
async fn sample_metrics(handle: &datum_cdc::CdcHandle) -> MetricSample {
let process = sample_process();
let lag = handle.lag().await.ok();
MetricSample {
time_ns: now_ns(),
cpu_seconds: process.0,
rss_kib: process.1,
threads: process.2,
retained_wal_bytes: lag.map(|lag| lag.retained_wal_bytes),
confirmed_lag_bytes: lag.map(|lag| lag.confirmed_lag_bytes),
slot_active: lag.map(|lag| lag.slot_active),
}
}
fn sample_process() -> (f64, u64, u64) {
let ticks_per_second = clk_tck();
let cpu_ticks = std::fs::read_to_string("/proc/self/stat")
.ok()
.and_then(|stat| {
let after_comm = stat.rsplit_once(") ")?.1;
let fields = after_comm.split_whitespace().collect::<Vec<_>>();
let utime = fields.get(11)?.parse::<u64>().ok()?;
let stime = fields.get(12)?.parse::<u64>().ok()?;
Some(utime + stime)
})
.unwrap_or(0);
let status = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
let rss_kib = status_value(&status, "VmRSS").unwrap_or(0);
let threads = status_value(&status, "Threads").unwrap_or(0);
(cpu_ticks as f64 / ticks_per_second, rss_kib, threads)
}
fn clk_tck() -> f64 {
std::process::Command::new("getconf")
.arg("CLK_TCK")
.output()
.ok()
.and_then(|output| String::from_utf8(output.stdout).ok())
.and_then(|value| value.trim().parse::<f64>().ok())
.unwrap_or(100.0)
}
fn status_value(status: &str, key: &str) -> Option<u64> {
status.lines().find_map(|line| {
let rest = line.strip_prefix(key)?.trim_start_matches(':').trim();
rest.split_whitespace().next()?.parse().ok()
})
}
fn build_summary(
run_id: &str,
expected: &[ExpectedEvent],
observed: &[ObservedEvent],
samples: &[MetricSample],
timed_out: bool,
) -> ObserverSummary {
let mut expected_by_key = BTreeMap::<i64, Vec<&ExpectedEvent>>::new();
let mut observed_by_key = BTreeMap::<i64, Vec<&ObservedEvent>>::new();
for row in expected {
expected_by_key.entry(row.key).or_default().push(row);
}
for row in observed {
observed_by_key.entry(row.key).or_default().push(row);
}
let mut ordered_per_key = true;
let mut duplicates = 0;
let mut missing = 0;
let mut mismatched = 0;
let mut latencies = Vec::new();
for (key, expected_rows) in &expected_by_key {
let observed_rows = observed_by_key.get(key).cloned().unwrap_or_default();
let expected_ops = expected_rows
.iter()
.map(|row| row.op.as_str())
.collect::<Vec<_>>();
let observed_ops = observed_rows
.iter()
.map(|row| row.op.as_str())
.collect::<Vec<_>>();
if expected_ops != observed_ops {
ordered_per_key = false;
}
duplicates += observed_rows.len().saturating_sub(expected_rows.len());
missing += expected_rows.len().saturating_sub(observed_rows.len());
for (expected_row, observed_row) in expected_rows.iter().zip(observed_rows.iter()) {
if expected_row.op != observed_row.op {
mismatched += 1;
} else {
latencies
.push((observed_row.observed_ns - expected_row.commit_ns) as f64 / 1_000_000.0);
}
}
}
for key in observed_by_key.keys() {
if !expected_by_key.contains_key(key) {
duplicates += observed_by_key[key].len();
ordered_per_key = false;
}
}
let first_seen = observed.iter().map(|row| row.observed_ns).min();
let last_seen = observed.iter().map(|row| row.observed_ns).max();
let throughput_events_per_sec = match (first_seen, last_seen) {
(Some(first), Some(last)) if last > first => {
observed.len() as f64 / ((last - first) as f64 / 1_000_000_000.0)
}
_ => 0.0,
};
ObserverSummary {
run_id: run_id.to_owned(),
expected_events: expected.len(),
observed_events: observed.len(),
throughput_events_per_sec,
latency_ms: latency_summary(&mut latencies),
counts: CountsSummary {
expected: counts(expected.iter().map(|row| row.op.as_str())),
observed: counts(observed.iter().map(|row| row.op.as_str())),
},
correctness: CorrectnessSummary {
events_observed_eq_committed: expected.len() == observed.len(),
ordered_per_key,
duplicates,
missing,
mismatched,
},
datum: process_summary(samples),
postgres: postgres_summary(samples),
timed_out_after_done: timed_out,
}
}
fn counts<'a>(ops: impl Iterator<Item = &'a str>) -> BTreeMap<String, usize> {
let mut counts = BTreeMap::new();
for op in ops {
*counts.entry(op.to_owned()).or_default() += 1;
}
counts
}
fn latency_summary(latencies: &mut [f64]) -> LatencySummary {
latencies.sort_by(|left, right| left.total_cmp(right));
let mean = if latencies.is_empty() {
None
} else {
Some(latencies.iter().sum::<f64>() / latencies.len() as f64)
};
LatencySummary {
count: latencies.len(),
min: latencies.first().copied(),
p50: percentile(latencies, 0.50),
p95: percentile(latencies, 0.95),
p99: percentile(latencies, 0.99),
max: latencies.last().copied(),
mean,
}
}
fn percentile(values: &[f64], p: f64) -> Option<f64> {
if values.is_empty() {
return None;
}
let rank = ((values.len() - 1) as f64 * p).round() as usize;
values.get(rank).copied()
}
fn process_summary(samples: &[MetricSample]) -> ProcessSummary {
let elapsed = samples
.first()
.zip(samples.last())
.map(|(first, last)| ((last.time_ns - first.time_ns) as f64 / 1_000_000_000.0).max(0.001))
.unwrap_or(0.001);
let cpu_core_seconds = samples
.first()
.zip(samples.last())
.map(|(first, last)| (last.cpu_seconds - first.cpu_seconds).max(0.0))
.unwrap_or(0.0);
ProcessSummary {
sample_count: samples.len(),
cpu_core_seconds,
avg_cpu_cores: cpu_core_seconds / elapsed,
peak_rss_mib: samples
.iter()
.map(|sample| sample.rss_kib)
.max()
.unwrap_or(0) as f64
/ 1024.0,
peak_threads: samples
.iter()
.map(|sample| sample.threads)
.max()
.unwrap_or(0),
}
}
fn postgres_summary(samples: &[MetricSample]) -> PostgresSummary {
let retained = samples
.iter()
.filter_map(|sample| sample.retained_wal_bytes)
.max();
let confirmed = samples
.iter()
.filter_map(|sample| sample.confirmed_lag_bytes)
.max();
PostgresSummary {
sample_count: samples
.iter()
.filter(|sample| sample.retained_wal_bytes.is_some())
.count(),
max_slot_retained_wal_bytes: retained,
max_confirmed_lag_bytes: confirmed,
slot_active_seen: samples
.iter()
.any(|sample| sample.slot_active.unwrap_or(false)),
}
}
fn now_ns() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos() as i64)
.unwrap_or_default()
}
#[allow(dead_code)]
fn _assert_column_value_send_sync(_: &ColumnValue, _: &BTreeSet<i32>) {}