mz_persist_client/internal/

metrics.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Prometheus monitoring metrics.

use async_stream::stream;
use mz_persist_types::stats::PartStatsMetrics;
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::time::{Duration, Instant};
use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore};

use async_trait::async_trait;
use bytes::Bytes;
use futures_util::StreamExt;
use mz_dyncfg::ConfigSet;
use mz_ore::bytes::SegmentedBytes;
use mz_ore::cast::{CastFrom, CastLossy};
use mz_ore::instrument;
use mz_ore::metric;
use mz_ore::metrics::aggregation::{AggregatedFamilies, AggregatedFamily, AggregationKey};
use mz_ore::metrics::{
    ComputedGauge, ComputedUIntGauge, Counter, DeleteOnDropCounter, DeleteOnDropGauge, IntCounter,
    MakeCollector, MetricsRegistry, PrometheusOpts, UIntGauge, UIntGaugeVec, raw,
};
use mz_ore::stats::histogram_seconds_buckets;
use mz_persist::location::{
    Blob, BlobMetadata, CaSResult, Consensus, ExternalError, ResultStream, SeqNo, VersionedData,
};
use mz_persist::metrics::{BlobHedgeMetrics, ColumnarMetrics, S3BlobMetrics};
use mz_persist::retry::RetryStream;
use mz_persist_types::Codec64;
use mz_postgres_client::metrics::PostgresClientMetrics;
use prometheus::core::{AtomicI64, AtomicU64, Collector, Desc, GenericGauge};
use prometheus::proto::MetricFamily;
use prometheus::{CounterVec, Gauge, GaugeVec, Histogram, HistogramVec, IntCounterVec};
use regex::Regex;
use timely::progress::Antichain;
use tokio_metrics::TaskMonitor;
use tracing::{Instrument, debug, info, info_span, warn};

use crate::cfg::{
    PER_SHARD_METRICS_ENABLE_REGEX, SHARD_INFO_METRIC, SHARD_METRICS, ShardMetricsExport,
};
use crate::fetch::{FETCH_SEMAPHORE_COST_ADJUSTMENT, FETCH_SEMAPHORE_PERMIT_ADJUSTMENT};
use crate::internal::paths::BlobKey;
use crate::{PersistConfig, ShardId};

/// Prometheus monitoring metrics.
///
/// Intentionally not Clone because we expect this to be passed around in an
/// Arc.
pub struct Metrics {
    _vecs: MetricsVecs,
    _uptime: ComputedGauge,

    /// Metrics for [Blob] usage.
    pub blob: BlobMetrics,
    /// Metrics for [Consensus] usage.
    pub consensus: ConsensusMetrics,
    /// Metrics of command evaluation.
    pub cmds: CmdsMetrics,
    /// Metrics for each retry loop.
    pub retries: RetriesMetrics,
    /// Metrics for batches written directly on behalf of a user (BatchBuilder
    /// or one of the sugar methods that use it).
    pub user: BatchWriteMetrics,
    /// Metrics for reading batch parts
    pub read: BatchPartReadMetrics,
    /// Metrics for compaction.
    pub compaction: CompactionMetrics,
    /// Metrics for garbage collection.
    pub gc: GcMetrics,
    /// Metrics for leasing and automatic lease expiry.
    pub lease: LeaseMetrics,
    /// Metrics for various encodings and decodings.
    pub codecs: CodecsMetrics,
    /// Metrics for (incremental) state updates and fetches.
    pub state: StateMetrics,
    /// Metrics for various per-shard measurements.
    pub shards: ShardsMetrics,
    /// Metrics for auditing persist usage
    pub audit: UsageAuditMetrics,
    /// Metrics for locking.
    pub locks: LocksMetrics,
    /// Metrics for StateWatch.
    pub watch: WatchMetrics,
    /// Metrics for PubSub client.
    pub pubsub_client: PubSubClientMetrics,
    /// Metrics for mfp/filter pushdown.
    pub pushdown: PushdownMetrics,
    /// Metrics for consolidation.
    pub consolidation: ConsolidationMetrics,
    /// Metrics for blob caching.
    pub blob_cache_mem: BlobMemCache,
    /// Metrics for tokio tasks.
    pub tasks: TasksMetrics,
    /// Metrics for columnar data encoding and decoding.
    pub columnar: ColumnarMetrics,
    /// Metrics for schemas and the schema registry.
    pub schema: SchemaMetrics,
    /// Metrics for inline writes.
    pub inline: InlineMetrics,
    /// Semaphore to limit memory/disk use by fetches.
    pub(crate) semaphore: SemaphoreMetrics,

    /// Metrics for the persist sink.
    pub sink: SinkMetrics,
    /// Metrics for the persist_source backpressure operator.
    pub backpressure: BackpressureMetrics,

    /// Metrics for S3-backed blob implementation
    pub s3_blob: S3BlobMetrics,
    /// Metrics for hedged blob gets
    pub blob_hedge: BlobHedgeMetrics,
    /// Metrics for Postgres-backed consensus implementation
    pub postgres_consensus: PostgresClientMetrics,

    #[allow(dead_code)]
    pub(crate) registry: MetricsRegistry,
}

impl std::fmt::Debug for Metrics {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Metrics").finish_non_exhaustive()
    }
}

impl Metrics {
    /// Returns a new [Metrics] instance connected to the given registry.
    pub fn new(cfg: &PersistConfig, registry: &MetricsRegistry) -> Self {
        let vecs = MetricsVecs::new(registry);
        let start = Instant::now();
        let uptime = registry.register_computed_gauge(
            metric!(
                name: "mz_persist_metadata_seconds",
                help: "server uptime, labels are build metadata",
                const_labels: {
                    "version" => cfg.build_version,
                    "build_type" => if cfg!(release) { "release" } else { "debug" }
                },
            ),
            move || start.elapsed().as_secs_f64(),
        );
        let s3_blob = S3BlobMetrics::new(registry);
        let columnar = ColumnarMetrics::new(registry);
        let mut per_shard_filter = PerShardMetricsFilter {
            configs: Arc::clone(&cfg.configs),
            mode: (String::new(), ShardMetricsExport::default()),
            invalid_mode: registry.register(metric!(
                name: "mz_persist_shard_metrics_mode_invalid",
                help: "count of metrics scrapes that found persist_shard_metrics unrecognized and so exported the default",
            )),
            allow: (String::new(), None),
            invalid_regex: registry.register(metric!(
                name: "mz_persist_per_shard_metrics_regex_invalid",
                help: "count of metrics scrapes that found persist_per_shard_metrics_enable_regex invalid and so kept no per-shard series",
            )),
        };
        registry.register_postprocessor(move |families| per_shard_filter.filter(families));
        Metrics {
            blob: vecs.blob_metrics(),
            consensus: vecs.consensus_metrics(),
            cmds: vecs.cmds_metrics(registry),
            retries: vecs.retries_metrics(),
            codecs: vecs.codecs_metrics(),
            user: BatchWriteMetrics::new(registry, "user"),
            read: vecs.batch_part_read_metrics(),
            compaction: CompactionMetrics::new(registry),
            gc: GcMetrics::new(registry),
            lease: LeaseMetrics::new(registry),
            state: StateMetrics::new(registry),
            shards: ShardsMetrics::new(registry, Arc::clone(&cfg.configs)),
            audit: UsageAuditMetrics::new(registry),
            locks: vecs.locks_metrics(),
            watch: WatchMetrics::new(registry),
            pubsub_client: PubSubClientMetrics::new(registry),
            pushdown: PushdownMetrics::new(registry),
            consolidation: ConsolidationMetrics::new(registry),
            blob_cache_mem: BlobMemCache::new(registry),
            tasks: TasksMetrics::new(registry),
            columnar,
            schema: SchemaMetrics::new(registry),
            inline: InlineMetrics::new(registry),
            semaphore: SemaphoreMetrics::new(cfg.clone(), registry.clone()),
            sink: SinkMetrics::new(registry),
            backpressure: BackpressureMetrics::new(registry),
            s3_blob,
            blob_hedge: BlobHedgeMetrics::new(registry),
            postgres_consensus: PostgresClientMetrics::new(registry, "mz_persist"),
            _vecs: vecs,
            _uptime: uptime,
            registry: registry.clone(),
        }
    }

    /// Returns the current lifetime write amplification reflected in these
    /// metrics.
    ///
    /// Only exposed for tests, persistcli, and benchmarks.
    pub fn write_amplification(&self) -> f64 {
        // This intentionally uses "bytes" for total and "goodbytes" for user so
        // that the overhead of our blob format is included.
        let total_written = self.blob.set.bytes.get();
        let user_written = self.user.goodbytes.get();
        #[allow(clippy::as_conversions)]
        {
            total_written as f64 / user_written as f64
        }
    }
}

#[derive(Debug)]
struct MetricsVecs {
    cmd_started: IntCounterVec,
    cmd_cas_mismatch: IntCounterVec,
    cmd_succeeded: IntCounterVec,
    cmd_failed: IntCounterVec,
    cmd_seconds: CounterVec,

    external_op_started: IntCounterVec,
    external_op_succeeded: IntCounterVec,
    external_op_failed: IntCounterVec,
    external_op_bytes: IntCounterVec,
    external_op_seconds: CounterVec,
    external_consensus_truncated_count: IntCounter,
    external_blob_delete_noop_count: IntCounter,
    external_blob_sizes: Histogram,
    external_rtt_latency: GaugeVec,
    external_op_latency: HistogramVec,

    retry_started: IntCounterVec,
    retry_finished: IntCounterVec,
    retry_retries: IntCounterVec,
    retry_sleep_seconds: CounterVec,

    encode_count: IntCounterVec,
    encode_seconds: CounterVec,
    decode_count: IntCounterVec,
    decode_seconds: CounterVec,

    read_part_bytes: IntCounterVec,
    read_part_goodbytes: IntCounterVec,
    read_part_count: IntCounterVec,
    read_part_seconds: CounterVec,
    read_ts_rewrite: IntCounterVec,

    lock_acquire_count: IntCounterVec,
    lock_blocking_acquire_count: IntCounterVec,
    lock_blocking_seconds: CounterVec,

    /// A minimal set of metrics imported into honeycomb for alerting.
    alerts_metrics: Arc<AlertsMetrics>,
}

impl MetricsVecs {
    fn new(registry: &MetricsRegistry) -> Self {
        MetricsVecs {
            cmd_started: registry.register(metric!(
                name: "mz_persist_cmd_started_count",
                help: "count of commands started",
                var_labels: ["cmd"],
            )),
            cmd_cas_mismatch: registry.register(metric!(
                name: "mz_persist_cmd_cas_mismatch_count",
                help: "count of command retries from CaS mismatch",
                var_labels: ["cmd"],
            )),
            cmd_succeeded: registry.register(metric!(
                name: "mz_persist_cmd_succeeded_count",
                help: "count of commands succeeded",
                var_labels: ["cmd"],
            )),
            cmd_failed: registry.register(metric!(
                name: "mz_persist_cmd_failed_count",
                help: "count of commands failed",
                var_labels: ["cmd"],
            )),
            cmd_seconds: registry.register(metric!(
                name: "mz_persist_cmd_seconds",
                help: "time spent applying commands",
                var_labels: ["cmd"],
            )),

            external_op_started: registry.register(metric!(
                name: "mz_persist_external_started_count",
                help: "count of external service calls started",
                var_labels: ["op"],
            )),
            external_op_succeeded: registry.register(metric!(
                name: "mz_persist_external_succeeded_count",
                help: "count of external service calls succeeded",
                var_labels: ["op"],
            )),
            external_op_failed: registry.register(metric!(
                name: "mz_persist_external_failed_count",
                help: "count of external service calls failed",
                var_labels: ["op"],
            )),
            external_op_bytes: registry.register(metric!(
                name: "mz_persist_external_bytes_count",
                help: "total size represented by external service calls",
                var_labels: ["op"],
            )),
            external_op_seconds: registry.register(metric!(
                name: "mz_persist_external_seconds",
                help: "time spent in external service calls",
                var_labels: ["op"],
            )),
            external_consensus_truncated_count: registry.register(metric!(
                name: "mz_persist_external_consensus_truncated_count",
                help: "count of versions deleted by consensus truncate calls",
            )),
            external_blob_delete_noop_count: registry.register(metric!(
                name: "mz_persist_external_blob_delete_noop_count",
                help: "count of blob delete calls that deleted a non-existent key",
            )),
            external_blob_sizes: registry.register(metric!(
                name: "mz_persist_external_blob_sizes",
                help: "histogram of blob sizes at put time",
                buckets: mz_ore::stats::HISTOGRAM_BYTE_BUCKETS.to_vec(),
            )),
            external_rtt_latency: registry.register(metric!(
                name: "mz_persist_external_rtt_latency",
                help: "roundtrip-time to external service as seen by this process",
                var_labels: ["external"],
            )),
            external_op_latency: registry.register(metric!(
                name: "mz_persist_external_op_latency",
                help: "rountrip latency observed by individual performance-critical operations",
                var_labels: ["op"],
                // NB: If we end up overrunning metrics quotas, we could plausibly cut this
                // down by switching to a factor of 4 between buckets (vs. the standard 2).
                buckets: histogram_seconds_buckets(0.000_500, 32.0),
            )),

            retry_started: registry.register(metric!(
                name: "mz_persist_retry_started_count",
                help: "count of retry loops started",
                var_labels: ["op"],
            )),
            retry_finished: registry.register(metric!(
                name: "mz_persist_retry_finished_count",
                help: "count of retry loops finished",
                var_labels: ["op"],
            )),
            retry_retries: registry.register(metric!(
                name: "mz_persist_retry_retries_count",
                help: "count of total attempts by retry loops",
                var_labels: ["op"],
            )),
            retry_sleep_seconds: registry.register(metric!(
                name: "mz_persist_retry_sleep_seconds",
                help: "time spent in retry loop backoff",
                var_labels: ["op"],
            )),

            encode_count: registry.register(metric!(
                name: "mz_persist_encode_count",
                help: "count of op encodes",
                var_labels: ["op"],
            )),
            encode_seconds: registry.register(metric!(
                name: "mz_persist_encode_seconds",
                help: "time spent in op encodes",
                var_labels: ["op"],
            )),
            decode_count: registry.register(metric!(
                name: "mz_persist_decode_count",
                help: "count of op decodes",
                var_labels: ["op"],
            )),
            decode_seconds: registry.register(metric!(
                name: "mz_persist_decode_seconds",
                help: "time spent in op decodes",
                var_labels: ["op"],
            )),

            read_part_bytes: registry.register(metric!(
                name: "mz_persist_read_batch_part_bytes",
                help: "total encoded size of batch parts read",
                var_labels: ["op"],
            )),
            read_part_goodbytes: registry.register(metric!(
                name: "mz_persist_read_batch_part_goodbytes",
                help: "total logical size of batch parts read",
                var_labels: ["op"],
            )),
            read_part_count: registry.register(metric!(
                name: "mz_persist_read_batch_part_count",
                help: "count of batch parts read",
                var_labels: ["op"],
            )),
            read_part_seconds: registry.register(metric!(
                name: "mz_persist_read_batch_part_seconds",
                help: "time spent reading batch parts",
                var_labels: ["op"],
            )),
            read_ts_rewrite: registry.register(metric!(
                name: "mz_persist_read_ts_rewite",
                help: "count of updates read with rewritten ts",
                var_labels: ["op"],
            )),

            lock_acquire_count: registry.register(metric!(
                name: "mz_persist_lock_acquire_count",
                help: "count of locks acquired",
                var_labels: ["op"],
            )),
            lock_blocking_acquire_count: registry.register(metric!(
                name: "mz_persist_lock_blocking_acquire_count",
                help: "count of locks acquired that required blocking",
                var_labels: ["op"],
            )),
            lock_blocking_seconds: registry.register(metric!(
                name: "mz_persist_lock_blocking_seconds",
                help: "time spent blocked for a lock",
                var_labels: ["op"],
            )),

            alerts_metrics: Arc::new(AlertsMetrics::new(registry)),
        }
    }

    fn cmds_metrics(&self, registry: &MetricsRegistry) -> CmdsMetrics {
        CmdsMetrics {
            init_state: self.cmd_metrics("init_state"),
            add_rollup: self.cmd_metrics("add_rollup"),
            remove_rollups: self.cmd_metrics("remove_rollups"),
            upgrade_version: self.cmd_metrics("upgrade_version"),
            register: self.cmd_metrics("register"),
            compare_and_append: self.cmd_metrics("compare_and_append"),
            compare_and_append_noop:             registry.register(metric!(
                name: "mz_persist_cmd_compare_and_append_noop",
                help: "count of compare_and_append retries that were discoverd to have already committed",
            )),
            compare_and_downgrade_since: self.cmd_metrics("compare_and_downgrade_since"),
            downgrade_since: self.cmd_metrics("downgrade_since"),
            expire_reader: self.cmd_metrics("expire_reader"),
            expire_writer: self.cmd_metrics("expire_writer"),
            merge_res: self.cmd_metrics("merge_res"),
            become_tombstone: self.cmd_metrics("become_tombstone"),
            compare_and_evolve_schema: self.cmd_metrics("compare_and_evolve_schema"),
            spine_exert: self.cmd_metrics("spine_exert"),
            fetch_upper_count: registry.register(metric!(
                name: "mz_persist_cmd_fetch_upper_count",
                help: "count of fetch_upper calls",
            ))
        }
    }

    fn cmd_metrics(&self, cmd: &str) -> CmdMetrics {
        CmdMetrics {
            name: cmd.to_owned(),
            started: self.cmd_started.with_label_values(&[cmd]),
            succeeded: self.cmd_succeeded.with_label_values(&[cmd]),
            cas_mismatch: self.cmd_cas_mismatch.with_label_values(&[cmd]),
            failed: self.cmd_failed.with_label_values(&[cmd]),
            seconds: self.cmd_seconds.with_label_values(&[cmd]),
        }
    }

    fn retries_metrics(&self) -> RetriesMetrics {
        RetriesMetrics {
            determinate: RetryDeterminate {
                apply_unbatched_cmd_cas: self.retry_metrics("apply_unbatched_cmd::cas"),
            },
            external: RetryExternal {
                batch_delete: Arc::new(self.retry_metrics("batch::delete")),
                batch_set: self.retry_metrics("batch::set"),
                blob_open: self.retry_metrics("blob::open"),
                compaction_noop_delete: Arc::new(self.retry_metrics("compaction_noop::delete")),
                consensus_open: self.retry_metrics("consensus::open"),
                fetch_batch_get: self.retry_metrics("fetch_batch::get"),
                fetch_state_scan: self.retry_metrics("fetch_state::scan"),
                gc_truncate: self.retry_metrics("gc::truncate"),
                maybe_init_cas: self.retry_metrics("maybe_init::cas"),
                rollup_delete: self.retry_metrics("rollup::delete"),
                rollup_get: self.retry_metrics("rollup::get"),
                rollup_set: self.retry_metrics("rollup::set"),
                hollow_run_get: self.retry_metrics("hollow_run::get"),
                hollow_run_set: self.retry_metrics("hollow_run::set"),
                storage_usage_shard_size: self.retry_metrics("storage_usage::shard_size"),
            },
            compare_and_append_idempotent: self.retry_metrics("compare_and_append_idempotent"),
            fetch_latest_state: self.retry_metrics("fetch_latest_state"),
            fetch_live_states: self.retry_metrics("fetch_live_states"),
            idempotent_cmd: self.retry_metrics("idempotent_cmd"),
            next_listen_batch: self.retry_metrics("next_listen_batch"),
            snapshot: self.retry_metrics("snapshot"),
        }
    }

    fn retry_metrics(&self, name: &str) -> RetryMetrics {
        RetryMetrics {
            name: name.to_owned(),
            started: self.retry_started.with_label_values(&[name]),
            finished: self.retry_finished.with_label_values(&[name]),
            retries: self.retry_retries.with_label_values(&[name]),
            sleep_seconds: self.retry_sleep_seconds.with_label_values(&[name]),
        }
    }

