fix: rest api calculates stage running time correctly (#1675)

* fix: rest api calculates stage running time correctly

* move current time method to utils

* update calculation

* add tests

* return 0 if stages not running

* address review comments

Author: milenkovicm

ballista/core/src/utils.rs

@@ -349,6 +349,14 @@ pub fn get_time_before(interval_seconds: u64) -> u64 {
         .as_secs()
 }
 
+/// current time since UNIX EPOCH
+pub fn get_current_time() -> u128 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_millis()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

ballista/scheduler/src/api/handlers.rs

@@ -13,6 +13,7 @@
 use crate::scheduler_server::event::QueryStageSchedulerEvent;
 use crate::state::execution_graph::ExecutionStage;
 use crate::state::execution_graph_dot::ExecutionGraphDot;
+use crate::state::execution_stage::TaskInfo;
 use crate::{api::SchedulerErrorResponse, scheduler_server::SchedulerServer};
 use axum::extract::Query;
 use axum::{
@@ -28,10 +29,11 @@ use ballista_core::serde::protobuf::{
 use ballista_core::serde::scheduler::{
     ExecutorOperatingSystemSpecification, ExecutorSpecification,
 };
+use ballista_core::utils::get_current_time;
 use datafusion::DATAFUSION_VERSION;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::displayable;
-use datafusion::physical_plan::metrics::{MetricValue, MetricsSet, Time};
+use datafusion::physical_plan::metrics::{MetricsSet, Time};
 use datafusion_proto::logical_plan::AsLogicalPlan;
 use datafusion_proto::physical_plan::AsExecutionPlan;
 #[cfg(feature = "graphviz-support")]
@@ -539,11 +541,9 @@ pub async fn get_query_stages<
                             .as_ref()
                             .map(|m| get_combined_count(m.as_slice(), "output_rows"))
                             .unwrap_or(0);
-                        summary.elapsed_compute = running_stage
-                            .stage_metrics
-                            .as_ref()
-                            .map(|m| get_elapsed_compute_nanos(m.as_slice()))
-                            .unwrap_or_default();
+                        summary.elapsed_compute = get_running_stage_time(&running_stage
+                            .task_infos, get_current_time())
+                            ;
                         summary.tasks = running_stage
                             .task_infos
                             .iter()
@@ -595,7 +595,7 @@ pub async fn get_query_stages<
                             "output_rows",
                         );
                         summary.elapsed_compute =
-                            get_elapsed_compute_nanos(&completed_stage.stage_metrics);
+                            get_finished_stage_time(&completed_stage.task_infos);
 
                         summary.tasks = completed_stage
                             .task_infos
@@ -722,19 +722,44 @@ fn format_job_status(status: &Option<Status>, elapsed_ms: u64) -> (String, Strin
     }
 }
 
-fn get_elapsed_compute_nanos(metrics: &[MetricsSet]) -> String {
-    let nanos: usize = metrics
+fn get_running_stage_time(task_infos: &[Option<TaskInfo>], current_time: u128) -> String {
+    let min_start = task_infos
         .iter()
-        .flat_map(|vec| {
-            vec.iter().map(|metric| match metric.as_ref().value() {
-                MetricValue::ElapsedCompute(time) => time.value(),
-                _ => 0,
-            })
-        })
-        .sum();
-    let t = Time::new();
-    t.add_duration(Duration::from_nanos(nanos as u64));
-    t.to_string()
+        .flat_map(|t| t.as_ref().map(|t| t.start_exec_time))
+        .filter(|t| *t > 0)
+        .min();
+
+    match (min_start, current_time) {
+        (Some(start), end) if end > start => {
+            let t = Time::new();
+            t.add_duration(Duration::from_millis((end - start) as u64));
+            t.to_string()
+        }
+        _ => "0".to_string(),
+    }
+}
+
+fn get_finished_stage_time(task_infos: &[TaskInfo]) -> String {
+    let min_start = task_infos
+        .iter()
+        .map(|t| t.start_exec_time)
+        .filter(|t| *t > 0)
+        .min();
+
+    let max_end = task_infos
+        .iter()
+        .map(|t| t.end_exec_time)
+        .filter(|t| *t > 0)
+        .max();
+
+    match (min_start, max_end) {
+        (Some(start), Some(end)) if end > start => {
+            let t = Time::new();
+            t.add_duration(Duration::from_millis((end - start) as u64));
+            t.to_string()
+        }
+        _ => "0".to_string(),
+    }
 }
 
 fn get_partition_counts(metrics: &[MetricsSet], partition_id: usize) -> (usize, usize) {
@@ -878,3 +903,102 @@ pub async fn get_scheduler_metrics<
             .unwrap(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::state::execution_stage::TaskInfo;
+    use ballista_core::serde::protobuf::task_status;
+
+    fn make_task_info(start: u128, end: u128) -> TaskInfo {
+        TaskInfo {
+            task_id: 0,
+            scheduled_time: 0,
+            launch_time: 0,
+            start_exec_time: start,
+            end_exec_time: end,
+            finish_time: 0,
+            task_status: task_status::Status::Running(Default::default()),
+        }
+    }
+
+    // --- get_finished_stage_time ---
+
+    #[test]
+    fn test_finished_empty_slice_returns_zero() {
+        assert_eq!(get_finished_stage_time(&[]), "0");
+    }
+
+    #[test]
+    fn test_finished_all_zero_timestamps_returns_zero() {
+        let tasks = vec![make_task_info(0, 0), make_task_info(0, 0)];
+        assert_eq!(get_finished_stage_time(&tasks), "0");
+    }
+
+    #[test]
+    fn test_finished_single_task_elapsed() {
+        // 600 - 100 = 500 ms → "500.00ms"
+        let tasks = vec![make_task_info(100, 600)];
+        assert_eq!(get_finished_stage_time(&tasks), "500.00ms");
+    }
+
+    #[test]
+    fn test_finished_picks_earliest_start_and_latest_end() {
+        // min start = 100, max end = 900 → 800 ms
+        let tasks = vec![
+            make_task_info(100, 500),
+            make_task_info(200, 900),
+            make_task_info(300, 700),
+        ];
+        assert_eq!(get_finished_stage_time(&tasks), "800.00ms");
+    }
+
+    #[test]
+    fn test_finished_end_before_start_returns_zero() {
+        let tasks = vec![make_task_info(900, 100)];
+        assert_eq!(get_finished_stage_time(&tasks), "0");
+    }
+
+    // --- get_running_stage_time ---
+
+    #[test]
+    fn test_running_empty_slice_returns_zero() {
+        assert_eq!(get_running_stage_time(&[], 1000), "0");
+    }
+
+    #[test]
+    fn test_running_all_none_returns_zero() {
+        let tasks: Vec<Option<TaskInfo>> = vec![None, None];
+        assert_eq!(get_running_stage_time(&tasks, 1000), "0");
+    }
+
+    #[test]
+    fn test_running_future_start_returns_zero() {
+        // start_exec_time beyond current time → elapsed clamped to 0
+        let tasks = vec![Some(make_task_info(u128::MAX, 0))];
+        assert_eq!(get_running_stage_time(&tasks, 1000), "0");
+    }
+
+    #[test]
+    fn test_running_past_start_returns_nonzero() {
+        let now = 4_000;
+        let start = 1_000;
+        let tasks = vec![Some(make_task_info(start, 0))];
+        assert_eq!(get_running_stage_time(&tasks, now), "3.00s");
+    }
+
+    #[test]
+    fn test_running_mixed_some_none_uses_earliest_some() {
+        let now = 3_000;
+        let earlier = 1_000;
+        let later = 2_000;
+        let tasks = vec![
+            None,
+            Some(make_task_info(later, 0)),
+            Some(make_task_info(earlier, 0)),
+            None,
+        ];
+        let result = get_running_stage_time(&tasks, now);
+        assert_eq!(result, "2.00s");
+    }
+}