fix(spider-scheduler): Cancel the runtime when the scheduler core exits on error. (#391)

Author: LinZhihao-723

components/spider-scheduler/src/core.rs

@@ -55,6 +55,9 @@ pub trait SchedulerCore: Send {
     /// and writes assignments to `sink`, repeating until `cancellation_token` is fired, at which
     /// point it returns.
     ///
+    /// The implementation does not need to fire the cancellation token when exiting on error: this
+    /// cancellation is handled by the runtime.
+    ///
     /// # Parameters
     ///
     /// * `storage_client` - The storage client used to poll the inbound queue and read state for

components/spider-scheduler/src/runtime.rs

@@ -148,12 +148,24 @@ pub async fn create_runtime<SchedulerStorageClientType: SchedulerStorageClient +
     let service = SchedulerServiceState::new(dispatch_queue_reader, registry, scheduler_id);
     let core = scheduler_config.make_core::<SchedulerStorageClientType, DispatchQueueWriter>();
 
-    let core_join_handle = tokio::spawn(core.run(
-        storage_client,
-        dispatch_queue_writer,
-        TaskAssignmentIdIssuer::new(),
-        cancellation_token.clone(),
-    ));
+    let core_cancellation_token = cancellation_token.clone();
+    let core_join_handle = tokio::spawn(async move {
+        core.run(
+            storage_client,
+            dispatch_queue_writer,
+            TaskAssignmentIdIssuer::new(),
+            core_cancellation_token.clone(),
+        )
+        .await
+        .inspect_err(|e| {
+            tracing::error!(
+                scheduler_id = % scheduler_id,
+                error = % e,
+                "Scheduler core exited on error. Cancelling runtime."
+            );
+            core_cancellation_token.cancel();
+        })
+    });
 
     let runtime = Runtime {
         core_join_handle,