diff --git a/components/spider-core/src/types.rs b/components/spider-core/src/types.rs index df9189789..d6f3baa54 100644 --- a/components/spider-core/src/types.rs +++ b/components/spider-core/src/types.rs @@ -1,2 +1,3 @@ pub mod id; pub mod io; +pub mod scheduler; diff --git a/components/spider-core/src/types/scheduler.rs b/components/spider-core/src/types/scheduler.rs new file mode 100644 index 000000000..b0cbcffea --- /dev/null +++ b/components/spider-core/src/types/scheduler.rs @@ -0,0 +1,11 @@ +use std::net::IpAddr; + +use crate::types::id::SchedulerId; + +/// The currently registered scheduler endpoint. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RegisteredScheduler { + pub id: SchedulerId, + pub ip_address: IpAddr, + pub port: u16, +} diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index 970dc1414..78e88fc4f 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -296,6 +296,64 @@ pub mod update_execution_manager_heartbeat_response { Error(super::ExecutionManagerLivenessError), } } +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RegisterSchedulerRequest { + #[prost(string, tag = "1")] + pub ip_address: ::prost::alloc::string::String, + #[prost(uint32, tag = "2")] + pub port: u32, +} +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct SchedulerRegistration { + #[prost(uint64, tag = "1")] + pub scheduler_id: u64, + #[prost(uint64, tag = "2")] + pub session_id: u64, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RegisterSchedulerResponse { + #[prost(oneof = "register_scheduler_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `RegisterSchedulerResponse`. +pub mod register_scheduler_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(message, tag = "1")] + Registration(super::SchedulerRegistration), + #[prost(message, tag = "2")] + Error(super::SchedulerRegistrationError), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct Scheduler { + #[prost(uint64, tag = "1")] + pub scheduler_id: u64, + #[prost(string, tag = "2")] + pub ip_address: ::prost::alloc::string::String, + #[prost(uint32, tag = "3")] + pub port: u32, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SchedulerRegistrations { + #[prost(message, repeated, tag = "1")] + pub schedulers: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct GetSchedulersResponse { + #[prost(oneof = "get_schedulers_response::Result", tags = "1, 2")] + pub result: ::core::option::Option, +} +/// Nested message and enum types in `GetSchedulersResponse`. +pub mod get_schedulers_response { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Result { + #[prost(message, tag = "1")] + Schedulers(super::SchedulerRegistrations), + #[prost(message, tag = "2")] + Error(super::SchedulerRegistrationError), + } +} #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct GetSessionResponse { #[prost(uint64, tag = "1")] @@ -637,6 +695,52 @@ pub mod execution_manager_liveness_error { } } } +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SchedulerRegistrationError { + #[prost(enumeration = "scheduler_registration_error::ErrCode", tag = "1")] + pub err_code: i32, + #[prost(string, tag = "2")] + pub message: ::prost::alloc::string::String, +} +/// Nested message and enum types in `SchedulerRegistrationError`. +pub mod scheduler_registration_error { + #[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + PartialOrd, + Ord, + ::prost::Enumeration + )] + #[repr(i32)] + pub enum ErrCode { + Unspecified = 0, + Server = 1, + } + impl ErrCode { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Unspecified => "ERR_CODE_UNSPECIFIED", + Self::Server => "SERVER", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "ERR_CODE_UNSPECIFIED" => Some(Self::Unspecified), + "SERVER" => Some(Self::Server), + _ => None, + } + } + } +} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum JobState { @@ -1679,6 +1783,159 @@ pub mod execution_manager_liveness_service_client { } } /// Generated client implementations. +pub mod scheduler_registration_service_client { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + use tonic::codegen::http::Uri; + #[derive(Debug, Clone)] + pub struct SchedulerRegistrationServiceClient { + inner: tonic::client::Grpc, + } + impl SchedulerRegistrationServiceClient { + /// Attempt to create a new client by connecting to a given endpoint. + pub async fn connect(dst: D) -> Result + where + D: TryInto, + D::Error: Into, + { + let conn = tonic::transport::Endpoint::new(dst)?.connect().await?; + Ok(Self::new(conn)) + } + } + impl SchedulerRegistrationServiceClient + where + T: tonic::client::GrpcService, + T::Error: Into, + T::ResponseBody: Body + std::marker::Send + 'static, + ::Error: Into + std::marker::Send, + { + pub fn new(inner: T) -> Self { + let inner = tonic::client::Grpc::new(inner); + Self { inner } + } + pub fn with_origin(inner: T, origin: Uri) -> Self { + let inner = tonic::client::Grpc::with_origin(inner, origin); + Self { inner } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> SchedulerRegistrationServiceClient> + where + F: tonic::service::Interceptor, + T::ResponseBody: Default, + T: tonic::codegen::Service< + http::Request, + Response = http::Response< + >::ResponseBody, + >, + >, + , + >>::Error: Into + std::marker::Send + std::marker::Sync, + { + SchedulerRegistrationServiceClient::new( + InterceptedService::new(inner, interceptor), + ) + } + /// Compress requests with the given encoding. + /// + /// This requires the server to support it otherwise it might respond with an + /// error. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.send_compressed(encoding); + self + } + /// Enable decompressing responses. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.inner = self.inner.accept_compressed(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_decoding_message_size(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.inner = self.inner.max_encoding_message_size(limit); + self + } + pub async fn register_scheduler( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.SchedulerRegistrationService/RegisterScheduler", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.SchedulerRegistrationService", + "RegisterScheduler", + ), + ); + self.inner.unary(req, path, codec).await + } + pub async fn get_schedulers( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.SchedulerRegistrationService/GetSchedulers", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new( + "storage.SchedulerRegistrationService", + "GetSchedulers", + ), + ); + self.inner.unary(req, path, codec).await + } + } +} +/// Generated client implementations. pub mod session_management_service_client { #![allow( unused_variables, @@ -3526,6 +3783,248 @@ pub mod execution_manager_liveness_service_server { } } /// Generated server implementations. +pub mod scheduler_registration_service_server { + #![allow( + unused_variables, + dead_code, + missing_docs, + clippy::wildcard_imports, + clippy::let_unit_value, + )] + use tonic::codegen::*; + /// Generated trait containing gRPC methods that should be implemented for use with SchedulerRegistrationServiceServer. + #[async_trait] + pub trait SchedulerRegistrationService: std::marker::Send + std::marker::Sync + 'static { + async fn register_scheduler( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn get_schedulers( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + } + #[derive(Debug)] + pub struct SchedulerRegistrationServiceServer { + inner: Arc, + accept_compression_encodings: EnabledCompressionEncodings, + send_compression_encodings: EnabledCompressionEncodings, + max_decoding_message_size: Option, + max_encoding_message_size: Option, + } + impl SchedulerRegistrationServiceServer { + pub fn new(inner: T) -> Self { + Self::from_arc(Arc::new(inner)) + } + pub fn from_arc(inner: Arc) -> Self { + Self { + inner, + accept_compression_encodings: Default::default(), + send_compression_encodings: Default::default(), + max_decoding_message_size: None, + max_encoding_message_size: None, + } + } + pub fn with_interceptor( + inner: T, + interceptor: F, + ) -> InterceptedService + where + F: tonic::service::Interceptor, + { + InterceptedService::new(Self::new(inner), interceptor) + } + /// Enable decompressing requests with the given encoding. + #[must_use] + pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.accept_compression_encodings.enable(encoding); + self + } + /// Compress responses with the given encoding, if the client supports it. + #[must_use] + pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self { + self.send_compression_encodings.enable(encoding); + self + } + /// Limits the maximum size of a decoded message. + /// + /// Default: `4MB` + #[must_use] + pub fn max_decoding_message_size(mut self, limit: usize) -> Self { + self.max_decoding_message_size = Some(limit); + self + } + /// Limits the maximum size of an encoded message. + /// + /// Default: `usize::MAX` + #[must_use] + pub fn max_encoding_message_size(mut self, limit: usize) -> Self { + self.max_encoding_message_size = Some(limit); + self + } + } + impl tonic::codegen::Service> + for SchedulerRegistrationServiceServer + where + T: SchedulerRegistrationService, + B: Body + std::marker::Send + 'static, + B::Error: Into + std::marker::Send + 'static, + { + type Response = http::Response; + type Error = std::convert::Infallible; + type Future = BoxFuture; + fn poll_ready( + &mut self, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + fn call(&mut self, req: http::Request) -> Self::Future { + match req.uri().path() { + "/storage.SchedulerRegistrationService/RegisterScheduler" => { + #[allow(non_camel_case_types)] + struct RegisterSchedulerSvc( + pub Arc, + ); + impl< + T: SchedulerRegistrationService, + > tonic::server::UnaryService + for RegisterSchedulerSvc { + type Response = super::RegisterSchedulerResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::register_scheduler( + &inner, + request, + ) + .await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = RegisterSchedulerSvc(inner); + let codec = tonic::codec::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } + "/storage.SchedulerRegistrationService/GetSchedulers" => { + #[allow(non_camel_case_types)] + struct GetSchedulersSvc(pub Arc); + impl< + T: SchedulerRegistrationService, + > tonic::server::UnaryService for GetSchedulersSvc { + type Response = super::GetSchedulersResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::get_schedulers( + &inner, + request, + ) + .await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = GetSchedulersSvc(inner); + let codec = tonic::codec::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } + _ => { + Box::pin(async move { + let mut response = http::Response::new(empty_body()); + let headers = response.headers_mut(); + headers + .insert( + tonic::Status::GRPC_STATUS, + (tonic::Code::Unimplemented as i32).into(), + ); + headers + .insert( + http::header::CONTENT_TYPE, + tonic::metadata::GRPC_CONTENT_TYPE, + ); + Ok(response) + }) + } + } + } + } + impl Clone for SchedulerRegistrationServiceServer { + fn clone(&self) -> Self { + let inner = self.inner.clone(); + Self { + inner, + accept_compression_encodings: self.accept_compression_encodings, + send_compression_encodings: self.send_compression_encodings, + max_decoding_message_size: self.max_decoding_message_size, + max_encoding_message_size: self.max_encoding_message_size, + } + } + } + /// Generated gRPC service name + pub const SERVICE_NAME: &str = "storage.SchedulerRegistrationService"; + impl tonic::server::NamedService for SchedulerRegistrationServiceServer { + const NAME: &'static str = SERVICE_NAME; + } +} +/// Generated server implementations. pub mod session_management_service_server { #![allow( unused_variables, diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index df771fb75..020ab5124 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -39,6 +39,11 @@ service ExecutionManagerLivenessService { returns (UpdateExecutionManagerHeartbeatResponse); } +service SchedulerRegistrationService { + rpc RegisterScheduler(RegisterSchedulerRequest) returns (RegisterSchedulerResponse); + rpc GetSchedulers(Void) returns (GetSchedulersResponse); +} + service SessionManagementService { rpc GetSession(Void) returns (GetSessionResponse); } @@ -211,6 +216,40 @@ message UpdateExecutionManagerHeartbeatResponse { } } +message RegisterSchedulerRequest { + string ip_address = 1; + uint32 port = 2; +} + +message SchedulerRegistration { + uint64 scheduler_id = 1; + uint64 session_id = 2; +} + +message RegisterSchedulerResponse { + oneof result { + SchedulerRegistration registration = 1; + SchedulerRegistrationError error = 2; + } +} + +message Scheduler { + uint64 scheduler_id = 1; + string ip_address = 2; + uint32 port = 3; +} + +message SchedulerRegistrations { + repeated Scheduler schedulers = 1; +} + +message GetSchedulersResponse { + oneof result { + SchedulerRegistrations schedulers = 1; + SchedulerRegistrationError error = 2; + } +} + message GetSessionResponse { uint64 session_id = 1; } @@ -321,3 +360,13 @@ message ExecutionManagerLivenessError { ErrCode err_code = 1; string message = 2; } + +message SchedulerRegistrationError { + enum ErrCode { + ERR_CODE_UNSPECIFIED = 0; + SERVER = 1; + } + + ErrCode err_code = 1; + string message = 2; +} diff --git a/components/spider-storage/src/db.rs b/components/spider-storage/src/db.rs index 453197015..d56c05005 100644 --- a/components/spider-storage/src/db.rs +++ b/components/spider-storage/src/db.rs @@ -11,5 +11,6 @@ pub use protocol::{ InternalJobOrchestration, RecoverableJobContext, ResourceGroupManagement, + SchedulerRegistrationManagement, SessionManagement, }; diff --git a/components/spider-storage/src/db/mariadb.rs b/components/spider-storage/src/db/mariadb.rs index 25d3e1814..230a04ac0 100644 --- a/components/spider-storage/src/db/mariadb.rs +++ b/components/spider-storage/src/db/mariadb.rs @@ -7,8 +7,9 @@ use spider_core::{ job::JobState, task::TaskGraph, types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, SessionId}, + id::{ExecutionManagerId, JobId, ResourceGroupId, SchedulerId, SessionId}, io::{TaskInput, TaskOutput}, + scheduler::RegisteredScheduler, }, }; use spider_derive::MySqlEnum; @@ -25,6 +26,7 @@ use crate::{ InternalJobOrchestration, RecoverableJobContext, ResourceGroupManagement, + SchedulerRegistrationManagement, SessionManagement, error::ExpectedStates, }, @@ -86,6 +88,9 @@ impl MariaDbStorageConnector { sqlx::query(execution_managers_creation_query()) .execute(&pool) .await?; + sqlx::query(schedulers_creation_query()) + .execute(&pool) + .await?; let session_id = sqlx::query_scalar::<_, SessionId>(BUMP_SESSION_ID_QUERY) .fetch_one(&pool) @@ -589,6 +594,57 @@ impl ExecutionManagerLivenessManagement for MariaDbStorageConnector { } } +#[async_trait] +impl SchedulerRegistrationManagement for MariaDbStorageConnector { + async fn register_scheduler( + &self, + ip_address: IpAddr, + port: u16, + ) -> Result { + const DELETE_QUERY: &str = + formatcp!("DELETE FROM `{table}`;", table = SCHEDULERS_TABLE_NAME,); + const INSERT_QUERY: &str = formatcp!( + "INSERT INTO `{table}` (`ip_address`, `port`) VALUES (?, ?) RETURNING `id`;", + table = SCHEDULERS_TABLE_NAME, + ); + + let mut tx = self.pool.begin().await?; + sqlx::query(DELETE_QUERY).execute(&mut *tx).await?; + let scheduler_id = sqlx::query_scalar(INSERT_QUERY) + .bind(ip_address.to_string()) + .bind(port) + .fetch_one(&mut *tx) + .await?; + tx.commit().await?; + Ok(scheduler_id) + } + + async fn get_schedulers(&self) -> Result, DbError> { + const QUERY: &str = formatcp!( + "SELECT `id`, `ip_address`, `port` FROM `{table}` ORDER BY `id` ASC;", + table = SCHEDULERS_TABLE_NAME, + ); + + let rows: Vec = sqlx::query_as(QUERY).fetch_all(&self.pool).await?; + rows.into_iter() + .map(SchedulerRowProjection::into_registered_scheduler) + .collect() + } + + async fn is_scheduler_registered(&self, scheduler_id: SchedulerId) -> Result { + const QUERY: &str = formatcp!( + "SELECT `id` FROM `{table}` WHERE `id` = ?;", + table = SCHEDULERS_TABLE_NAME, + ); + + let registered_scheduler_id: Option = sqlx::query_scalar(QUERY) + .bind(scheduler_id) + .fetch_optional(&self.pool) + .await?; + Ok(registered_scheduler_id.is_some()) + } +} + impl SessionManagement for MariaDbStorageConnector { fn session_id(&self) -> SessionId { self.session_id @@ -606,6 +662,7 @@ const MYSQL_ER_DUP_ENTRY: u16 = 1062; const RESOURCE_GROUPS_TABLE_NAME: &str = "resource_groups"; const JOBS_TABLE_NAME: &str = "jobs"; const EXECUTION_MANAGERS_TABLE_NAME: &str = "execution_managers"; +const SCHEDULERS_TABLE_NAME: &str = "schedulers"; const SESSIONS_TABLE_NAME: &str = "sessions"; const UPDATE_JOB_STATE: &str = formatcp!( @@ -674,6 +731,20 @@ CREATE TABLE IF NOT EXISTS `{EXECUTION_MANAGERS_TABLE_NAME}` ( ) } +#[must_use] +const fn schedulers_creation_query() -> &'static str { + formatcp!( + r" +CREATE TABLE IF NOT EXISTS `{SCHEDULERS_TABLE_NAME}` ( + `id` BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, + `ip_address` VARCHAR(45) NOT NULL, + `port` SMALLINT UNSIGNED NOT NULL, + `created_at` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (`id`) +);" + ) +} + #[must_use] const fn sessions_creation_query() -> &'static str { formatcp!( @@ -732,6 +803,41 @@ impl RecoverableJobRowProjection { } } +/// A raw row selected from the schedulers table. +#[derive(sqlx::FromRow)] +struct SchedulerRowProjection { + id: SchedulerId, + ip_address: String, + port: u16, +} + +impl SchedulerRowProjection { + /// Converts the row projection into [`RegisteredScheduler`]. + /// + /// # Returns + /// + /// The registered scheduler on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`DbError::CorruptedDbState`] if the scheduler IP address is invalid. + fn into_registered_scheduler(self) -> Result { + let ip_address = self.ip_address.parse().map_err(|error| { + DbError::CorruptedDbState(format!( + "scheduler `{}` has invalid IP address `{}`: {error}", + self.id, self.ip_address + )) + })?; + Ok(RegisteredScheduler { + id: self.id, + ip_address, + port: self.port, + }) + } +} + /// Gets the job state with exclusive lock on the row. /// /// # Returns diff --git a/components/spider-storage/src/db/protocol.rs b/components/spider-storage/src/db/protocol.rs index 4ef56b2b6..3d1ec60a7 100644 --- a/components/spider-storage/src/db/protocol.rs +++ b/components/spider-storage/src/db/protocol.rs @@ -4,8 +4,9 @@ use async_trait::async_trait; use spider_core::{ job::JobState, types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, SessionId}, + id::{ExecutionManagerId, JobId, ResourceGroupId, SchedulerId, SessionId}, io::TaskOutput, + scheduler::RegisteredScheduler, }, }; @@ -28,6 +29,7 @@ pub struct RecoverableJobContext { /// * [`InternalJobOrchestration`] /// * [`ResourceGroupManagement`] /// * [`ExecutionManagerLivenessManagement`] +/// * [`SchedulerRegistrationManagement`] /// * [`SessionManagement`] #[async_trait] pub trait DbStorage: @@ -35,6 +37,7 @@ pub trait DbStorage: + InternalJobOrchestration + ResourceGroupManagement + ExecutionManagerLivenessManagement + + SchedulerRegistrationManagement + SessionManagement { } @@ -426,6 +429,65 @@ pub trait ExecutionManagerLivenessManagement: Clone + Send + Sync { ) -> Result, DbError>; } +/// Defines the storage interface for scheduler registration in the database. +#[async_trait] +pub trait SchedulerRegistrationManagement: Clone + Send + Sync { + /// Registers the scheduler in the database. + /// + /// For now, only one scheduler can be registered at a time. Registering a new scheduler removes + /// any previously registered scheduler before allocating the new scheduler ID. + /// + /// # Parameters + /// + /// * `ip_address` - The scheduler IP address. + /// * `port` - The scheduler port. + /// + /// # Returns + /// + /// The ID of the registered scheduler on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`sqlx::error::Error`] on DB operation failure. + async fn register_scheduler( + &self, + ip_address: IpAddr, + port: u16, + ) -> Result; + + /// Gets registered schedulers. + /// + /// # Returns + /// + /// The registered schedulers on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`sqlx::error::Error`] on DB operation failure. + async fn get_schedulers(&self) -> Result, DbError>; + + /// Checks whether the scheduler with the given ID is registered. + /// + /// # Parameters + /// + /// * `scheduler_id` - The scheduler ID to check. + /// + /// # Returns + /// + /// Whether the scheduler is registered on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`sqlx::error::Error`] on DB operation failure. + async fn is_scheduler_registered(&self, scheduler_id: SchedulerId) -> Result; +} + /// Defines the storage interface for session management. /// /// A session ID is a monotonically increasing value that bumps each time the storage layer diff --git a/components/spider-storage/src/state/service.rs b/components/spider-storage/src/state/service.rs index eb2e26a8b..a72d615d6 100644 --- a/components/spider-storage/src/state/service.rs +++ b/components/spider-storage/src/state/service.rs @@ -4,8 +4,17 @@ use spider_core::{ job::JobState, task::{TaskGraph, TaskIndex}, types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, SessionId, TaskId, TaskInstanceId}, + id::{ + ExecutionManagerId, + JobId, + ResourceGroupId, + SchedulerId, + SessionId, + TaskId, + TaskInstanceId, + }, io::{ExecutionContext, TaskInput, TaskOutput}, + scheduler::RegisteredScheduler, }, }; use spider_tdl::{ @@ -653,6 +662,54 @@ impl< Ok(()) } + /// Registers the scheduler. + /// + /// Registering a scheduler invalidates any previously registered scheduler. + /// + /// # Returns + /// + /// The ID of the registered scheduler on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`SchedulerRegistrationManagement::register_scheduler`]'s return values on + /// failure. + pub async fn register_scheduler( + &self, + ip_address: IpAddr, + port: u16, + ) -> Result { + let scheduler_id = self.inner.db.register_scheduler(ip_address, port).await?; + tracing::info!( + scheduler_id = ? scheduler_id, + ip = ? ip_address, + port, + "Scheduler registered.", + ); + Ok(scheduler_id) + } + + /// Gets registered schedulers. + /// + /// # Returns + /// + /// The registered schedulers on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`SchedulerRegistrationManagement::get_schedulers`]'s return values on failure. + pub async fn get_schedulers(&self) -> Result, StorageServerError> { + self.inner + .db + .get_schedulers() + .await + .map_err(StorageServerError::from) + } + /// Validates that the given `session_id` matches the session ID captured at service creation /// time. /// @@ -699,6 +756,7 @@ struct ServiceStateInner< #[cfg(test)] mod tests { + use spider_core::{ job::JobState, task::{ diff --git a/components/spider-storage/src/state/test_utils.rs b/components/spider-storage/src/state/test_utils.rs index b931b34c0..e5aa511af 100644 --- a/components/spider-storage/src/state/test_utils.rs +++ b/components/spider-storage/src/state/test_utils.rs @@ -10,8 +10,9 @@ use dashmap::DashMap; use spider_core::{ job::JobState, types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, SessionId, TaskInstanceId}, + id::{ExecutionManagerId, JobId, ResourceGroupId, SchedulerId, SessionId, TaskInstanceId}, io::TaskOutput, + scheduler::RegisteredScheduler, }, }; @@ -29,6 +30,7 @@ use crate::{ InternalJobOrchestration, RecoverableJobContext, ResourceGroupManagement, + SchedulerRegistrationManagement, SessionManagement, }, ready_queue::ReadyQueueSender, @@ -251,6 +253,25 @@ impl ExecutionManagerLivenessManagement for MockDbConnector { } } +#[async_trait::async_trait] +impl SchedulerRegistrationManagement for MockDbConnector { + async fn register_scheduler( + &self, + _ip_address: IpAddr, + _port: u16, + ) -> Result { + unreachable!("not implemented for mock connector") + } + + async fn get_schedulers(&self) -> Result, DbError> { + unreachable!("not implemented for mock connector") + } + + async fn is_scheduler_registered(&self, _scheduler_id: SchedulerId) -> Result { + unreachable!("not implemented for mock connector") + } +} + impl SessionManagement for MockDbConnector { fn session_id(&self) -> SessionId { self.session_id diff --git a/components/spider-storage/tests/mariadb_test.rs b/components/spider-storage/tests/mariadb_test.rs index f58a020f4..029825610 100644 --- a/components/spider-storage/tests/mariadb_test.rs +++ b/components/spider-storage/tests/mariadb_test.rs @@ -6,7 +6,7 @@ use std::{ use spider_core::{ job::JobState, types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId}, + id::{ExecutionManagerId, JobId, ResourceGroupId, SchedulerId}, io::TaskInput, }, }; @@ -19,6 +19,7 @@ use spider_storage::{ InternalJobOrchestration, MariaDbStorageConnector, ResourceGroupManagement, + SchedulerRegistrationManagement, SessionManagement, }, }; @@ -33,6 +34,8 @@ const TEST_INPUT_PAYLOAD_SIZE: usize = 128; /// Number of execution managers to register in multi-EM tests. const TEST_NUM_EMS: usize = 3; +const TEST_SCHEDULER_PORT: u16 = 5678; +const TEST_UPDATED_SCHEDULER_PORT: u16 = 6789; /// Builds a task graph with a single task for DB-layer tests. /// @@ -55,6 +58,19 @@ async fn register_test_em(storage: &MariaDbStorageConnector) -> ExecutionManager .expect("register_execution_manager should succeed") } +/// # Returns +/// +/// Whether the scheduler is registered. +async fn is_scheduler_registered( + storage: &MariaDbStorageConnector, + scheduler_id: SchedulerId, +) -> bool { + storage + .is_scheduler_registered(scheduler_id) + .await + .expect("is_scheduler_registered should succeed") +} + #[tokio::test] #[ignore = "requires MariaDB"] async fn test_register_job() { @@ -1005,6 +1021,49 @@ async fn test_get_dead_execution_managers_multiple() { } } +#[tokio::test] +#[ignore = "requires MariaDB"] +#[serial_test::file_serial] +async fn test_register_scheduler_replaces_previous_scheduler() { + let storage = create_mariadb_connector().await; + let scheduler_ip_address = IpAddr::V4(Ipv4Addr::LOCALHOST); + let updated_scheduler_ip_address = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)); + + let first_scheduler_id = storage + .register_scheduler(scheduler_ip_address, TEST_SCHEDULER_PORT) + .await + .expect("first register_scheduler should succeed"); + let second_scheduler_id = storage + .register_scheduler(updated_scheduler_ip_address, TEST_UPDATED_SCHEDULER_PORT) + .await + .expect("second register_scheduler should succeed"); + let schedulers = storage + .get_schedulers() + .await + .expect("get_schedulers should succeed"); + + assert_ne!( + first_scheduler_id, second_scheduler_id, + "new registration should allocate a fresh scheduler ID" + ); + assert!( + !is_scheduler_registered(&storage, first_scheduler_id).await, + "old scheduler should be removed after a new registration" + ); + assert!( + is_scheduler_registered(&storage, second_scheduler_id).await, + "new scheduler should remain registered" + ); + assert_eq!( + schedulers.len(), + 1, + "only the latest scheduler should remain" + ); + assert_eq!(schedulers[0].id, second_scheduler_id); + assert_eq!(schedulers[0].ip_address, updated_scheduler_ip_address); + assert_eq!(schedulers[0].port, TEST_UPDATED_SCHEDULER_PORT); +} + #[tokio::test] #[ignore = "requires MariaDB"] async fn test_session_id_returned_after_connect() {