Skip to main content

mz_storage/
statistics.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Helpers for managing storage statistics.
11//!
12//!
13//! This module collects statistics related to sources and sinks. Statistics, as exposed
14//! to their respective system tables have strong semantics, defined within the
15//! `mz_storage_types::statistics` module. This module collects and aggregates metrics
16//! across workers according to those semantics.
17//!
18//! Note that it _simultaneously_ collect prometheus metrics for the given statistics. Those
19//! metrics _do not have the same strong semantics_, which is _by design_ to ensure we
20//! are able to categorically debug sources and sinks during complex failures (or bugs
21//! with statistics collection itself). Prometheus metrics are
22//! - Never dropped or reset until a source/sink is dropped.
23//! - Entirely independent across workers.
24
25use std::cell::RefCell;
26use std::collections::BTreeMap;
27use std::rc::Rc;
28use std::time::Instant;
29
30use mz_ore::metric;
31use mz_ore::metrics::{
32    DeleteOnDropCounter, DeleteOnDropGauge, IntCounterVec, IntGaugeVec, MetricTag,
33    MetricVisibility, MetricsRegistry, UIntGaugeVec,
34};
35use mz_repr::{GlobalId, Timestamp};
36use mz_storage_client::statistics::{Gauge, SinkStatisticsUpdate, SourceStatisticsUpdate};
37use mz_storage_types::sources::SourceEnvelope;
38use prometheus::core::{AtomicI64, AtomicU64};
39use serde::{Deserialize, Serialize};
40use timely::PartialOrder;
41use timely::progress::frontier::Antichain;
42
43// Note(guswynn): ordinarily these metric structs would be in the `metrics` modules, but we
44// put them here so they can be near the user-facing definitions as well.
45
46#[derive(Clone, Debug)]
47pub(crate) struct SourceStatisticsMetricDefs {
48    // Counters
49    pub(crate) messages_received: IntCounterVec,
50    pub(crate) updates_staged: IntCounterVec,
51    pub(crate) updates_committed: IntCounterVec,
52    pub(crate) bytes_received: IntCounterVec,
53
54    // Gauges
55    pub(crate) snapshot_committed: UIntGaugeVec,
56    pub(crate) bytes_indexed: UIntGaugeVec,
57    pub(crate) records_indexed: UIntGaugeVec,
58    pub(crate) rehydration_latency_ms: IntGaugeVec,
59
60    pub(crate) offset_known: UIntGaugeVec,
61    pub(crate) offset_committed: UIntGaugeVec,
62    pub(crate) snapshot_records_known: UIntGaugeVec,
63    pub(crate) snapshot_records_staged: UIntGaugeVec,
64
65    // Just prometheus.
66    pub(crate) envelope_state_tombstones: UIntGaugeVec,
67}
68
69impl SourceStatisticsMetricDefs {
70    pub(crate) fn register_with(registry: &MetricsRegistry) -> Self {
71        Self {
72            snapshot_committed: registry.register(metric!(
73                name: "mz_source_snapshot_committed",
74                help: "Whether or not the worker has committed the initial snapshot for a source.",
75                var_labels: ["source_id", "worker_id", "parent_source_id", "shard_id"],
76            )),
77            messages_received: registry.register(metric!(
78                name: "mz_source_messages_received",
79                help: "The number of raw messages the worker has received from upstream.",
80                var_labels: ["source_id", "worker_id", "parent_source_id"],
81                visibility: MetricVisibility::Public,
82                tags: [MetricTag::Source],
83            )),
84            updates_staged: registry.register(metric!(
85                name: "mz_source_updates_staged",
86                help: "The number of updates (inserts + deletes) the worker has written but not yet committed to the storage layer.",
87                var_labels: ["source_id", "worker_id", "parent_source_id", "shard_id"],
88            )),
89            updates_committed: registry.register(metric!(
90                name: "mz_source_updates_committed",
91                help: "The number of updates (inserts + deletes) the worker has committed into the storage layer.",
92                var_labels: ["source_id", "worker_id", "parent_source_id", "shard_id"],
93            )),
94            bytes_received: registry.register(metric!(
95                name: "mz_source_bytes_received",
96                help: "The number of bytes worth of messages the worker has received from upstream. The way the bytes are counted is source-specific.",
97                var_labels: ["source_id", "worker_id", "parent_source_id"],
98                visibility: MetricVisibility::Public,
99                tags: [MetricTag::Source],
100            )),
101            bytes_indexed: registry.register(metric!(
102                name: "mz_source_bytes_indexed",
103                help: "The number of bytes of the source envelope state kept. This will be specific to the envelope in use.",
104                var_labels: ["source_id", "worker_id", "parent_source_id", "shard_id"],
105            )),
106            records_indexed: registry.register(metric!(
107                name: "mz_source_records_indexed",
108                help: "The number of records in the source envelope state. This will be specific to the envelope in use",
109                var_labels: ["source_id", "worker_id", "parent_source_id", "shard_id"],
110            )),
111            envelope_state_tombstones: registry.register(metric!(
112                name: "mz_source_envelope_state_tombstones",
113                help: "The number of outstanding tombstones in the source envelope state. This will be specific to the envelope in use",
114                var_labels: ["source_id", "worker_id", "parent_source_id", "shard_id"],
115            )),
116            rehydration_latency_ms: registry.register(metric!(
117                name: "mz_source_rehydration_latency_ms",
118                help: "The amount of time in milliseconds it took for the worker to rehydrate the source envelope state. This will be specific to the envelope in use.",
119                var_labels: ["source_id", "worker_id", "parent_source_id", "shard_id", "envelope"],
120            )),
121            offset_known: registry.register(metric!(
122                name: "mz_source_offset_known",
123                help: "The total number of _values_ (source-defined unit) present in upstream.",
124                var_labels: ["source_id", "worker_id", "shard_id"],
125                visibility: MetricVisibility::Public,
126                tags: [MetricTag::Source],
127            )),
128            offset_committed: registry.register(metric!(
129                name: "mz_source_offset_committed",
130                help: "The total number of _values_ (source-defined unit) we have fully processed, and storage and committed.",
131                var_labels: ["source_id", "worker_id", "shard_id"],
132                visibility: MetricVisibility::Public,
133                tags: [MetricTag::Source],
134            )),
135            snapshot_records_known: registry.register(metric!(
136                name: "mz_source_snapshot_records_known",
137                help: "The total number of records in the source's snapshot",
138                var_labels: ["source_id", "worker_id", "shard_id", "parent_source_id"],
139            )),
140            snapshot_records_staged: registry.register(metric!(
141                name: "mz_source_snapshot_records_staged",
142                help: "The total number of records read from the source's snapshot",
143                var_labels: ["source_id", "worker_id", "shard_id", "parent_source_id"],
144            )),
145        }
146    }
147}
148
149/// Prometheus metrics for user-facing source metrics.
150#[derive(Debug)]
151pub struct SourceStatisticsMetrics {
152    // Counters
153    pub(crate) messages_received: DeleteOnDropCounter<AtomicU64, Vec<String>>,
154    pub(crate) updates_staged: DeleteOnDropCounter<AtomicU64, Vec<String>>,
155    pub(crate) updates_committed: DeleteOnDropCounter<AtomicU64, Vec<String>>,
156    pub(crate) bytes_received: DeleteOnDropCounter<AtomicU64, Vec<String>>,
157
158    // Gauges
159    pub(crate) snapshot_committed: DeleteOnDropGauge<AtomicU64, Vec<String>>,
160    pub(crate) bytes_indexed: DeleteOnDropGauge<AtomicU64, Vec<String>>,
161    pub(crate) records_indexed: DeleteOnDropGauge<AtomicU64, Vec<String>>,
162    pub(crate) rehydration_latency_ms: DeleteOnDropGauge<AtomicI64, Vec<String>>,
163
164    pub(crate) offset_known: DeleteOnDropGauge<AtomicU64, Vec<String>>,
165    pub(crate) offset_committed: DeleteOnDropGauge<AtomicU64, Vec<String>>,
166    pub(crate) snapshot_records_known: DeleteOnDropGauge<AtomicU64, Vec<String>>,
167    pub(crate) snapshot_records_staged: DeleteOnDropGauge<AtomicU64, Vec<String>>,
168
169    // Just prometheus.
170    pub(crate) envelope_state_tombstones: DeleteOnDropGauge<AtomicU64, Vec<String>>,
171}
172
173impl SourceStatisticsMetrics {
174    pub(crate) fn new(
175        defs: &SourceStatisticsMetricDefs,
176        id: GlobalId,
177        worker_id: usize,
178        parent_source_id: GlobalId,
179        shard_id: &mz_persist_client::ShardId,
180        envelope: SourceEnvelope,
181    ) -> SourceStatisticsMetrics {
182        let shard = shard_id.to_string();
183        let envelope = match envelope {
184            SourceEnvelope::None(_) => "none",
185            SourceEnvelope::Upsert(_) => "upsert",
186            SourceEnvelope::CdcV2 => "cdcv2",
187        };
188
189        SourceStatisticsMetrics {
190            snapshot_committed: defs.snapshot_committed.get_delete_on_drop_metric(vec![
191                id.to_string(),
192                worker_id.to_string(),
193                parent_source_id.to_string(),
194                shard.clone(),
195            ]),
196            messages_received: defs.messages_received.get_delete_on_drop_metric(vec![
197                id.to_string(),
198                worker_id.to_string(),
199                parent_source_id.to_string(),
200            ]),
201            updates_staged: defs.updates_staged.get_delete_on_drop_metric(vec![
202                id.to_string(),
203                worker_id.to_string(),
204                parent_source_id.to_string(),
205                shard.clone(),
206            ]),
207            updates_committed: defs.updates_committed.get_delete_on_drop_metric(vec![
208                id.to_string(),
209                worker_id.to_string(),
210                parent_source_id.to_string(),
211                shard.clone(),
212            ]),
213            bytes_received: defs.bytes_received.get_delete_on_drop_metric(vec![
214                id.to_string(),
215                worker_id.to_string(),
216                parent_source_id.to_string(),
217            ]),
218            bytes_indexed: defs.bytes_indexed.get_delete_on_drop_metric(vec![
219                id.to_string(),
220                worker_id.to_string(),
221                parent_source_id.to_string(),
222                shard.clone(),
223            ]),
224            records_indexed: defs.records_indexed.get_delete_on_drop_metric(vec![
225                id.to_string(),
226                worker_id.to_string(),
227                parent_source_id.to_string(),
228                shard.clone(),
229            ]),
230            envelope_state_tombstones: defs.envelope_state_tombstones.get_delete_on_drop_metric(
231                vec![
232                    id.to_string(),
233                    worker_id.to_string(),
234                    parent_source_id.to_string(),
235                    shard.clone(),
236                ],
237            ),
238            rehydration_latency_ms: defs.rehydration_latency_ms.get_delete_on_drop_metric(vec![
239                id.to_string(),
240                worker_id.to_string(),
241                parent_source_id.to_string(),
242                shard.clone(),
243                envelope.to_string(),
244            ]),
245            offset_known: defs.offset_known.get_delete_on_drop_metric(vec![
246                id.to_string(),
247                worker_id.to_string(),
248                shard.clone(),
249            ]),
250            offset_committed: defs.offset_committed.get_delete_on_drop_metric(vec![
251                id.to_string(),
252                worker_id.to_string(),
253                shard.clone(),
254            ]),
255            snapshot_records_known: defs.snapshot_records_known.get_delete_on_drop_metric(vec![
256                id.to_string(),
257                worker_id.to_string(),
258                shard.clone(),
259                parent_source_id.to_string(),
260            ]),
261            snapshot_records_staged: defs.snapshot_records_staged.get_delete_on_drop_metric(vec![
262                id.to_string(),
263                worker_id.to_string(),
264                shard.clone(),
265                parent_source_id.to_string(),
266            ]),
267        }
268    }
269}
270
271#[derive(Clone, Debug)]
272pub(crate) struct SinkStatisticsMetricDefs {
273    // Counters
274    pub(crate) messages_staged: IntCounterVec,
275    pub(crate) messages_committed: IntCounterVec,
276    pub(crate) bytes_staged: IntCounterVec,
277    pub(crate) bytes_committed: IntCounterVec,
278}
279
280impl SinkStatisticsMetricDefs {
281    pub(crate) fn register_with(registry: &MetricsRegistry) -> Self {
282        Self {
283            messages_staged: registry.register(metric!(
284                name: "mz_sink_messages_staged",
285                help: "The number of messages staged but possibly not committed to the sink.",
286                var_labels: ["sink_id", "worker_id"],
287            )),
288            messages_committed: registry.register(metric!(
289                name: "mz_sink_messages_committed",
290                help: "The number of messages committed to the sink.",
291                var_labels: ["sink_id", "worker_id"],
292            )),
293            bytes_staged: registry.register(metric!(
294                name: "mz_sink_bytes_staged",
295                help: "The number of bytes staged but possibly not committed to the sink.",
296                var_labels: ["sink_id", "worker_id"],
297                visibility: MetricVisibility::Public,
298                tags: [MetricTag::Sink],
299            )),
300            bytes_committed: registry.register(metric!(
301                name: "mz_sink_bytes_committed",
302                help: "The number of bytes committed to the sink.",
303                var_labels: ["sink_id", "worker_id"],
304                visibility: MetricVisibility::Public,
305                tags: [MetricTag::Sink],
306            )),
307        }
308    }
309}
310
311/// Prometheus metrics for user-facing sink metrics.
312#[derive(Debug)]
313pub struct SinkStatisticsMetrics {
314    // Counters
315    pub(crate) messages_staged: DeleteOnDropCounter<AtomicU64, Vec<String>>,
316    pub(crate) messages_committed: DeleteOnDropCounter<AtomicU64, Vec<String>>,
317    pub(crate) bytes_staged: DeleteOnDropCounter<AtomicU64, Vec<String>>,
318    pub(crate) bytes_committed: DeleteOnDropCounter<AtomicU64, Vec<String>>,
319}
320
321impl SinkStatisticsMetrics {
322    pub(crate) fn new(
323        defs: &SinkStatisticsMetricDefs,
324        id: GlobalId,
325        worker_id: usize,
326    ) -> SinkStatisticsMetrics {
327        SinkStatisticsMetrics {
328            messages_staged: defs
329                .messages_staged
330                .get_delete_on_drop_metric(vec![id.to_string(), worker_id.to_string()]),
331            messages_committed: defs
332                .messages_committed
333                .get_delete_on_drop_metric(vec![id.to_string(), worker_id.to_string()]),
334            bytes_staged: defs
335                .bytes_staged
336                .get_delete_on_drop_metric(vec![id.to_string(), worker_id.to_string()]),
337            bytes_committed: defs
338                .bytes_committed
339                .get_delete_on_drop_metric(vec![id.to_string(), worker_id.to_string()]),
340        }
341    }
342}
343
344/// Meta data needed to maintain source statistics.
345#[derive(Debug, Clone)]
346pub struct SourceStatisticsMetadata {
347    /// The resumption upper of the source.
348    resume_upper: Antichain<Timestamp>,
349    /// The as_of of the ingestion dataflow, where an export that snapshots lands its snapshot.
350    as_of: Antichain<Timestamp>,
351    /// Time at which the source (more precisely: this metadata object) was created.
352    created_at: Instant,
353}
354
355impl SourceStatisticsMetadata {
356    /// Create a new `SourceStatisticsMetadata` object.
357    pub fn new(resume_upper: Antichain<Timestamp>, as_of: Antichain<Timestamp>) -> Self {
358        Self {
359            resume_upper,
360            as_of,
361            created_at: Instant::now(),
362        }
363    }
364}
365
366#[derive(Debug)]
367struct StatsInner<Stats, Metrics> {
368    stats: Stats,
369    prom: Metrics,
370}
371
372/// A helper struct designed to make it easy for operators to update user-facing metrics.
373/// This struct also ensures that each stack is also incremented in prometheus.
374///
375/// Caveats:
376/// - There is one Prometheus timeseries-per-worker, and we label it with the source id and the
377/// source id of the parent source (if there is no parent, these labels have the same value).
378///     - Some metrics also have the shard id we are writing metrics for.
379/// - The prometheus metrics do not have the same timestamps as the ones exposed in sql, because
380/// they are written at different times.
381///     - This may be fixed in the future when we write the metrics from storaged directly.
382///     - The value also eventually converge to the same value.
383#[derive(Debug)]
384pub struct StorageStatistics<Stats, Metrics, Meta> {
385    // Note also that the `DeleteOnDropCounter`'s in the `SourceStatisticsMetrics`
386    // already are in an `Arc`, so this is a bit of extra wrapping, but the cost
387    // shouldn't cost too much.
388    stats: Rc<RefCell<StatsInner<Stats, Metrics>>>,
389    /// Meta data needed to maintain statistics.
390    meta: Meta,
391}
392
393impl<Stats, Metrics, Meta: Clone> Clone for StorageStatistics<Stats, Metrics, Meta> {
394    fn clone(&self) -> Self {
395        Self {
396            stats: Rc::clone(&self.stats),
397            meta: self.meta.clone(),
398        }
399    }
400}
401
402/// Statistics maintained for sources. Gauges can be uninitialized.
403#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
404pub struct SourceStatisticsRecord {
405    id: GlobalId,
406    worker_id: usize,
407    // Counters
408    messages_received: u64,
409    bytes_received: u64,
410    updates_staged: u64,
411    updates_committed: u64,
412
413    // Gauges are always wrapped in an `Option` that represents if that gauge has been
414    // initialized by that worker.
415    records_indexed: Option<u64>,
416    bytes_indexed: Option<u64>,
417
418    // This field is nullable, so its value is an `Option`.
419    rehydration_latency_ms: Option<Option<i64>>,
420
421    // The following fields are able to be unset when shipped to the controller, so their
422    // values are `Option`'s
423    snapshot_records_known: Option<Option<u64>>,
424    snapshot_records_staged: Option<Option<u64>>,
425    snapshot_committed: Option<bool>,
426    offset_known: Option<Option<u64>>,
427    offset_committed: Option<Option<u64>>,
428
429    // Just prometheus.
430    envelope_state_tombstones: u64,
431}
432
433impl SourceStatisticsRecord {
434    fn reset_gauges(&mut self) {
435        // These gauges MUST be initialized across all workers before we aggregate and
436        // report statistics.
437        self.rehydration_latency_ms = None;
438        self.snapshot_committed = None;
439
440        // We consider these gauges always initialized
441        self.bytes_indexed = Some(0);
442        self.records_indexed = Some(0);
443
444        // We don't yet populate these, so we consider the initialized (with an empty value).
445        self.snapshot_records_known = Some(None);
446        self.snapshot_records_staged = Some(None);
447        self.offset_known = Some(None);
448        self.offset_committed = Some(None);
449
450        self.envelope_state_tombstones = 0;
451    }
452
453    /// Reset counters so that we continue to ship diffs to the controller.
454    fn reset_counters(&mut self) {
455        self.messages_received = 0;
456        self.bytes_received = 0;
457        self.updates_staged = 0;
458        self.updates_committed = 0;
459    }
460
461    /// Convert this record into an `SourceStatisticsUpdate` to be merged
462    /// across workers. Gauges must be initialized.
463    fn as_update(&self) -> SourceStatisticsUpdate {
464        let SourceStatisticsRecord {
465            id,
466            worker_id: _,
467            messages_received,
468            bytes_received,
469            updates_staged,
470            updates_committed,
471            records_indexed,
472            bytes_indexed,
473            rehydration_latency_ms,
474            snapshot_records_known,
475            snapshot_records_staged,
476            snapshot_committed,
477            offset_known,
478            offset_committed,
479            envelope_state_tombstones: _,
480        } = self.clone();
481
482        SourceStatisticsUpdate {
483            id,
484            messages_received: messages_received.into(),
485            bytes_received: bytes_received.into(),
486            updates_staged: updates_staged.into(),
487            updates_committed: updates_committed.into(),
488            records_indexed: Gauge::gauge(records_indexed.unwrap()),
489            bytes_indexed: Gauge::gauge(bytes_indexed.unwrap()),
490            rehydration_latency_ms: Gauge::gauge(rehydration_latency_ms.unwrap()),
491            snapshot_records_known: Gauge::gauge(snapshot_records_known.unwrap()),
492            snapshot_records_staged: Gauge::gauge(snapshot_records_staged.unwrap()),
493            snapshot_committed: Gauge::gauge(snapshot_committed.unwrap()),
494            offset_known: Gauge::gauge(offset_known.unwrap()),
495            offset_committed: Gauge::gauge(offset_committed.unwrap()),
496        }
497    }
498}
499
500/// Statistics maintained for sinks. Gauges can be uninitialized.
501#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
502pub struct SinkStatisticsRecord {
503    id: GlobalId,
504    worker_id: usize,
505    // Counters
506    messages_staged: u64,
507    messages_committed: u64,
508    bytes_staged: u64,
509    bytes_committed: u64,
510}
511
512impl SinkStatisticsRecord {
513    fn reset_gauges(&self) {}
514
515    /// Reset counters so that we continue to ship diffs to the controller.
516    fn reset_counters(&mut self) {
517        self.messages_staged = 0;
518        self.messages_committed = 0;
519        self.bytes_staged = 0;
520        self.bytes_committed = 0;
521    }
522
523    /// Convert this record into an `SinkStatisticsUpdate` to be merged
524    /// across workers. Gauges must be initialized.
525    fn as_update(&self) -> SinkStatisticsUpdate {
526        let SinkStatisticsRecord {
527            id,
528            worker_id: _,
529            messages_staged,
530            messages_committed,
531            bytes_staged,
532            bytes_committed,
533        } = self.clone();
534
535        SinkStatisticsUpdate {
536            id,
537            messages_staged: messages_staged.into(),
538            messages_committed: messages_committed.into(),
539            bytes_staged: bytes_staged.into(),
540            bytes_committed: bytes_committed.into(),
541        }
542    }
543}
544
545/// Statistics maintained for sources.
546pub type SourceStatistics =
547    StorageStatistics<SourceStatisticsRecord, SourceStatisticsMetrics, SourceStatisticsMetadata>;
548
549/// Statistics maintained for sinks.
550pub type SinkStatistics = StorageStatistics<SinkStatisticsRecord, SinkStatisticsMetrics, ()>;
551
552impl SourceStatistics {
553    pub(crate) fn new(
554        id: GlobalId,
555        worker_id: usize,
556        metrics: &SourceStatisticsMetricDefs,
557        parent_source_id: GlobalId,
558        shard_id: &mz_persist_client::ShardId,
559        envelope: SourceEnvelope,
560        resume_upper: Antichain<Timestamp>,
561        as_of: Antichain<Timestamp>,
562    ) -> Self {
563        Self {
564            stats: Rc::new(RefCell::new(StatsInner {
565                stats: SourceStatisticsRecord {
566                    id,
567                    worker_id,
568                    messages_received: 0,
569                    updates_staged: 0,
570                    updates_committed: 0,
571                    bytes_received: 0,
572                    records_indexed: Some(0),
573                    bytes_indexed: Some(0),
574                    rehydration_latency_ms: None,
575                    snapshot_records_staged: Some(None),
576                    snapshot_records_known: Some(None),
577                    snapshot_committed: None,
578                    offset_known: Some(None),
579                    offset_committed: Some(None),
580                    envelope_state_tombstones: 0,
581                },
582                prom: SourceStatisticsMetrics::new(
583                    metrics,
584                    id,
585                    worker_id,
586                    parent_source_id,
587                    shard_id,
588                    envelope,
589                ),
590            })),
591            meta: SourceStatisticsMetadata::new(resume_upper, as_of),
592        }
593    }
594
595    /// Reset gauges to uninitialized state.
596    /// This does not reset prometheus gauges, which will be overwritten with new values.
597    pub fn reset_gauges(&self) {
598        let mut cur = self.stats.borrow_mut();
599        cur.stats.reset_gauges();
600    }
601
602    /// Get a snapshot of the data, returning `None` if all gauges are not initialized.
603    ///
604    /// This also resets counters, so that we continue to move diffs around.
605    pub fn snapshot(&self) -> Option<SourceStatisticsRecord> {
606        let mut cur = self.stats.borrow_mut();
607
608        match &cur.stats {
609            SourceStatisticsRecord {
610                records_indexed: Some(_),
611                bytes_indexed: Some(_),
612                rehydration_latency_ms: Some(_),
613                snapshot_records_known: Some(_),
614                snapshot_records_staged: Some(_),
615                snapshot_committed: Some(_),
616                offset_known: Some(_),
617                offset_committed: Some(_),
618                ..
619            } => {
620                let ret = Some(cur.stats.clone());
621                cur.stats.reset_counters();
622                ret
623            }
624            _ => None,
625        }
626    }
627
628    /// Set the `snapshot_committed` stat based on the reported upper, and
629    /// mark the stats as initialized.
630    ///
631    /// - In sql, we ensure that we _never_ reset `snapshot_committed` to `false`, but gauges and
632    /// counters are ordinarily reset to 0 in Prometheus, so on restarts this value may be inconsistent.
633    // TODO(guswynn): Actually test that this initialization logic works.
634    pub fn initialize_snapshot_committed(&self, upper: &Antichain<Timestamp>) {
635        self.update_snapshot_committed(upper);
636    }
637
638    /// Set the `snapshot_committed` stat based on the reported upper.
639    pub fn update_snapshot_committed(&self, upper: &Antichain<Timestamp>) {
640        // An upper past the minimum is not enough. When the as_of is past the minimum, the sink
641        // commits the empty interval below it before the snapshot lands at the as_of.
642        let value = PartialOrder::less_than(&self.meta.as_of, upper);
643        let mut cur = self.stats.borrow_mut();
644        cur.stats.snapshot_committed = Some(value);
645        cur.prom.snapshot_committed.set(if value { 1 } else { 0 });
646    }
647
648    /// Increment the `messages_received` stat.
649    pub fn inc_messages_received_by(&self, value: u64) {
650        let mut cur = self.stats.borrow_mut();
651        cur.stats.messages_received = cur.stats.messages_received + value;
652        cur.prom.messages_received.inc_by(value);
653    }
654
655    /// Increment the `updates` stat.
656    pub fn inc_updates_staged_by(&self, value: u64) {
657        let mut cur = self.stats.borrow_mut();
658        cur.stats.updates_staged = cur.stats.updates_staged + value;
659        cur.prom.updates_staged.inc_by(value);
660    }
661
662    /// Increment the `messages_committed` stat.
663    pub fn inc_updates_committed_by(&self, value: u64) {
664        let mut cur = self.stats.borrow_mut();
665        cur.stats.updates_committed = cur.stats.updates_committed + value;
666        cur.prom.updates_committed.inc_by(value);
667    }
668
669    /// Increment the `bytes_received` stat.
670    pub fn inc_bytes_received_by(&self, value: u64) {
671        let mut cur = self.stats.borrow_mut();
672        cur.stats.bytes_received = cur.stats.bytes_received + value;
673        cur.prom.bytes_received.inc_by(value);
674    }
675
676    /// Update the `bytes_indexed` stat.
677    /// A positive value will add and a negative value will subtract.
678    pub fn update_bytes_indexed_by(&self, value: i64) {
679        let mut cur = self.stats.borrow_mut();
680        if let Some(updated) = cur
681            .stats
682            .bytes_indexed
683            .unwrap_or(0)
684            .checked_add_signed(value)
685        {
686            cur.stats.bytes_indexed = Some(updated);
687            cur.prom.bytes_indexed.set(updated);
688        } else {
689            let bytes_indexed = cur.stats.bytes_indexed.unwrap_or(0);
690            tracing::warn!(
691                "Unexpected u64 overflow while updating bytes_indexed value {} with {}",
692                bytes_indexed,
693                value
694            );
695            cur.stats.bytes_indexed = Some(0);
696            cur.prom.bytes_indexed.set(0);
697        }
698    }
699
700    /// Set the `bytes_indexed` to the given value
701    pub fn set_bytes_indexed(&self, value: i64) {
702        let mut cur = self.stats.borrow_mut();
703        let value = if value < 0 {
704            tracing::warn!("Unexpected negative value for bytes_indexed {}", value);
705            0
706        } else {
707            value.unsigned_abs()
708        };
709        cur.stats.bytes_indexed = Some(value);
710        cur.prom.bytes_indexed.set(value);
711    }
712
713    /// Update the `records_indexed` stat.
714    /// A positive value will add and a negative value will subtract.
715    pub fn update_records_indexed_by(&self, value: i64) {
716        let mut cur = self.stats.borrow_mut();
717        if let Some(updated) = cur
718            .stats
719            .records_indexed
720            .unwrap_or(0)
721            .checked_add_signed(value)
722        {
723            cur.stats.records_indexed = Some(updated);
724            cur.prom.records_indexed.set(updated);
725        } else {
726            let records_indexed = cur.stats.records_indexed.unwrap_or(0);
727            tracing::warn!(
728                "Unexpected u64 overflow while updating records_indexed value {} with {}",
729                records_indexed,
730                value
731            );
732            cur.stats.records_indexed = Some(0);
733            cur.prom.records_indexed.set(0);
734        }
735    }
736
737    /// Set the `records_indexed` to the given value
738    pub fn set_records_indexed(&self, value: i64) {
739        let mut cur = self.stats.borrow_mut();
740        let value = if value < 0 {
741            tracing::warn!("Unexpected negative value for records_indexed {}", value);
742            0
743        } else {
744            value.unsigned_abs()
745        };
746        cur.stats.records_indexed = Some(value);
747        cur.prom.records_indexed.set(value);
748    }
749
750    /// Initialize the `rehydration_latency_ms` stat as `NULL`.
751    pub fn initialize_rehydration_latency_ms(&self) {
752        let mut cur = self.stats.borrow_mut();
753        cur.stats.rehydration_latency_ms = Some(None);
754    }
755
756    /// Update the `envelope_state_tombstones` stat.
757    /// A positive value will add and a negative value will subtract.
758    // TODO(guswynn): consider exposing this to users
759    pub fn update_envelope_state_tombstones_by(&self, value: i64) {
760        let mut cur = self.stats.borrow_mut();
761        if let Some(updated) = cur
762            .stats
763            .envelope_state_tombstones
764            .checked_add_signed(value)
765        {
766            cur.stats.envelope_state_tombstones = updated;
767            cur.prom.envelope_state_tombstones.set(updated);
768        } else {
769            let envelope_state_tombstones = cur.stats.envelope_state_tombstones;
770            tracing::warn!(
771                "Unexpected u64 overflow while updating envelope_state_tombstones value {} with {}",
772                envelope_state_tombstones,
773                value
774            );
775            cur.stats.envelope_state_tombstones = 0;
776            cur.prom.envelope_state_tombstones.set(0);
777        }
778    }
779
780    /// Set the `rehydration_latency_ms` stat based on the reported upper.
781    pub fn update_rehydration_latency_ms(&self, upper: &Antichain<Timestamp>) {
782        let mut cur = self.stats.borrow_mut();
783
784        if matches!(cur.stats.rehydration_latency_ms, Some(Some(_))) {
785            return; // source was already hydrated before
786        }
787        // An export that snapshots has a resume upper at or below the as_of, and passes the
788        // resume upper before its snapshot commits (see `update_snapshot_committed`).
789        if !PartialOrder::less_than(&self.meta.resume_upper, upper)
790            || !PartialOrder::less_than(&self.meta.as_of, upper)
791        {
792            return; // source is not yet hydrated
793        }
794
795        let elapsed = self.meta.created_at.elapsed();
796        let value = elapsed
797            .as_millis()
798            .try_into()
799            .expect("Rehydration took more than ~584 million years!");
800        cur.stats.rehydration_latency_ms = Some(Some(value));
801        cur.prom.rehydration_latency_ms.set(value);
802    }
803
804    /// Set the `offset_known` stat to the given value.
805    pub fn set_offset_known(&self, value: u64) {
806        let mut cur = self.stats.borrow_mut();
807        cur.stats.offset_known = Some(Some(value));
808        cur.prom.offset_known.set(value);
809    }
810
811    /// Set the `offset_committed` stat to the given value.
812    pub fn set_offset_committed(&self, value: u64) {
813        let mut cur = self.stats.borrow_mut();
814        cur.stats.offset_committed = Some(Some(value));
815        cur.prom.offset_committed.set(value);
816    }
817
818    /// Set the `snapshot_records_known` stat to the given value.
819    pub fn set_snapshot_records_known(&self, value: u64) {
820        let mut cur = self.stats.borrow_mut();
821        cur.stats.snapshot_records_known = Some(Some(value));
822        cur.prom.snapshot_records_known.set(value);
823    }
824
825    /// Set the `snapshot_records_known` stat to the given value.
826    pub fn set_snapshot_records_staged(&self, value: u64) {
827        let mut cur = self.stats.borrow_mut();
828        cur.stats.snapshot_records_staged = Some(Some(value));
829        cur.prom.snapshot_records_staged.set(value);
830    }
831}
832
833impl SinkStatistics {
834    pub(crate) fn new(id: GlobalId, worker_id: usize, metrics: &SinkStatisticsMetricDefs) -> Self {
835        Self {
836            stats: Rc::new(RefCell::new(StatsInner {
837                stats: SinkStatisticsRecord {
838                    id,
839                    worker_id,
840                    messages_staged: 0,
841                    messages_committed: 0,
842                    bytes_staged: 0,
843                    bytes_committed: 0,
844                },
845                prom: SinkStatisticsMetrics::new(metrics, id, worker_id),
846            })),
847            meta: (),
848        }
849    }
850
851    /// Reset gauges to uninitialized state.
852    /// This does not reset prometheus gauges, which will be overwritten with new values.
853    pub fn reset_gauges(&self) {
854        let cur = self.stats.borrow_mut();
855        cur.stats.reset_gauges();
856    }
857
858    /// Get a snapshot of the data, returning `None` if all gauges are not initialized.
859    ///
860    /// This also resets counters, so that we continue to move diffs around.
861    pub fn snapshot(&self) -> Option<SinkStatisticsRecord> {
862        let mut cur = self.stats.borrow_mut();
863
864        match &cur.stats {
865            SinkStatisticsRecord { .. } => {
866                let ret = Some(cur.stats.clone());
867                cur.stats.reset_counters();
868                ret
869            }
870        }
871    }
872
873    /// Increment the `messages_staged` stat.
874    pub fn inc_messages_staged_by(&self, value: u64) {
875        let mut cur = self.stats.borrow_mut();
876        cur.stats.messages_staged = cur.stats.messages_staged + value;
877        cur.prom.messages_staged.inc_by(value);
878    }
879
880    /// Increment the `bytes_received` stat.
881    pub fn inc_bytes_staged_by(&self, value: u64) {
882        let mut cur = self.stats.borrow_mut();
883        cur.stats.bytes_staged = cur.stats.bytes_staged + value;
884        cur.prom.bytes_staged.inc_by(value);
885    }
886
887    /// Increment the `messages_committed` stat.
888    pub fn inc_messages_committed_by(&self, value: u64) {
889        let mut cur = self.stats.borrow_mut();
890        cur.stats.messages_committed = cur.stats.messages_committed + value;
891        cur.prom.messages_committed.inc_by(value);
892    }
893
894    /// Increment the `bytes_committed` stat.
895    pub fn inc_bytes_committed_by(&self, value: u64) {
896        let mut cur = self.stats.borrow_mut();
897        cur.stats.bytes_committed = cur.stats.bytes_committed + value;
898        cur.prom.bytes_committed.inc_by(value);
899    }
900}
901
902/// A structure that keeps track of _local_ statistics, as well as aggregating
903/// statistics into a single worker (currently worker 0).
904///
905/// This is because we ONLY want to emit statistics for _gauge-style-statistics_ when
906/// ALL workers have caught up to the _currently running instance of a source/sink_ and have
907/// reported an accurate gauge. This prevents issues like 1 worker reporting
908/// `snapshot_committed=true` before the others workers have caught up to report that they haven't
909/// finished snapshotting.
910///
911/// The API flow is:
912/// - Initialize sources/sinks with `initialize_source` or `initialize_sink`. This advances the
913/// local epoch.
914/// - Advance the _global epoch_ with `advance_global_epoch`. This should always be called strictly
915/// before reinitializing objects.
916/// - Emit a snapshot of local data with `emit_local`, which can be broadcasted to other workers.
917/// - Ingest data from other workers with `ingest`.
918/// - Emit a `snapshot` of _global data_ (i.e. skipping objects that don't have all workers caught
919/// up) with `snapshot`.
920///
921/// All functions can and should be called from ALL workers.
922///
923/// Note that we hand-roll a statistics epoch here, but in the future, when the storage
924/// _controller_ has top-level support for _dataflow epochs_, this logic can be entirely moved
925/// there.
926pub struct AggregatedStatistics {
927    worker_id: usize,
928    worker_count: usize,
929    local_source_statistics: BTreeMap<GlobalId, (usize, GlobalId, SourceStatistics)>,
930    local_sink_statistics: BTreeMap<GlobalId, (usize, SinkStatistics)>,
931
932    global_source_statistics:
933        BTreeMap<GlobalId, (usize, GlobalId, Vec<Option<SourceStatisticsUpdate>>)>,
934    global_sink_statistics: BTreeMap<GlobalId, (usize, Vec<Option<SinkStatisticsUpdate>>)>,
935}
936
937impl AggregatedStatistics {
938    /// Initialize a new `AggregatedStatistics`.
939    pub fn new(worker_id: usize, worker_count: usize) -> Self {
940        AggregatedStatistics {
941            worker_id,
942            worker_count,
943            local_source_statistics: Default::default(),
944            local_sink_statistics: Default::default(),
945            global_source_statistics: Default::default(),
946            global_sink_statistics: Default::default(),
947        }
948    }
949
950    /// Get a collection of `SourceStatistics` for the ingestion `ingestion_id`.
951    pub fn get_ingestion_stats(
952        &self,
953        ingestion_id: &GlobalId,
954    ) -> BTreeMap<GlobalId, SourceStatistics> {
955        let mut ingestion_stats = BTreeMap::new();
956        for (id, (_epoch, ingestion_id2, stats)) in self.local_source_statistics.iter() {
957            if ingestion_id == ingestion_id2 {
958                ingestion_stats.insert(id.clone(), stats.clone());
959            }
960        }
961        ingestion_stats
962    }
963
964    /// Get a `SourceStatistics` for an id, if it exists.
965    pub fn get_source(&self, id: &GlobalId) -> Option<&SourceStatistics> {
966        self.local_source_statistics.get(id).map(|(_, _, s)| s)
967    }
968
969    /// Get a `SinkStatistics` for an id, if it exists.
970    pub fn get_sink(&self, id: &GlobalId) -> Option<&SinkStatistics> {
971        self.local_sink_statistics.get(id).map(|(_, s)| s)
972    }
973
974    /// Deinitialize an object. Other methods other than `initialize_source` and `initialize_sink`
975    /// will never overwrite this.
976    pub fn deinitialize(&mut self, id: GlobalId) {
977        self.local_source_statistics.remove(&id);
978        self.local_sink_statistics.remove(&id);
979        self.global_source_statistics.remove(&id);
980        self.global_sink_statistics.remove(&id);
981    }
982
983    /// Advance the _global epoch_ for statistics.
984    ///
985    /// Gauge values from previous epochs will be ignored. Counter values from previous epochs will
986    /// still be applied as usual.
987    pub fn advance_global_epoch(&mut self, id: GlobalId) {
988        if let Some((epoch, _ingestion_id, stats)) = self.global_source_statistics.get_mut(&id) {
989            *epoch += 1;
990            for worker_stats in stats {
991                if let Some(update) = worker_stats {
992                    update.reset_gauges();
993                }
994            }
995        }
996        if let Some((epoch, stats)) = self.global_sink_statistics.get_mut(&id) {
997            *epoch += 1;
998            for worker_stats in stats {
999                if let Some(update) = worker_stats {
1000                    update.reset_gauges();
1001                }
1002            }
1003        }
1004    }
1005
1006    /// Re-initialize a source. If it already exists, then its local epoch is advanced.
1007    pub fn initialize_source<F: FnOnce() -> SourceStatistics>(
1008        &mut self,
1009        id: GlobalId,
1010        ingestion_id: GlobalId,
1011        resume_upper: Antichain<Timestamp>,
1012        as_of: Antichain<Timestamp>,
1013        stats: F,
1014    ) {
1015        self.local_source_statistics
1016            .entry(id)
1017            .and_modify(|(epoch, ingestion_id2, stats)| {
1018                assert_eq!(ingestion_id, *ingestion_id2);
1019                *epoch += 1;
1020                stats.reset_gauges();
1021                stats.meta = SourceStatisticsMetadata::new(resume_upper, as_of);
1022            })
1023            .or_insert_with(|| (0, ingestion_id, stats()));
1024
1025        if self.worker_id == 0 {
1026            self.global_source_statistics
1027                .entry(id)
1028                .or_insert_with(|| (0, ingestion_id, vec![None; self.worker_count]));
1029        }
1030    }
1031
1032    /// Re-initialize a sink. If it already exists, then its local epoch is advanced.
1033    pub fn initialize_sink<F: FnOnce() -> SinkStatistics>(&mut self, id: GlobalId, stats: F) {
1034        self.local_sink_statistics
1035            .entry(id)
1036            .and_modify(|(epoch, stats)| {
1037                *epoch += 1;
1038                stats.reset_gauges();
1039            })
1040            .or_insert_with(|| (0, stats()));
1041        if self.worker_id == 0 {
1042            self.global_sink_statistics
1043                .entry(id)
1044                .or_insert_with(|| (0, vec![None; self.worker_count]));
1045        }
1046    }
1047
1048    /// Ingest data from other workers.
1049    pub fn ingest(
1050        &mut self,
1051        source_statistics: Vec<(usize, SourceStatisticsRecord)>,
1052        sink_statistics: Vec<(usize, SinkStatisticsRecord)>,
1053    ) {
1054        // Currently, only worker 0 ingest data from other workers.
1055        if self.worker_id != 0 {
1056            return;
1057        }
1058
1059        for (epoch, stat) in source_statistics {
1060            if let Some((global_epoch, _, stats)) = self.global_source_statistics.get_mut(&stat.id)
1061            {
1062                // The record might be from a previous incarnation of the source. If that's the
1063                // case, we only incorporate its counter values and ignore its gauge values.
1064                let epoch_match = epoch >= *global_epoch;
1065                let mut update = stat.as_update();
1066                match (&mut stats[stat.worker_id], epoch_match) {
1067                    (None, true) => stats[stat.worker_id] = Some(update),
1068                    (None, false) => {
1069                        update.reset_gauges();
1070                        stats[stat.worker_id] = Some(update);
1071                    }
1072                    (Some(occupied), true) => occupied.incorporate(update),
1073                    (Some(occupied), false) => occupied.incorporate_counters(update),
1074                }
1075            }
1076        }
1077
1078        for (epoch, stat) in sink_statistics {
1079            if let Some((global_epoch, stats)) = self.global_sink_statistics.get_mut(&stat.id) {
1080                // The record might be from a previous incarnation of the sink. If that's the
1081                // case, we only incorporate its counter values and ignore its gauge values.
1082                let epoch_match = epoch >= *global_epoch;
1083                let update = stat.as_update();
1084                match (&mut stats[stat.worker_id], epoch_match) {
1085                    (None, true) => stats[stat.worker_id] = Some(update),
1086                    (None, false) => {
1087                        update.reset_gauges();
1088                        stats[stat.worker_id] = Some(update);
1089                    }
1090                    (Some(occupied), true) => occupied.incorporate(update),
1091                    (Some(occupied), false) => occupied.incorporate_counters(update),
1092                }
1093            }
1094        }
1095    }
1096
1097    fn _emit_local(
1098        &mut self,
1099    ) -> (
1100        Vec<(usize, SourceStatisticsRecord)>,
1101        Vec<(usize, SinkStatisticsRecord)>,
1102    ) {
1103        let sources = self
1104            .local_source_statistics
1105            .values_mut()
1106            .flat_map(|(epoch, _, s)| s.snapshot().map(|v| (*epoch, v)))
1107            .collect();
1108
1109        let sinks = self
1110            .local_sink_statistics
1111            .values_mut()
1112            .flat_map(|(epoch, s)| s.snapshot().map(|v| (*epoch, v)))
1113            .collect();
1114
1115        (sources, sinks)
1116    }
1117
1118    /// Emit a snapshot of this workers local data.
1119    pub fn emit_local(
1120        &mut self,
1121    ) -> (
1122        Vec<(usize, SourceStatisticsRecord)>,
1123        Vec<(usize, SinkStatisticsRecord)>,
1124    ) {
1125        // As an optimization, worker 0 does not broadcast it data. It ingests
1126        // it in `snapshot`.
1127        if self.worker_id == 0 {
1128            return (Vec::new(), Vec::new());
1129        }
1130
1131        self._emit_local()
1132    }
1133
1134    /// Emit a _global_ snapshot of data. This does not include objects whose workers have not
1135    /// initialized gauges for the current epoch.
1136    pub fn snapshot(&mut self) -> (Vec<SourceStatisticsUpdate>, Vec<SinkStatisticsUpdate>) {
1137        if !self.worker_id == 0 {
1138            return (Vec::new(), Vec::new());
1139        }
1140
1141        let (sources, sinks) = self._emit_local();
1142        self.ingest(sources, sinks);
1143
1144        let sources = self
1145            .global_source_statistics
1146            .iter_mut()
1147            .filter_map(|(_, (_, _, s))| {
1148                if s.iter().all(|s| s.is_some()) {
1149                    let ret = Some(SourceStatisticsUpdate::summarize(|| {
1150                        s.iter().filter_map(Option::as_ref)
1151                    }));
1152
1153                    // Reset the counters so we only report diffs.
1154                    s.iter_mut().for_each(|s| {
1155                        if let Some(s) = s {
1156                            s.reset_counters();
1157                        }
1158                    });
1159                    ret
1160                } else {
1161                    None
1162                }
1163            })
1164            .collect();
1165
1166        let sinks = self
1167            .global_sink_statistics
1168            .iter_mut()
1169            .filter_map(|(_, (_, s))| {
1170                if s.iter().all(|s| s.is_some()) {
1171                    let ret = Some(SinkStatisticsUpdate::summarize(|| {
1172                        s.iter().filter_map(Option::as_ref)
1173                    }));
1174
1175                    // Reset the counters so we only report diffs.
1176                    s.iter_mut().for_each(|s| {
1177                        if let Some(s) = s {
1178                            s.reset_counters();
1179                        }
1180                    });
1181                    ret
1182                } else {
1183                    None
1184                }
1185            })
1186            .collect();
1187
1188        (sources, sinks)
1189    }
1190}