feat(spider-storage): Implement scheduler registration gRPC service and re-enqueue ready tasks when a scheduler reconnects.

[LinZhihao-723]

Description

This PR implements scheduler registration in the storage service end-to-end. ServiceState::register_scheduler is fleshed out and wired into the previously-stubbed SchedulerRegistrationService::register_scheduler gRPC handler, so a scheduler can register with storage and receive its assigned SchedulerId and the current storage session.

The core behavior added here is recovery on scheduler reconnect: registration is mutually exclusive, and when a new registration replaces a previously-connected scheduler, storage re-enqueues every cached job's ready tasks in a background task so the new scheduler's inbound queue is repopulated. To support cancelling the service if that background resend fails, the runtime's CancellationToken is now threaded into ServiceState. Along the way, ServiceState::new's growing argument list is bundled into a ServiceStateParams struct to silence clippy::too_many_arguments.

The scope is limited to SchedulerRegistrationService. get_schedulers remains unimplemented (it now returns an explicit UNIMPLEMENTED status instead of panicking).

Service layer (state/service.rs)

• register_scheduler now guards registration with a per-service async Mutex<bool> (has_previous_scheduler_connection), making registrations mutually exclusive: only one can be in flight at a time. It calls db.register_scheduler, and on success, if a scheduler was previously registered, it spawns a background task that calls JobCache::resend_ready_tasks(). If that resend fails, the task logs the error and cancels the service via the CancellationToken. The flag is then set to true. The resend runs off the critical path (in a spawned task), so the RPC returns promptly.
• ServiceStateInner gains two fields: has_previous_scheduler_connection: tokio::sync::Mutex<bool> and cancellation_token: CancellationToken.
• Because register_scheduler moves clones of job_cache and the cancellation token into a tokio::spawned future, the three service type parameters now carry a 'static bound (ServiceState, ServiceStateInner, and — transitively — Runtime and GrpcServiceState).

Constructor bundle (ServiceStateParams, state/service.rs + state.rs)

• Added ServiceStateParams, a pub struct holding the eight fields ServiceState::new needs (db, session_id, job_cache, ready_queue_sender, ready_queue_receiver, task_instance_pool_connector, job_cache_gc_handle, cancellation_token). new now takes a single ServiceStateParams and destructures it, resolving the clippy::too_many_arguments warning that the new cancellation_token argument would otherwise trip.
• Re-exported ServiceStateParams from the state module so runtime.rs (and other constructors) can build it with struct-literal syntax.

gRPC layer (grpc.rs)

• Implemented SchedulerRegistrationService::register_scheduler (was todo!("Not implemented")): it unpacks the request into (ip_addr, port), calls ServiceState::register_scheduler, and returns a RegisterSchedulerResponse carrying SchedulerRegistration { scheduler_id, session_id }.
• get_schedulers now returns Status::unimplemented("not implemented") instead of todo! — a still-unimplemented RPC no longer panics (and takes down) the service if called.
• Added scheduler_registration_service_error_handler, the single chokepoint mapping a StorageServerError from this service to a Status. Every error maps to an opaque INTERNAL via default_error_handler (non-strict), mirroring inbound_queue_service_error_handler.
• GrpcServiceState's type parameters gain the 'static bound to match ServiceState.

Runtime wiring (state/runtime.rs)

• create_runtime now passes cancellation_token.clone() into ServiceStateParams when constructing the service state.
• Runtime's type parameters gain the 'static bound (propagated from ServiceState).
• Both the production and test call sites were migrated to the ServiceStateParams { .. } struct-literal form.

Test mock (state/test_utils.rs)

• MockDbConnector::register_scheduler previously did unreachable!("not implemented for mock connector"), which would panic the moment the service called into it. It now returns an incrementing SchedulerId from a new next_scheduler_id: Arc<AtomicUsize> counter, mirroring the existing register_execution_manager mock. get_schedulers / is_scheduler_registered remain unreachable!.

Tests

• Added register_scheduler_resends_ready_tasks_only_when_replacing_previous_scheduler. It builds a service backed by a real ready queue, submits and starts a single-task job, and drains the ready task that start_job enqueues so the queue is empty before probing. It then asserts that the first register_scheduler does not resend (queue stays empty), and that the second register_scheduler does resend the job's ready task (one entry, matching job_id). Because the resend runs in a spawned background task, the test uses tokio::task::yield_now() to let it run before polling.
• The existing constructor-touching tests (cancel_job_/succeed_task_instance_/fail_task_instance_enqueues_terminal_job_for_cache_gc, and the two service factories) were updated to the ServiceStateParams form.

Checklist

• The PR satisfies the contribution guidelines.
• This is a breaking change and that has been indicated in the PR title, OR this isn't a
  breaking change.
• Necessary docs have been updated, OR no docs need to be updated.

Validation performed

• Ensure all workflows pass.

Summary by CodeRabbit

• New Features

  • Added scheduler registration support in storage runtime gRPC.
  • Returning scheduler registration details now includes the scheduler ID and current session information.