    fn codecs_metrics(&self) -> CodecsMetrics {
        CodecsMetrics {
            state: self.codec_metrics("state"),
            state_diff: self.codec_metrics("state_diff"),
            batch: self.codec_metrics("batch"),
            key: self.codec_metrics("key"),
            val: self.codec_metrics("val"),
        }
    }

    fn codec_metrics(&self, op: &str) -> CodecMetrics {
        CodecMetrics {
            encode_count: self.encode_count.with_label_values(&[op]),
            encode_seconds: self.encode_seconds.with_label_values(&[op]),
            decode_count: self.decode_count.with_label_values(&[op]),
            decode_seconds: self.decode_seconds.with_label_values(&[op]),
        }
    }

    fn blob_metrics(&self) -> BlobMetrics {
        BlobMetrics {
            set: self.external_op_metrics("blob_set", true),
            get: self.external_op_metrics("blob_get", true),
            list_keys: self.external_op_metrics("blob_list_keys", false),
            delete: self.external_op_metrics("blob_delete", false),
            restore: self.external_op_metrics("restore", false),
            delete_noop: self.external_blob_delete_noop_count.clone(),
            blob_sizes: self.external_blob_sizes.clone(),
            rtt_latency: self.external_rtt_latency.with_label_values(&["blob"]),
        }
    }

    fn consensus_metrics(&self) -> ConsensusMetrics {
        ConsensusMetrics {
            list_keys: self.external_op_metrics("consensus_list_keys", false),
            head: self.external_op_metrics("consensus_head", false),
            compare_and_set: self.external_op_metrics("consensus_cas", true),
            scan: self.external_op_metrics("consensus_scan", false),
            truncate: self.external_op_metrics("consensus_truncate", false),
            truncated_count: self.external_consensus_truncated_count.clone(),
            rtt_latency: self.external_rtt_latency.with_label_values(&["consensus"]),
        }
    }

    fn external_op_metrics(&self, op: &str, latency_histogram: bool) -> ExternalOpMetrics {
        ExternalOpMetrics {
            started: self.external_op_started.with_label_values(&[op]),
            succeeded: self.external_op_succeeded.with_label_values(&[op]),
            failed: self.external_op_failed.with_label_values(&[op]),
            bytes: self.external_op_bytes.with_label_values(&[op]),
            seconds: self.external_op_seconds.with_label_values(&[op]),
            seconds_histogram: if latency_histogram {
                Some(self.external_op_latency.with_label_values(&[op]))
            } else {
                None
            },
            alerts_metrics: Arc::clone(&self.alerts_metrics),
        }
    }

    fn batch_part_read_metrics(&self) -> BatchPartReadMetrics {
        BatchPartReadMetrics {
            listen: self.read_metrics("listen"),
            snapshot: self.read_metrics("snapshot"),
            batch_fetcher: self.read_metrics("batch_fetcher"),
            compaction: self.read_metrics("compaction"),
            unindexed: self.read_metrics("unindexed"),
        }
    }

    fn read_metrics(&self, op: &str) -> ReadMetrics {
        ReadMetrics {
            part_bytes: self.read_part_bytes.with_label_values(&[op]),
            part_goodbytes: self.read_part_goodbytes.with_label_values(&[op]),
            part_count: self.read_part_count.with_label_values(&[op]),
            seconds: self.read_part_seconds.with_label_values(&[op]),
            ts_rewrite: self.read_ts_rewrite.with_label_values(&[op]),
        }
    }

    fn locks_metrics(&self) -> LocksMetrics {
        LocksMetrics {
            applier_read_cacheable: self.lock_metrics("applier_read_cacheable"),
            applier_read_noncacheable: self.lock_metrics("applier_read_noncacheable"),
            applier_write: self.lock_metrics("applier_write"),
            watch: self.lock_metrics("watch"),
        }
    }

    fn lock_metrics(&self, op: &str) -> LockMetrics {
        LockMetrics {
            acquire_count: self.lock_acquire_count.with_label_values(&[op]),
            blocking_acquire_count: self.lock_blocking_acquire_count.with_label_values(&[op]),
            blocking_seconds: self.lock_blocking_seconds.with_label_values(&[op]),
        }
    }
}

#[derive(Debug)]
pub struct CmdMetrics {
    pub(crate) name: String,
    pub(crate) started: IntCounter,
    pub(crate) cas_mismatch: IntCounter,
    pub(crate) succeeded: IntCounter,
    pub(crate) failed: IntCounter,
    pub(crate) seconds: Counter,
}

impl CmdMetrics {
    pub async fn run_cmd<R, E, F, CmdFn>(
        &self,
        shard_metrics: &ShardMetrics,
        cmd_fn: CmdFn,
    ) -> Result<R, E>
    where
        F: std::future::Future<Output = Result<R, E>>,
        CmdFn: FnOnce() -> F,
    {
        self.started.inc();
        let start = Instant::now();
        let res = cmd_fn().await;
        self.seconds.inc_by(start.elapsed().as_secs_f64());
        match res.as_ref() {
            Ok(_) => {
                self.succeeded.inc();
                shard_metrics.cmd_succeeded.inc();
            }
            Err(_) => self.failed.inc(),
        };
        res
    }
}

#[derive(Debug)]
pub struct CmdsMetrics {
    pub(crate) init_state: CmdMetrics,
    pub(crate) add_rollup: CmdMetrics,
    pub(crate) remove_rollups: CmdMetrics,
    pub(crate) upgrade_version: CmdMetrics,
    pub(crate) register: CmdMetrics,
    pub(crate) compare_and_append: CmdMetrics,
    pub(crate) compare_and_append_noop: IntCounter,
    pub(crate) compare_and_downgrade_since: CmdMetrics,
    pub(crate) downgrade_since: CmdMetrics,
    pub(crate) expire_reader: CmdMetrics,
    pub(crate) expire_writer: CmdMetrics,
    pub(crate) merge_res: CmdMetrics,
    pub(crate) become_tombstone: CmdMetrics,
    pub(crate) compare_and_evolve_schema: CmdMetrics,
    pub(crate) spine_exert: CmdMetrics,
    pub(crate) fetch_upper_count: IntCounter,
}

#[derive(Debug)]
pub struct RetryMetrics {
    pub(crate) name: String,
    pub(crate) started: IntCounter,
    pub(crate) finished: IntCounter,
    pub(crate) retries: IntCounter,
    pub(crate) sleep_seconds: Counter,
}

impl RetryMetrics {
    pub(crate) fn stream(&self, retry: RetryStream) -> MetricsRetryStream {
        MetricsRetryStream::new(retry, self)
    }
}

#[derive(Debug)]
pub struct RetryDeterminate {
    pub(crate) apply_unbatched_cmd_cas: RetryMetrics,
}

#[derive(Debug)]
pub struct RetryExternal {
    pub(crate) batch_delete: Arc<RetryMetrics>,
    pub(crate) batch_set: RetryMetrics,
    pub(crate) blob_open: RetryMetrics,
    pub(crate) compaction_noop_delete: Arc<RetryMetrics>,
    pub(crate) consensus_open: RetryMetrics,
    pub(crate) fetch_batch_get: RetryMetrics,
    pub(crate) fetch_state_scan: RetryMetrics,
    pub(crate) gc_truncate: RetryMetrics,
    pub(crate) maybe_init_cas: RetryMetrics,
    pub(crate) rollup_delete: RetryMetrics,
    pub(crate) rollup_get: RetryMetrics,
    pub(crate) rollup_set: RetryMetrics,
    pub(crate) hollow_run_get: RetryMetrics,
    pub(crate) hollow_run_set: RetryMetrics,
    pub(crate) storage_usage_shard_size: RetryMetrics,
}

#[derive(Debug)]
pub struct RetriesMetrics {
    pub(crate) determinate: RetryDeterminate,
    pub(crate) external: RetryExternal,

    pub(crate) compare_and_append_idempotent: RetryMetrics,
    pub(crate) fetch_latest_state: RetryMetrics,
    pub(crate) fetch_live_states: RetryMetrics,
    pub(crate) idempotent_cmd: RetryMetrics,
    pub(crate) next_listen_batch: RetryMetrics,
    pub(crate) snapshot: RetryMetrics,
}

#[derive(Debug)]
pub struct BatchPartReadMetrics {
    pub(crate) listen: ReadMetrics,
    pub(crate) snapshot: ReadMetrics,
    pub(crate) batch_fetcher: ReadMetrics,
    pub(crate) compaction: ReadMetrics,
    pub(crate) unindexed: ReadMetrics,
}

#[derive(Debug, Clone)]
pub struct ReadMetrics {
    pub(crate) part_bytes: IntCounter,
    pub(crate) part_goodbytes: IntCounter,
    pub(crate) part_count: IntCounter,
    pub(crate) seconds: Counter,
    pub(crate) ts_rewrite: IntCounter,
}

// This one is Clone in contrast to the others because it has to get moved into
// a task.
#[derive(Debug, Clone)]
pub struct BatchWriteMetrics {
    pub(crate) bytes: IntCounter,
    pub(crate) goodbytes: IntCounter,
    pub(crate) seconds: Counter,
    pub(crate) write_stalls: IntCounter,
    pub(crate) key_lower_too_big: IntCounter,

    pub(crate) unordered: IntCounter,
    pub(crate) codec_order: IntCounter,
    pub(crate) structured_order: IntCounter,
    _order_counts: IntCounterVec,

    pub(crate) step_stats: Counter,
    pub(crate) step_part_writing: Counter,
    pub(crate) step_inline: Counter,
}

impl BatchWriteMetrics {
    fn new(registry: &MetricsRegistry, name: &str) -> Self {
        let order_counts: IntCounterVec = registry.register(metric!(
                name: format!("mz_persist_{}_write_batch_order", name),
                help: "count of batches by the data ordering",
                var_labels: ["order"],
        ));
        let unordered = order_counts.with_label_values(&["unordered"]);
        let codec_order = order_counts.with_label_values(&["codec"]);
        let structured_order = order_counts.with_label_values(&["structured"]);

        BatchWriteMetrics {
            bytes: registry.register(metric!(
                name: format!("mz_persist_{}_bytes", name),
                help: format!("total encoded size of {} batches written", name),
            )),
            goodbytes: registry.register(metric!(
                name: format!("mz_persist_{}_goodbytes", name),
                help: format!("total logical size of {} batches written", name),
            )),
            seconds: registry.register(metric!(
                name: format!("mz_persist_{}_write_batch_part_seconds", name),
                help: format!("time spent writing {} batches", name),
            )),
            write_stalls: registry.register(metric!(
                name: format!("mz_persist_{}_write_stall_count", name),
                help: format!(
                    "count of {} writes stalling to await max outstanding reqs",
                    name
                ),
            )),
            key_lower_too_big: registry.register(metric!(
                name: format!("mz_persist_{}_key_lower_too_big", name),
                help: format!(
                    "count of {} writes that were unable to write a key lower, because the size threshold was too low",
                    name
                ),
            )),
            unordered,
            codec_order,
            structured_order,
            _order_counts: order_counts,
            step_stats: registry.register(metric!(
                name: format!("mz_persist_{}_step_stats", name),
                help: format!("time spent computing {} update stats", name),
            )),
            step_part_writing: registry.register(metric!(
                name: format!("mz_persist_{}_step_part_writing", name),
                help: format!("blocking time spent writing parts for {} updates", name),
            )),
            step_inline: registry.register(metric!(
                name: format!("mz_persist_{}_step_inline", name),
                help: format!("time spent encoding {} inline batches", name)
            )),
        }
    }
}

#[derive(Debug)]
pub struct CompactionMetrics {
    pub(crate) requested: IntCounter,
    pub(crate) dropped: IntCounter,
    pub(crate) disabled: IntCounter,
    pub(crate) skipped: IntCounter,
    pub(crate) started: IntCounter,
    pub(crate) applied: IntCounter,
    pub(crate) timed_out: IntCounter,
    pub(crate) failed: IntCounter,
    pub(crate) noop: IntCounter,
    pub(crate) seconds: Counter,
    pub(crate) concurrency_waits: IntCounter,
    pub(crate) queued_seconds: Counter,
    pub(crate) memory_violations: IntCounter,
    pub(crate) runs_compacted: IntCounter,
    pub(crate) chunks_compacted: IntCounter,
    pub(crate) not_all_prefetched: IntCounter,
    pub(crate) parts_prefetched: IntCounter,
    pub(crate) parts_waited: IntCounter,
    pub(crate) fast_path_eligible: IntCounter,
    pub(crate) admin_count: IntCounter,

    pub(crate) applied_exact_match: IntCounter,
    pub(crate) applied_subset_match: IntCounter,
    pub(crate) not_applied_too_many_updates: IntCounter,

    pub(crate) batch: BatchWriteMetrics,
    pub(crate) steps: CompactionStepTimings,
    pub(crate) schema_selection: CompactionSchemaSelection,

    pub(crate) _steps_vec: CounterVec,
}

impl CompactionMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        let step_timings: CounterVec = registry.register(metric!(
                name: "mz_persist_compaction_step_seconds",
                help: "time spent on individual steps of compaction",
                var_labels: ["step"],
        ));
        let schema_selection: CounterVec = registry.register(metric!(
            name: "mz_persist_compaction_schema_selection",
            help: "count of compactions and how we did schema selection",
            var_labels: ["selection"],
        ));

        CompactionMetrics {
            requested: registry.register(metric!(
                name: "mz_persist_compaction_requested",
                help: "count of total compaction requests",
            )),
            dropped: registry.register(metric!(
                name: "mz_persist_compaction_dropped",
                help: "count of total compaction requests dropped due to a full queue",
            )),
            disabled: registry.register(metric!(
                name: "mz_persist_compaction_disabled",
                help: "count of total compaction requests dropped because compaction was disabled",
            )),
            skipped: registry.register(metric!(
                name: "mz_persist_compaction_skipped",
                help: "count of compactions skipped due to heuristics",
            )),
            started: registry.register(metric!(
                name: "mz_persist_compaction_started",
                help: "count of compactions started",
            )),
            failed: registry.register(metric!(
                name: "mz_persist_compaction_failed",
                help: "count of compactions failed",
            )),
            applied: registry.register(metric!(
                name: "mz_persist_compaction_applied",
                help: "count of compactions applied to state",
            )),
            timed_out: registry.register(metric!(
                name: "mz_persist_compaction_timed_out",
                help: "count of compactions that timed out",
            )),
            noop: registry.register(metric!(
                name: "mz_persist_compaction_noop",
                help: "count of compactions discarded (obsolete)",
            )),
            seconds: registry.register(metric!(
                name: "mz_persist_compaction_seconds",
                help: "time spent in compaction",
            )),
            concurrency_waits: registry.register(metric!(
                name: "mz_persist_compaction_concurrency_waits",
                help: "count of compaction requests that ever blocked due to concurrency limit",
            )),
            queued_seconds: registry.register(metric!(
                name: "mz_persist_compaction_queued_seconds",
                help: "time that compaction requests spent queued",
            )),
            memory_violations: registry.register(metric!(
                name: "mz_persist_compaction_memory_violations",
                help: "count of compaction memory requirement violations",
            )),
            runs_compacted: registry.register(metric!(
                name: "mz_persist_compaction_runs_compacted",
                help: "count of runs compacted",
            )),
            chunks_compacted: registry.register(metric!(
                name: "mz_persist_compaction_chunks_compacted",
                help: "count of run chunks compacted",
            )),
            not_all_prefetched: registry.register(metric!(
                name: "mz_persist_compaction_not_all_prefetched",
                help: "count of compactions where not all inputs were prefetched",
            )),
            parts_prefetched: registry.register(metric!(
                name: "mz_persist_compaction_parts_prefetched",
                help: "count of compaction parts completely prefetched by the time they're needed",
            )),
            parts_waited: registry.register(metric!(
                name: "mz_persist_compaction_parts_waited",
                help: "count of compaction parts that had to be waited on",
            )),
            fast_path_eligible: registry.register(metric!(
                name: "mz_persist_compaction_fast_path_eligible",
                help: "count of compaction requests that could have used the fast-path optimization",
            )),
            admin_count: registry.register(metric!(
                name: "mz_persist_compaction_admin_count",
                help: "count of compaction requests that were performed by admin tooling",
            )),
            applied_exact_match: registry.register(metric!(
                name: "mz_persist_compaction_applied_exact_match",
                help: "count of merge results that exactly replaced a SpineBatch",
            )),
            applied_subset_match: registry.register(metric!(
                name: "mz_persist_compaction_applied_subset_match",
                help: "count of merge results that replaced a subset of a SpineBatch",
            )),
            not_applied_too_many_updates: registry.register(metric!(
                name: "mz_persist_compaction_not_applied_too_many_updates",
                help: "count of merge results that did not apply due to too many updates",
            )),
            batch: BatchWriteMetrics::new(registry, "compaction"),
            steps: CompactionStepTimings::new(step_timings.clone()),
            schema_selection: CompactionSchemaSelection::new(schema_selection.clone()),
            _steps_vec: step_timings,
        }
    }
}

#[derive(Debug)]
pub struct CompactionStepTimings {
    pub(crate) part_fetch_seconds: Counter,
    pub(crate) heap_population_seconds: Counter,
}

impl CompactionStepTimings {
    fn new(step_timings: CounterVec) -> CompactionStepTimings {
        CompactionStepTimings {
            part_fetch_seconds: step_timings.with_label_values(&["part_fetch"]),
            heap_population_seconds: step_timings.with_label_values(&["heap_population"]),
        }
    }
}

#[derive(Debug)]
pub struct CompactionSchemaSelection {
    pub(crate) recent_schema: Counter,
    pub(crate) no_schema: Counter,
}

impl CompactionSchemaSelection {
    fn new(schema_selection: CounterVec) -> CompactionSchemaSelection {
        CompactionSchemaSelection {
            recent_schema: schema_selection.with_label_values(&["recent"]),
            no_schema: schema_selection.with_label_values(&["none"]),
        }
    }
}

#[derive(Debug)]
pub struct GcMetrics {
    pub(crate) noop: IntCounter,
    pub(crate) started: IntCounter,
    pub(crate) finished: IntCounter,
    pub(crate) merged: IntCounter,
    pub(crate) seconds: Counter,
    pub(crate) steps: GcStepTimings,
}

#[derive(Debug)]
pub struct GcStepTimings {
    pub(crate) find_removable_rollups: Counter,
    pub(crate) fetch_seconds: Counter,
    pub(crate) find_deletable_blobs_seconds: Counter,
    pub(crate) delete_rollup_seconds: Counter,
    pub(crate) delete_batch_part_seconds: Counter,
    pub(crate) truncate_diff_seconds: Counter,
    pub(crate) remove_rollups_from_state: Counter,
    pub(crate) post_gc_calculations_seconds: Counter,
}

impl GcStepTimings {
    fn new(step_timings: CounterVec) -> Self {
        Self {
            find_removable_rollups: step_timings.with_label_values(&["find_removable_rollups"]),
            fetch_seconds: step_timings.with_label_values(&["fetch"]),
            find_deletable_blobs_seconds: step_timings.with_label_values(&["find_deletable_blobs"]),
            delete_rollup_seconds: step_timings.with_label_values(&["delete_rollup"]),
            delete_batch_part_seconds: step_timings.with_label_values(&["delete_batch_part"]),
            truncate_diff_seconds: step_timings.with_label_values(&["truncate_diff"]),
            remove_rollups_from_state: step_timings
                .with_label_values(&["remove_rollups_from_state"]),
            post_gc_calculations_seconds: step_timings.with_label_values(&["post_gc_calculations"]),
        }
    }
}

