review: Address comments

Author: peasee

ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -247,13 +247,18 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                 // subscribers (e.g. SpiceAI's QueryHandle) who receive the
                 // Completed event can immediately read the Successful status
                 // without hitting a retry/polling loop.
-                if let Err(e) = self.state.task_manager.succeed_job(&job_id).await {
-                    error!("Fail to invoke succeed_job for job {job_id} due to {e:?}");
+                match self.state.task_manager.succeed_job(&job_id).await {
+                    Ok(()) => {
+                        // Broadcast job completed state after status is persisted
+                        self.broadcast_job_state(JobStateEvent::completed(&job_id));
+                    }
+                    Err(e) => {
+                        error!(
+                            "Fail to invoke succeed_job for job {job_id} due to {e:?}"
+                        );
+                    }
                 }
 
-                // Broadcast job completed state after status is persisted
-                self.broadcast_job_state(JobStateEvent::completed(&job_id));
-
                 self.state.clean_up_successful_job(job_id);
             }
             QueryStageSchedulerEvent::JobRunningFailed {
@@ -283,15 +288,17 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                                 ))
                                 .await?;
                         }
+                        // Broadcast job failed state after status is persisted
+                        self.broadcast_job_state(JobStateEvent::failed(
+                            &job_id,
+                            &fail_message,
+                        ));
                     }
                     Err(e) => {
                         error!("Fail to invoke abort_job for job {job_id} due to {e:?}");
                     }
                 }
 
-                // Broadcast job failed state after status is persisted
-                self.broadcast_job_state(JobStateEvent::failed(&job_id, &fail_message));
-
                 self.state.clean_up_failed_job(job_id);
             }
             QueryStageSchedulerEvent::JobUpdated(job_id) => {
@@ -306,23 +313,23 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                 info!("Job {job_id} Cancelled");
 
                 // Persist terminal status before broadcasting so subscribers
-                // can immediately read the Cancelled status on receipt of the event.
+                // can immediately read the terminal status on receipt of the event.
+                // Note: cancel_job routes to abort_job, persisting a Failed status.
                 match self.state.task_manager.cancel_job(&job_id).await {
                     Ok((running_tasks, _pending_tasks)) => {
                         event_sender
                             .post_event(QueryStageSchedulerEvent::CancelTasks(
                                 running_tasks,
                             ))
                             .await?;
+                        // Broadcast cancelled state after status is persisted
+                        self.broadcast_job_state(JobStateEvent::cancelled(&job_id));
                     }
                     Err(e) => {
                         error!("Fail to invoke cancel_job for job {job_id} due to {e:?}");
                     }
                 }
 
-                // Broadcast cancelled state after status is persisted
-                self.broadcast_job_state(JobStateEvent::cancelled(&job_id));
-
                 self.state.clean_up_failed_job(job_id);
             }
             QueryStageSchedulerEvent::TaskUpdating(executor_id, tasks_status) => {

ballista/scheduler/src/state/aqe/mod.rs

@@ -919,9 +919,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph {
 
     /// Return all currently running tasks along with the executor ID on which they are assigned
     fn running_tasks(&self) -> Vec<RunningTaskInfo> {
-        self.stages
-            .iter()
-            .flat_map(|(_, stage)| {
+        self.stages.values().flat_map(|stage| {
                 if let ExecutionStage::Running(stage) = stage {
                     stage
                         .running_tasks()

ballista/scheduler/src/state/execution_graph.rs

@@ -1319,9 +1319,7 @@ impl ExecutionGraph for StaticExecutionGraph {
 
     /// Return all currently running tasks along with the executor ID on which they are assigned
     fn running_tasks(&self) -> Vec<RunningTaskInfo> {
-        self.stages
-            .iter()
-            .flat_map(|(_, stage)| {
+        self.stages.values().flat_map(|stage| {
                 if let ExecutionStage::Running(stage) = stage {
                     stage
                         .running_tasks()

ballista/scheduler/src/state/executor_manager.rs

@@ -438,9 +438,7 @@ impl ExecutorManager {
             get_time_before(self.config.executor_termination_grace_period);
 
         self.cluster_state
-            .executor_heartbeats()
-            .iter()
-            .filter_map(|(_exec, heartbeat)| {
+            .executor_heartbeats().values().filter_map(|heartbeat| {
                 let terminating = matches!(
                     heartbeat
                         .status