From 4cd95846e0ce0210bcc3c49f3239100c52c8e6cf Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Sun, 2 Aug 2026 00:57:31 -0700 Subject: [PATCH] [adapters] Have output endpoints report their own checkpoint totals A checkpoint needs each output endpoint's `transmitted_records` as of the moment the endpoint has transmitted the output the checkpoint covers, since that is the state the pipeline resumes from. It cannot read that itself: the circuit thread takes the checkpoint while the endpoints are still working, and it must not block waiting for them. So it guessed using this formula: ```rust transmitted_records: snapshot.transmitted_records + snapshot.buffered_records + snapshot.queued_records, ``` There are several issues with this: 1. The second and third terms are the number of records output by the circuit, not the number of records send by the connector. They can have non-unit multiplicities and canceling out weights, making them an approximation to what the connector will output by the time the checkpoint completes 2. There was a window of double counting when the same records were accounted by both `queued_records` and `transmitted_records`. The endpoint thread raises `transmitted_records` from inside `push_batch_to_encoder`, while the batch those records came from stays in `queued_records` until `output_batch` releases it afterwards. 3. This aside, each counter is implemented as a separate atomic variable, meaning that they can be in an inconsistent state for short periods of time. The second issue above caused the suspend_barrier4 test to fail. Having tried several solutions, I concluded that the only correct way to do this is to have the endpoint report the actual transmitted number of records when it's finished processing records associated with the checkpoint, which is what this commit implements. As a bonums, `transmitted_bytes` can now be tracked as well, and no longer has to be left out of the checkpoint. Signed-off-by: Leonid Ryzhyk --- crates/adapters/src/controller.rs | 21 ++ crates/adapters/src/controller/checkpoint.rs | 16 +- crates/adapters/src/controller/stats.rs | 121 ++++++++- crates/adapters/src/controller/test.rs | 259 +++++++++++++++++++ 4 files changed, 406 insertions(+), 11 deletions(-) diff --git a/crates/adapters/src/controller.rs b/crates/adapters/src/controller.rs index f0235b8df6e..619a1653884 100644 --- a/crates/adapters/src/controller.rs +++ b/crates/adapters/src/controller.rs @@ -9141,6 +9141,14 @@ impl RunningCheckpoint { ) }) .collect(); + // Ask the output endpoints to record what they have transmitted once + // they catch up with this checkpoint. `CheckpointThread` waits for + // them and collects the readings. + circuit + .controller + .status + .arm_output_checkpoint(processed_records); + let output_statistics = { let outputs_by_name: HashMap = circuit .controller @@ -9349,6 +9357,19 @@ impl CheckpointThread { }); } + // Every endpoint has now transmitted the output this checkpoint covers + // and recorded what that came to. Those readings are the endpoint + // statistics the checkpoint stores: taken by the endpoint itself at + // the moment it caught up, they neither miss what it was still + // transmitting when the checkpoint started nor count what it went on + // to transmit afterwards. + for (endpoint_name, transmitted) in self.status.take_output_checkpoint_totals() { + if let Some(metrics) = self.checkpoint.output_statistics.get_mut(&endpoint_name) { + metrics.transmitted_records = transmitted.records; + metrics.transmitted_bytes = transmitted.bytes; + } + } + // Finalize the checkpoint on storage. // // [Checkpoint::write] commits to stable storage. diff --git a/crates/adapters/src/controller/checkpoint.rs b/crates/adapters/src/controller/checkpoint.rs index 8cc9e7f2add..79e316cf477 100644 --- a/crates/adapters/src/controller/checkpoint.rs +++ b/crates/adapters/src/controller/checkpoint.rs @@ -237,16 +237,12 @@ impl CheckpointOutputEndpointMetrics { pub fn from_endpoint_status(status: &OutputEndpointStatus, snapshot_sent: bool) -> Self { let snapshot = status.metrics.snapshot(); Self { - // Includes all the records that have been transmitted plus all - // of the records that will be transmitted by the time we commit - // the checkpoint. - transmitted_records: snapshot.transmitted_records - + snapshot.buffered_records - + snapshot.queued_records, - - // Only the bytes and errors already transmitted, not including - // those that will be transmitted by the time we commit the - // checkpoint (we don't have proper statistics for those). + // What the endpoint had transmitted when the checkpoint started. + // `CheckpointThread` replaces both figures with what it had + // transmitted once it caught up with the checkpoint, which is what + // the checkpoint needs; these stand only for an endpoint that goes + // away before it gets there. + transmitted_records: snapshot.transmitted_records, transmitted_bytes: snapshot.transmitted_bytes, // We can't predict how many errors there will be by the time we diff --git a/crates/adapters/src/controller/stats.rs b/crates/adapters/src/controller/stats.rs index 9138364acef..b9cf8f068ab 100644 --- a/crates/adapters/src/controller/stats.rs +++ b/crates/adapters/src/controller/stats.rs @@ -73,7 +73,7 @@ use parking_lot::{RwLock, RwLockReadGuard}; use serde::{Deserialize, Serialize}; use size_of::HumanBytes; use std::{ - collections::{BTreeMap, BTreeSet, VecDeque}, + collections::{BTreeMap, BTreeSet, HashMap, VecDeque}, fmt::Display, sync::{ Arc, Mutex, @@ -1295,6 +1295,36 @@ impl ControllerStatus { }; } + /// Ask every output endpoint to record what it has transmitted once it has + /// transmitted everything derived from the first `threshold` records the + /// pipeline ingested. + /// + /// Call from the circuit thread when a checkpoint starts; + /// [`Self::take_output_checkpoint_totals`] collects the readings once the + /// endpoints have got there. + pub fn arm_output_checkpoint(&self, threshold: u64) { + for endpoint_stats in self.output_status().values() { + endpoint_stats.arm_checkpoint(threshold); + } + } + + /// Collect what each output endpoint transmitted for the armed checkpoint, + /// keyed by endpoint name, and disarm them all. + /// + /// An endpoint that never caught up, because it was removed first, is + /// absent from the result. + pub fn take_output_checkpoint_totals(&self) -> HashMap { + self.output_status() + .values() + .filter_map(|endpoint_stats| { + Some(( + endpoint_stats.endpoint_name.clone(), + endpoint_stats.take_checkpoint_totals()?, + )) + }) + .collect() + } + pub fn output_buffers_full(&self) -> bool { self.output_status() .values() @@ -2702,6 +2732,42 @@ pub struct OutputEndpointStatus { /// Connector-specific metrics for Prometheus export. pub custom_metrics: Option>, + + /// The checkpoint waiting on this endpoint, if one is in progress. + /// + /// A checkpoint needs each endpoint's [`Self::transmitted_records`] as of + /// the moment the endpoint has transmitted the output the checkpoint + /// covers, since that is the state the pipeline resumes from. The circuit + /// thread cannot read it at that moment: it takes the checkpoint while the endpoints + /// are still working, and it must not block waiting for them. So it arms + /// this latch instead, and the endpoint takes the reading itself the + /// moment it gets there. + checkpoint_latch: Mutex>, +} + +/// What an output endpoint had put on the transport at a particular moment. +#[derive(Clone, Copy, Debug)] +pub struct TransmittedTotals { + pub records: u64, + pub bytes: u64, +} + +/// A checkpoint waiting for an output endpoint to catch up with it. +struct CheckpointLatch { + /// Pipeline records the checkpoint covers. The endpoint has caught up once + /// its [`OutputEndpointMetrics::total_processed_input_records`] reaches + /// this, i.e., the endpoint has transmitted every batch derived from + /// `threshold` records. + /// + /// The endpoint lands on this figure exactly, because the circuit hands it + /// one batch per step and the checkpoint covers whole steps. The exception + /// is output buffering: a flush can span the checkpoint, carrying the + /// counter past it in one write, and the reading then covers the whole + /// flush. + threshold: u64, + + /// What the endpoint had transmitted when it caught up, once it does. + transmitted_totals: Option, } impl OutputEndpointStatus { @@ -2756,6 +2822,54 @@ impl OutputEndpointStatus { pub fn transmitted_records(&self) -> u64 { self.metrics.transmitted_records.load(Ordering::Acquire) } + + /// What this endpoint has put on the transport so far. + fn transmitted_totals(&self) -> TransmittedTotals { + TransmittedTotals { + records: self.metrics.transmitted_records.load(Ordering::Acquire), + bytes: self.metrics.transmitted_bytes.load(Ordering::Acquire), + } + } + + /// Ask this endpoint to record what it has transmitted once it has + /// transmitted everything derived from the first `threshold` records the + /// pipeline ingested. See [`Self::checkpoint_latch`]. + /// + /// An endpoint that is already there records it now: nothing more will + /// reach it on an idle pipeline, so it would never take the reading + /// otherwise. + fn arm_checkpoint(&self, threshold: u64) { + let mut latch = self.checkpoint_latch.lock().unwrap(); + let transmitted_totals = (self.num_total_processed_input_records() >= threshold) + .then(|| self.transmitted_totals()); + *latch = Some(CheckpointLatch { + threshold, + transmitted_totals, + }); + } + + /// Record what this endpoint has transmitted if it has just caught up with + /// the checkpoint waiting on it. Call after advancing + /// [`OutputEndpointMetrics::total_processed_input_records`]. + fn note_checkpoint_progress(&self) { + let mut latch = self.checkpoint_latch.lock().unwrap(); + if let Some(latch) = latch.as_mut() + && latch.transmitted_totals.is_none() + && self.num_total_processed_input_records() >= latch.threshold + { + latch.transmitted_totals = Some(self.transmitted_totals()); + } + } + + /// Take the reading this endpoint recorded for the checkpoint waiting on + /// it, and disarm. `None` if it never caught up. + fn take_checkpoint_totals(&self) -> Option { + self.checkpoint_latch + .lock() + .unwrap() + .take()? + .transmitted_totals + } } impl OutputEndpointStatus { @@ -2782,6 +2896,7 @@ impl OutputEndpointStatus { transport_errors: Mutex::new(transport_errors), health: Mutex::new(None), custom_metrics: None, + checkpoint_latch: Mutex::new(None), } } @@ -2802,6 +2917,7 @@ impl OutputEndpointStatus { self.metrics .total_processed_steps .store(processed.total_processed_steps, Ordering::Release); + self.note_checkpoint_progress(); } let old = self @@ -2848,8 +2964,11 @@ impl OutputEndpointStatus { self.metrics .total_processed_steps .store(processed.total_processed_steps, Ordering::Release); + self.note_checkpoint_progress(); } + /// A chunk of the batch the endpoint is transmitting has reached the + /// transport. fn output_buffer(&self, num_bytes: usize, num_records: usize) { self.metrics .transmitted_bytes diff --git a/crates/adapters/src/controller/test.rs b/crates/adapters/src/controller/test.rs index e839a41b260..17e48b9f681 100644 --- a/crates/adapters/src/controller/test.rs +++ b/crates/adapters/src/controller/test.rs @@ -5196,6 +5196,265 @@ fn test_removing_a_lagging_output_endpoint_republishes_completion() { ); } +/// A bare [`ControllerStatus`] for the output-endpoint accounting tests. +fn test_controller_status(name: &str) -> crate::ControllerStatus { + crate::ControllerStatus::new( + serde_json::from_value(json!({ "name": name, "workers": 1 })).unwrap(), + 0, + None, + uuid::Uuid::nil(), + ) +} + +/// A file sink writing CSV, which transmits one record per tuple. +fn test_output_config() -> OutputEndpointConfig { + serde_json::from_value(json!({ + "stream": "v1", + "transport": { "name": "file_output", "config": { "path": "/dev/null" } }, + "format": { "name": "csv", "config": {} } + })) + .unwrap() +} + +/// An output endpoint must report what it had transmitted at the moment it +/// caught up with a checkpoint: not less, because the endpoint was still +/// writing when the checkpoint started, and not more, because it keeps working +/// while the checkpoint commits. +/// +/// The checkpoint stores this as the endpoint's `transmitted_records`, so +/// getting it wrong makes a resumed connector misreport for the rest of the +/// pipeline's life. Adding up the live counters cannot produce it: a batch +/// stays queued until the endpoint has finished writing it, by which time +/// `transmitted_records` already covers what went out, so the two overlap. +#[test] +fn test_output_endpoint_reports_totals_at_the_checkpoint() { + use super::stats::ProcessedRecords; + + // Two steps of output that the circuit queued before the endpoint thread + // got to either. These are the batch sizes from the `suspend_barrier4` CI + // failure that exposed the double count. + const FIRST_BATCH: usize = 630; + const SECOND_BATCH: usize = 370; + const TOTAL: u64 = (FIRST_BATCH + SECOND_BATCH) as u64; + + let status = test_controller_status("test_checkpoint_output_stats"); + let output_config = test_output_config(); + status.add_output(&0, "out", &output_config, None, true); + + let parker = Parker::new(); + let unparker = parker.unparker().clone(); + let processed = |records: u64, steps| ProcessedRecords { + total_processed_input_records: records, + total_processed_steps: steps, + }; + let transmitted = || { + status + .output_status() + .get(&0) + .unwrap() + .transmitted_records() + }; + let reported = || { + status + .take_output_checkpoint_totals() + .remove("out") + .map(|totals| totals.records) + }; + + // The circuit queues two steps of output, then takes a checkpoint covering + // every record it has ingested. + status.enqueue_batch(0, FIRST_BATCH); + status.enqueue_batch(0, SECOND_BATCH); + status.arm_output_checkpoint(TOTAL); + + // The endpoint writes the first batch to the transport. That write does not + // take the batch off the queue; the endpoint thread does that once it is + // done with the whole batch. The endpoint owes the checkpoint the rest, so + // it has nothing to report yet. + status.output_buffer(0, 4096, FIRST_BATCH); + status.output_batch( + 0, + Some(processed(FIRST_BATCH as u64, 1)), + FIRST_BATCH, + &unparker, + ); + + // The second batch takes it to the checkpoint. + status.output_buffer(0, 4096, SECOND_BATCH); + status.output_batch(0, Some(processed(TOTAL, 2)), SECOND_BATCH, &unparker); + assert_eq!(transmitted(), TOTAL); + + // The circuit keeps running while the checkpoint commits. What the endpoint + // transmits now belongs to the steps after the checkpoint, and the + // pipeline re-emits it on resume, so the checkpoint must not count it. + const AFTERWARDS: usize = 500; + status.enqueue_batch(0, AFTERWARDS); + status.output_buffer(0, 4096, AFTERWARDS); + status.output_batch( + 0, + Some(processed(TOTAL + AFTERWARDS as u64, 3)), + AFTERWARDS, + &unparker, + ); + assert_eq!(transmitted(), TOTAL + AFTERWARDS as u64); + + assert_eq!(reported(), Some(TOTAL)); + assert_eq!(reported(), None, "collecting the reading disarms the latch"); +} + +/// An endpoint that is already caught up when the checkpoint starts must +/// report right away. +/// +/// Nothing more reaches it on an idle pipeline, so it would never take the +/// reading otherwise, and the checkpoint would wait on it forever. +#[test] +fn test_idle_output_endpoint_reports_at_once() { + use super::stats::ProcessedRecords; + + const RECORDS: usize = 100; + + let status = test_controller_status("test_checkpoint_idle_endpoint"); + let output_config = test_output_config(); + status.add_output(&0, "out", &output_config, None, true); + + let parker = Parker::new(); + let unparker = parker.unparker().clone(); + + status.enqueue_batch(0, RECORDS); + status.output_buffer(0, 4096, RECORDS); + status.output_batch( + 0, + Some(ProcessedRecords { + total_processed_input_records: RECORDS as u64, + total_processed_steps: 1, + }), + RECORDS, + &unparker, + ); + + status.arm_output_checkpoint(RECORDS as u64); + assert_eq!( + status + .take_output_checkpoint_totals() + .remove("out") + .map(|totals| totals.records), + Some(RECORDS as u64) + ); +} + +/// An endpoint removed before it catches up reports nothing, and the +/// checkpoint keeps the figures it captured when it started. +#[test] +fn test_removed_output_endpoint_reports_nothing() { + const RECORDS: usize = 100; + + let status = test_controller_status("test_checkpoint_removed_endpoint"); + let output_config = test_output_config(); + status.add_output(&0, "out", &output_config, None, true); + + // The endpoint is still holding the batch when the checkpoint starts. + status.enqueue_batch(0, RECORDS); + status.arm_output_checkpoint(RECORDS as u64); + + status.remove_output(&0); + assert!(status.take_output_checkpoint_totals().is_empty()); +} + +/// The reading an endpoint reports must be exact however the checkpoint +/// interleaves with the endpoint's own work. +/// +/// Each round arms the checkpoint at a different point, so it lands before the +/// endpoint has caught up, at the moment it does, and after. The endpoint +/// transmits one record per record the circuit ingested, so whatever the +/// checkpoint covers is exactly what it must report. +#[test] +fn test_output_endpoint_reports_totals_under_concurrency() { + use super::stats::ProcessedRecords; + use std::sync::mpsc::channel; + use std::sync::{Arc, Mutex}; + use std::thread; + + const BATCHES: usize = 400; + // Transmitting a batch in two chunks widens the window in which the + // endpoint has put part of a batch on the transport but has not yet + // released it. + const FIRST_CHUNK: usize = 630; + const SECOND_CHUNK: usize = 370; + const BATCH: usize = FIRST_CHUNK + SECOND_CHUNK; + const ROUNDS: usize = 16; + + for round in 0..ROUNDS { + let status = Arc::new(test_controller_status("test_checkpoint_concurrency")); + status.add_output(&0, "out", &test_output_config(), None, true); + + let (sender, receiver) = channel(); + // Stands in for the circuit thread. The thread that produces output is + // the thread that arms the checkpoint, so no output can appear while a + // checkpoint is being armed, and the endpoint can never be past what + // the checkpoint covers. + let ingested = Arc::new(Mutex::new(0u64)); + + let circuit = { + let status = status.clone(); + let ingested = ingested.clone(); + thread::spawn(move || { + for _ in 0..BATCHES { + let mut ingested = ingested.lock().unwrap(); + status.enqueue_batch(0, BATCH); + *ingested += BATCH as u64; + let processed = *ingested; + drop(ingested); + sender.send(processed).unwrap(); + } + }) + }; + + let endpoint = { + let status = status.clone(); + thread::spawn(move || { + let parker = Parker::new(); + let unparker = parker.unparker().clone(); + while let Ok(processed) = receiver.recv() { + status.output_buffer(0, 4096, FIRST_CHUNK); + thread::yield_now(); + status.output_buffer(0, 4096, SECOND_CHUNK); + status.output_batch( + 0, + Some(ProcessedRecords { + total_processed_input_records: processed, + total_processed_steps: 0, + }), + BATCH, + &unparker, + ); + } + }) + }; + + let arm_after = ((round * BATCHES) / ROUNDS * BATCH) as u64; + while *ingested.lock().unwrap() < arm_after { + thread::yield_now(); + } + let threshold = { + let ingested = ingested.lock().unwrap(); + status.arm_output_checkpoint(*ingested); + *ingested + }; + + circuit.join().unwrap(); + endpoint.join().unwrap(); + + assert_eq!( + status + .take_output_checkpoint_totals() + .remove("out") + .map(|totals| totals.records), + Some(threshold), + "round {round}: the checkpoint covers {threshold} records" + ); + } +} + /// Only a changed relation makes an output endpoint fall behind. A connector whose /// own definition changed still emits nothing for inputs already processed, so it /// stays caught up and keeps its seeded progress counter.