impl GcMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        let step_timings: CounterVec = registry.register(metric!(
                name: "mz_persist_gc_step_seconds",
                help: "time spent on individual steps of gc",
                var_labels: ["step"],
        ));
        GcMetrics {
            noop: registry.register(metric!(
                name: "mz_persist_gc_noop",
                help: "count of garbage collections skipped because they were already done",
            )),
            started: registry.register(metric!(
                name: "mz_persist_gc_started",
                help: "count of garbage collections started",
            )),
            finished: registry.register(metric!(
                name: "mz_persist_gc_finished",
                help: "count of garbage collections finished",
            )),
            merged: registry.register(metric!(
                name: "mz_persist_gc_merged_reqs",
                help: "count of garbage collection requests merged",
            )),
            seconds: registry.register(metric!(
                name: "mz_persist_gc_seconds",
                help: "time spent in garbage collections",
            )),
            steps: GcStepTimings::new(step_timings),
        }
    }
}

#[derive(Debug)]
pub struct LeaseMetrics {
    pub(crate) timeout_read: IntCounter,
    /// Counted where `shard_source` reports a lost lease. Only observable where
    /// the error handler survives it (storage); compute halts right after.
    pub(crate) expired_live_read: IntCounter,
    pub(crate) dropped_part: IntCounter,
}

impl LeaseMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        LeaseMetrics {
            timeout_read: registry.register(metric!(
                name: "mz_persist_lease_timeout_read",
                help: "count of readers whose lease timed out",
            )),
            expired_live_read: registry.register(metric!(
                name: "mz_persist_lease_expired_live_read",
                help: "count of fetches that found their reader's lease had timed out",
            )),
            dropped_part: registry.register(metric!(
                name: "mz_persist_lease_dropped_part",
                help: "count of LeasedBatchParts that were dropped without being politely returned",
            )),
        }
    }
}

struct IncOnDrop(IntCounter);

impl Drop for IncOnDrop {
    fn drop(&mut self) {
        self.0.inc()
    }
}

pub struct MetricsRetryStream {
    retry: RetryStream,
    pub(crate) retries: IntCounter,
    sleep_seconds: Counter,
    _finished: IncOnDrop,
}

impl MetricsRetryStream {
    pub fn new(retry: RetryStream, metrics: &RetryMetrics) -> Self {
        metrics.started.inc();
        MetricsRetryStream {
            retry,
            retries: metrics.retries.clone(),
            sleep_seconds: metrics.sleep_seconds.clone(),
            _finished: IncOnDrop(metrics.finished.clone()),
        }
    }

    /// How many times [Self::sleep] has been called.
    pub fn attempt(&self) -> usize {
        self.retry.attempt()
    }

    /// The next sleep (without jitter for easy printing in logs).
    pub fn next_sleep(&self) -> Duration {
        self.retry.next_sleep()
    }

    /// Executes the next sleep in the series.
    ///
    /// This isn't cancel-safe, so it consumes and returns self, to prevent
    /// accidental mis-use.
    pub async fn sleep(self) -> Self {
        self.retries.inc();
        self.sleep_seconds
            .inc_by(self.retry.next_sleep().as_secs_f64());
        let retry = self.retry.sleep().await;
        MetricsRetryStream {
            retry,
            retries: self.retries,
            sleep_seconds: self.sleep_seconds,
            _finished: self._finished,
        }
    }
}

#[derive(Debug)]
pub struct CodecsMetrics {
    pub(crate) state: CodecMetrics,
    pub(crate) state_diff: CodecMetrics,
    pub(crate) batch: CodecMetrics,
    pub(crate) key: CodecMetrics,
    pub(crate) val: CodecMetrics,
    // Intentionally not adding time and diff because they're just
    // `{to,from}_le_bytes`.
}

#[derive(Debug)]
pub struct CodecMetrics {
    pub(crate) encode_count: IntCounter,
    pub(crate) encode_seconds: Counter,
    pub(crate) decode_count: IntCounter,
    pub(crate) decode_seconds: Counter,
}

impl CodecMetrics {
    pub(crate) fn encode<R, F: FnOnce() -> R>(&self, f: F) -> R {
        let now = Instant::now();
        let r = f();
        self.encode_count.inc();
        self.encode_seconds.inc_by(now.elapsed().as_secs_f64());
        r
    }

    pub(crate) fn decode<R, F: FnOnce() -> R>(&self, f: F) -> R {
        let now = Instant::now();
        let r = f();
        self.decode_count.inc();
        self.decode_seconds.inc_by(now.elapsed().as_secs_f64());
        r
    }
}

#[derive(Debug)]
pub struct StateMetrics {
    pub(crate) apply_spine_fast_path: IntCounter,
    pub(crate) apply_spine_slow_path: IntCounter,
    pub(crate) apply_spine_slow_path_lenient: IntCounter,
    pub(crate) apply_spine_slow_path_lenient_adjustment: IntCounter,
    pub(crate) apply_spine_slow_path_with_reconstruction: IntCounter,
    pub(crate) apply_spine_flattened: IntCounter,
    pub(crate) update_state_noop_path: IntCounter,
    pub(crate) update_state_empty_path: IntCounter,
    pub(crate) update_state_fast_path: IntCounter,
    pub(crate) update_state_slow_path: IntCounter,
    pub(crate) rollup_at_seqno_migration: IntCounter,
    pub(crate) fetch_recent_live_diffs_fast_path: IntCounter,
    pub(crate) fetch_recent_live_diffs_slow_path: IntCounter,
    pub(crate) writer_added: IntCounter,
    pub(crate) writer_removed: IntCounter,
    pub(crate) force_apply_hostname: IntCounter,
    pub(crate) rollup_write_success: IntCounter,
    pub(crate) rollup_write_noop_latest: IntCounter,
    pub(crate) rollup_write_noop_truncated: IntCounter,
}

impl StateMetrics {
    pub(crate) fn new(registry: &MetricsRegistry) -> Self {
        let rollup_write_noop: IntCounterVec = registry.register(metric!(
                name: "mz_persist_state_rollup_write_noop",
                help: "count of no-op rollup writes",
                var_labels: ["reason"],
        ));

        StateMetrics {
            apply_spine_fast_path: registry.register(metric!(
                name: "mz_persist_state_apply_spine_fast_path",
                help: "count of spine diff applications that hit the fast path",
            )),
            apply_spine_slow_path: registry.register(metric!(
                name: "mz_persist_state_apply_spine_slow_path",
                help: "count of spine diff applications that hit the slow path",
            )),
            apply_spine_slow_path_lenient: registry.register(metric!(
                name: "mz_persist_state_apply_spine_slow_path_lenient",
                help: "count of spine diff applications that hit the lenient compaction apply path",
            )),
            apply_spine_slow_path_lenient_adjustment: registry.register(metric!(
                name: "mz_persist_state_apply_spine_slow_path_lenient_adjustment",
                help: "count of adjustments made by the lenient compaction apply path",
            )),
            apply_spine_slow_path_with_reconstruction: registry.register(metric!(
                name: "mz_persist_state_apply_spine_slow_path_with_reconstruction",
                help: "count of spine diff applications that hit the slow path with extra spine reconstruction step",
            )),
            apply_spine_flattened: registry.register(metric!(
                name: "mz_persist_state_apply_spine_flattened",
                help: "count of spine diff applications that flatten the trace",
            )),
            update_state_noop_path: registry.register(metric!(
                name: "mz_persist_state_update_state_noop_path",
                help: "count of state update applications that no-oped due to shared state",
            )),
            update_state_empty_path: registry.register(metric!(
                name: "mz_persist_state_update_state_empty_path",
                help: "count of state update applications that found no new updates",
            )),
            update_state_fast_path: registry.register(metric!(
                name: "mz_persist_state_update_state_fast_path",
                help: "count of state update applications that hit the fast path",
            )),
            update_state_slow_path: registry.register(metric!(
                name: "mz_persist_state_update_state_slow_path",
                help: "count of state update applications that hit the slow path",
            )),
            rollup_at_seqno_migration: registry.register(metric!(
                name: "mz_persist_state_rollup_at_seqno_migration",
                help: "count of fetch_rollup_at_seqno calls that only worked because of the migration",
            )),
            fetch_recent_live_diffs_fast_path: registry.register(metric!(
                name: "mz_persist_state_fetch_recent_live_diffs_fast_path",
                help: "count of fetch_recent_live_diffs that hit the fast path",
            )),
            fetch_recent_live_diffs_slow_path: registry.register(metric!(
                name: "mz_persist_state_fetch_recent_live_diffs_slow_path",
                help: "count of fetch_recent_live_diffs that hit the slow path",
            )),
            writer_added: registry.register(metric!(
                name: "mz_persist_state_writer_added",
                help: "count of writers added to the state",
            )),
            writer_removed: registry.register(metric!(
                name: "mz_persist_state_writer_removed",
                help: "count of writers removed from the state",
            )),
            force_apply_hostname: registry.register(metric!(
                name: "mz_persist_state_force_applied_hostname",
                help: "count of when hostname diffs needed to be force applied",
            )),
            rollup_write_success: registry.register(metric!(
                name: "mz_persist_state_rollup_write_success",
                help: "count of rollups written successful (may not be linked in to state)",
            )),
            rollup_write_noop_latest: rollup_write_noop.with_label_values(&["latest"]),
            rollup_write_noop_truncated: rollup_write_noop.with_label_values(&["truncated"]),
        }
    }
}

#[derive(Debug)]
pub struct ShardsMetrics {
    // Unlike all the other metrics in here, ShardsMetrics intentionally uses
    // the DeleteOnDrop wrappers. A process might stop using a shard (drop all
    // handles to it) but e.g. the set of commands never changes.
    //
    // Every per-shard family also has a bounded aggregate whose cardinality
    // does not grow with the shard count. `PER_SHARD_FAMILIES` lists which.
    //
    // The scrape-time aggregates, plus `mz_persist_shard_count` and
    // `mz_persist_stale_shard_count`, are not fields here: they live in a
    // `ShardsAggregateMetrics` collector that the registry owns, sharing this
    // struct's `shards` map so one scrape walk feeds all of them.
    encoded_rollup_size: mz_ore::metrics::UIntGaugeVec,
    encoded_diff_size: mz_ore::metrics::IntCounterVec,
    hollow_batch_count: mz_ore::metrics::UIntGaugeVec,
    spine_batch_count: mz_ore::metrics::UIntGaugeVec,
    batch_part_count: mz_ore::metrics::UIntGaugeVec,
    batch_part_version_count: mz_ore::metrics::UIntGaugeVec,
    update_count: mz_ore::metrics::UIntGaugeVec,
    rollup_count: mz_ore::metrics::UIntGaugeVec,
    largest_batch_size: mz_ore::metrics::UIntGaugeVec,
    seqnos_held: mz_ore::metrics::UIntGaugeVec,
    seqnos_since_last_rollup: mz_ore::metrics::UIntGaugeVec,
    gc_seqno_held_parts: mz_ore::metrics::UIntGaugeVec,
    gc_live_diffs: mz_ore::metrics::UIntGaugeVec,
    gc_finished: mz_ore::metrics::IntCounterVec,
    compaction_applied: mz_ore::metrics::IntCounterVec,
    cmd_succeeded: mz_ore::metrics::IntCounterVec,
    usage_current_state_batches_bytes: mz_ore::metrics::UIntGaugeVec,
    usage_current_state_rollups_bytes: mz_ore::metrics::UIntGaugeVec,
    usage_referenced_not_current_state_bytes: mz_ore::metrics::UIntGaugeVec,
    usage_not_leaked_not_referenced_bytes: mz_ore::metrics::UIntGaugeVec,
    usage_leaked_bytes: mz_ore::metrics::UIntGaugeVec,
    blob_gets: mz_ore::metrics::IntCounterVec,
    blob_sets: mz_ore::metrics::IntCounterVec,
    unconsolidated_snapshot: mz_ore::metrics::IntCounterVec,
    inline_part_count: UIntGaugeVec,
    compact_batches: UIntGaugeVec,
    compacting_batches: UIntGaugeVec,
    noncompact_batches: UIntGaugeVec,
    process_blob_gets: IntCounter,
    process_blob_sets: IntCounter,
    process_encoded_diff_size: IntCounter,
    process_unconsolidated_snapshot: IntCounter,
    // We hand out `Arc<ShardMetrics>` to read and write handles, but store it
    // here as `Weak`. This allows us to discover if it's no longer in use and
    // so we can remove it from the map.
    shards: Arc<Mutex<BTreeMap<ShardId, Weak<ShardMetrics>>>>,
}

/// Removes the per-shard series this process should not export from gathered
/// families, as configured by [`SHARD_METRICS`] and
/// [`PER_SHARD_METRICS_ENABLE_REGEX`].
///
/// It matches the names in `PER_SHARD_FAMILIES`, not a name prefix, because
/// `mz_persist_shard_count`, `mz_persist_stale_shard_count` and the
/// `mz_persist_shard_*_percentile` / `_sum` / `_topk` aggregates and
/// `mz_persist_shard_info` share the prefix.
///
/// NOTE: Postprocessors run after collection, so the `invalid_*` counters land
/// in the scrape after the one that bumped them.
struct PerShardMetricsFilter {
    configs: Arc<ConfigSet>,
    /// The last seen value of `SHARD_METRICS` and what it parsed to.
    mode: (String, ShardMetricsExport),
    invalid_mode: IntCounter,
    /// The last seen value of `PER_SHARD_METRICS_ENABLE_REGEX`, and its
    /// compiled form if it is non-empty and valid.
    allow: (String, Option<Regex>),
    invalid_regex: IntCounter,
}

impl PerShardMetricsFilter {
    fn filter(&mut self, families: &mut Vec<MetricFamily>) {
        if self.mode().per_shard() {
            return;
        }
        let allow = self.allow_regex();
        families.retain_mut(|f| {
            if !PER_SHARD_FAMILIES.iter().any(|p| p.name == f.name()) {
                return true;
            }
            let Some(allow) = allow else {
                return false;
            };
            f.mut_metric().retain(|m| {
                m.get_label()
                    .iter()
                    .any(|l| matches!(l.name(), "shard" | "name") && allow.is_match(l.value()))
            });
            !f.get_metric().is_empty()
        });
    }

    /// Returns the regex of series to keep, or `None` to keep none.
    ///
    /// Recompiles only when the config value changes, so an invalid value is
    /// logged once, but it is counted on every scrape. It fails closed: an
    /// invalid regex keeps nothing rather than everything, which would bring
    /// back the full per-shard cardinality.
    fn allow_regex(&mut self) -> Option<&Regex> {
        let value = PER_SHARD_METRICS_ENABLE_REGEX.get(&self.configs);
        if value != self.allow.0 {
            let compiled = if value.is_empty() {
                None
            } else {
                Regex::new(&value)
                    .inspect_err(|err| {
                        warn!(
                            "invalid {}, keeping no per-shard series: {err}",
                            PER_SHARD_METRICS_ENABLE_REGEX.name()
                        )
                    })
                    .ok()
            };
            self.allow = (value, compiled);
        }
        if !self.allow.0.is_empty() && self.allow.1.is_none() {
            self.invalid_regex.inc();
        }
        self.allow.1.as_ref()
    }

    /// The export mode, or the default if the config value is unrecognized.
    ///
    /// Reparses only on change, so an unrecognized value logs once but counts
    /// every scrape.
    fn mode(&mut self) -> ShardMetricsExport {
        let value = SHARD_METRICS.get(&self.configs);
        if value != self.mode.0 {
            let parsed = ShardMetricsExport::parse(&value).unwrap_or_else(|| {
                let default = ShardMetricsExport::default();
                warn!(
                    "unrecognized {}: '{value}', exporting '{}'",
                    SHARD_METRICS.name(),
                    default.as_str()
                );
                default
            });
            self.mode = (value, parsed);
        }
        if self.mode.1.as_str() != self.mode.0 {
            self.invalid_mode.inc();
        }
        self.mode.1
    }
}

impl ShardsMetrics {
    fn new(registry: &MetricsRegistry, configs: Arc<ConfigSet>) -> Self {
        let shards = Arc::new(Mutex::new(BTreeMap::new()));
        registry.register_collector(ShardsAggregateMetrics::new(Arc::clone(&shards), configs));
        ShardsMetrics {
            encoded_rollup_size: registry.register(metric!(
                name: "mz_persist_shard_rollup_size_bytes",
                help: "total encoded rollup size by shard",
                var_labels: ["shard", "name"],
            )),
            encoded_diff_size: registry.register(metric!(
                name: "mz_persist_shard_diff_size_bytes",
                help: "total encoded diff size by shard",
                var_labels: ["shard", "name"],
            )),
            hollow_batch_count: registry.register(metric!(
                name: "mz_persist_shard_hollow_batch_count",
                help: "count of hollow batches by shard",
                var_labels: ["shard", "name"],
            )),
            spine_batch_count: registry.register(metric!(
                name: "mz_persist_shard_spine_batch_count",
                help: "count of spine batches by shard",
                var_labels: ["shard", "name"],
            )),
            batch_part_count: registry.register(metric!(
                name: "mz_persist_shard_batch_part_count",
                help: "count of batch parts by shard",
                var_labels: ["shard", "name"],
            )),
            batch_part_version_count: registry.register(metric!(
                name: "mz_persist_shard_batch_part_version_count",
                help: "count of batch parts by shard and version",
                var_labels: ["shard", "name", "version"],
            )),
            update_count: registry.register(metric!(
                name: "mz_persist_shard_update_count",
                help: "count of updates by shard",
                var_labels: ["shard", "name"],
            )),
            rollup_count: registry.register(metric!(
                name: "mz_persist_shard_rollup_count",
                help: "count of rollups by shard",
                var_labels: ["shard", "name"],
            )),
            largest_batch_size: registry.register(metric!(
                name: "mz_persist_shard_largest_batch_size",
                help: "largest encoded batch size by shard",
                var_labels: ["shard", "name"],
            )),
            seqnos_held: registry.register(metric!(
                name: "mz_persist_shard_seqnos_held",
                help: "maximum count of gc-ineligible states by shard",
                var_labels: ["shard", "name"],
            )),
            seqnos_since_last_rollup: registry.register(metric!(
                name: "mz_persist_shard_seqnos_since_last_rollup",
                help: "count of seqnos since last rollup",
                var_labels: ["shard", "name"],
            )),
            gc_seqno_held_parts: registry.register(metric!(
                name: "mz_persist_shard_gc_seqno_held_parts",
                help: "count of parts referenced by some live state but not the current state (ie. parts kept only to satisfy seqno holds) at GC time",
                var_labels: ["shard", "name"],
            )),
            gc_live_diffs: registry.register(metric!(
                name: "mz_persist_shard_gc_live_diffs",
                help: "the number of diffs (or, alternatively, the number of seqnos) present in consensus state at GC time",
                var_labels: ["shard", "name"],
            )),
            gc_finished: registry.register(metric!(
                name: "mz_persist_shard_gc_finished",
                help: "count of garbage collections finished by shard",
                var_labels: ["shard", "name"],
            )),
            compaction_applied: registry.register(metric!(
                name: "mz_persist_shard_compaction_applied",
                help: "count of compactions applied to state by shard",
                var_labels: ["shard", "name"],
            )),
            cmd_succeeded: registry.register(metric!(
                name: "mz_persist_shard_cmd_succeeded",
                help: "count of commands succeeded by shard",
                var_labels: ["shard", "name"],
            )),
            usage_current_state_batches_bytes: registry.register(metric!(
                name: "mz_persist_shard_usage_current_state_batches_bytes",
                help: "data in batches/parts referenced by current version of state",
                var_labels: ["shard", "name"],
            )),
            usage_current_state_rollups_bytes: registry.register(metric!(
                name: "mz_persist_shard_usage_current_state_rollups_bytes",
                help: "data in rollups referenced by current version of state",
                var_labels: ["shard", "name"],
            )),
            usage_referenced_not_current_state_bytes: registry.register(metric!(
                name: "mz_persist_shard_usage_referenced_not_current_state_bytes",
                help: "data referenced only by a previous version of state",
                var_labels: ["shard", "name"],
            )),
            usage_not_leaked_not_referenced_bytes: registry.register(metric!(
                name: "mz_persist_shard_usage_not_leaked_not_referenced_bytes",
                help: "data written by an active writer but not referenced by any version of state",
                var_labels: ["shard", "name"],
            )),
            usage_leaked_bytes: registry.register(metric!(
                name: "mz_persist_shard_usage_leaked_bytes",
                help: "data reclaimable by a leaked blob detector",
                var_labels: ["shard", "name"],
            )),
            blob_gets: registry.register(metric!(
                name: "mz_persist_shard_blob_gets",
                help: "number of Blob::get calls for this shard",
                var_labels: ["shard", "name"],
            )),
            blob_sets: registry.register(metric!(
                name: "mz_persist_shard_blob_sets",
                help: "number of Blob::set calls for this shard",
                var_labels: ["shard", "name"],
            )),
            unconsolidated_snapshot: registry.register(metric!(
                name: "mz_persist_shard_unconsolidated_snapshot",
                help: "in snapshot_and_read, the number of times consolidating the raw data wasn't enough to produce consolidated output",
                var_labels: ["shard", "name"],
            )),
            inline_part_count: registry.register(metric!(
                name: "mz_persist_shard_inline_part_count",
                help: "count of parts inline in shard metadata",
                var_labels: ["shard", "name"],
            )),
            compact_batches: registry.register(metric!(
                name: "mz_persist_shard_compact_batches",
                help: "number of fully compact batches in the shard",
                var_labels: ["shard", "name"],
            )),
            compacting_batches: registry.register(metric!(
                name: "mz_persist_shard_compacting_batches",
                help: "number of batches in the shard with compactions in progress",
                var_labels: ["shard", "name"],
            )),
            noncompact_batches: registry.register(metric!(
                name: "mz_persist_shard_noncompact_batches",
                help: "number of batches in the shard that aren't compact and have no ongoing compaction",
                var_labels: ["shard", "name"],
            )),
            process_blob_gets: registry.register(metric!(
                name: "mz_persist_blob_gets",
                help: "number of Blob::get calls for batch parts, over all shards",
            )),
            process_blob_sets: registry.register(metric!(
                name: "mz_persist_blob_sets",
                help: "number of Blob::set calls for batch parts, over all shards",
            )),
            process_encoded_diff_size: registry.register(metric!(
                name: "mz_persist_diff_size_bytes",
                help: "total encoded diff size, over all shards",
            )),
            process_unconsolidated_snapshot: registry.register(metric!(
                name: "mz_persist_unconsolidated_snapshot",
                help: "in snapshot_and_read, the number of times consolidating the raw data wasn't enough to produce consolidated output, over all shards",
            )),
            shards,
        }
    }

