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.