from<ActionStage> trait

Author: MarcusSorealheis

nativelink-scheduler/src/memory_awaited_action_db.rs

@@ -647,57 +647,17 @@ impl<I: InstantWrapper, NowFn: Fn() -> I + Clone + Send + Sync> AwaitedActionDbI
                 metrics.execution_stage_transitions.add(1, &base_attrs);
 
                 // Update active count for old stage
-                let old_stage_attrs = match old_stage {
-                    ActionStage::Unknown => vec![opentelemetry::KeyValue::new(
-                        nativelink_util::metrics::EXECUTION_STAGE,
-                        ExecutionStage::Unknown,
-                    )],
-                    ActionStage::CacheCheck => vec![opentelemetry::KeyValue::new(
-                        nativelink_util::metrics::EXECUTION_STAGE,
-                        ExecutionStage::CacheCheck,
-                    )],
-                    ActionStage::Queued => vec![opentelemetry::KeyValue::new(
-                        nativelink_util::metrics::EXECUTION_STAGE,
-                        ExecutionStage::Queued,
-                    )],
-                    ActionStage::Executing => vec![opentelemetry::KeyValue::new(
-                        nativelink_util::metrics::EXECUTION_STAGE,
-                        ExecutionStage::Executing,
-                    )],
-                    ActionStage::Completed(_) | ActionStage::CompletedFromCache(_) => {
-                        vec![opentelemetry::KeyValue::new(
-                            nativelink_util::metrics::EXECUTION_STAGE,
-                            ExecutionStage::Completed,
-                        )]
-                    }
-                };
+                let old_stage_attrs = vec![opentelemetry::KeyValue::new(
+                    nativelink_util::metrics::EXECUTION_STAGE,
+                    ExecutionStage::from(old_stage.clone()),
+                )];
                 metrics.execution_active_count.add(-1, &old_stage_attrs);
 
                 // Update active count for new stage
-                let new_stage_attrs = match new_stage {
-                    ActionStage::Unknown => vec![opentelemetry::KeyValue::new(
-                        nativelink_util::metrics::EXECUTION_STAGE,
-                        ExecutionStage::Unknown,
-                    )],
-                    ActionStage::CacheCheck => vec![opentelemetry::KeyValue::new(
-                        nativelink_util::metrics::EXECUTION_STAGE,
-                        ExecutionStage::CacheCheck,
-                    )],
-                    ActionStage::Queued => vec![opentelemetry::KeyValue::new(
-                        nativelink_util::metrics::EXECUTION_STAGE,
-                        ExecutionStage::Queued,
-                    )],
-                    ActionStage::Executing => vec![opentelemetry::KeyValue::new(
-                        nativelink_util::metrics::EXECUTION_STAGE,
-                        ExecutionStage::Executing,
-                    )],
-                    ActionStage::Completed(_) | ActionStage::CompletedFromCache(_) => {
-                        vec![opentelemetry::KeyValue::new(
-                            nativelink_util::metrics::EXECUTION_STAGE,
-                            ExecutionStage::Completed,
-                        )]
-                    }
-                };
+                let new_stage_attrs = vec![opentelemetry::KeyValue::new(
+                    nativelink_util::metrics::EXECUTION_STAGE,
+                    ExecutionStage::from(new_stage.clone()),
+                )];
                 metrics.execution_active_count.add(1, &new_stage_attrs);
 
                 // Record completion metrics

nativelink-util/src/metrics.rs

@@ -16,6 +16,8 @@ use std::sync::LazyLock;
 
 use opentelemetry::{InstrumentationScope, KeyValue, Value, global, metrics};
 
+use crate::action_messages::ActionStage;
+
 // Metric attribute keys for cache operations.
 pub const CACHE_TYPE: &str = "cache.type";
 pub const CACHE_OPERATION: &str = "cache.operation.name";
@@ -111,6 +113,20 @@ impl From<ExecutionStage> for Value {
     }
 }
 
+impl From<ActionStage> for ExecutionStage {
+    fn from(stage: ActionStage) -> Self {
+        match stage {
+            ActionStage::Unknown => ExecutionStage::Unknown,
+            ActionStage::CacheCheck => ExecutionStage::CacheCheck,
+            ActionStage::Queued => ExecutionStage::Queued,
+            ActionStage::Executing => ExecutionStage::Executing,
+            ActionStage::Completed(_) | ActionStage::CompletedFromCache(_) => {
+                ExecutionStage::Completed
+            }
+        }
+    }
+}
+
 /// Results of remote execution operations.
 #[derive(Debug, Clone, Copy)]
 pub enum ExecutionResult {