Skip to content

feat(spider-scheduler): Reschedule task assignments returned by lost execution managers. #395

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LinZhihao-723 merged 4 commits into from
Jul 14, 2026
Merged

feat(spider-scheduler): Reschedule task assignments returned by lost execution managers. #395

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7138f36
Done with implementation.
LinZhihao-723 Jul 13, 2026
7749d22
Update test cases.
LinZhihao-723 Jul 13, 2026
2c02a97
Merge branch 'main' into reschedule-queue-wiring
LinZhihao-723 Jul 13, 2026
0cc7b99
Merge branch 'main' into reschedule-queue-wiring
sitaowang1998 Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

9 changes: 4 additions & 5 deletions components/spider-core/src/types/scheduler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,9 @@ pub struct TaskAssignment {

/// The task to dispatch.
pub task_id: TaskId,

/// The scheduler's view of storage's session when the assignment was produced.
pub session_id: SessionId,
}

/// A record of a task assignment previously produced by the scheduler.
Expand All @@ -50,8 +53,7 @@ impl TaskAssignmentRecord {

/// A task assignment handed to the execution manager by the scheduler.
///
/// `session_id` is the scheduler's view of storage's session at the moment the assignment was
/// produced. The execution manager pins this exact value on every subsequent storage call for the
/// The execution manager pins `task_assignment.session_id` on every subsequent storage call for the
/// attempt.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SchedulerResponse {
Expand All @@ -60,7 +62,4 @@ pub struct SchedulerResponse {

/// The scheduler that produced the assignment.
pub scheduler_id: SchedulerId,

/// The scheduler's view of storage's session when the assignment was produced.
pub session_id: SessionId,
}
20 changes: 10 additions & 10 deletions components/spider-execution-manager/src/runtime.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,16 +302,16 @@ impl<
};

tracing::info!(
bundle_session = response.session_id,
bundle_session = response.task_assignment.session_id,
job_id = ? response.task_assignment.job_id,
task_id = ? response.task_assignment.task_id,
"Received a new task assignment from the scheduler."
);

let current_session = self.session_tracker.current();
if response.session_id < current_session {
if response.task_assignment.session_id < current_session {
tracing::warn!(
bundle_session = response.session_id,
bundle_session = response.task_assignment.session_id,
current_session,
job_id = ? response.task_assignment.job_id,
task_id = ? response.task_assignment.task_id,
Expand All @@ -320,9 +320,9 @@ impl<
self.mark_consume(&response);
continue;
}
if response.session_id > current_session {
if response.task_assignment.session_id > current_session {
tracing::info!(
new_session = response.session_id,
new_session = response.task_assignment.session_id,
"Observed a newer session via the scheduler. Refreshing liveness."
);
self.liveness_handle.refresh().await;
Expand Down Expand Up @@ -355,9 +355,9 @@ impl<
})?;

let current_session = self.session_tracker.current();
if response.session_id < current_session {
if response.task_assignment.session_id < current_session {
tracing::warn!(
bundle_session = response.session_id,
bundle_session = response.task_assignment.session_id,
current_session,
job_id = ? response.task_assignment.job_id,
task_id = ? response.task_assignment.task_id,
Expand All @@ -375,7 +375,7 @@ impl<
job: response.task_assignment.job_id,
task: response.task_assignment.task_id,
task_instance_id,
session: response.session_id,
session: response.task_assignment.session_id,
},
outcome,
));
Expand Down Expand Up @@ -413,7 +413,7 @@ impl<
response.task_assignment.job_id,
response.task_assignment.task_id,
self.em_id,
response.session_id,
response.task_assignment.session_id,
) => result,
};

Expand All @@ -425,7 +425,7 @@ impl<
Err(err) => match &err {
StorageResponseError::StaleSession(message) => {
tracing::warn!(
bundle_session = response.session_id,
bundle_session = response.task_assignment.session_id,
error = % message,
job_id = ? response.task_assignment.job_id,
task_id = ? response.task_assignment.task_id,
Expand Down
4 changes: 2 additions & 2 deletions components/spider-proto-rust/src/assignment.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,9 @@ impl TryFrom<NextTaskResponse> for Option<SchedulerResponse> {
resource_group_id: ResourceGroupId::from(assignment.resource_group_id),
job_id: JobId::from(assignment.job_id),
task_id,
session_id: assignment.session_id,
},
scheduler_id: SchedulerId::from(assignment.scheduler_id),
session_id: assignment.session_id,
}))
}
Some(next_task_response::Result::NoTask(common::Void {})) => Ok(None),
Expand Down Expand Up @@ -117,7 +117,7 @@ mod tests {
assert_eq!(scheduler_response.task_assignment.job_id, JobId::from(3));
assert_eq!(scheduler_response.task_assignment.task_id, TaskId::Index(7));
assert_eq!(scheduler_response.scheduler_id, SchedulerId::from(4));
assert_eq!(scheduler_response.session_id, 5);
assert_eq!(scheduler_response.task_assignment.session_id, 5);
}