    pub fn shard(&self, shard_id: &ShardId, name: &str) -> Arc<ShardMetrics> {
        let mut shards = self.shards.lock().expect("mutex poisoned");
        if let Some(shard) = shards.get(shard_id) {
            if let Some(shard) = shard.upgrade() {
                return Arc::clone(&shard);
            } else {
                assert!(shards.remove(shard_id).is_some());
            }
        }
        let shard = Arc::new(ShardMetrics::new(shard_id, name, self));
        assert!(
            shards
                .insert(shard_id.clone(), Arc::downgrade(&shard))
                .is_none()
        );
        shard
    }

    /// Returns the shards that still have handles, and forgets the rest.
    ///
    /// Holds the `shards` lock, which `Self::shard` needs, so it does nothing
    /// but upgrade each entry.
    fn live_shards(
        shards: &Arc<Mutex<BTreeMap<ShardId, Weak<ShardMetrics>>>>,
    ) -> Vec<Arc<ShardMetrics>> {
        let mut shards = shards.lock().expect("mutex poisoned");
        let mut live = Vec::with_capacity(shards.len());
        shards.retain(|_, metrics| match metrics.upgrade() {
            Some(metrics) => {
                live.push(metrics);
                true
            }
            None => false,
        });
        live
    }
}

/// A family that [`ShardsMetrics::new`] registers with one series per shard,
/// and the bounded aggregate that covers it.
struct PerShardFamily {
    name: &'static str,
    aggregate: ShardAggregate,
}

enum ShardAggregate {
    /// Folded into `<name>_percentile` and `<name>_topk` at scrape time.
    Gauge(ShardGauge),
    /// Summed over shards into `mz_persist_batch_part_version_count`.
    VersionSum,
    /// Counted in full by the named process-level counter.
    ///
    /// NOTE: For `gc_finished`, `compaction_applied` and `cmd_succeeded` the
    /// process counter is kept equal to the sum over shards by a paired
    /// `.inc()` at each call site, not by construction like [`ShardCounter`].
    /// Add one without the other and the total drifts, with no test to catch it.
    Counter {
        // Only read by tests: that it names a registered counter, and that
        // every export mode keeps that counter.
        #[allow(dead_code)]
        total: &'static str,
    },
}

struct ShardGauge {
    /// The registered help of the per-shard family, which the aggregate
    /// extends. `per_shard_families_match_registrations` checks they match.
    help: &'static str,
    /// How many of the largest shards `<name>_topk` names.
    top_k: usize,
    /// Whether to emit `<name>_sum`. See [`AggregatedFamily::sum`].
    sum: bool,
    get: fn(&ShardMetrics) -> u64,
}

/// The default `ShardGauge::top_k`.
const SHARD_TOP_K: usize = 10;

/// Every per-shard family. Must list exactly the families with a `shard`
/// label that `ShardsMetrics::new` registers.
const PER_SHARD_FAMILIES: &[PerShardFamily] = &[
    PerShardFamily {
        name: "mz_persist_shard_rollup_size_bytes",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "total encoded rollup size by shard",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.latest_rollup_size.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_diff_size_bytes",
        aggregate: ShardAggregate::Counter {
            total: "mz_persist_diff_size_bytes",
        },
    },
    PerShardFamily {
        name: "mz_persist_shard_hollow_batch_count",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "count of hollow batches by shard",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.hollow_batch_count.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_spine_batch_count",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "count of spine batches by shard",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.spine_batch_count.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_batch_part_count",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "count of batch parts by shard",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.batch_part_count.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_batch_part_version_count",
        aggregate: ShardAggregate::VersionSum,
    },
    PerShardFamily {
        name: "mz_persist_shard_update_count",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "count of updates by shard",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.update_count.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_rollup_count",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "count of rollups by shard",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.rollup_count.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_largest_batch_size",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "largest encoded batch size by shard",
            top_k: SHARD_TOP_K,
            sum: false,
            get: |m| m.largest_batch_size.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_seqnos_held",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "maximum count of gc-ineligible states by shard",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.seqnos_held.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_seqnos_since_last_rollup",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "count of seqnos since last rollup",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.seqnos_since_last_rollup.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_gc_seqno_held_parts",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "count of parts referenced by some live state but not the current state (ie. parts kept only to satisfy seqno holds) at GC time",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.gc_seqno_held_parts.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_gc_live_diffs",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "the number of diffs (or, alternatively, the number of seqnos) present in consensus state at GC time",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.gc_live_diffs.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_gc_finished",
        aggregate: ShardAggregate::Counter {
            total: "mz_persist_gc_finished",
        },
    },
    PerShardFamily {
        name: "mz_persist_shard_compaction_applied",
        aggregate: ShardAggregate::Counter {
            total: "mz_persist_compaction_applied",
        },
    },
    PerShardFamily {
        name: "mz_persist_shard_cmd_succeeded",
        aggregate: ShardAggregate::Counter {
            total: "mz_persist_cmd_succeeded_count",
        },
    },
    PerShardFamily {
        name: "mz_persist_shard_usage_current_state_batches_bytes",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "data in batches/parts referenced by current version of state",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.usage_current_state_batches_bytes.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_usage_current_state_rollups_bytes",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "data in rollups referenced by current version of state",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.usage_current_state_rollups_bytes.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_usage_referenced_not_current_state_bytes",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "data referenced only by a previous version of state",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.usage_referenced_not_current_state_bytes.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_usage_not_leaked_not_referenced_bytes",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "data written by an active writer but not referenced by any version of state",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.usage_not_leaked_not_referenced_bytes.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_usage_leaked_bytes",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "data reclaimable by a leaked blob detector",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.usage_leaked_bytes.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_blob_gets",
        aggregate: ShardAggregate::Counter {
            total: "mz_persist_blob_gets",
        },
    },
    PerShardFamily {
        name: "mz_persist_shard_blob_sets",
        aggregate: ShardAggregate::Counter {
            total: "mz_persist_blob_sets",
        },
    },
    PerShardFamily {
        name: "mz_persist_shard_unconsolidated_snapshot",
        aggregate: ShardAggregate::Counter {
            total: "mz_persist_unconsolidated_snapshot",
        },
    },
    PerShardFamily {
        name: "mz_persist_shard_inline_part_count",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "count of parts inline in shard metadata",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.inline_part_count.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_compact_batches",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "number of fully compact batches in the shard",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.compact_batches.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_compacting_batches",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "number of batches in the shard with compactions in progress",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.compacting_batches.get(),
        }),
    },
    PerShardFamily {
        name: "mz_persist_shard_noncompact_batches",
        aggregate: ShardAggregate::Gauge(ShardGauge {
            help: "number of batches in the shard that aren't compact and have no ongoing compaction",
            top_k: SHARD_TOP_K,
            sum: true,
            get: |m| m.noncompact_batches.get(),
        }),
    },
];

/// The per-shard gauges and their names, in `PER_SHARD_FAMILIES` order.
fn shard_gauges() -> impl Iterator<Item = (&'static str, &'static ShardGauge)> {
    PER_SHARD_FAMILIES
        .iter()
        .filter_map(|f| match &f.aggregate {
            ShardAggregate::Gauge(g) => Some((f.name, g)),
            _ => None,
        })
}

const SHARD_INFO_NAME: &str = "mz_persist_shard_info";
const BATCH_PART_VERSION_COUNT_NAME: &str = "mz_persist_batch_part_version_count";
const BATCH_PART_VERSION_COUNT_HELP: &str = "count of batch parts by version, over all shards";

/// Identifies a shard in the `_topk` families, with the same labels as the
/// per-shard series.
#[derive(Debug)]
struct ShardKey(Arc<ShardMetrics>);

impl AggregationKey for ShardKey {
    const LABEL_NAMES: &'static [&'static str] = &["shard", "name"];

    fn label_values(&self) -> Vec<String> {
        vec![self.0.shard_id.to_string(), self.0.name.clone()]
    }
}

/// Process-level metrics derived from the shards map, collected in a single
/// walk: the shard counts, percentiles, a top-K and (where meaningful) a sum
/// per per-shard gauge, the per-version batch part count summed over shards,
/// and, if [`SHARD_INFO_METRIC`] is set, `mz_persist_shard_info`.
///
/// All but `mz_persist_shard_info` have a cardinality independent of the
/// number of shards.
#[derive(Debug)]
struct ShardsAggregateMetrics {
    shards: Arc<Mutex<BTreeMap<ShardId, Weak<ShardMetrics>>>>,
    configs: Arc<ConfigSet>,
    count: GenericGauge<AtomicI64>,
    stale_count: GenericGauge<AtomicI64>,
    gauges: AggregatedFamilies<ShardKey>,
    /// The getters behind `gauges`, same order, so the per-shard walk reads a
    /// slice instead of re-deriving them per shard.
    gauge_gets: Vec<fn(&ShardMetrics) -> u64>,
    // Only used for `desc`. `collect` builds a fresh vec each scrape so that
    // versions no shard has anymore are not exported.
    batch_part_version_count: raw::UIntGaugeVec,
    // Only used for `desc`, like `batch_part_version_count`.
    shard_info: raw::UIntGaugeVec,
}

impl ShardsAggregateMetrics {
    fn new(
        shards: Arc<Mutex<BTreeMap<ShardId, Weak<ShardMetrics>>>>,
        configs: Arc<ConfigSet>,
    ) -> Self {
        let gauges = shard_gauges()
            .map(|(name, g)| AggregatedFamily {
                name: name.to_string(),
                help: g.help.to_string(),
                sum: g.sum,
                top_k: Some(g.top_k),
            })
            .collect();
        ShardsAggregateMetrics {
            shards,
            configs,
            count: MakeCollector::make_collector(metric!(
                name: "mz_persist_shard_count",
                help: "count of all active shards on this process",
            )),
            stale_count: MakeCollector::make_collector(metric!(
                name: "mz_persist_stale_shard_count",
                help: "count of shards on this process whose persisted state version \
                       is behind this process's build version; per-process, so summing \
                       across processes counts (shard, process) pairs, not distinct shards",
            )),
            gauges: AggregatedFamilies::new(gauges),
            gauge_gets: shard_gauges().map(|(_, g)| g.get).collect(),
            batch_part_version_count: new_batch_part_version_count(),
            shard_info: new_shard_info(),
        }
    }

    /// The bounded aggregates of the per-shard families over `live`.
    fn summaries(&self, live: Vec<Arc<ShardMetrics>>) -> Vec<MetricFamily> {
        let mut versions = BTreeMap::<String, u64>::new();
        let mut values = Vec::new();
        let mut snapshot = self.gauges.snapshot();
        for m in live {
            values.clear();
            values.extend(self.gauge_gets.iter().map(|get| get(&m)));
            let version_map = m
                .batch_part_version_map
                .lock()
                .expect("mutex should not be poisoned");
            for (version, metrics) in version_map.iter() {
                let n = metrics.batch_part_version_count.get();
                if n == 0 {
                    continue;
                }
                match versions.get_mut(version) {
                    Some(total) => *total += n,
                    None => {
                        versions.insert(version.clone(), n);
                    }
                }
            }
            drop(version_map);
            snapshot.push(ShardKey(m), &values);
        }

        let mut families = snapshot.finish();
        let batch_part_version_count = new_batch_part_version_count();
        for (version, n) in versions {
            batch_part_version_count
                .with_label_values(&[version.as_str()])
                .set(n);
        }
        families.extend(batch_part_version_count.collect());
        families
    }
}

/// Lists the shards, which the `_topk` families name only the largest of. One
/// series per shard, so unlike the summaries it grows with the shard count.
fn new_shard_info() -> raw::UIntGaugeVec {
    raw::UIntGaugeVec::new(
        PrometheusOpts::new(SHARD_INFO_NAME, "1 for each active shard on this process"),
        ShardKey::LABEL_NAMES,
    )
    .expect("valid gauge vec")
}

fn new_batch_part_version_count() -> raw::UIntGaugeVec {
    raw::UIntGaugeVec::new(
        PrometheusOpts::new(BATCH_PART_VERSION_COUNT_NAME, BATCH_PART_VERSION_COUNT_HELP),
        &["version"],
    )
    .expect("valid gauge vec")
}

impl Collector for ShardsAggregateMetrics {
    fn desc(&self) -> Vec<&Desc> {
        let mut descs = self.count.desc();
        descs.extend(self.stale_count.desc());
        descs.extend(self.gauges.descs());
        descs.extend(self.batch_part_version_count.desc());
        descs.extend(self.shard_info.desc());
        descs
    }

    fn collect(&self) -> Vec<MetricFamily> {
        let live = ShardsMetrics::live_shards(&self.shards);
        let mut count = 0;
        let mut stale_count = 0;
        for m in &live {
            count += 1;
            if m.stale.load(Ordering::Relaxed) {
                stale_count += 1;
            }
        }
        self.count.set(count);
        self.stale_count.set(stale_count);

        let mut families = Vec::new();
        if SHARD_INFO_METRIC.get(&self.configs) {
            let shard_info = new_shard_info();
            for m in &live {
                shard_info
                    .with_label_values(&[m.shard_id.to_string(), m.name.clone()])
                    .set(1);
            }
            families.extend(shard_info.collect());
        }
        if ShardMetricsExport::get(&self.configs).summary() {
            families.extend(self.summaries(live));
        }
        families.extend(self.count.collect());
        families.extend(self.stale_count.collect());
        families
    }
}

#[derive(Debug)]
pub struct ShardMetrics {
    pub shard_id: ShardId,
    pub name: String,
    pub largest_batch_size: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub latest_rollup_size: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub encoded_diff_size: ShardCounter,
    pub hollow_batch_count: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub spine_batch_count: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub batch_part_count: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    batch_part_version_count: mz_ore::metrics::UIntGaugeVec,
    batch_part_version_map: Mutex<BTreeMap<String, BatchPartVersionMetrics>>,
    pub update_count: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub rollup_count: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub seqnos_held: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub seqnos_since_last_rollup: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub gc_seqno_held_parts: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub gc_live_diffs: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub usage_current_state_batches_bytes: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub usage_current_state_rollups_bytes: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub usage_referenced_not_current_state_bytes: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub usage_not_leaked_not_referenced_bytes: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub usage_leaked_bytes: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub gc_finished: DeleteOnDropCounter<AtomicU64, Vec<String>>,
    pub compaction_applied: DeleteOnDropCounter<AtomicU64, Vec<String>>,
    pub cmd_succeeded: DeleteOnDropCounter<AtomicU64, Vec<String>>,
    pub blob_gets: ShardCounter,
    pub blob_sets: ShardCounter,
    pub unconsolidated_snapshot: ShardCounter,
    pub inline_part_count: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub compact_batches: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub compacting_batches: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    pub noncompact_batches: DeleteOnDropGauge<AtomicU64, Vec<String>>,
    // Not a registered per-shard series: this feeds the process-level
    // `mz_persist_stale_shard_count` computed gauge, which sums it across shards
    // instead of exporting one series per shard.
    pub stale: AtomicBool,
}

impl ShardMetrics {
    pub fn new(shard_id: &ShardId, name: &str, shards_metrics: &ShardsMetrics) -> Self {
        let shard = shard_id.to_string();
        ShardMetrics {
            shard_id: *shard_id,
            name: name.to_string(),
            latest_rollup_size: shards_metrics
                .encoded_rollup_size
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            encoded_diff_size: ShardCounter {
                shard: shards_metrics
                    .encoded_diff_size
                    .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
                process: shards_metrics.process_encoded_diff_size.clone(),
            },
            hollow_batch_count: shards_metrics
                .hollow_batch_count
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            spine_batch_count: shards_metrics
                .spine_batch_count
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            batch_part_count: shards_metrics
                .batch_part_count
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            batch_part_version_count: shards_metrics.batch_part_version_count.clone(),
            batch_part_version_map: Mutex::new(BTreeMap::new()),
            update_count: shards_metrics
                .update_count
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            rollup_count: shards_metrics
                .rollup_count
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            largest_batch_size: shards_metrics
                .largest_batch_size
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            seqnos_held: shards_metrics
                .seqnos_held
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            seqnos_since_last_rollup: shards_metrics
                .seqnos_since_last_rollup
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            gc_seqno_held_parts: shards_metrics
                .gc_seqno_held_parts
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            gc_live_diffs: shards_metrics
                .gc_live_diffs
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            gc_finished: shards_metrics
                .gc_finished
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            compaction_applied: shards_metrics
                .compaction_applied
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            cmd_succeeded: shards_metrics
                .cmd_succeeded
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            usage_current_state_batches_bytes: shards_metrics
                .usage_current_state_batches_bytes
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            usage_current_state_rollups_bytes: shards_metrics
                .usage_current_state_rollups_bytes
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            usage_referenced_not_current_state_bytes: shards_metrics
                .usage_referenced_not_current_state_bytes
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            usage_not_leaked_not_referenced_bytes: shards_metrics
                .usage_not_leaked_not_referenced_bytes
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            usage_leaked_bytes: shards_metrics
                .usage_leaked_bytes
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            blob_gets: ShardCounter {
                shard: shards_metrics
                    .blob_gets
                    .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
                process: shards_metrics.process_blob_gets.clone(),
            },
            blob_sets: ShardCounter {
                shard: shards_metrics
                    .blob_sets
                    .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
                process: shards_metrics.process_blob_sets.clone(),
            },
            unconsolidated_snapshot: ShardCounter {
                shard: shards_metrics
                    .unconsolidated_snapshot
                    .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
                process: shards_metrics.process_unconsolidated_snapshot.clone(),
            },
            inline_part_count: shards_metrics
                .inline_part_count
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            compact_batches: shards_metrics
                .compact_batches
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            compacting_batches: shards_metrics
                .compacting_batches
                .get_delete_on_drop_metric(vec![shard.clone(), name.to_string()]),
            noncompact_batches: shards_metrics
                .noncompact_batches
                .get_delete_on_drop_metric(vec![shard, name.to_string()]),
            stale: AtomicBool::new(false),
        }
    }

