From 1399d9153e565ebbe14a0adaa1f2f4acf7ee0247 Mon Sep 17 00:00:00 2001 From: Jack Luo Date: Thu, 6 Aug 2026 03:50:19 -0400 Subject: [PATCH 1/3] fix(spider-execution-manager): Send `prev_assignment` at most once per `next_task` call (fixes #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 --- .../src/client/grpc/scheduler.rs | 142 +++++++++++++++++- 1 file changed, 140 insertions(+), 2 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/scheduler.rs b/components/spider-execution-manager/src/client/grpc/scheduler.rs index c018717b3..442c12664 100644 --- a/components/spider-execution-manager/src/client/grpc/scheduler.rs +++ b/components/spider-execution-manager/src/client/grpc/scheduler.rs @@ -54,16 +54,18 @@ impl SchedulerClient for GrpcSchedulerClient { async fn next_task( &self, em_id: ExecutionManagerId, - prev_assignment: Option, + mut prev_assignment: Option, wait_time_ms: u64, ) -> Result { + // The scheduler completes `prev_assignment` by removing it from its registry, so it must be + // sent at most once; every later poll iteration sends `None`. loop { let response = self .connection_pool .get_client() .next_task(scheduler::NextTaskRequest { execution_manager_id: em_id.get(), - prev_assignment: prev_assignment.map(Into::into), + prev_assignment: prev_assignment.take().map(Into::into), wait_time_ms, }) .await @@ -147,8 +149,144 @@ fn to_protocol_error(error: impl std::fmt::Display) -> SchedulerError { #[cfg(test)] mod tests { + 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 { + 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>>>, + } + + #[async_trait] + impl SchedulerService for FakeScheduler { + async fn next_task( + &self, + request: Request, + ) -> Result, 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, + ) -> Result, Status> { + Ok(Response::new(common::Void {})) + } + + async fn shutdown( + &self, + _request: Request, + ) -> Result, 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(()) + } + #[test] fn status_maps_unavailable_to_transport() { const MESSAGE: &str = "connection lost"; From 892efa66fefaba3385a62e364ee0813d90d729cf Mon Sep 17 00:00:00 2001 From: Jack Luo Date: Thu, 6 Aug 2026 11:26:20 -0400 Subject: [PATCH 2/3] test(spider-execution-manager): Drop the `next_task` long-poll regression 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 --- .../src/client/grpc/scheduler.rs | 136 ------------------ 1 file changed, 136 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/scheduler.rs b/components/spider-execution-manager/src/client/grpc/scheduler.rs index 442c12664..b819dcb2c 100644 --- a/components/spider-execution-manager/src/client/grpc/scheduler.rs +++ b/components/spider-execution-manager/src/client/grpc/scheduler.rs @@ -149,144 +149,8 @@ fn to_protocol_error(error: impl std::fmt::Display) -> SchedulerError { #[cfg(test)] mod tests { - 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 { - 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>>>, - } - - #[async_trait] - impl SchedulerService for FakeScheduler { - async fn next_task( - &self, - request: Request, - ) -> Result, 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, - ) -> Result, Status> { - Ok(Response::new(common::Void {})) - } - - async fn shutdown( - &self, - _request: Request, - ) -> Result, 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(()) - } - #[test] fn status_maps_unavailable_to_transport() { const MESSAGE: &str = "connection lost"; From 60203b6b9c51725ab28479ea1fede5d4b91d772c Mon Sep 17 00:00:00 2001 From: Jack Luo Date: Thu, 6 Aug 2026 13:06:03 -0400 Subject: [PATCH 3/3] docs(spider-execution-manager): Name `prev_assignment` explicitly in 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. --- .../spider-execution-manager/src/client/grpc/scheduler.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/components/spider-execution-manager/src/client/grpc/scheduler.rs b/components/spider-execution-manager/src/client/grpc/scheduler.rs index b819dcb2c..006a12396 100644 --- a/components/spider-execution-manager/src/client/grpc/scheduler.rs +++ b/components/spider-execution-manager/src/client/grpc/scheduler.rs @@ -57,8 +57,8 @@ impl SchedulerClient for GrpcSchedulerClient { mut prev_assignment: Option, wait_time_ms: u64, ) -> Result { - // The scheduler completes `prev_assignment` by removing it from its registry, so it must be - // sent at most once; every later poll iteration sends `None`. + // Completing an assignment removes it from the scheduler's registry, so `prev_assignment` + // must be sent at most once; every later poll iteration sends `None`. loop { let response = self .connection_pool