fix(spider-execution-manager): Send prev_assignment at most once per next_task call (fixes #429). - #430
Merged
LinZhihao-723 merged 3 commits intoAug 7, 2026
Conversation
…r `next_task` call (fixes y-scope#429). `GrpcSchedulerClient::next_task` long-polls the scheduler in an internal loop until a task becomes available, but re-sent the same `prev_assignment` on every iteration. The scheduler completes an assignment by removing it from its registry, so only the first iteration succeeded; each subsequent one made the scheduler fail to complete an assignment it had already dropped, logging a warning on every poll for as long as the execution manager stayed idle. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Contributor
WalkthroughThe scheduler now consumes ChangesScheduler polling
Estimated code review effort: 1 (Trivial) | ~5 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
LinZhihao-723
left a comment
LinZhihao-723
left a comment
Member
There was a problem hiding this comment.
The fix is correct. Let's remove the unit tests as explained in my comment.
Comment on lines
+152
to
+289
| use std::net::TcpListener as StdTcpListener; | ||
| use std::sync::Arc; | ||
| use std::sync::Mutex; | ||
| use std::time::Duration; | ||
|
|
||
| use spider_core::types::id::SchedulerId; | ||
| use spider_core::types::id::TaskAssignmentId; | ||
| use spider_proto_rust::common; | ||
| use spider_proto_rust::common::TaskId as ProtoTaskId; | ||
| use spider_proto_rust::common::task_id::Kind as ProtoTaskIdKind; | ||
| use spider_proto_rust::scheduler::SchedulerAssignment; | ||
| use spider_proto_rust::scheduler::SchedulerService; | ||
| use spider_proto_rust::scheduler::SchedulerServiceServer; | ||
| use spider_proto_rust::scheduler::next_task_response; | ||
| use tonic::Request; | ||
| use tonic::Response; | ||
| use tonic::transport::Server; | ||
|
|
||
| use super::*; | ||
|
|
||
| /// The number of `NoTask` replies the fake scheduler sends before handing out an assignment. | ||
| const NUM_NO_TASK_REPLIES: usize = 3; | ||
|
|
||
| /// The maximum time spent waiting for the test server to accept connections. | ||
| const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); | ||
|
|
||
| /// The delay between connection attempts while the test server is still binding. | ||
| const CONNECT_RETRY_DELAY: Duration = Duration::from_millis(50); | ||
|
|
||
| /// Connects to `endpoint`, retrying for up to [`CONNECT_TIMEOUT`] while the test server binds | ||
| /// its listener. | ||
| /// | ||
| /// # Returns | ||
| /// | ||
| /// A connected [`GrpcSchedulerClient`] on success. | ||
| /// | ||
| /// # Errors | ||
| /// | ||
| /// Forwards [`GrpcSchedulerClient::connect`]'s failure once [`CONNECT_TIMEOUT`] elapses. | ||
| async fn connect_with_retries(endpoint: &Endpoint) -> anyhow::Result<GrpcSchedulerClient> { | ||
| let pool_size = NonZeroUsize::new(1).expect("1 is non-zero"); | ||
| let deadline = tokio::time::Instant::now() + CONNECT_TIMEOUT; | ||
| loop { | ||
| match GrpcSchedulerClient::connect(endpoint.clone(), pool_size).await { | ||
| Ok(client) => return Ok(client), | ||
| Err(error) if tokio::time::Instant::now() >= deadline => return Err(error.into()), | ||
| Err(_) => tokio::time::sleep(CONNECT_RETRY_DELAY).await, | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// A fake scheduler that records every request's `prev_assignment`, replying `NoTask` | ||
| /// [`NUM_NO_TASK_REPLIES`] times before handing out an assignment. | ||
| struct FakeScheduler { | ||
| observed: Arc<Mutex<Vec<Option<scheduler::TaskAssignmentRecord>>>>, | ||
| } | ||
|
|
||
| #[async_trait] | ||
| impl SchedulerService for FakeScheduler { | ||
| async fn next_task( | ||
| &self, | ||
| request: Request<scheduler::NextTaskRequest>, | ||
| ) -> Result<Response<scheduler::NextTaskResponse>, Status> { | ||
| let num_requests = { | ||
| let mut observed = self.observed.lock().expect("lock shouldn't be poisoned"); | ||
| observed.push(request.into_inner().prev_assignment); | ||
| observed.len() | ||
| }; | ||
| let result = if num_requests > NUM_NO_TASK_REPLIES { | ||
| next_task_response::Result::Assignment(SchedulerAssignment { | ||
| id: 1, | ||
| resource_group_id: 2, | ||
| job_id: 3, | ||
| task_id: Some(ProtoTaskId { | ||
| kind: Some(ProtoTaskIdKind::Index(0)), | ||
| }), | ||
| scheduler_id: 4, | ||
| session_id: 5, | ||
| }) | ||
| } else { | ||
| next_task_response::Result::NoTask(common::Void {}) | ||
| }; | ||
| Ok(Response::new(scheduler::NextTaskResponse { | ||
| result: Some(result), | ||
| })) | ||
| } | ||
|
|
||
| async fn heartbeat( | ||
| &self, | ||
| _request: Request<scheduler::HeartbeatRequest>, | ||
| ) -> Result<Response<common::Void>, Status> { | ||
| Ok(Response::new(common::Void {})) | ||
| } | ||
|
|
||
| async fn shutdown( | ||
| &self, | ||
| _request: Request<scheduler::ShutdownRequest>, | ||
| ) -> Result<Response<common::Void>, Status> { | ||
| Ok(Response::new(common::Void {})) | ||
| } | ||
| } | ||
|
|
||
| /// Tests that `prev_assignment` is sent exactly once even when the client long-polls several | ||
| /// times before a task becomes available. The scheduler completes it by removing it from its | ||
| /// registry, so re-sending it makes the scheduler fail to complete an already-dropped | ||
| /// assignment. | ||
| #[tokio::test] | ||
| async fn prev_assignment_is_sent_once_across_long_poll_iterations() -> anyhow::Result<()> { | ||
| let observed = Arc::new(Mutex::new(Vec::new())); | ||
| // Reserve an ephemeral port, then release it so the server can bind it. | ||
| let address = StdTcpListener::bind("127.0.0.1:0")?.local_addr()?; | ||
| let service = SchedulerServiceServer::new(FakeScheduler { | ||
| observed: Arc::clone(&observed), | ||
| }); | ||
| let server = tokio::spawn(Server::builder().add_service(service).serve(address)); | ||
|
|
||
| let endpoint = Endpoint::from_shared(format!("http://{address}"))?; | ||
| let client = connect_with_retries(&endpoint).await?; | ||
| let prev_assignment = | ||
| TaskAssignmentRecord::new(TaskAssignmentId::from(7), SchedulerId::from(4)); | ||
| client | ||
| .next_task(ExecutionManagerId::from(2), Some(prev_assignment), 0) | ||
| .await?; | ||
| server.abort(); | ||
|
|
||
| let observed = observed.lock().expect("lock shouldn't be poisoned").clone(); | ||
| assert_eq!(observed.len(), NUM_NO_TASK_REPLIES + 1); | ||
| assert_eq!( | ||
| observed[0], | ||
| Some(scheduler::TaskAssignmentRecord { id: 7, from: 4 }) | ||
| ); | ||
| assert!( | ||
| observed[1..].iter().all(Option::is_none), | ||
| "later poll iterations resent the assignment: {observed:?}" | ||
| ); | ||
| Ok(()) | ||
| } | ||
|
|
Member
There was a problem hiding this comment.
Prefer not to add this unit test since it adds too much unnecessary scaffolding to test this simple behavior. This increases the lines of code changes needed when we need to make any change to the scheduler gRPC trait.
…sion test. Per review: the fake `SchedulerService` is more scaffolding than the one-line behaviour warrants, and it couples this file to every future change to the scheduler gRPC trait. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…the `next_task` loop comment. The comment used three pronouns for two referents, so `it must be sent at most once` read as if the scheduler's registry were the thing being sent. Naming the field removes the ambiguity.
LinZhihao-723
approved these changes
Aug 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes #429.
GrpcSchedulerClient::next_taskre-sentprev_assignmenton every iteration of its long-poll loop, so the scheduler logged a warning for each iteration after the first. This callsOption::takeon the field, which sends it on the first iteration andNoneon every later one. #429 carries the mechanism and the log output.Not addressed here
When the first RPC fails at the transport layer, the completion is lost and the assignment stays in the scheduler's registry until the execution manager's liveness expires. That predates this change and is worth handling separately.
Checklist
Validation performed
Cluster reproduction. A cluster of 1 storage, 1 scheduler, and 4 execution managers, backed by MariaDB 10.11.16, emitted a steady 3 warnings/sec once every job had reached
Succeeded. Details in #429.Suites.
cargo test -p spider-execution-managerandtask lint:check-rustpass on this branch.