    pub(crate) fn set_batch_part_versions<'a>(
        &self,
        batch_parts_by_version: impl Iterator<Item = &'a str>,
    ) {
        let mut map = self
            .batch_part_version_map
            .lock()
            .expect("mutex should not be poisoned");
        // NB: It's a bit sus that the below assumes that no one else is
        // concurrently modifying the atomics in the gauges, but we're holding
        // the mutex this whole time, so it should be true.

        // We want to do this in a way that avoids allocating (e.g. summing up a
        // map). First reset everything.
        for x in map.values() {
            x.batch_part_version_count.set(0);
        }

        // Then go through the iterator, creating new entries as necessary and
        // adding.
        for key in batch_parts_by_version {
            if !map.contains_key(key) {
                map.insert(
                    key.to_owned(),
                    BatchPartVersionMetrics {
                        batch_part_version_count: self
                            .batch_part_version_count
                            .get_delete_on_drop_metric(vec![
                                self.shard_id.to_string(),
                                self.name.clone(),
                                key.to_owned(),
                            ]),
                    },
                );
            }
            let value = map.get(key).expect("inserted above");
            value.batch_part_version_count.inc();
        }
    }
}

/// A per-shard counter that also bumps a process-level total, so the total
/// stays monotonic when shards are dropped.
#[derive(Debug)]
pub struct ShardCounter {
    shard: DeleteOnDropCounter<AtomicU64, Vec<String>>,
    process: IntCounter,
}

impl ShardCounter {
    pub fn inc(&self) {
        self.inc_by(1);
    }

    pub fn inc_by(&self, v: u64) {
        self.shard.inc_by(v);
        self.process.inc_by(v);
    }
}

#[derive(Debug)]
pub struct BatchPartVersionMetrics {
    pub batch_part_version_count: DeleteOnDropGauge<AtomicU64, Vec<String>>,
}

/// Metrics recorded by audits of persist usage
#[derive(Debug)]
pub struct UsageAuditMetrics {
    /// Size of all batch parts stored in Blob
    pub blob_batch_part_bytes: UIntGauge,
    /// Count of batch parts stored in Blob
    pub blob_batch_part_count: UIntGauge,
    /// Size of all state rollups stored in Blob
    pub blob_rollup_bytes: UIntGauge,
    /// Count of state rollups stored in Blob
    pub blob_rollup_count: UIntGauge,
    /// Size of Blob
    pub blob_bytes: UIntGauge,
    /// Count of all blobs
    pub blob_count: UIntGauge,
    /// Time spent fetching blob metadata
    pub step_blob_metadata: Counter,
    /// Time spent fetching state versions
    pub step_state: Counter,
    /// Time spent doing math
    pub step_math: Counter,
}

impl UsageAuditMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        let step_timings: CounterVec = registry.register(metric!(
                name: "mz_persist_audit_step_seconds",
                help: "time spent on individual steps of audit",
                var_labels: ["step"],
        ));
        UsageAuditMetrics {
            blob_batch_part_bytes: registry.register(metric!(
                name: "mz_persist_audit_blob_batch_part_bytes",
                help: "total size of batch parts in blob",
            )),
            blob_batch_part_count: registry.register(metric!(
                name: "mz_persist_audit_blob_batch_part_count",
                help: "count of batch parts in blob",
            )),
            blob_rollup_bytes: registry.register(metric!(
                name: "mz_persist_audit_blob_rollup_bytes",
                help: "total size of state rollups stored in blob",
            )),
            blob_rollup_count: registry.register(metric!(
                name: "mz_persist_audit_blob_rollup_count",
                help: "count of all state rollups in blob",
            )),
            blob_bytes: registry.register(metric!(
                name: "mz_persist_audit_blob_bytes",
                help: "total size of blob",
            )),
            blob_count: registry.register(metric!(
                name: "mz_persist_audit_blob_count",
                help: "count of all blobs",
            )),
            step_blob_metadata: step_timings.with_label_values(&["blob_metadata"]),
            step_state: step_timings.with_label_values(&["state"]),
            step_math: step_timings.with_label_values(&["math"]),
        }
    }
}

/// Represents a change in a number of updates kept in a data structure
/// (e.g., a buffer length or capacity change).
#[derive(Debug)]
pub enum UpdateDelta {
    /// A negative delta in the number of updates.
    Negative(u64),
    /// A non-negative delta in the number of updates.
    NonNegative(u64),
}

impl UpdateDelta {
    /// Creates a new `UpdateDelta` from the difference between a new value
    /// for a number of updates and the corresponding old value.
    pub fn new(new: usize, old: usize) -> Self {
        if new < old {
            UpdateDelta::Negative(CastFrom::cast_from(old - new))
        } else {
            UpdateDelta::NonNegative(CastFrom::cast_from(new - old))
        }
    }
}

/// Metrics for the persist sink. (While this lies slightly outside the usual
/// abstraction boundary of the client, it's convenient to manage them together.
#[derive(Debug, Clone)]
pub struct SinkMetrics {
    /// Cumulative record insertions made to the correction buffer across workers
    correction_insertions_total: IntCounter,
    /// Cumulative record deletions made to the correction buffer across workers
    correction_deletions_total: IntCounter,
    /// Cumulative capacity increases made to the correction buffer across workers
    correction_capacity_increases_total: IntCounter,
    /// Cumulative capacity decreases made to the correction buffer across workers
    correction_capacity_decreases_total: IntCounter,
    /// Maximum length observed for any one correction buffer per worker
    correction_max_per_sink_worker_len_updates: raw::UIntGaugeVec,
    /// Maximum capacity observed for any one correction buffer per worker
    correction_max_per_sink_worker_capacity_updates: raw::UIntGaugeVec,
}

impl SinkMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        SinkMetrics {
            correction_insertions_total: registry.register(metric!(
                name: "mz_persist_sink_correction_insertions_total",
                help: "The cumulative insertions observed on the correction buffer across workers and persist sinks.",
            )),
            correction_deletions_total: registry.register(metric!(
                name: "mz_persist_sink_correction_deletions_total",
                help: "The cumulative deletions observed on the correction buffer across workers and persist sinks.",
            )),
            correction_capacity_increases_total: registry.register(metric!(
                name: "mz_persist_sink_correction_capacity_increases_total",
                help: "The cumulative capacity increases observed on the correction buffer across workers and persist sinks.",
            )),
            correction_capacity_decreases_total: registry.register(metric!(
                name: "mz_persist_sink_correction_capacity_decreases_total",
                help: "The cumulative capacity decreases observed on the correction buffer across workers and persist sinks.",
            )),
            correction_max_per_sink_worker_len_updates: registry.register(metric!(
                name: "mz_persist_sink_correction_max_per_sink_worker_len_updates",
                help: "The maximum length observed for the correction buffer of any single persist sink per worker.",
                var_labels: ["worker_id"],
            )),
            correction_max_per_sink_worker_capacity_updates: registry.register(metric!(
                name: "mz_persist_sink_correction_max_per_sink_worker_capacity_updates",
                help: "The maximum capacity observed for the correction buffer of any single persist sink per worker.",
                var_labels: ["worker_id"],
            )),
        }
    }

    /// Obtains a `SinkWorkerMetrics` instance, which allows for metric reporting
    /// from a specific `persist_sink` instance for a given worker. The reports will
    /// update metrics shared across workers, but provide per-worker contributions
    /// to them.
    pub fn for_worker(&self, worker_id: usize) -> SinkWorkerMetrics {
        let worker = worker_id.to_string();
        let correction_max_per_sink_worker_len_updates = self
            .correction_max_per_sink_worker_len_updates
            .with_label_values(&[&worker]);
        let correction_max_per_sink_worker_capacity_updates = self
            .correction_max_per_sink_worker_capacity_updates
            .with_label_values(&[&worker]);
        SinkWorkerMetrics {
            correction_max_per_sink_worker_len_updates,
            correction_max_per_sink_worker_capacity_updates,
        }
    }

    /// Reports updates to the length and capacity of the correction buffer in the
    /// `write_batches` operator of a `persist_sink`.
    ///
    /// This method updates monotonic metrics based on the deltas and thus can be
    /// called across workers and instances of `persist_sink`.
    pub fn report_correction_update_deltas(
        &self,
        correction_len_delta: UpdateDelta,
        correction_cap_delta: UpdateDelta,
    ) {
        // Report insertions or deletions.
        match correction_len_delta {
            UpdateDelta::NonNegative(delta) => {
                if delta > 0 {
                    self.correction_insertions_total.inc_by(delta)
                }
            }
            UpdateDelta::Negative(delta) => self.correction_deletions_total.inc_by(delta),
        }
        // Report capacity increases or decreases.
        match correction_cap_delta {
            UpdateDelta::NonNegative(delta) => {
                if delta > 0 {
                    self.correction_capacity_increases_total.inc_by(delta)
                }
            }
            UpdateDelta::Negative(delta) => self.correction_capacity_decreases_total.inc_by(delta),
        }
    }
}

/// Metrics for the persist sink that are labeled per-worker.
#[derive(Clone, Debug)]
pub struct SinkWorkerMetrics {
    correction_max_per_sink_worker_len_updates: UIntGauge,
    correction_max_per_sink_worker_capacity_updates: UIntGauge,
}

impl SinkWorkerMetrics {
    /// Reports the length and capacity of the correction buffer in the `write_batches`
    /// operator of `persist_sink`.
    ///
    /// This method is used to update metrics that are kept per worker.
    pub fn report_correction_update_totals(&self, correction_len: usize, correction_cap: usize) {
        // Maintain per-worker peaks.
        let correction_len = CastFrom::cast_from(correction_len);
        if correction_len > self.correction_max_per_sink_worker_len_updates.get() {
            self.correction_max_per_sink_worker_len_updates
                .set(correction_len);
        }
        let correction_cap = CastFrom::cast_from(correction_cap);
        if correction_cap > self.correction_max_per_sink_worker_capacity_updates.get() {
            self.correction_max_per_sink_worker_capacity_updates
                .set(correction_cap);
        }
    }
}

/// Metrics for the `persist_source` backpressure operator, summed over every
/// instance of the operator in the process. Like [SinkMetrics], these belong
/// to a dataflow operator rather than the client, but the client owns the only
/// registry the operator can reach.
#[derive(Debug)]
pub struct BackpressureMetrics {
    /// Bytes emitted by backpressure operators.
    pub emitted_bytes: IntCounter,
    /// Sum over live operator instances of the inflight bytes each one most
    /// recently stalled on. Instances contribute deltas, so the sum stays exact
    /// as operators start and stop.
    pub last_backpressured_bytes: UIntGauge,
    /// Bytes retired by processing downstream of backpressure operators.
    pub retired_bytes: IntCounter,
}

impl BackpressureMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        BackpressureMetrics {
            emitted_bytes: registry.register(metric!(
                name: "mz_persist_backpressure_emitted_bytes",
                help: "bytes emitted by backpressure operators",
            )),
            last_backpressured_bytes: registry.register(metric!(
                name: "mz_persist_backpressure_last_backpressured_bytes",
                help: "sum over backpressure operators of the inflight bytes each last \
                    stalled on; not comparable to retired bytes, but nonzero growth \
                    indicates backpressure is happening",
            )),
            retired_bytes: registry.register(metric!(
                name: "mz_persist_backpressure_retired_bytes",
                help: "bytes retired by processing downstream of backpressure operators",
            )),
        }
    }
}

/// A minimal set of metrics imported into honeycomb for alerting.
#[derive(Debug)]
pub struct AlertsMetrics {
    pub(crate) blob_failures: IntCounter,
    pub(crate) consensus_failures: IntCounter,
}

impl AlertsMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        AlertsMetrics {
            blob_failures: registry.register(metric!(
                name: "mz_persist_blob_failures",
                help: "count of all blob operation failures",
                const_labels: {"honeycomb" => "import"},
            )),
            consensus_failures: registry.register(metric!(
                name: "mz_persist_consensus_failures",
                help: "count of determinate consensus operation failures",
                const_labels: {"honeycomb" => "import"},
            )),
        }
    }
}

/// Metrics for the PubSubServer implementation.
#[derive(Debug)]
pub struct PubSubServerMetrics {
    pub(crate) active_connections: UIntGauge,
    pub(crate) broadcasted_diff_count: IntCounter,
    pub(crate) broadcasted_diff_bytes: IntCounter,
    pub(crate) broadcasted_diff_dropped_channel_full: IntCounter,

    pub(crate) push_seconds: Counter,
    pub(crate) subscribe_seconds: Counter,
    pub(crate) unsubscribe_seconds: Counter,
    pub(crate) connection_cleanup_seconds: Counter,

    pub(crate) push_call_count: IntCounter,
    pub(crate) subscribe_call_count: IntCounter,
    pub(crate) unsubscribe_call_count: IntCounter,
}

impl PubSubServerMetrics {
    pub(crate) fn new(registry: &MetricsRegistry) -> Self {
        let op_timings: CounterVec = registry.register(metric!(
                name: "mz_persist_pubsub_server_operation_seconds",
                help: "time spent in pubsub server performing each operation",
                var_labels: ["op"],
        ));
        let call_count: IntCounterVec = registry.register(metric!(
                name: "mz_persist_pubsub_server_call_count",
                help: "count of each pubsub server message received",
                var_labels: ["call"],
        ));

        Self {
            active_connections: registry.register(metric!(
                    name: "mz_persist_pubsub_server_active_connections",
                    help: "number of active connections to server",
            )),
            broadcasted_diff_count: registry.register(metric!(
                    name: "mz_persist_pubsub_server_broadcasted_diff_count",
                    help: "count of total broadcast diff messages sent",
            )),
            broadcasted_diff_bytes: registry.register(metric!(
                    name: "mz_persist_pubsub_server_broadcasted_diff_bytes",
                    help: "count of total broadcast diff bytes sent",
            )),
            broadcasted_diff_dropped_channel_full: registry.register(metric!(
                    name: "mz_persist_pubsub_server_broadcasted_diff_dropped_channel_full",
                    help: "count of diffs dropped due to full connection channel",
            )),

            push_seconds: op_timings.with_label_values(&["push"]),
            subscribe_seconds: op_timings.with_label_values(&["subscribe"]),
            unsubscribe_seconds: op_timings.with_label_values(&["unsubscribe"]),
            connection_cleanup_seconds: op_timings.with_label_values(&["cleanup"]),

            push_call_count: call_count.with_label_values(&["push"]),
            subscribe_call_count: call_count.with_label_values(&["subscribe"]),
            unsubscribe_call_count: call_count.with_label_values(&["unsubscribe"]),
        }
    }
}

/// Metrics for the PubSubClient implementation.
#[derive(Debug)]
pub struct PubSubClientMetrics {
    pub sender: PubSubClientSenderMetrics,
    pub receiver: PubSubClientReceiverMetrics,
    pub grpc_connection: PubSubGrpcClientConnectionMetrics,
}

impl PubSubClientMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        PubSubClientMetrics {
            sender: PubSubClientSenderMetrics::new(registry),
            receiver: PubSubClientReceiverMetrics::new(registry),
            grpc_connection: PubSubGrpcClientConnectionMetrics::new(registry),
        }
    }
}

#[derive(Debug)]
pub struct PubSubGrpcClientConnectionMetrics {
    pub(crate) connected: UIntGauge,
    pub(crate) connection_established_count: IntCounter,
    pub(crate) connect_call_attempt_count: IntCounter,
    pub(crate) broadcast_recv_lagged_count: IntCounter,
    pub(crate) grpc_error_count: IntCounter,
}

impl PubSubGrpcClientConnectionMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        Self {
            connected: registry.register(metric!(
                    name: "mz_persist_pubsub_client_grpc_connected",
                    help: "whether the grpc client is currently connected",
            )),
            connection_established_count: registry.register(metric!(
                    name: "mz_persist_pubsub_client_grpc_connection_established_count",
                    help: "count of grpc connection establishments to pubsub server",
            )),
            connect_call_attempt_count: registry.register(metric!(
                    name: "mz_persist_pubsub_client_grpc_connect_call_attempt_count",
                    help: "count of connection call attempts (including retries) to pubsub server",
            )),
            broadcast_recv_lagged_count: registry.register(metric!(
                    name: "mz_persist_pubsub_client_grpc_broadcast_recv_lagged_count",
                    help: "times a message was missed by broadcast receiver due to lag",
            )),
            grpc_error_count: registry.register(metric!(
                    name: "mz_persist_pubsub_client_grpc_error_count",
                    help: "count of grpc errors received",
            )),
        }
    }
}

#[derive(Clone, Debug)]
pub struct PubSubClientReceiverMetrics {
    pub(crate) push_received: IntCounter,
    pub(crate) unknown_message_received: IntCounter,
    pub(crate) approx_diff_latency_seconds: Histogram,

    pub(crate) state_pushed_diff_fast_path: IntCounter,
    pub(crate) state_pushed_diff_slow_path_succeeded: IntCounter,
    pub(crate) state_pushed_diff_slow_path_failed: IntCounter,

    pub(crate) diff_applied: IntCounter,
    pub(crate) diff_not_applied_stale: IntCounter,
    pub(crate) diff_not_applied_out_of_order: IntCounter,
}

impl PubSubClientReceiverMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        let call_received: IntCounterVec = registry.register(metric!(
                name: "mz_persist_pubsub_client_call_received",
                help: "times a pubsub client call was received",
                var_labels: ["call"],
        ));

        Self {
            push_received: call_received.with_label_values(&["push"]),
            unknown_message_received: call_received.with_label_values(&["unknown"]),
            approx_diff_latency_seconds: registry.register(metric!(
                name: "mz_persist_pubsub_client_approx_diff_apply_latency_seconds",
                help: "histogram of (approximate) latency between sending a diff and applying it",
                buckets: prometheus::exponential_buckets(0.001, 2.0, 13).expect("buckets"),
            )),

            state_pushed_diff_fast_path: registry.register(metric!(
                name: "mz_persist_pubsub_client_receiver_state_push_diff_fast_path",
                help: "count fast-path state push_diff calls",
            )),
            state_pushed_diff_slow_path_succeeded: registry.register(metric!(
                name: "mz_persist_pubsub_client_receiver_state_push_diff_slow_path_succeeded",
                help: "count of successful slow-path state push_diff calls",
            )),
            state_pushed_diff_slow_path_failed: registry.register(metric!(
                name: "mz_persist_pubsub_client_receiver_state_push_diff_slow_path_failed",
                help: "count of unsuccessful slow-path state push_diff calls",
            )),
            diff_applied: registry.register(metric!(
                name: "mz_persist_pubsub_client_receiver_diff_applied",
                help: "number of diffs received via pubsub that applied",
            )),
            diff_not_applied_stale: registry.register(metric!(
                name: "mz_persist_pubsub_client_receiver_diff_not_applied_stale",
                help: "number of diffs received via pubsub that did not apply due to staleness",
            )),
            diff_not_applied_out_of_order: registry.register(metric!(
                name: "mz_persist_pubsub_client_receiver_diff_not_applied_out_of_order",
                help: "number of diffs received via pubsub that did not apply due to out-of-order delivery",
            )),
        }
    }
}

