compute,storage: annotate wallclock lag histogram with workload_class

teskje · teskje · commit 308563aae0c2 · 2025-03-27T11:33:23.000+01:00
This commit adds a "workload_class" label to measurements in
`mz_wallclock_global_lag_histogram`. The value of the label is the
workload class of the cluster maintaining the respective collection.
diff --git a/src/compute-client/src/controller/instance.rs b/src/compute-client/src/controller/instance.rs
@@ -39,7 +39,7 @@ use mz_ore::{soft_assert_or_log, soft_panic_or_log};
 use mz_repr::adt::interval::Interval;
 use mz_repr::adt::timestamp::CheckedTimestamp;
 use mz_repr::refresh_schedule::RefreshSchedule;
-use mz_repr::{Datum, DatumMap, Diff, GlobalId, Row};
+use mz_repr::{Datum, Diff, GlobalId, Row};
 use mz_storage_client::controller::{IntrospectionType, WallclockLagHistogramPeriod};
 use mz_storage_types::read_holds::{self, ReadHold};
 use mz_storage_types::read_policy::ReadPolicy;
@@ -228,6 +228,10 @@ pub(super) struct Instance<T: ComputeControllerTimestamp> {
     /// controlled by it are allowed to affect changes to external systems
     /// (largely persist).
     read_only: bool,
+    /// The workload class of this instance.
+    ///
+    /// This is currently only used to annotate metrics.
+    workload_class: Option<String>,
     /// The replicas of this compute instance.
     replicas: BTreeMap<ReplicaId, ReplicaState<T>>,
     /// Currently installed compute collections.
@@ -608,8 +612,13 @@ impl<T: ComputeControllerTimestamp> Instance<T> {
         }
 
         let histogram_period = WallclockLagHistogramPeriod::from_epoch_millis(now_ms, &self.dyncfg);
+        let histogram_labels = match &self.workload_class {
+            Some(wc) => [("workload_class", wc.clone())].into(),
+            None => BTreeMap::new(),
+        };
 
         let mut histogram_updates = Vec::new();
+        let mut row_buf = Row::default();
         for (collection_id, collection) in &mut self.collections {
             if !ENABLE_WALLCLOCK_LAG_HISTOGRAM_COLLECTION.get(&self.dyncfg) {
                 continue;
@@ -619,19 +628,22 @@ impl<T: ComputeControllerTimestamp> Instance<T> {
                 let lag = collection.shared.lock_write_frontier(|f| frontier_lag(f));
                 let bucket = lag.as_secs().next_power_of_two();
 
-                let key = (histogram_period, bucket);
+                let key = (histogram_period, bucket, histogram_labels.clone());
                 *stash.entry(key).or_default() += 1;
 
                 if refresh_histogram {
-                    for ((period, lag), count) in std::mem::take(stash) {
-                        let row = Row::pack_slice(&[
+                    for ((period, lag, labels), count) in std::mem::take(stash) {
+                        let mut packer = row_buf.packer();
+                        packer.extend([
                             Datum::TimestampTz(period.start),
                             Datum::TimestampTz(period.end),
                             Datum::String(&collection_id.to_string()),
                             Datum::UInt64(lag),
-                            Datum::Map(DatumMap::empty()),
                         ]);
-                        histogram_updates.push((row, count));
+                        let labels = labels.iter().map(|(k, v)| (*k, Datum::String(v)));
+                        packer.push_dict(labels);
+
+                        histogram_updates.push((row_buf.clone(), count));
                     }
                 }
             }
@@ -846,6 +858,7 @@ impl<T: ComputeControllerTimestamp> Instance<T> {
             storage_collections: _,
             initialized,
             read_only,
+            workload_class,
             replicas,
             collections,
             log_sources: _,
@@ -905,6 +918,7 @@ impl<T: ComputeControllerTimestamp> Instance<T> {
         let map = serde_json::Map::from_iter([
             field("initialized", initialized)?,
             field("read_only", read_only)?,
+            field("workload_class", workload_class)?,
             field("replicas", replicas)?,
             field("collections", collections)?,
             field("peeks", peeks)?,
@@ -968,6 +982,7 @@ where
             storage_collections: storage,
             initialized: false,
             read_only: true,
+            workload_class: None,
             replicas: Default::default(),
             collections,
             log_sources,
@@ -1022,6 +1037,10 @@ where
     /// Update instance configuration.
     #[mz_ore::instrument(level = "debug")]
     pub fn update_configuration(&mut self, config_params: ComputeParameters) {
+        if let Some(workload_class) = &config_params.workload_class {
+            self.workload_class = workload_class.clone();
+        }
+
         self.send(ComputeCommand::UpdateConfiguration(config_params));
     }
 
@@ -2211,8 +2230,19 @@ struct CollectionState<T: ComputeControllerTimestamp> {
     /// Frontier wallclock lag measurements stashed until the next `WallclockLagHistogram`
     /// introspection update.
     ///
+    /// Keys are `(period, lag, labels)` triples, values are counts.
+    ///
     /// If this is `None`, wallclock lag is not tracked for this collection.
-    wallclock_lag_histogram_stash: Option<BTreeMap<(WallclockLagHistogramPeriod, u64), i64>>,
+    wallclock_lag_histogram_stash: Option<
+        BTreeMap<
+            (
+                WallclockLagHistogramPeriod,
+                u64,
+                BTreeMap<&'static str, String>,
+            ),
+            i64,
+        >,
+    >,
 }
 
 impl<T: ComputeControllerTimestamp> CollectionState<T> {
diff --git a/src/controller/src/clusters.rs b/src/controller/src/clusters.rs
@@ -364,6 +364,8 @@ where
         id: ClusterId,
         workload_class: Option<String>,
     ) -> Result<(), anyhow::Error> {
+        self.storage
+            .update_instance_workload_class(id, workload_class.clone());
         self.compute
             .update_instance_workload_class(id, workload_class)?;
         Ok(())
diff --git a/src/storage-client/src/controller.rs b/src/storage-client/src/controller.rs
@@ -383,6 +383,13 @@ pub trait StorageController: Debug {
     /// Panics if a storage instance with the given ID does not exist.
     fn drop_instance(&mut self, id: StorageInstanceId);
 
+    /// Updates a storage instance's workload class.
+    fn update_instance_workload_class(
+        &mut self,
+        id: StorageInstanceId,
+        workload_class: Option<String>,
+    );
+
     /// Connects the storage instance to the specified replica.
     ///
     /// If the storage instance is already attached to a replica, communication
diff --git a/src/storage-controller/src/instance.rs b/src/storage-controller/src/instance.rs
@@ -55,6 +55,10 @@ use crate::history::CommandHistory;
 /// lead to panics.
 #[derive(Debug)]
 pub(crate) struct Instance<T> {
+    /// The workload class of this instance.
+    ///
+    /// This is currently only used to annotate metrics.
+    pub workload_class: Option<String>,
     /// The replicas connected to this storage instance.
     replicas: BTreeMap<ReplicaId, Replica<T>>,
     /// The ingestions currently running on this instance.
@@ -114,6 +118,7 @@ where
         let epoch = ClusterStartupEpoch::new(envd_epoch, 0);
 
         let mut instance = Self {
+            workload_class: None,
             replicas: Default::default(),
             active_ingestions: Default::default(),
             ingestion_exports: Default::default(),
diff --git a/src/storage-controller/src/lib.rs b/src/storage-controller/src/lib.rs
@@ -48,9 +48,7 @@ use mz_persist_types::Codec64;
 use mz_proto::RustType;
 use mz_repr::adt::interval::Interval;
 use mz_repr::adt::timestamp::CheckedTimestamp;
-use mz_repr::{
-    Datum, DatumMap, Diff, GlobalId, RelationDesc, RelationVersion, Row, TimestampManipulation,
-};
+use mz_repr::{Datum, Diff, GlobalId, RelationDesc, RelationVersion, Row, TimestampManipulation};
 use mz_storage_client::client::{
     ProtoStorageCommand, ProtoStorageResponse, RunIngestionCommand, RunOneshotIngestion,
     RunSinkCommand, Status, StatusUpdate, StorageCommand, StorageResponse, TableData,
@@ -536,6 +534,19 @@ where
         assert!(instance.is_some(), "storage instance {id} does not exist");
     }
 
+    fn update_instance_workload_class(
+        &mut self,
+        id: StorageInstanceId,
+        workload_class: Option<String>,
+    ) {
+        let instance = self
+            .instances
+            .get_mut(&id)
+            .unwrap_or_else(|| panic!("instance {id} does not exist"));
+
+        instance.workload_class = workload_class;
+    }
+
     fn connect_replica(
         &mut self,
         instance_id: StorageInstanceId,
@@ -3478,6 +3489,7 @@ where
 
         let mut history_updates = Vec::new();
         let mut histogram_updates = Vec::new();
+        let mut row_buf = Row::default();
         for frontiers in self.storage_collections.active_collection_frontiers() {
             let id = frontiers.id;
             let Some(collection) = self.collections.get_mut(&id) else {
@@ -3507,19 +3519,36 @@ where
 
             if let Some(stash) = &mut collection.wallclock_lag_histogram_stash {
                 let bucket = lag.as_secs().next_power_of_two();
-                let key = (histogram_period, bucket);
+
+                let instance_id = match &collection.extra_state {
+                    CollectionStateExtra::Ingestion(i) => Some(i.instance_id),
+                    CollectionStateExtra::Export(e) => Some(e.cluster_id()),
+                    CollectionStateExtra::None => None,
+                };
+                let workload_class = instance_id
+                    .and_then(|id| self.instances.get(&id))
+                    .and_then(|i| i.workload_class.clone());
+                let labels = match workload_class {
+                    Some(wc) => [("workload_class", wc.clone())].into(),
+                    None => BTreeMap::new(),
+                };
+
+                let key = (histogram_period, bucket, labels);
                 *stash.entry(key).or_default() += 1;
 
                 if refresh_histogram {
-                    for ((period, lag), count) in std::mem::take(stash) {
-                        let row = Row::pack_slice(&[
+                    for ((period, lag, labels), count) in std::mem::take(stash) {
+                        let mut packer = row_buf.packer();
+                        packer.extend([
                             Datum::TimestampTz(period.start),
                             Datum::TimestampTz(period.end),
                             Datum::String(&id.to_string()),
                             Datum::UInt64(lag),
-                            Datum::Map(DatumMap::empty()),
                         ]);
-                        histogram_updates.push((row, count));
+                        let labels = labels.iter().map(|(k, v)| (*k, Datum::String(v)));
+                        packer.push_dict(labels);
+
+                        histogram_updates.push((row_buf.clone(), count));
                     }
                 }
             }
@@ -3666,8 +3695,19 @@ struct CollectionState<T: TimelyTimestamp> {
     /// Frontier wallclock lag measurements stashed until the next `WallclockLagHistogram`
     /// introspection update.
     ///
+    /// Keys are `(period, lag, labels)` triples, values are counts.
+    ///
     /// If this is `None`, wallclock lag is not tracked for this collection.
-    wallclock_lag_histogram_stash: Option<BTreeMap<(WallclockLagHistogramPeriod, u64), i64>>,
+    wallclock_lag_histogram_stash: Option<
+        BTreeMap<
+            (
+                WallclockLagHistogramPeriod,
+                u64,
+                BTreeMap<&'static str, String>,
+            ),
+            i64,
+        >,
+    >,
     /// Frontier wallclock lag metrics tracked for this collection.
     wallclock_lag_metrics: WallclockLagMetrics,
 }