• Bug Fixes

  • Improved scheduler replacement handling so ready tasks are resent when a new scheduler replaces an existing one.
  • If task resending fails, the service now shuts down cleanly instead of continuing in a bad state.
  • Scheduler lookup still returns an explicit “unimplemented” response rather than failing unexpectedly.

[coderabbitai]

Walkthrough

This PR implements the previously stubbed SchedulerRegistrationService gRPC endpoints (register_scheduler, get_schedulers), refactors ServiceState construction to use a new ServiceStateParams struct with a CancellationToken, adds mutex-guarded resend-on-scheduler-replacement logic, and updates test mocks/utilities accordingly.

Changes

Scheduler registration and service state refactor

Layer / File(s)	Summary
Scheduler registration gRPC endpoints components/spider-storage/src/grpc.rs	Implements register_scheduler and get_schedulers gRPC handlers, adds a dedicated error handler, imports SchedulerRegistration, and tightens GrpcServiceState generic bounds with 'static.
ServiceStateParams introduction and re-export components/spider-storage/src/state.rs, components/spider-storage/src/state/service.rs	Introduces ServiceStateParams bundling constructor inputs plus a cancellation_token, refactors ServiceState::new to destructure it, adds has_previous_scheduler_connection and cancellation_token to ServiceStateInner, and re-exports the new type.
Mutually exclusive scheduler registration with resend logic components/spider-storage/src/state/service.rs	Reworks register_scheduler to lock a mutex flag and conditionally spawn a background resend of ready tasks on replacement, cancelling the service via CancellationToken on failure.
Runtime and service construction wiring components/spider-storage/src/state/runtime.rs	Adds 'static bounds to Runtime generics and updates create_runtime/create_test_runtime to build ServiceState via ServiceStateParams literals.
Test helper and mock updates components/spider-storage/src/state/service.rs, components/spider-storage/src/state/test_utils.rs	Updates unit tests to use ServiceStateParams, expands the scheduler resend test, and updates MockDbConnector to allocate incrementing scheduler IDs instead of panicking.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant SchedulerRegistrationService
  participant ServiceState
  participant JobCache

  Client->>SchedulerRegistrationService: register_scheduler(ip_addr, port)
  SchedulerRegistrationService->>ServiceState: register_scheduler(ip_addr, port)
  ServiceState->>ServiceState: lock has_previous_scheduler_connection
  alt previous scheduler existed
    ServiceState->>JobCache: spawn resend_ready_tasks()
    JobCache-->>ServiceState: success or error
    ServiceState->>ServiceState: cancel() via CancellationToken on error
  end
  ServiceState-->>SchedulerRegistrationService: scheduler_id or error
  SchedulerRegistrationService-->>Client: RegisterSchedulerResponse or error status

Loading

Possibly related PRs

• y-scope/spider#342: Updates the same get_schedulers gRPC handler signature in grpc.rs.
• y-scope/spider#347: Adds scheduler registration protocol types and storage support that align with the newly implemented endpoints.
• y-scope/spider#368: Adds a SchedulerStorageClient::register call flow that directly consumes the registration API implemented here.

Suggested reviewers: sitaowang1998

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: scheduler registration gRPC implementation plus reconnect-triggered ready-task re-enqueueing.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

components/spider-storage/src/state/service.rs (1)

715-745: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value

Detached resend task: consider overlapping resends and shutdown behaviour.

The mutex correctly serializes registrations and gates the resend. Two edge cases worth confirming:

• The resend is fire-and-forget and only guarded by the flag, not the lock, so rapid re-registrations can spawn multiple resend_ready_tasks() executions that run concurrently against the ready queue.
• The spawned handle isn't tracked, so it can't be awaited/aborted during Runtime::stop.

If both are acceptable for the reconnect flow (resend is idempotent-enough and shutdown-race tolerable), this is fine as-is.

🤖 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-storage/src/state/service.rs` around lines 715 - 745, The
resend path in register_scheduler can spawn multiple concurrent
resend_ready_tasks() jobs on rapid re-registration, and the detached
tokio::spawn handle is never tracked for shutdown. Serialize or deduplicate the
resend by moving the guard/state to cover the spawned work, and store the
JoinHandle somewhere in the service so Runtime::stop can abort or await it using
the existing cancellation_token and job_cache flow.

🤖 Prompt for all review comments with 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.

Nitpick comments:
In `@components/spider-storage/src/state/service.rs`:
- Around line 715-745: The resend path in register_scheduler can spawn multiple
concurrent resend_ready_tasks() jobs on rapid re-registration, and the detached
tokio::spawn handle is never tracked for shutdown. Serialize or deduplicate the
resend by moving the guard/state to cover the spawned work, and store the
JoinHandle somewhere in the service so Runtime::stop can abort or await it using
the existing cancellation_token and job_cache flow.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 2cafb32c-269f-427f-a3c2-d1c5db4ee5c5

📥 Commits

Reviewing files that changed from the base of the PR and between 99b43d0 and 4290c0a.

📒 Files selected for processing (5)

• components/spider-storage/src/grpc.rs
• components/spider-storage/src/state.rs
• components/spider-storage/src/state/runtime.rs
• components/spider-storage/src/state/service.rs
• components/spider-storage/src/state/test_utils.rs