#[derive(Debug)]
pub struct PubSubClientSenderMetrics {
    pub push: PubSubClientCallMetrics,
    pub subscribe: PubSubClientCallMetrics,
    pub unsubscribe: PubSubClientCallMetrics,
}

#[derive(Debug)]
pub struct PubSubClientCallMetrics {
    pub(crate) succeeded: IntCounter,
    pub(crate) bytes_sent: IntCounter,
    pub(crate) failed: IntCounter,
}

impl PubSubClientSenderMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        let call_bytes_sent: IntCounterVec = registry.register(metric!(
                name: "mz_persist_pubsub_client_call_bytes_sent",
                help: "number of bytes sent for a given pubsub client call",
                var_labels: ["call"],
        ));
        let call_succeeded: IntCounterVec = registry.register(metric!(
                name: "mz_persist_pubsub_client_call_succeeded",
                help: "times a pubsub client call succeeded",
                var_labels: ["call"],
        ));
        let call_failed: IntCounterVec = registry.register(metric!(
                name: "mz_persist_pubsub_client_call_failed",
                help: "times a pubsub client call failed",
                var_labels: ["call"],
        ));

        Self {
            push: PubSubClientCallMetrics {
                succeeded: call_succeeded.with_label_values(&["push"]),
                failed: call_failed.with_label_values(&["push"]),
                bytes_sent: call_bytes_sent.with_label_values(&["push"]),
            },
            subscribe: PubSubClientCallMetrics {
                succeeded: call_succeeded.with_label_values(&["subscribe"]),
                failed: call_failed.with_label_values(&["subscribe"]),
                bytes_sent: call_bytes_sent.with_label_values(&["subscribe"]),
            },
            unsubscribe: PubSubClientCallMetrics {
                succeeded: call_succeeded.with_label_values(&["unsubscribe"]),
                failed: call_failed.with_label_values(&["unsubscribe"]),
                bytes_sent: call_bytes_sent.with_label_values(&["unsubscribe"]),
            },
        }
    }
}

#[derive(Debug)]
pub struct LocksMetrics {
    pub(crate) applier_read_cacheable: LockMetrics,
    pub(crate) applier_read_noncacheable: LockMetrics,
    pub(crate) applier_write: LockMetrics,
    pub(crate) watch: LockMetrics,
}

#[derive(Debug, Clone)]
pub struct LockMetrics {
    pub(crate) acquire_count: IntCounter,
    pub(crate) blocking_acquire_count: IntCounter,
    pub(crate) blocking_seconds: Counter,
}

#[derive(Debug)]
pub struct WatchMetrics {
    pub(crate) wait_woken_via_watch: IntCounter,
    pub(crate) wait_woken_via_sleep: IntCounter,
    pub(crate) wait_resolved_via_watch: IntCounter,
    pub(crate) wait_resolved_via_sleep: IntCounter,
    pub(crate) notify_sent: IntCounter,
    pub(crate) notify_upper_sent: IntCounter,
    pub(crate) notify_noop: IntCounter,
    pub(crate) notify_recv: IntCounter,
    pub(crate) notify_lagged: IntCounter,
    pub(crate) notify_wait_started: IntCounter,
    pub(crate) notify_wait_finished: IntCounter,
}

impl WatchMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        WatchMetrics {
            wait_woken_via_watch: registry.register(metric!(
                name: "mz_persist_wait_woken_via_watch",
                help: "count of wait-for-uppers wakes via watch notify",
            )),
            wait_woken_via_sleep: registry.register(metric!(
                name: "mz_persist_wait_woken_via_sleep",
                help: "count of wait-for-uppers wakes via sleep",
            )),
            wait_resolved_via_watch: registry.register(metric!(
                name: "mz_persist_wait_resolved_via_watch",
                help: "count of wait-for-uppers resolved via watch notify",
            )),
            wait_resolved_via_sleep: registry.register(metric!(
                name: "mz_persist_wait_resolved_via_sleep",
                help: "count of wait-for-uppers resolved via sleep",
            )),
            notify_sent: registry.register(metric!(
                name: "mz_persist_watch_notify_sent",
                help: "count of watch notifications sent to a non-empty broadcast channel",
            )),
            notify_upper_sent: registry.register(metric!(
                name: "mz_persist_watch_notify_upper_sent",
                help: "count of strict shard upper advances signaled to upper waiters",
            )),
            notify_noop: registry.register(metric!(
                name: "mz_persist_watch_notify_noop",
                help: "count of watch notifications sent to an broadcast channel",
            )),
            notify_recv: registry.register(metric!(
                name: "mz_persist_watch_notify_recv",
                help: "count of watch notifications received from the broadcast channel",
            )),
            notify_lagged: registry.register(metric!(
                name: "mz_persist_watch_notify_lagged",
                help: "count of lagged events in the watch notification broadcast channel",
            )),
            notify_wait_started: registry.register(metric!(
                name: "mz_persist_watch_notify_wait_started",
                help: "count of watch wait calls started",
            )),
            notify_wait_finished: registry.register(metric!(
                name: "mz_persist_watch_notify_wait_finished",
                help: "count of watch wait calls resolved",
            )),
        }
    }
}

#[derive(Debug)]
pub struct PushdownMetrics {
    pub(crate) parts_filtered_count: IntCounter,
    pub(crate) parts_filtered_bytes: IntCounter,
    pub(crate) parts_fetched_count: IntCounter,
    pub(crate) parts_fetched_bytes: IntCounter,
    pub(crate) parts_audited_count: IntCounter,
    pub(crate) parts_audited_bytes: IntCounter,
    pub(crate) parts_inline_count: IntCounter,
    pub(crate) parts_inline_bytes: IntCounter,
    pub(crate) parts_faked_count: IntCounter,
    pub(crate) parts_faked_bytes: IntCounter,
    pub(crate) parts_stats_trimmed_count: IntCounter,
    pub(crate) parts_stats_trimmed_bytes: IntCounter,
    pub(crate) parts_projection_trimmed_bytes: IntCounter,
    pub part_stats: PartStatsMetrics,
}

impl PushdownMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        PushdownMetrics {
            parts_filtered_count: registry.register(metric!(
                name: "mz_persist_pushdown_parts_filtered_count",
                help: "count of parts filtered by pushdown",
            )),
            parts_filtered_bytes: registry.register(metric!(
                name: "mz_persist_pushdown_parts_filtered_bytes",
                help: "total size of parts filtered by pushdown in bytes",
            )),
            parts_fetched_count: registry.register(metric!(
                name: "mz_persist_pushdown_parts_fetched_count",
                help: "count of parts not filtered by pushdown",
            )),
            parts_fetched_bytes: registry.register(metric!(
                name: "mz_persist_pushdown_parts_fetched_bytes",
                help: "total size of parts not filtered by pushdown in bytes",
            )),
            parts_audited_count: registry.register(metric!(
                name: "mz_persist_pushdown_parts_audited_count",
                help: "count of parts fetched only for pushdown audit",
            )),
            parts_audited_bytes: registry.register(metric!(
                name: "mz_persist_pushdown_parts_audited_bytes",
                help: "total size of parts fetched only for pushdown audit",
            )),
            parts_inline_count: registry.register(metric!(
                name: "mz_persist_pushdown_parts_inline_count",
                help: "count of parts not fetched because they were inline",
            )),
            parts_inline_bytes: registry.register(metric!(
                name: "mz_persist_pushdown_parts_inline_bytes",
                help: "total size of parts not fetched because they were inline",
            )),
            parts_faked_count: registry.register(metric!(
                name: "mz_persist_pushdown_parts_faked_count",
                help: "count of parts faked because of aggressive projection pushdown",
            )),
            parts_faked_bytes: registry.register(metric!(
                name: "mz_persist_pushdown_parts_faked_bytes",
                help: "total size of parts replaced with fakes by aggressive projection pushdown",
            )),
            parts_stats_trimmed_count: registry.register(metric!(
                name: "mz_persist_pushdown_parts_stats_trimmed_count",
                help: "count of trimmed part stats",
            )),
            parts_stats_trimmed_bytes: registry.register(metric!(
                name: "mz_persist_pushdown_parts_stats_trimmed_bytes",
                help: "total bytes trimmed from part stats",
            )),
            parts_projection_trimmed_bytes: registry.register(metric!(
                name: "mz_persist_pushdown_parts_projection_trimmed_bytes",
                help: "total bytes trimmed from columnar data because of projection pushdown",
            )),
            part_stats: PartStatsMetrics::new(registry),
        }
    }
}

#[derive(Debug)]
pub struct ConsolidationMetrics {
    pub(crate) parts_fetched: IntCounter,
    pub(crate) parts_skipped: IntCounter,
    pub(crate) parts_wasted: IntCounter,
    pub(crate) wrong_sort: IntCounter,
}

impl ConsolidationMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        ConsolidationMetrics {
            parts_fetched: registry.register(metric!(
                name: "mz_persist_consolidation_parts_fetched_count",
                help: "count of parts that were fetched and used during consolidation",
            )),
            parts_skipped: registry.register(metric!(
                name: "mz_persist_consolidation_parts_skipped_count",
                help: "count of parts that were never needed during consolidation",
            )),
            parts_wasted: registry.register(metric!(
                name: "mz_persist_consolidation_parts_wasted_count",
                help: "count of parts that were fetched but not needed during consolidation",
            )),
            wrong_sort: registry.register(metric!(
                name: "mz_persist_consolidation_wrong_sort_count",
                help: "count of runs that were sorted using the wrong ordering for the current consolidation",
            )),
        }
    }
}

#[derive(Debug)]
pub struct BlobMemCache {
    pub(crate) size_blobs: UIntGauge,
    pub(crate) size_bytes: UIntGauge,
    pub(crate) hits_blobs: IntCounter,
    pub(crate) hits_bytes: IntCounter,
    pub(crate) evictions: IntCounter,
}

impl BlobMemCache {
    fn new(registry: &MetricsRegistry) -> Self {
        BlobMemCache {
            size_blobs: registry.register(metric!(
                name: "mz_persist_blob_cache_size_blobs",
                help: "count of blobs in the cache",
                const_labels: {"cache" => "mem"},
            )),
            size_bytes: registry.register(metric!(
                name: "mz_persist_blob_cache_size_bytes",
                help: "total size of blobs in the cache",
                const_labels: {"cache" => "mem"},
            )),
            hits_blobs: registry.register(metric!(
                name: "mz_persist_blob_cache_hits_blobs",
                help: "count of blobs served via cache instead of s3",
                const_labels: {"cache" => "mem"},
            )),
            hits_bytes: registry.register(metric!(
                name: "mz_persist_blob_cache_hits_bytes",
                help: "total size of blobs served via cache instead of s3",
                const_labels: {"cache" => "mem"},
            )),
            evictions: registry.register(metric!(
                name: "mz_persist_blob_cache_evictions",
                help: "count of capacity-based cache evictions",
                const_labels: {"cache" => "mem"},
            )),
        }
    }
}

#[derive(Debug)]
pub struct SemaphoreMetrics {
    cfg: PersistConfig,
    registry: MetricsRegistry,
    fetch: OnceCell<MetricsSemaphore>,
}

impl SemaphoreMetrics {
    fn new(cfg: PersistConfig, registry: MetricsRegistry) -> Self {
        SemaphoreMetrics {
            cfg,
            registry,
            fetch: OnceCell::new(),
        }
    }

    /// We can't easily change the number of permits, and the dyncfgs are all
    /// set to defaults on process start, so make sure we only initialize the
    /// semaphore once we've synced dyncfgs at least once.
    async fn fetch(&self) -> &MetricsSemaphore {
        if let Some(x) = self.fetch.get() {
            // Common case of already initialized avoids the cloning below.
            return x;
        }
        let cfg = self.cfg.clone();
        let registry = self.registry.clone();
        let init = async move {
            let total_permits = match cfg.announce_memory_limit {
                // Non-cc replicas have the old physical flow control mechanism,
                // so only apply this one on cc replicas.
                Some(mem) if cfg.is_cc_active => {
                    // We can't easily adjust the number of permits later, so
                    // make sure we've synced dyncfg values at least once.
                    info!("fetch semaphore awaiting first dyncfg values");
                    let () = cfg.configs_synced_once().await;
                    let total_permits = usize::cast_lossy(
                        f64::cast_lossy(mem) * FETCH_SEMAPHORE_PERMIT_ADJUSTMENT.get(&cfg),
                    );
                    info!("fetch_semaphore got first dyncfg values");
                    total_permits
                }
                Some(_) | None => Semaphore::MAX_PERMITS,
            };
            MetricsSemaphore::new(&registry, "fetch", total_permits)
        };
        self.fetch.get_or_init(|| init).await
    }

    pub(crate) async fn acquire_fetch_permits(&self, encoded_size_bytes: usize) -> MetricsPermits {
        // Adjust the requested permits to account for the difference between
        // encoded_size_bytes and the decoded size in lgalloc.
        let requested_permits = f64::cast_lossy(encoded_size_bytes);
        let requested_permits = requested_permits * FETCH_SEMAPHORE_COST_ADJUSTMENT.get(&self.cfg);
        let requested_permits = usize::cast_lossy(requested_permits);
        self.fetch().await.acquire_permits(requested_permits).await
    }
}

#[derive(Debug)]
pub struct MetricsSemaphore {
    name: &'static str,
    semaphore: Arc<Semaphore>,
    total_permits: usize,
    acquire_count: IntCounter,
    blocking_count: IntCounter,
    blocking_seconds: Counter,
    acquired_permits: IntCounter,
    released_permits: IntCounter,
    _available_permits: ComputedUIntGauge,
}

impl MetricsSemaphore {
    pub fn new(registry: &MetricsRegistry, name: &'static str, total_permits: usize) -> Self {
        let total_permits = std::cmp::min(total_permits, Semaphore::MAX_PERMITS);
        // TODO: Sadly, tokio::sync::Semaphore makes it difficult to have a
        // dynamic total_permits count.
        let semaphore = Arc::new(Semaphore::new(total_permits));
        MetricsSemaphore {
            name,
            total_permits,
            acquire_count: registry.register(metric!(
                name: "mz_persist_semaphore_acquire_count",
                help: "count of acquire calls (not acquired permits count)",
                const_labels: {"name" => name},
            )),
            blocking_count: registry.register(metric!(
                name: "mz_persist_semaphore_blocking_count",
                help: "count of acquire calls that had to block",
                const_labels: {"name" => name},
            )),
            blocking_seconds: registry.register(metric!(
                name: "mz_persist_semaphore_blocking_seconds",
                help: "total time spent blocking on permit acquisition",
                const_labels: {"name" => name},
            )),
            acquired_permits: registry.register(metric!(
                name: "mz_persist_semaphore_acquired_permits",
                help: "total sum of acquired permits",
                const_labels: {"name" => name},
            )),
            released_permits: registry.register(metric!(
                name: "mz_persist_semaphore_released_permits",
                help: "total sum of released permits",
                const_labels: {"name" => name},
            )),
            _available_permits: registry.register_computed_gauge(
                metric!(
                    name: "mz_persist_semaphore_available_permits",
                    help: "currently available permits according to the semaphore",
                ),
                {
                    let semaphore = Arc::clone(&semaphore);
                    move || u64::cast_from(semaphore.available_permits())
                },
            ),
            semaphore,
        }
    }

    pub async fn acquire_permits(&self, requested_permits: usize) -> MetricsPermits {
        // HACK: Cap the request at the total permit count. This prevents
        // deadlock, even if the cfg gets set to some small value.
        let total_permits = u32::try_from(self.total_permits).unwrap_or(u32::MAX);
        let requested_permits = u32::try_from(requested_permits).unwrap_or(u32::MAX);
        let requested_permits = std::cmp::min(requested_permits, total_permits);
        let wrap = |_permit| {
            self.acquired_permits.inc_by(u64::from(requested_permits));
            MetricsPermits {
                _permit,
                released_metric: self.released_permits.clone(),
                count: requested_permits,
            }
        };

        // Special-case non-blocking happy path.
        self.acquire_count.inc();
        match Arc::clone(&self.semaphore).try_acquire_many_owned(requested_permits) {
            Ok(x) => return wrap(x),
            Err(_) => {}
        };

        // Sad path, gotta block.
        self.blocking_count.inc();
        let start = Instant::now();
        let ret = Arc::clone(&self.semaphore)
            .acquire_many_owned(requested_permits)
            .instrument(info_span!("acquire_permits"))
            .await;
        let elapsed = start.elapsed();
        self.blocking_seconds.inc_by(elapsed.as_secs_f64());
        debug!(
            "acquisition of {} {} permits blocked for {:?}",
            self.name, requested_permits, elapsed
        );
        wrap(ret.expect("semaphore is never closed"))
    }
}

#[derive(Debug)]
pub struct MetricsPermits {
    _permit: OwnedSemaphorePermit,
    released_metric: IntCounter,
    count: u32,
}

impl Drop for MetricsPermits {
    fn drop(&mut self) {
        self.released_metric.inc_by(u64::from(self.count))
    }
}

#[derive(Debug)]
pub struct ExternalOpMetrics {
    started: IntCounter,
    succeeded: IntCounter,
    failed: IntCounter,
    bytes: IntCounter,
    seconds: Counter,
    seconds_histogram: Option<Histogram>,
    alerts_metrics: Arc<AlertsMetrics>,
}

impl ExternalOpMetrics {
    async fn run_op<R, F, OpFn, ErrFn>(
        &self,
        op_fn: OpFn,
        on_err_fn: ErrFn,
    ) -> Result<R, ExternalError>
    where
        F: std::future::Future<Output = Result<R, ExternalError>>,
        OpFn: FnOnce() -> F,
        ErrFn: FnOnce(&AlertsMetrics, &ExternalError),
    {
        self.started.inc();
        let start = Instant::now();
        let res = op_fn().await;
        let elapsed_seconds = start.elapsed().as_secs_f64();
        self.seconds.inc_by(elapsed_seconds);
        if let Some(h) = &self.seconds_histogram {
            h.observe(elapsed_seconds);
        }
        match res.as_ref() {
            Ok(_) => self.succeeded.inc(),
            Err(err) => {
                self.failed.inc();
                on_err_fn(&self.alerts_metrics, err);
            }
        };
        res
    }

    fn run_stream<'a, R: 'a, S, OpFn, ErrFn>(
        &'a self,
        op_fn: OpFn,
        mut on_err_fn: ErrFn,
    ) -> impl futures::Stream<Item = Result<R, ExternalError>> + 'a
    where
        S: futures::Stream<Item = Result<R, ExternalError>> + Unpin + 'a,
        OpFn: FnOnce() -> S,
        ErrFn: FnMut(&AlertsMetrics, &ExternalError) + 'a,
    {
        self.started.inc();
        let start = Instant::now();
        let mut stream = op_fn();
        stream! {
            let mut succeeded = true;
            while let Some(res) = stream.next().await {
                if let Err(err) = res.as_ref() {
                    on_err_fn(&self.alerts_metrics, err);
                    succeeded = false;
                }
                yield res;
            }
            if succeeded {
                self.succeeded.inc()
            } else {
                self.failed.inc()
            }
            let elapsed_seconds = start.elapsed().as_secs_f64();
            self.seconds.inc_by(elapsed_seconds);
            if let Some(h) = &self.seconds_histogram {
                h.observe(elapsed_seconds);
            }
        }
    }
}

#[derive(Debug)]
pub struct BlobMetrics {
    set: ExternalOpMetrics,
    get: ExternalOpMetrics,
    list_keys: ExternalOpMetrics,
    delete: ExternalOpMetrics,
    restore: ExternalOpMetrics,
    delete_noop: IntCounter,
    blob_sizes: Histogram,
    pub rtt_latency: Gauge,
}