#[test]
Expand Down
1 change: 1 addition & 0 deletions components/spider-scheduler/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,5 +31,6 @@ tracing = { version = "0.1.41", default-features = false, features = ["std"] }
[dev-dependencies]
anyhow = "1.0.102"
dashmap = "6.1.0"
rand = "0.9.1"

@coderabbitai coderabbitai Bot Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check if rand 0.9.1 exists on crates.io
curl -s https://crates.io/api/v1/crates/rand/0.9.1 | jq '.version.num, .version.yanked'

# Check for security advisories
gh api graphql -f query='
{
  securityVulnerabilities(first: 5, ecosystem: RUST, package: "rand") {
    nodes {
      advisory { summary severity publishedAt }
      vulnerableVersionRange
      firstPatchedVersion { identifier }
    }
  }
}'

Repository: y-scope/spider

Length of output: 880

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "Raw crates.io response for rand 0.9.1:"
curl -s https://crates.io/api/v1/crates/rand/0.9.1 | jq .

echo
echo "Searching published rand versions for 0.9.1:"
curl -s "https://crates.io/api/v1/crates/rand" | jq -r '.versions[].num' | grep -x '0.9.1' || true

Repository: y-scope/spider

Length of output: 599

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "Trying crates.io sparse index entry for rand:"
curl -s https://index.crates.io/ra/nd/rand | sed -n '1,120p'

Repository: y-scope/spider

Length of output: 50371

🏁 Script executed:

#!/bin/bash
set -euo pipefail

curl -s https://index.crates.io/ra/nd/rand | grep -m1 '"vers":"0.9.1"' || true

Repository: y-scope/spider

Length of output: 1542

