feat(spider-storage): Add scheduler registration protocols and MariaDB implementation.#347
Merged
Merged
Conversation
Contributor
WalkthroughAdds full-stack scheduler registration support: a new ChangesScheduler Registration Storage Support
Sequence Diagram(s)sequenceDiagram
participant Client as gRPC Client
participant StorageServer as StorageServer (SchedulerRegistrationService)
participant ServiceState
participant MariaDbStorageConnector
participant DB as MariaDB
Client->>StorageServer: RegisterScheduler(ip_address, port)
StorageServer->>ServiceState: register_scheduler(ip_address, port)
ServiceState->>MariaDbStorageConnector: register_scheduler(ip_address, port)
MariaDbStorageConnector->>DB: BEGIN TRANSACTION
MariaDbStorageConnector->>DB: DELETE FROM schedulers
MariaDbStorageConnector->>DB: INSERT (ip_address, port)
DB-->>MariaDbStorageConnector: SchedulerId
MariaDbStorageConnector->>DB: COMMIT
MariaDbStorageConnector-->>ServiceState: Ok(SchedulerId)
ServiceState-->>StorageServer: Ok(SchedulerId)
StorageServer-->>Client: RegisterSchedulerResponse { scheduler }
Client->>StorageServer: GetSchedulers()
StorageServer->>ServiceState: get_schedulers()
ServiceState->>MariaDbStorageConnector: get_schedulers()
MariaDbStorageConnector->>DB: SELECT * FROM schedulers ORDER BY id
DB-->>MariaDbStorageConnector: Vec<RegisteredScheduler>
MariaDbStorageConnector-->>ServiceState: Ok(Vec<RegisteredScheduler>)
ServiceState-->>StorageServer: Ok(Vec<RegisteredScheduler>)
StorageServer-->>Client: GetSchedulersResponse { scheduler_registrations }
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@components/spider-storage/src/db/mariadb.rs`:
- Around line 599-619: The register_scheduler function performs a DELETE
followed by INSERT within a transaction, but concurrent invocations can
interleave causing multiple schedulers to be registered simultaneously,
violating the single-active-scheduler invariant. Fix this by ensuring the DELETE
and INSERT operations are atomic and exclusive: either use a higher transaction
isolation level (SERIALIZABLE), add row-level locking with SELECT FOR UPDATE on
a singleton control row, or modify the logic to use an atomic REPLACE/UPDATE
statement that naturally enforces the single-active constraint. Additionally,
the same concurrency issue exists in the code referenced at line 738-744 (Also
applies to), so apply the same singleton-enforcing mechanism there to ensure
both locations consistently prevent multiple active schedulers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 7edf5154-2274-4510-af3c-c0e9fe4ba3d4
⛔ Files ignored due to path filters (1)
components/spider-proto-rust/src/generated/storage.rsis excluded by!**/generated/**
📒 Files selected for processing (7)
components/spider-proto/storage/storage.protocomponents/spider-storage/src/db.rscomponents/spider-storage/src/db/mariadb.rscomponents/spider-storage/src/db/protocol.rscomponents/spider-storage/src/state/service.rscomponents/spider-storage/src/state/test_utils.rscomponents/spider-storage/tests/mariadb_test.rs
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
components/spider-storage/src/state/test_utils.rs (1)
256-273: ⚡ Quick winImplement scheduler registration methods for testability.
All three
SchedulerRegistrationManagementmethods are stubbed withunreachable!(), which will panic if called by any test. This prevents unit testing ofServiceStatescheduler methods and is inconsistent with other mock trait implementations in this file (e.g.,ExecutionManagerLivenessManagementat lines 217-254).Consider implementing a simple in-memory mock similar to the execution manager implementation using
Arc<DashMap<SchedulerId, RegisteredScheduler>>and an atomic counter for ID generation.🤖 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/test_utils.rs` around lines 256 - 273, The three methods register_scheduler, get_schedulers, and is_scheduler_registered in the SchedulerRegistrationManagement trait implementation for MockDbConnector are stubbed with unreachable!() macros, which will panic if called during tests. Replace these stubs with actual in-memory mock implementations using Arc<DashMap> to store registered schedulers and an atomic counter for generating unique SchedulerIds (similar to the ExecutionManagerLivenessManagement implementation pattern already in the file). The register_scheduler method should generate a new ID and store the scheduler entry, get_schedulers should return all stored scheduler entries, and is_scheduler_registered should check if a given ID exists in the map.components/spider-storage/src/db/protocol.rs (1)
432-471: ⚡ Quick winConsider whether
get_schedulersreturn type matches the singleton constraint.The documentation at line 437 states that "only one scheduler can be registered at a time," and the PR objectives confirm this design assumption. However,
get_schedulersreturnsVec<RegisteredScheduler>, which suggests support for multiple schedulers.If the singleton constraint is a temporary limitation and multiple schedulers may be supported in the future, the
Vecreturn type provides forward compatibility. Otherwise, consider returningOption<RegisteredScheduler>to make the single-scheduler constraint explicit in the API signature.🤖 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/db/protocol.rs` around lines 432 - 471, The SchedulerRegistrationManagement trait's get_schedulers method returns Vec<RegisteredScheduler>, which contradicts the documented singleton constraint that states only one scheduler can be registered at a time. To align the API signature with the documented design constraint, change the return type of the get_schedulers method from Vec<RegisteredScheduler> to Option<RegisteredScheduler>. This makes the single-scheduler limitation explicit in the type system and clarifies to API consumers that only zero or one scheduler can be registered.
🤖 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/db/protocol.rs`:
- Around line 432-471: The SchedulerRegistrationManagement trait's
get_schedulers method returns Vec<RegisteredScheduler>, which contradicts the
documented singleton constraint that states only one scheduler can be registered
at a time. To align the API signature with the documented design constraint,
change the return type of the get_schedulers method from
Vec<RegisteredScheduler> to Option<RegisteredScheduler>. This makes the
single-scheduler limitation explicit in the type system and clarifies to API
consumers that only zero or one scheduler can be registered.
In `@components/spider-storage/src/state/test_utils.rs`:
- Around line 256-273: The three methods register_scheduler, get_schedulers, and
is_scheduler_registered in the SchedulerRegistrationManagement trait
implementation for MockDbConnector are stubbed with unreachable!() macros, which
will panic if called during tests. Replace these stubs with actual in-memory
mock implementations using Arc<DashMap> to store registered schedulers and an
atomic counter for generating unique SchedulerIds (similar to the
ExecutionManagerLivenessManagement implementation pattern already in the file).
The register_scheduler method should generate a new ID and store the scheduler
entry, get_schedulers should return all stored scheduler entries, and
is_scheduler_registered should check if a given ID exists in the map.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 5a457c28-75d3-441c-9cfd-5fbe45efd336
📒 Files selected for processing (7)
components/spider-core/src/types.rscomponents/spider-core/src/types/scheduler.rscomponents/spider-storage/src/db.rscomponents/spider-storage/src/db/mariadb.rscomponents/spider-storage/src/db/protocol.rscomponents/spider-storage/src/state/service.rscomponents/spider-storage/src/state/test_utils.rs
💤 Files with no reviewable changes (1)
- components/spider-storage/src/db.rs
✅ Files skipped from review due to trivial changes (1)
- components/spider-core/src/types.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- components/spider-storage/src/db/mariadb.rs
LinZhihao-723
requested changes
Jun 17, 2026
LinZhihao-723
left a comment
LinZhihao-723
left a comment
Member
There was a problem hiding this comment.
Changes I made:
- Move
RegisteredSchedulerintospider-core. This type should be shared across other crates. And it's not the best idea to depend onspider-storagefor type sharing. - The mock DB connector for the service doesn't make sense: the correctness of this service purely depends on the DB layer implementation; the service doesn't do anything but forward the request. These test cases are not actually testing the service, but testing wheteher ur mock DB connector is "mocking" the intentional behavior, lol. Removed them all for cleaning.
|
|
||
| message GetSchedulersResponse { | ||
| oneof result { | ||
| SchedulerRegistrations schedulers = 1; |
Member
There was a problem hiding this comment.
I understand why we need session ID in SchedulerRegistrations, but this is why also needed for GetSchedulersResponse?
Collaborator
Author
There was a problem hiding this comment.
This is likely the first function the execution manager calls. It can get the session id the first chance it has.
Member
There was a problem hiding this comment.
The first call is EM registration, right?
We should probably design this API to be session-independent since it might be called multiple times through the lifetime of an EM.
LinZhihao-723
changed the title
LinZhihao-723
approved these changes
Jun 18, 2026
LinZhihao-723
left a comment
LinZhihao-723
left a comment
Member
There was a problem hiding this comment.
Directly modified the PR title.
This was referenced Jun 25, 2026
4 tasks
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
This PR:
Note
We assume at any given time, there is only one scheduler. The current registration removes all previous scheduler registration.
Checklist
breaking change.
Validation performed
Summary by CodeRabbit
Release Notes
New Features
Tests