#[derive(Debug)]
pub struct MetricsBlob {
    blob: Arc<dyn Blob>,
    metrics: Arc<Metrics>,
}

impl MetricsBlob {
    pub fn new(blob: Arc<dyn Blob>, metrics: Arc<Metrics>) -> Self {
        MetricsBlob { blob, metrics }
    }

    fn on_err(alerts_metrics: &AlertsMetrics, _err: &ExternalError) {
        alerts_metrics.blob_failures.inc()
    }
}

#[async_trait]
impl Blob for MetricsBlob {
    #[instrument(name = "blob::get", fields(shard=blob_key_shard_id(key)))]
    async fn get(&self, key: &str) -> Result<Option<SegmentedBytes>, ExternalError> {
        let res = self
            .metrics
            .blob
            .get
            .run_op(|| self.blob.get(key), Self::on_err)
            .await;
        if let Ok(Some(value)) = res.as_ref() {
            self.metrics
                .blob
                .get
                .bytes
                .inc_by(u64::cast_from(value.len()));
        }
        res
    }

    #[instrument(name = "blob::list_keys_and_metadata", fields(shard=blob_key_shard_id(key_prefix)))]
    async fn list_keys_and_metadata(
        &self,
        key_prefix: &str,
        f: &mut (dyn FnMut(BlobMetadata) + Send + Sync),
    ) -> Result<(), ExternalError> {
        let mut byte_total = 0;
        let mut instrumented = |blob_metadata: BlobMetadata| {
            // Track the size of the _keys_, not the blobs, so that we get a
            // sense for how much network bandwidth these calls are using.
            byte_total += blob_metadata.key.len();
            f(blob_metadata)
        };

        let res = self
            .metrics
            .blob
            .list_keys
            .run_op(
                || {
                    self.blob
                        .list_keys_and_metadata(key_prefix, &mut instrumented)
                },
                Self::on_err,
            )
            .await;

        self.metrics
            .blob
            .list_keys
            .bytes
            .inc_by(u64::cast_from(byte_total));

        res
    }

    #[instrument(name = "blob::set", fields(shard=blob_key_shard_id(key),size_bytes=value.len()))]
    async fn set(&self, key: &str, value: Bytes) -> Result<(), ExternalError> {
        let bytes = value.len();
        let res = self
            .metrics
            .blob
            .set
            .run_op(|| self.blob.set(key, value), Self::on_err)
            .await;
        if res.is_ok() {
            self.metrics.blob.set.bytes.inc_by(u64::cast_from(bytes));
            self.metrics.blob.blob_sizes.observe(f64::cast_lossy(bytes));
        }
        res
    }

    #[instrument(name = "blob::delete", fields(shard=blob_key_shard_id(key)))]
    async fn delete(&self, key: &str) -> Result<Option<usize>, ExternalError> {
        let bytes = self
            .metrics
            .blob
            .delete
            .run_op(|| self.blob.delete(key), Self::on_err)
            .await?;
        if let Some(bytes) = bytes {
            self.metrics.blob.delete.bytes.inc_by(u64::cast_from(bytes));
        } else {
            self.metrics.blob.delete_noop.inc();
        }
        Ok(bytes)
    }

    async fn restore(&self, key: &str) -> Result<(), ExternalError> {
        self.metrics
            .blob
            .restore
            .run_op(|| self.blob.restore(key), Self::on_err)
            .await
    }
}

#[derive(Debug)]
pub struct ConsensusMetrics {
    list_keys: ExternalOpMetrics,
    head: ExternalOpMetrics,
    compare_and_set: ExternalOpMetrics,
    scan: ExternalOpMetrics,
    truncate: ExternalOpMetrics,
    truncated_count: IntCounter,
    pub rtt_latency: Gauge,
}

#[derive(Debug)]
pub struct MetricsConsensus {
    consensus: Arc<dyn Consensus>,
    metrics: Arc<Metrics>,
}

impl MetricsConsensus {
    pub fn new(consensus: Arc<dyn Consensus>, metrics: Arc<Metrics>) -> Self {
        MetricsConsensus { consensus, metrics }
    }

    fn on_err(alerts_metrics: &AlertsMetrics, err: &ExternalError) {
        // As of 2022-09-06, regular determinate errors are expected in
        // Consensus (i.e. "txn conflict, please retry"), so only count the
        // indeterminate ones.
        if let ExternalError::Indeterminate(_) = err {
            alerts_metrics.consensus_failures.inc()
        }
    }
}

#[async_trait]
impl Consensus for MetricsConsensus {
    fn list_keys(&self) -> ResultStream<'_, String> {
        Box::pin(
            self.metrics
                .consensus
                .list_keys
                .run_stream(|| self.consensus.list_keys(), Self::on_err),
        )
    }

    #[instrument(name = "consensus::head", fields(shard=key))]
    async fn head(&self, key: &str) -> Result<Option<VersionedData>, ExternalError> {
        let res = self
            .metrics
            .consensus
            .head
            .run_op(|| self.consensus.head(key), Self::on_err)
            .await;
        if let Ok(Some(data)) = res.as_ref() {
            self.metrics
                .consensus
                .head
                .bytes
                .inc_by(u64::cast_from(data.data.len()));
        }
        res
    }

    #[instrument(name = "consensus::compare_and_set", fields(shard=key,size_bytes=new.data.len()))]
    async fn compare_and_set(
        &self,
        key: &str,
        new: VersionedData,
    ) -> Result<CaSResult, ExternalError> {
        let bytes = new.data.len();
        let res = self
            .metrics
            .consensus
            .compare_and_set
            .run_op(|| self.consensus.compare_and_set(key, new), Self::on_err)
            .await;
        match res.as_ref() {
            Ok(CaSResult::Committed) => self
                .metrics
                .consensus
                .compare_and_set
                .bytes
                .inc_by(u64::cast_from(bytes)),
            Ok(CaSResult::ExpectationMismatch) | Err(_) => {}
        }
        res
    }

    #[instrument(name = "consensus::scan", fields(shard=key))]
    async fn scan(
        &self,
        key: &str,
        from: SeqNo,
        limit: usize,
    ) -> Result<Vec<VersionedData>, ExternalError> {
        let res = self
            .metrics
            .consensus
            .scan
            .run_op(|| self.consensus.scan(key, from, limit), Self::on_err)
            .await;
        if let Ok(dataz) = res.as_ref() {
            let bytes: usize = dataz.iter().map(|x| x.data.len()).sum();
            self.metrics
                .consensus
                .scan
                .bytes
                .inc_by(u64::cast_from(bytes));
        }
        res
    }

    #[instrument(name = "consensus::truncate", fields(shard=key))]
    async fn truncate(&self, key: &str, seqno: SeqNo) -> Result<Option<usize>, ExternalError> {
        let metrics = &self.metrics.consensus;
        let deleted = metrics
            .truncate
            .run_op(|| self.consensus.truncate(key, seqno), Self::on_err)
            .await?;
        if let Some(deleted) = deleted {
            metrics.truncated_count.inc_by(u64::cast_from(deleted));
        }
        Ok(deleted)
    }
}

/// A standard set of metrics for an async task. Call [TaskMetrics::instrument_task] to instrument
/// a future and report its metrics for this task type.
#[derive(Debug, Clone)]
pub struct TaskMetrics {
    f64_gauges: Vec<(Gauge, fn(&tokio_metrics::TaskMetrics) -> f64)>,
    u64_gauges: Vec<(
        GenericGauge<AtomicU64>,
        fn(&tokio_metrics::TaskMetrics) -> u64,
    )>,
    monitor: TaskMonitor,
}

impl TaskMetrics {
    pub fn new(name: &str) -> Self {
        let monitor = TaskMonitor::new();
        Self {
            f64_gauges: vec![
                (
                    Gauge::make_collector(metric!(
                        name: "mz_persist_task_total_idle_duration",
                        help: "Seconds of time spent idling, ie. waiting for a task to be woken up.",
                        const_labels: {"name" => name}
                    )),
                    |m| m.total_idle_duration.as_secs_f64(),
                ),
                (
                    Gauge::make_collector(metric!(
                        name: "mz_persist_task_total_scheduled_duration",
                        help: "Seconds of time spent scheduled, ie. ready to poll but not yet polled.",
                        const_labels: {"name" => name}
                    )),
                    |m| m.total_scheduled_duration.as_secs_f64(),
                ),
            ],
            u64_gauges: vec![
                (
                    MakeCollector::make_collector(metric!(
                        name: "mz_persist_task_total_scheduled_count",
                        help: "The total number of task schedules. Useful for computing the average scheduled time.",
                        const_labels: {"name" => name}
                    )),
                    |m| m.total_scheduled_count,
                ),
                (
                    MakeCollector::make_collector(metric!(
                        name: "mz_persist_task_total_idled_count",
                        help: "The total number of task idles. Useful for computing the average idle time.",
                        const_labels: {"name" => name}
                    ,
                    )),
                    |m| m.total_idled_count,
                ),
            ],
            monitor,
        }
    }

    /// Instrument the provided future. The expectation is that the result will be executed
    /// as a task. (See [TaskMonitor::instrument] for more context.)
    pub fn instrument_task<F>(&self, task: F) -> tokio_metrics::Instrumented<F> {
        TaskMonitor::instrument(&self.monitor, task)
    }
}

impl Collector for TaskMetrics {
    fn desc(&self) -> Vec<&Desc> {
        let mut descs = Vec::with_capacity(self.f64_gauges.len() + self.u64_gauges.len());
        for (g, _) in &self.f64_gauges {
            descs.extend(g.desc());
        }
        for (g, _) in &self.u64_gauges {
            descs.extend(g.desc());
        }
        descs
    }

    fn collect(&self) -> Vec<MetricFamily> {
        let mut families = Vec::with_capacity(self.f64_gauges.len() + self.u64_gauges.len());
        let metrics = self.monitor.cumulative();
        for (g, metrics_fn) in &self.f64_gauges {
            g.set(metrics_fn(&metrics));
            families.extend(g.collect());
        }
        for (g, metrics_fn) in &self.u64_gauges {
            g.set(metrics_fn(&metrics));
            families.extend(g.collect());
        }
        families
    }
}

#[derive(Debug)]
pub struct TasksMetrics {
    pub heartbeat_read: TaskMetrics,
}

impl TasksMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        let heartbeat_read = TaskMetrics::new("heartbeat_read");
        registry.register_collector(heartbeat_read.clone());
        TasksMetrics { heartbeat_read }
    }
}

#[derive(Debug)]
pub struct SchemaMetrics {
    pub(crate) cache_fetch_state_count: IntCounter,
    pub(crate) cache_schema: SchemaCacheMetrics,
    pub(crate) cache_migration: SchemaCacheMetrics,
    pub(crate) migration_count_same: IntCounter,
    pub(crate) migration_count_codec: IntCounter,
    pub(crate) migration_count_either: IntCounter,
    pub(crate) migration_len_legacy_codec: IntCounter,
    pub(crate) migration_len_either_codec: IntCounter,
    pub(crate) migration_len_either_arrow: IntCounter,
    pub(crate) migration_new_count: IntCounter,
    pub(crate) migration_new_seconds: Counter,
    pub(crate) migration_migrate_seconds: Counter,
}

impl SchemaMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        let cached: IntCounterVec = registry.register(metric!(
            name: "mz_persist_schema_cache_cached_count",
            help: "count of schema cache entries served from cache",
            var_labels: ["op"],
        ));
        let computed: IntCounterVec = registry.register(metric!(
            name: "mz_persist_schema_cache_computed_count",
            help: "count of schema cache entries computed",
            var_labels: ["op"],
        ));
        let unavailable: IntCounterVec = registry.register(metric!(
            name: "mz_persist_schema_cache_unavailable_count",
            help: "count of schema cache entries unavailable at current state",
            var_labels: ["op"],
        ));
        let added: IntCounterVec = registry.register(metric!(
            name: "mz_persist_schema_cache_added_count",
            help: "count of schema cache entries added",
            var_labels: ["op"],
        ));
        let dropped: IntCounterVec = registry.register(metric!(
            name: "mz_persist_schema_cache_dropped_count",
            help: "count of schema cache entries dropped",
            var_labels: ["op"],
        ));
        let cache = |name| SchemaCacheMetrics {
            cached_count: cached.with_label_values(&[name]),
            computed_count: computed.with_label_values(&[name]),
            unavailable_count: unavailable.with_label_values(&[name]),
            added_count: added.with_label_values(&[name]),
            dropped_count: dropped.with_label_values(&[name]),
        };
        let migration_count: IntCounterVec = registry.register(metric!(
            name: "mz_persist_schema_migration_count",
            help: "count of fetch part migrations",
            var_labels: ["op"],
        ));
        let migration_len: IntCounterVec = registry.register(metric!(
            name: "mz_persist_schema_migration_len",
            help: "count of migrated update records",
            var_labels: ["op"],
        ));
        SchemaMetrics {
            cache_fetch_state_count: registry.register(metric!(
                name: "mz_persist_schema_cache_fetch_state_count",
                help: "count of state fetches by the schema cache",
            )),
            cache_schema: cache("schema"),
            cache_migration: cache("migration"),
            migration_count_same: migration_count.with_label_values(&["same"]),
            migration_count_codec: migration_count.with_label_values(&["codec"]),
            migration_count_either: migration_count.with_label_values(&["either"]),
            migration_len_legacy_codec: migration_len.with_label_values(&["legacy_codec"]),
            migration_len_either_codec: migration_len.with_label_values(&["either_codec"]),
            migration_len_either_arrow: migration_len.with_label_values(&["either_arrow"]),
            migration_new_count: registry.register(metric!(
                name: "mz_persist_schema_migration_new_count",
                help: "count of migrations constructed",
            )),
            migration_new_seconds: registry.register(metric!(
                name: "mz_persist_schema_migration_new_seconds",
                help: "seconds spent constructing migration logic",
            )),
            migration_migrate_seconds: registry.register(metric!(
                name: "mz_persist_schema_migration_migrate_seconds",
                help: "seconds spent applying migration logic",
            )),
        }
    }
}

#[derive(Debug, Clone)]
pub struct SchemaCacheMetrics {
    pub(crate) cached_count: IntCounter,
    pub(crate) computed_count: IntCounter,
    pub(crate) unavailable_count: IntCounter,
    pub(crate) added_count: IntCounter,
    pub(crate) dropped_count: IntCounter,
}

#[derive(Debug)]
pub struct InlineMetrics {
    pub(crate) part_commit_count: IntCounter,
    pub(crate) part_commit_bytes: IntCounter,
    pub(crate) backpressure: BatchWriteMetrics,
}

impl InlineMetrics {
    fn new(registry: &MetricsRegistry) -> Self {
        InlineMetrics {
            part_commit_count: registry.register(metric!(
                name: "mz_persist_inline_part_commit_count",
                help: "count of inline parts committed to state",
            )),
            part_commit_bytes: registry.register(metric!(
                name: "mz_persist_inline_part_commit_bytes",
                help: "total size of of inline parts committed to state",
            )),
            backpressure: BatchWriteMetrics::new(registry, "inline_backpressure"),
        }
    }
}

fn blob_key_shard_id(key: &str) -> Option<String> {
    let (shard_id, _) = BlobKey::parse_ids(key).ok()?;
    Some(shard_id.to_string())
}