Bump rand to 0.9.3 or later rand 0.9.1 is published, but it falls in the advisory range >= 0.9.0, < 0.9.3. components/spider-scheduler/Cargo.toml:34

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/spider-scheduler/Cargo.toml` at line 34, Update the rand
dependency declaration in Cargo.toml from 0.9.1 to version 0.9.3 or later,
ensuring the resolved dependency is outside the advisory range.

tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread"] }
tokio-util = { version = "0.7.18", features = ["rt"] }
4 changes: 4 additions & 0 deletions components/spider-scheduler/src/core.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use spider_core::types::id::TaskAssignmentId;
use crate::dispatch_queue::DispatchQueueSink;
use crate::error::SchedulerError;
use crate::storage_client::SchedulerStorageClient;
use crate::types::TaskAssignment;

/// Single-source ID issuer for creating globally unique IDs for task assignments.
pub struct TaskAssignmentIdIssuer {
Expand Down Expand Up @@ -63,6 +64,8 @@ pub trait SchedulerCore: Send {
/// * `storage_client` - The storage client used to poll the inbound queue and read state for
/// placement.
/// * `sink` - The dispatch sink that assignments are written to.
/// * `reschedule_queue_reader` - The reader side of the re-schedule queue, delivering task
/// assignments returned for re-placement when an execution manager is lost.
/// * `id_issuer` - The single-source ID issuer for creating globally unique IDs for task
/// assignments.
/// * `cancellation_token` - The token to signal the scheduling loop to stop.
Expand All @@ -74,6 +77,7 @@ pub trait SchedulerCore: Send {
self: Box<Self>,
storage_client: Self::StorageClient,
sink: Self::Sink,
reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver<TaskAssignment>,
id_issuer: TaskAssignmentIdIssuer,
cancellation_token: tokio_util::sync::CancellationToken,
) -> Result<(), SchedulerError>;
Expand Down
89 changes: 87 additions & 2 deletions components/spider-scheduler/src/core_impl/round_robin/implementation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,13 +114,15 @@ impl<
self: Box<Self>,
storage_client: Self::StorageClient,
sink: Self::Sink,
reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver<TaskAssignment>,
id_issuer: TaskAssignmentIdIssuer,
cancellation_token: CancellationToken,
) -> Result<(), SchedulerError> {
RoundRobin::new(
SessionId::default(),
storage_client,
sink,
reschedule_queue_reader,
id_issuer,
cancellation_token,
self.config,
Expand Down Expand Up @@ -202,6 +204,7 @@ pub(super) struct RoundRobin<
pub(super) finalizing_job_queue: VecDeque<(JobId, Instant)>,

pub(super) inbound_queue_reader: AsyncInboundQueueReader<SchedulerStorageClientType>,
pub(super) reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver<TaskAssignment>,
}

impl<
Expand All @@ -220,6 +223,7 @@ impl<
storage_session_id: SessionId,
storage_client: SchedulerStorageClientType,
sink: DispatchQueueSinkType,
reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver<TaskAssignment>,
id_issuer: TaskAssignmentIdIssuer,
cancellation_token: CancellationToken,
config: RoundRobinConfig,
Expand Down Expand Up @@ -254,21 +258,25 @@ impl<
finalizing_jobs,
finalizing_job_queue,
inbound_queue_reader,
reschedule_queue_reader,
}
}

/// Executes a single scheduling tick: consumes any completed inbound poll, then makes
/// scheduling decisions to fill the dispatch queue.
/// Executes a single scheduling tick: consumes any completed inbound poll, re-injects
/// assignments returned by lost execution managers, then makes scheduling decisions to fill the
/// dispatch queue.
///
/// # Errors
///
/// Returns an error if:
///
/// * Forwards [`Self::consume_inbound_poll_result`]'s return values on failure.
/// * Forwards [`Self::reschedule`]'s return values on failure.
/// * Forwards [`Self::make_schedule_decisions`]'s return values on failure.
pub(super) async fn tick(&mut self) -> Result<(), SchedulerError> {
tracing::info!("Starting scheduling tick.");
self.consume_inbound_poll_result().await?;
self.reschedule()?;
self.make_schedule_decisions().await?;
self.retire_expired_finalizing_jobs();
Ok(())
Expand Down Expand Up @@ -488,6 +496,78 @@ impl<
Ok(())
}

/// Re-injects assignments recovered from lost execution managers into the scheduler.
///
/// The assignments are grouped by task kind and processed through the same
/// commit-ready, cleanup-ready, and ready paths as inbound entries. This ensures that
/// finalization semantics and buffered-task deduplication are applied consistently.
///
/// Rescheduled assignments are enqueued without respecting the internal buffer's capacity
/// limit. As a result, the buffer may temporarily contain more assignments than its configured
/// capacity.
///
/// # Errors
///
/// Returns an error if:
///
/// * [`SchedulerError::Internal`] if the re-schedule queue reader fails to receive a buffered
/// assignment.
/// * Forwards [`Self::enqueue_commit_ready_entries`]'s return values on failure.
/// * Forwards [`Self::enqueue_cleanup_ready_entries`]'s return values on failure.
fn reschedule(&mut self) -> Result<(), SchedulerError> {
if self.reschedule_queue_reader.is_empty() {
return Ok(());
}

// Snapshot the size so assignments pushed concurrently during the drain are deferred to the
// next tick instead of extending this loop unboundedly.
let mut reschedule_queue_size = self.reschedule_queue_reader.len();
let mut ready_entries = Vec::new();
let mut commit_ready_entries = Vec::new();
let mut cleanup_ready_entries = Vec::new();
while reschedule_queue_size > 0 {
let assignment = self.reschedule_queue_reader.try_recv().map_err(|e| {
tracing::error!(
err = % e,
"Reschedule queue reader failed to receive a message."
);
SchedulerError::Internal(
"reschedule queue reader failed to receive a message".to_string(),
)
})?;
reschedule_queue_size -= 1;

if assignment.session_id < self.storage_session_id {
continue;
}

let entry = InboundEntry {
resource_group_id: assignment.resource_group_id,
job_id: assignment.job_id,
task_id: assignment.task_id,
};
match &assignment.task_id {
TaskId::Index(_) => {
ready_entries.push(entry);
}
TaskId::Commit => {
commit_ready_entries.push(entry);
}
TaskId::Cleanup => {
cleanup_ready_entries.push(entry);
}
}
}

// Load commit-ready tasks and cleanup-ready tasks first to avoid loading a job that is
// already finalizing.
self.enqueue_commit_ready_entries(commit_ready_entries)?;
self.enqueue_cleanup_ready_entries(cleanup_ready_entries)?;
self.enqueue_ready_entries(ready_entries);

Ok(())
}

/// Enqueues polled commit-ready entries: each entry's job is marked finalizing, queued for a
/// commit-task assignment, and removed from the active or pending set.
///
Expand Down Expand Up @@ -732,6 +812,7 @@ impl<
job_id,
resource_group_id,
task_id: TaskId::Cleanup,
session_id: self.storage_session_id,
})
.await?;
self.buffered_tasks.remove(&(job_id, TaskId::Cleanup));
Expand All @@ -752,6 +833,7 @@ impl<
job_id,
resource_group_id,
task_id: TaskId::Commit,
session_id: self.storage_session_id,
})
.await?;
self.buffered_tasks.remove(&(job_id, TaskId::Commit));
Expand All @@ -771,6 +853,7 @@ impl<
job_id,
resource_group_id: job_entry.resource_group_id,
task_id,
session_id: self.storage_session_id,
})
.await?;
self.buffered_tasks.remove(&(job_id, task_id));
Expand Down Expand Up @@ -800,6 +883,8 @@ impl<
///
/// * Forwards [`AsyncInboundQueueReader::start`]'s return values on failure.
fn start_inbound_poll(&mut self) -> Result<(), SchedulerError> {
// The reschedule path can load the ready buffers beyond their configured capacity, so
// `saturating_sub` floors the remaining poll budget at zero instead of underflowing.
let num_commit_ready_tasks = self.commit_ready_jobs.len();
let num_cleanup_ready_tasks = self.cleanup_ready_jobs.len();
let max_commit_ready_entries = self
Expand Down
Loading
Loading