/// Encode a frontier into an i64 acceptable for use in metrics.
pub fn encode_ts_metric<T: Codec64>(ts: &Antichain<T>) -> i64 {
    // We have two problems in mapping a persist frontier into a metric.
    // First is that we only have a `T: Timestamp+Codec64`. Second, is
    // mapping an antichain to a single counter value. We solve both by
    // taking advantage of the fact that in practice, timestamps in mz are
    // currently always a u64 (and if we switch them, it will be to an i64).
    // This means that for all values that mz would actually produce,
    // interpreting the encoded bytes as a little-endian i64 will work.
    // Both of them impl PartialOrder, so in practice, there will always be
    // zero or one elements in the antichain.
    match ts.elements().first() {
        Some(ts) => i64::from_le_bytes(Codec64::encode(ts)),
        None => i64::MAX,
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeSet;

    use mz_dyncfg::ConfigUpdates;
    use mz_ore::metrics::aggregation::PERCENTILES;

    use crate::cfg::SHARD_METRICS;

    use super::*;

    #[mz_ore::test]
    fn shards_aggregate_metrics_one_pass() {
        let registry = MetricsRegistry::new();
        let configs = Arc::clone(&PersistConfig::new_for_tests().configs);
        let shards = ShardsMetrics::new(&registry, Arc::clone(&configs));
        let agg = ShardsAggregateMetrics::new(Arc::clone(&shards.shards), configs);
        let a = shards.shard(&ShardId::new(), "a");
        let b = shards.shard(&ShardId::new(), "b");

        // collect() recomputes both gauges from a single walk of the map.
        agg.collect();
        assert_eq!(agg.count.get(), 2);
        assert_eq!(agg.stale_count.get(), 0);

        a.stale.store(true, Ordering::Relaxed);
        agg.collect();
        assert_eq!(agg.count.get(), 2);
        assert_eq!(agg.stale_count.get(), 1);

        b.stale.store(true, Ordering::Relaxed);
        agg.collect();
        assert_eq!(agg.stale_count.get(), 2);

        // Flipping back and dropping a handle both lower the counts: the gauge
        // is recomputed, not ticked.
        a.stale.store(false, Ordering::Relaxed);
        drop(b);
        agg.collect();
        assert_eq!(agg.count.get(), 1);
        assert_eq!(agg.stale_count.get(), 0);
    }

    #[mz_ore::test]
    fn per_shard_families_match_registrations() {
        let cfg = PersistConfig::new_for_tests();
        set_shard_metrics(&cfg, ShardMetricsExport::Both);
        let registry = MetricsRegistry::new();
        let metrics = Metrics::new(&cfg, &registry);
        let shard = metrics.shards.shard(&ShardId::new(), "s");
        shard.set_batch_part_versions(["v1"].into_iter());
        let gathered = registry.gather();

        let registered: BTreeSet<_> = gathered
            .iter()
            .filter(|f| !f.name().ends_with("_topk"))
            .filter(|f| {
                f.get_metric()
                    .iter()
                    .any(|m| m.get_label().iter().any(|l| l.name() == "shard"))
            })
            .map(|f| f.name())
            .collect();
        let listed: BTreeSet<_> = PER_SHARD_FAMILIES.iter().map(|f| f.name).collect();
        assert_eq!(registered, listed);

        // `ShardGauge::help` is a second copy of the registered help, and
        // becomes the `_percentile` / `_topk` help. Keep the wording the same.
        for (name, gauge) in shard_gauges() {
            let registered = find(&gathered, name)
                .unwrap_or_else(|| panic!("{name} not gathered"))
                .help();
            assert_eq!(gauge.help, registered, "{name}");
        }

        let gathered: BTreeSet<_> = gathered.iter().map(|f| f.name()).collect();
        for family in PER_SHARD_FAMILIES {
            if let ShardAggregate::Counter { total } = family.aggregate {
                assert!(gathered.contains(total), "{total} not gathered");
            }
        }
    }

    #[mz_ore::test]
    fn every_shard_gauge_has_its_aggregates() {
        let registry = MetricsRegistry::new();
        let configs = Arc::clone(&PersistConfig::new_for_tests().configs);
        let shards = ShardsMetrics::new(&registry, Arc::clone(&configs));
        let agg = ShardsAggregateMetrics::new(Arc::clone(&shards.shards), configs);
        let declared: BTreeSet<_> = agg.desc().iter().map(|d| d.fq_name.clone()).collect();
        for (gauge, g) in shard_gauges() {
            for suffix in ["_percentile", "_topk"] {
                let name = format!("{gauge}{suffix}");
                assert!(declared.contains(&name), "{name} not declared");
            }
            let sum = format!("{gauge}_sum");
            assert_eq!(declared.contains(&sum), g.sum, "{sum}");
        }
    }

    fn find<'a>(families: &'a [MetricFamily], name: &str) -> Option<&'a MetricFamily> {
        families.iter().find(|f| f.name() == name)
    }

    fn label<'a>(m: &'a prometheus::proto::Metric, name: &str) -> &'a str {
        m.get_label()
            .iter()
            .find(|l| l.name() == name)
            .map(|l| l.value())
            .unwrap_or_else(|| panic!("no label {name}"))
    }

    /// The `(label, value)` pairs of a gauge family, as a set since gather
    /// order is not insertion order.
    fn gauge_set(
        families: &[MetricFamily],
        name: &str,
        label_name: &str,
    ) -> BTreeSet<(String, u64)> {
        find(families, name).map_or_else(BTreeSet::new, |f| {
            f.get_metric()
                .iter()
                .map(|m| {
                    (
                        label(m, label_name).to_string(),
                        u64::cast_lossy(m.get_gauge().value()),
                    )
                })
                .collect()
        })
    }

    fn scalar(families: &[MetricFamily], name: &str) -> f64 {
        let f = find(families, name).unwrap_or_else(|| panic!("{name} not gathered"));
        let m = &f.get_metric()[0];
        if f.get_field_type() == prometheus::proto::MetricType::COUNTER {
            m.get_counter().value()
        } else {
            m.get_gauge().value()
        }
    }

    /// Builds `Metrics` on a fresh registry with `n` shards that populate
    /// every per-shard family.
    fn metrics_with_shards(n: u64) -> (PersistConfig, MetricsRegistry, Vec<Arc<ShardMetrics>>) {
        let cfg = PersistConfig::new_for_tests();
        let registry = MetricsRegistry::new();
        let metrics = Metrics::new(&cfg, &registry);
        let shards = (1..=n)
            .map(|i| {
                let shard = metrics.shards.shard(&ShardId::new(), &format!("s{i}"));
                shard.usage_current_state_batches_bytes.set(i);
                shard.set_batch_part_versions(["v1"].into_iter());
                shard
            })
            .collect();
        (cfg, registry, shards)
    }

    fn set_shard_metrics(cfg: &PersistConfig, mode: ShardMetricsExport) {
        let mut updates = ConfigUpdates::default();
        updates.add(&SHARD_METRICS, mode.as_str());
        cfg.apply_from(&updates);
    }

    fn family_names(families: &[MetricFamily]) -> BTreeSet<String> {
        families.iter().map(|f| f.name().to_string()).collect()
    }

    /// The families that stay in every mode: the shard counts and the
    /// process-level counters the per-shard counters roll up into.
    fn always_exported() -> BTreeSet<String> {
        let mut names: BTreeSet<_> = ["mz_persist_shard_count", "mz_persist_stale_shard_count"]
            .into_iter()
            .map(String::from)
            .collect();
        names.extend(PER_SHARD_FAMILIES.iter().filter_map(|f| match f.aggregate {
            ShardAggregate::Counter { total } => Some(total.to_string()),
            _ => None,
        }));
        names
    }

    #[mz_ore::test]
    fn shard_metrics_modes_round_trip() {
        for mode in ShardMetricsExport::ALL {
            assert_eq!(ShardMetricsExport::parse(mode.as_str()), Some(mode));
        }
        assert_eq!(
            ShardMetricsExport::parse(SHARD_METRICS.default()),
            Some(ShardMetricsExport::default())
        );
        assert_eq!(ShardMetricsExport::parse("not_a_mode"), None);
    }

    #[mz_ore::test]
    fn unknown_shard_metrics_mode_exports_the_default() {
        let (cfg, registry, _shards) = metrics_with_shards(3);
        let summary = family_names(&registry.gather());
        let invalid =
            |families: &[MetricFamily]| scalar(families, "mz_persist_shard_metrics_mode_invalid");

        let mut updates = ConfigUpdates::default();
        updates.add(&SHARD_METRICS, "summry");
        cfg.apply_from(&updates);

        let gathered = registry.gather();
        assert_eq!(family_names(&gathered), summary);
        // Counted on every scrape while it stays unrecognized, one scrape late.
        assert_eq!(invalid(&gathered), 0.0);
        assert_eq!(invalid(&registry.gather()), 1.0);

        set_shard_metrics(&cfg, ShardMetricsExport::Both);
        let gathered = registry.gather();
        assert!(summary.is_subset(&family_names(&gathered)));
        assert_eq!(invalid(&gathered), 2.0);
        assert_eq!(invalid(&registry.gather()), 2.0);
    }

    #[mz_ore::test]
    fn counter_totals_are_unchanged_by_the_mode() {
        let (cfg, registry, shards) = metrics_with_shards(3);
        for (i, shard) in shards.iter().enumerate() {
            shard.blob_gets.inc_by(u64::cast_from(i) + 1);
            shard.blob_sets.inc();
        }
        let totals = |families: &[MetricFamily]| -> BTreeMap<&'static str, f64> {
            PER_SHARD_FAMILIES
                .iter()
                .filter_map(|f| match f.aggregate {
                    ShardAggregate::Counter { total } => Some((total, scalar(families, total))),
                    _ => None,
                })
                .collect()
        };
        let both = registry.gather();
        assert_eq!(totals(&both)["mz_persist_blob_gets"], 6.0);

        // The modes drop the per-shard families. The totals that replace
        // them don't move.
        for mode in ShardMetricsExport::ALL {
            set_shard_metrics(&cfg, mode);
            assert_eq!(totals(&registry.gather()), totals(&both), "{mode:?}");
        }
    }

    #[mz_ore::test]
    fn each_shard_metrics_mode_exports_its_families() {
        let (cfg, registry, _shards) = metrics_with_shards(3);
        set_shard_metrics(&cfg, ShardMetricsExport::Both);
        let both = family_names(&registry.gather());
        let per_shard: BTreeSet<_> = PER_SHARD_FAMILIES
            .iter()
            .map(|f| f.name.to_string())
            .collect();
        let summary: BTreeSet<_> = both
            .iter()
            .filter(|n| {
                n.ends_with("_percentile")
                    || n.ends_with("_sum")
                    || n.ends_with("_topk")
                    || *n == "mz_persist_batch_part_version_count"
            })
            .cloned()
            .collect();
        let always = always_exported();
        assert!(both.is_superset(&per_shard));
        assert!(both.is_superset(&always));
        assert!(!summary.is_empty());
        assert!(always.is_disjoint(&per_shard) && always.is_disjoint(&summary));

        // The mode is read on every gather, so switching it needs no new
        // `Metrics`. `ALL` ends with `Both`, which restores everything.
        for mode in ShardMetricsExport::ALL {
            set_shard_metrics(&cfg, mode);
            let mut expected = both.clone();
            if !mode.per_shard() {
                expected.retain(|n| !per_shard.contains(n));
            }
            if !mode.summary() {
                expected.retain(|n| !summary.contains(n));
            }
            assert!(expected.is_superset(&always), "{mode:?}");
            assert_eq!(family_names(&registry.gather()), expected, "{mode:?}");
        }
    }

    #[mz_ore::test]
    #[cfg_attr(miri, ignore)] // too slow
    fn summary_scrape_is_flat_in_shard_count() {
        let scrape = |n: u64, mode: ShardMetricsExport| {
            let (cfg, registry, _shards) = metrics_with_shards(n);
            set_shard_metrics(&cfg, mode);
            let families = registry.gather();
            let series: usize = families.iter().map(|f| f.get_metric().len()).sum();
            let body = prometheus::TextEncoder::new()
                .encode_to_string(&families)
                .expect("encodable");
            (series, body.len())
        };
        let (few_series, few_bytes) = scrape(10, ShardMetricsExport::Summary);
        let (many_series, many_bytes) = scrape(10_000, ShardMetricsExport::Summary);
        let (both_series, both_bytes) = scrape(10_000, ShardMetricsExport::Both);
        assert_eq!(few_series, many_series);
        assert!(
            both_series >= many_series + PER_SHARD_FAMILIES.len() * 10_000,
            "{both_series} vs {many_series}"
        );
        // Only the digits of the aggregated values grow with the shard count.
        assert!(many_bytes < few_bytes + 4096, "{many_bytes} vs {few_bytes}");
        assert!(
            both_bytes > 100 * many_bytes,
            "{both_bytes} vs {many_bytes}"
        );
    }

    #[mz_ore::test]
    fn shard_info_lists_live_shards_when_enabled() {
        let (cfg, registry, mut shards) = metrics_with_shards(3);
        let info =
            |registry: &MetricsRegistry| gauge_set(&registry.gather(), SHARD_INFO_NAME, "name");
        let mut updates = ConfigUpdates::default();
        updates.add(&SHARD_INFO_METRIC, true);
        cfg.apply_from(&updates);
        // Independent of the mode, including the one that drops everything else.
        for mode in ShardMetricsExport::ALL {
            set_shard_metrics(&cfg, mode);
            let expected = ["s1", "s2", "s3"].map(|n| (n.to_string(), 1)).into();
            assert_eq!(info(&registry), expected, "{mode:?}");
        }

        shards.pop();
        assert_eq!(info(&registry).len(), 2);
    }

    fn set_enable_regex(cfg: &PersistConfig, regex: &str) {
        let mut updates = ConfigUpdates::default();
        updates.add(&PER_SHARD_METRICS_ENABLE_REGEX, regex);
        cfg.apply_from(&updates);
    }

    /// The per-shard series in `families`, keyed by family name and labels.
    fn per_shard_series(
        families: &[MetricFamily],
    ) -> BTreeMap<(String, Vec<(String, String)>), prometheus::proto::Metric> {
        families
            .iter()
            .filter(|f| PER_SHARD_FAMILIES.iter().any(|p| p.name == f.name()))
            .flat_map(|f| {
                f.get_metric().iter().map(|m| {
                    let labels = m
                        .get_label()
                        .iter()
                        .map(|l| (l.name().to_string(), l.value().to_string()))
                        .collect();
                    ((f.name().to_string(), labels), m.clone())
                })
            })
            .collect()
    }

    /// The `name` labels of the shards with per-shard series in `families`.
    fn exported_shard_names(families: &[MetricFamily]) -> BTreeSet<String> {
        per_shard_series(families)
            .into_keys()
            .map(|(_, labels)| {
                labels
                    .into_iter()
                    .find(|(k, _)| k == "name")
                    .expect("per-shard series have a name label")
                    .1
            })
            .collect()
    }

    #[mz_ore::test]
    fn enable_regex_keeps_matching_shards_unchanged() {
        let (cfg, registry, shards) = metrics_with_shards(5);
        set_shard_metrics(&cfg, ShardMetricsExport::Both);
        let both = registry.gather();
        set_shard_metrics(&cfg, ShardMetricsExport::Summary);
        let summary = registry.gather();
        assert!(exported_shard_names(&summary).is_empty());

        // By shard id: exactly that shard's series come back, byte for byte
        // as they are in `both`, so dashboards need no change.
        let shard_id = shards[2].shard_id.to_string();
        set_enable_regex(&cfg, &shard_id);
        let enabled = registry.gather();
        let expected: BTreeMap<_, _> = per_shard_series(&both)
            .into_iter()
            .filter(|((_, labels), _)| labels.contains(&("shard".into(), shard_id.clone())))
            .collect();
        assert_eq!(expected.len(), PER_SHARD_FAMILIES.len());
        assert_eq!(per_shard_series(&enabled), expected);
        assert_eq!(family_names(&enabled), {
            let mut names = family_names(&summary);
            names.extend(PER_SHARD_FAMILIES.iter().map(|f| f.name.to_string()));
            names
        });

        // By `name` label, and several at once.
        set_enable_regex(&cfg, "^s2$");
        assert_eq!(
            exported_shard_names(&registry.gather()),
            BTreeSet::from(["s2".to_string()])
        );
        set_enable_regex(&cfg, "^s(1|4)$");
        assert_eq!(
            exported_shard_names(&registry.gather()),
            BTreeSet::from(["s1".to_string(), "s4".to_string()])
        );

        // Empty keeps none again.
        set_enable_regex(&cfg, "");
        assert_eq!(family_names(&registry.gather()), family_names(&summary));

        // It applies in `none` too, and not while per-shard series are on.
        set_enable_regex(&cfg, "^s2$");
        set_shard_metrics(&cfg, ShardMetricsExport::None);
        let gathered = registry.gather();
        assert_eq!(
            exported_shard_names(&gathered),
            BTreeSet::from(["s2".to_string()])
        );
        assert!(gathered.iter().all(|f| !f.name().ends_with("_percentile")));
        set_shard_metrics(&cfg, ShardMetricsExport::Both);
        assert_eq!(exported_shard_names(&registry.gather()).len(), 5);
    }

    #[mz_ore::test]
    fn unanchored_enable_regex_over_matches() {
        let cfg = PersistConfig::new_for_tests();
        let registry = MetricsRegistry::new();
        let metrics = Metrics::new(&cfg, &registry);
        let shard = |name: &str| metrics.shards.shard(&ShardId::new(), name);
        // `r` is not a hex digit, so these names can't collide with the
        // random shard ids the second half matches against.
        let _named: Vec<_> = (1..=10).map(|i| shard(&format!("src{i}"))).collect();
        set_shard_metrics(&cfg, ShardMetricsExport::Summary);

        // Unanchored, so a pattern meant for one shard keeps every name it
        // is a substring of.
        set_enable_regex(&cfg, "src1");
        assert_eq!(
            exported_shard_names(&registry.gather()),
            BTreeSet::from(["src1".to_string(), "src10".to_string()])
        );

        // It is tried against `shard` and `name`, so a pattern scoped to one
        // shard's id also keeps any shard whose name contains that fragment.
        // The naming here is contrived, but a pattern loose enough to hit
        // both labels is what brings the shed cardinality back.
        let target = shard("by-id");
        let fragment = target.shard_id.to_string()[3..9].to_string();
        let decoy = shard(&format!("t-{fragment}-t"));
        assert!(!decoy.shard_id.to_string().contains(&fragment));
        set_enable_regex(&cfg, &fragment);
        let kept = exported_shard_names(&registry.gather());
        assert!(kept.contains(&target.name), "{kept:?}");
        assert!(kept.contains(&decoy.name), "{kept:?}");
    }

    #[mz_ore::test]
    fn invalid_enable_regex_fails_closed() {
        let (cfg, registry, _shards) = metrics_with_shards(3);
        set_shard_metrics(&cfg, ShardMetricsExport::Summary);
        let summary = family_names(&registry.gather());
        let invalid = |families: &[MetricFamily]| {
            scalar(families, "mz_persist_per_shard_metrics_regex_invalid")
        };

        set_enable_regex(&cfg, "(");
        let gathered = registry.gather();
        assert_eq!(family_names(&gathered), summary);
        // Counted on every scrape while it stays invalid, one scrape late.
        assert_eq!(invalid(&gathered), 0.0);
        assert_eq!(invalid(&registry.gather()), 1.0);

        set_enable_regex(&cfg, "^s1$");
        let gathered = registry.gather();
        assert_eq!(
            exported_shard_names(&gathered),
            BTreeSet::from(["s1".to_string()])
        );
        assert_eq!(invalid(&gathered), 2.0);
        assert_eq!(invalid(&registry.gather()), 2.0);
    }

    fn gauge_count(families: &[MetricFamily], name: &str) -> usize {
        find(families, name).map_or(0, |f| f.get_metric().len())
    }

    /// The value of one series of a `_percentile` family.
    fn percentile(families: &[MetricFamily], name: &str, percentile: &str) -> u64 {
        gauge_set(families, name, "percentile")
            .into_iter()
            .find_map(|(p, v)| (p == percentile).then_some(v))
            .unwrap_or_else(|| panic!("no {percentile} in {name}"))
    }

    #[mz_ore::test]
    fn shard_aggregates_fold_live_shards() {
        let registry = MetricsRegistry::new();
        let metrics = Metrics::new(&PersistConfig::new_for_tests(), &registry);
        let mut shards: Vec<_> = (1..=12u64)
            .map(|i| {
                let shard = metrics.shards.shard(&ShardId::new(), &format!("s{i}"));
                shard.usage_current_state_batches_bytes.set(i * 100);
                shard
            })
            .collect();
        shards[0].set_batch_part_versions(["v1", "v1", "v2"].into_iter());
        shards[1].set_batch_part_versions(["v1"].into_iter());

        let gathered = registry.gather();
        assert_eq!(scalar(&gathered, "mz_persist_shard_count"), 12.0);
        let percentile_families = gathered
            .iter()
            .filter(|f| {
                f.name().starts_with("mz_persist_shard_") && f.name().ends_with("_percentile")
            })
            .count();
        assert_eq!(percentile_families, shard_gauges().count());
        for (gauge, _) in shard_gauges() {
            let name = format!("{gauge}_percentile");
            assert_eq!(gauge_count(&gathered, &name), PERCENTILES.len(), "{name}");
        }
        let percentiles = "mz_persist_shard_usage_current_state_batches_bytes_percentile";
        assert_eq!(percentile(&gathered, percentiles, "min"), 100);
        assert_eq!(percentile(&gathered, percentiles, "p50"), 600);
        assert_eq!(percentile(&gathered, percentiles, "max"), 1200);
        let top_k = "mz_persist_shard_usage_current_state_batches_bytes_topk";
        let expected: BTreeSet<_> = shards[2..]
            .iter()
            .map(|s| {
                (
                    s.shard_id.to_string(),
                    s.usage_current_state_batches_bytes.get(),
                )
            })
            .collect();
        assert_eq!(gauge_set(&gathered, top_k, "shard"), expected);
        assert_eq!(
            gauge_set(&gathered, "mz_persist_batch_part_version_count", "version"),
            BTreeSet::from([("v1".to_string(), 3), ("v2".to_string(), 1)])
        );

        // Dropping the largest shard removes it from the next gather, and a
        // version no shard has anymore is not exported.
        let largest = shards.pop().expect("non-empty");
        let largest_id = largest.shard_id.to_string();
        drop(largest);
        shards[0].set_batch_part_versions(["v1"].into_iter());
        let gathered = registry.gather();
        assert_eq!(scalar(&gathered, "mz_persist_shard_count"), 11.0);
        assert_eq!(percentile(&gathered, percentiles, "max"), 1100);
        let top_k_shards = gauge_set(&gathered, top_k, "shard");
        assert_eq!(top_k_shards.len(), SHARD_TOP_K);
        assert!(top_k_shards.iter().all(|(shard, _)| *shard != largest_id));
        assert_eq!(
            gauge_set(&gathered, "mz_persist_batch_part_version_count", "version"),
            BTreeSet::from([("v1".to_string(), 2)])
        );
    }

    #[mz_ore::test]
    fn process_counters_survive_shard_drop() {
        let cfg = PersistConfig::new_for_tests();
        set_shard_metrics(&cfg, ShardMetricsExport::Both);
        let registry = MetricsRegistry::new();
        let metrics = Metrics::new(&cfg, &registry);
        let a = metrics.shards.shard(&ShardId::new(), "a");
        let b = metrics.shards.shard(&ShardId::new(), "b");
        a.blob_gets.inc_by(5);
        b.blob_gets.inc();
        a.encoded_diff_size.inc_by(100);

        let gathered = registry.gather();
        assert_eq!(scalar(&gathered, "mz_persist_blob_gets"), 6.0);
        assert_eq!(gauge_count(&gathered, "mz_persist_shard_blob_gets"), 2);

        drop(a);
        let gathered = registry.gather();
        assert_eq!(scalar(&gathered, "mz_persist_blob_gets"), 6.0);
        assert_eq!(scalar(&gathered, "mz_persist_diff_size_bytes"), 100.0);
        assert_eq!(gauge_count(&gathered, "mz_persist_shard_blob_gets"), 1);
    }
}