diff --git a/Cargo.lock b/Cargo.lock index 69c950eddd..03b01ad06e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1020,6 +1020,7 @@ dependencies = [ "const_format", "non-empty-string", "rmp-serde", + "secrecy", "serde", "spider-client", "spider-core", @@ -1027,6 +1028,8 @@ dependencies = [ "strsim", "thiserror", "tokio", + "tokio-util", + "tonic", "tracing", ] diff --git a/components/api-server/src/error.rs b/components/api-server/src/error.rs index d983890940..162290a5c3 100644 --- a/components/api-server/src/error.rs +++ b/components/api-server/src/error.rs @@ -80,9 +80,9 @@ impl From for ClientError { impl From for ClientError { fn from(value: clp_rust_utils::Error) -> Self { match value { - clp_rust_utils::Error::MsgpackEncode(_) | clp_rust_utils::Error::SerdeYaml(_) => { - Self::MalformedData - } + clp_rust_utils::Error::MsgpackEncode(_) + | clp_rust_utils::Error::MsgpackDecode(_) + | clp_rust_utils::Error::SerdeYaml(_) => Self::MalformedData, clp_rust_utils::Error::UnsupportedS3Endpoint(_) => { Self::InvalidInput(value.to_string()) } diff --git a/components/clp-rust-utils/src/clp_config/package/config.rs b/components/clp-rust-utils/src/clp_config/package/config.rs index 9f2a481f56..f911728a27 100644 --- a/components/clp-rust-utils/src/clp_config/package/config.rs +++ b/components/clp-rust-utils/src/clp_config/package/config.rs @@ -1,5 +1,9 @@ -use std::path::{Path, PathBuf}; +use std::{ + num::{NonZeroU32, NonZeroU64}, + path::{Path, PathBuf}, +}; +use non_empty_string::NonEmptyString; use serde::Deserialize; use crate::{ @@ -27,6 +31,8 @@ pub struct Config { pub logs_input: LogsInput, pub archive_output: ArchiveOutput, pub telemetry: Telemetry, + pub spider: Option, + pub compression_coordinator: Option, } impl Default for Config { @@ -44,6 +50,8 @@ impl Default for Config { }, archive_output: ArchiveOutput::default(), telemetry: Telemetry::default(), + spider: None, + compression_coordinator: None, } } } @@ -471,6 +479,70 @@ impl Default for Telemetry { } } +/// Compression coordinator configuration. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq)] +#[serde(default)] +pub struct CompressionCoordinator { + pub resource_group: SpiderResourceGroup, + pub job_polling_interval_millisecs: NonZeroU64, + pub result_polling: PollingBackoff, + pub compression_task_max_retry: u32, + pub commit_task_max_retry: u32, + pub database_connection_pool_size: NonZeroU32, + pub termination_timeout_secs: NonZeroU64, + pub commit_task_soft_timeout_secs: NonZeroU64, + pub commit_task_hard_timeout_secs: NonZeroU64, +} + +impl Default for CompressionCoordinator { + fn default() -> Self { + Self { + resource_group: SpiderResourceGroup { + name: NonEmptyString::new("compression-coordinator".to_owned()) + .expect("default resource group name should not be empty"), + }, + job_polling_interval_millisecs: NonZeroU64::new(100) + .expect("default jobs poll delay should not be zero"), + result_polling: PollingBackoff { + init_backoff_millisecs: NonZeroU64::new(100) + .expect("default result polling init backoff should not be zero"), + max_backoff_millisecs: NonZeroU64::new(1000) + .expect("default result polling max backoff should not be zero"), + }, + compression_task_max_retry: 1, + commit_task_max_retry: 1, + database_connection_pool_size: NonZeroU32::new(10) + .expect("default database connection pool size should not be zero"), + termination_timeout_secs: NonZeroU64::new(30) + .expect("default termination timeout should not be zero"), + commit_task_soft_timeout_secs: NonZeroU64::new(45) + .expect("default commit task soft timeout should not be zero"), + commit_task_hard_timeout_secs: NonZeroU64::new(60) + .expect("default commit task hard timeout should not be zero"), + } + } +} + +/// Spider configuration. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq)] +pub struct Spider { + pub host: NonEmptyString, + pub port: u16, +} + +/// Spider resource group configuration. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq)] +pub struct SpiderResourceGroup { + pub name: NonEmptyString, +} + +/// Polling backoff configuration. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq)] +pub struct PollingBackoff { + pub init_backoff_millisecs: NonZeroU64, + pub max_backoff_millisecs: NonZeroU64, +} + /// # Returns /// /// `path` unchanged if it is already absolute, otherwise joined with `root`. diff --git a/components/clp-rust-utils/src/error.rs b/components/clp-rust-utils/src/error.rs index 2b565dc8d9..cd84a432f1 100644 --- a/components/clp-rust-utils/src/error.rs +++ b/components/clp-rust-utils/src/error.rs @@ -5,6 +5,9 @@ pub enum Error { #[error("`rmp_serde::encode::Error`: {0}")] MsgpackEncode(#[from] rmp_serde::encode::Error), + #[error("`rmp_serde::decode::Error`: {0}")] + MsgpackDecode(#[from] rmp_serde::decode::Error), + #[error("`std::io::Error`: {0}")] Io(#[from] std::io::Error), diff --git a/components/clp-rust-utils/src/job_config/clp_io_config.rs b/components/clp-rust-utils/src/job_config/clp_io_config.rs index 7f4e19cb51..a9d54a70c2 100644 --- a/components/clp-rust-utils/src/job_config/clp_io_config.rs +++ b/components/clp-rust-utils/src/job_config/clp_io_config.rs @@ -1,5 +1,5 @@ use non_empty_string::NonEmptyString; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use crate::{ clp_config::S3Config, @@ -8,14 +8,14 @@ use crate::{ }; /// Represents CLP IO config. -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ClpIoConfig { pub input: InputConfig, pub output: OutputConfig, } /// An enum representing CLP input config. -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(tag = "type")] pub enum InputConfig { #[serde(rename = "s3")] @@ -32,7 +32,7 @@ pub enum InputConfig { } /// Represents S3 input config. -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct S3InputConfig { #[serde(flatten)] pub s3_config: S3Config, @@ -43,7 +43,7 @@ pub struct S3InputConfig { pub unstructured: bool, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct S3ObjectMetadataInputConfig { #[serde(flatten)] pub s3_config: S3Config, @@ -56,7 +56,7 @@ pub struct S3ObjectMetadataInputConfig { } /// Represents CLP output config. -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct OutputConfig { pub target_archive_size: u64, pub target_dictionaries_size: u64, diff --git a/components/clp-rust-utils/src/serde/brotli_msgpack.rs b/components/clp-rust-utils/src/serde/brotli_msgpack.rs index a85afc15a0..0ac1beae90 100644 --- a/components/clp-rust-utils/src/serde/brotli_msgpack.rs +++ b/components/clp-rust-utils/src/serde/brotli_msgpack.rs @@ -1,7 +1,7 @@ -use std::io::Write; +use std::io::{Read, Write}; -use brotli::CompressorWriter; -use serde::Serialize; +use brotli::{CompressorWriter, Decompressor}; +use serde::{Serialize, de::DeserializeOwned}; use crate::Error; @@ -11,7 +11,8 @@ pub struct BrotliMsgpack {} impl BrotliMsgpack { /// Serialize a value to a Brotli-compressed `MessagePack` byte sequence. /// - /// # Return + /// # Returns + /// /// A vector of bytes containing the serialized byte sequence. /// /// # Errors @@ -26,4 +27,22 @@ impl BrotliMsgpack { brotli_compressor.write_all(&msgpack_data)?; Ok(brotli_compressor.into_inner()) } + + /// Deserialize an owned value from a Brotli-compressed `MessagePack` byte sequence. + /// + /// # Returns + /// + /// The deserialized value. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`rmp_serde::from_slice`]'s errors on failure. + /// * Forwards [`std::io::Read::read_to_end`]'s errors on failure. + pub fn deserialize(data: &[u8]) -> Result { + let mut msgpack_data = Vec::new(); + Decompressor::new(data, 4096).read_to_end(&mut msgpack_data)?; + rmp_serde::from_slice(&msgpack_data).map_err(Into::into) + } } diff --git a/components/compression-coordinator/Cargo.toml b/components/compression-coordinator/Cargo.toml index be6952f983..df48c7ecdc 100644 --- a/components/compression-coordinator/Cargo.toml +++ b/components/compression-coordinator/Cargo.toml @@ -19,11 +19,14 @@ clp-rust-utils = { path = "../clp-rust-utils" } const_format = "0.2.35" non-empty-string = "0.2.6" rmp-serde = "1.3.1" +secrecy = "0.10.3" serde = { version = "1.0.228", features = ["derive"] } spider-client = { git = "https://github.com/y-scope/spider.git", branch = "main" } spider-core = { git = "https://github.com/y-scope/spider.git", branch = "main" } sqlx = { version = "0.8.6", features = ["runtime-tokio", "mysql"] } strsim = "0.11.1" thiserror = "2.0.18" -tokio = { version = "1.52.3", features = ["time", "rt-multi-thread"] } +tokio = { version = "1.52.3", features = ["time", "rt-multi-thread", "signal"] } +tokio-util = "0.7.18" +tonic = "0.14.6" tracing = "0.1.44" diff --git a/components/compression-coordinator/src/bin/compression_coordinator.rs b/components/compression-coordinator/src/bin/compression_coordinator.rs index 8e18f0072c..cf069567df 100644 --- a/components/compression-coordinator/src/bin/compression_coordinator.rs +++ b/components/compression-coordinator/src/bin/compression_coordinator.rs @@ -1,6 +1,11 @@ -use std::path::PathBuf; +use std::{path::PathBuf, time::Duration}; use clap::Parser; +use clp_rust_utils::{ + clp_config::package::{self}, + database::mysql::create_clp_db_mysql_pool, + serde::yaml, +}; /// Command-line arguments for the compression coordinator. #[derive(Debug, Parser)] @@ -11,9 +16,141 @@ struct Cli { config: PathBuf, } +#[allow(clippy::too_many_lines)] #[tokio::main] async fn main() -> anyhow::Result<()> { - let _cli = Cli::parse(); + let args = Cli::parse(); - Ok(()) + let _guard = clp_rust_utils::logging::set_up_logging("compression_coordinator.log"); + + let config: package::config::Config = yaml::from_path(args.config).inspect_err(|e| { + tracing::error!(error = % e, "Failed to load the configuration file."); + })?; + + if !matches!(config.logs_input, package::config::LogsInput::S3 { .. }) { + const ERROR_MESSAGE: &str = "The compression coordinator only supports S3 logs inputs."; + tracing::error!(ERROR_MESSAGE); + anyhow::bail!(ERROR_MESSAGE); + } + + if !matches!( + config.archive_output.storage, + package::config::ArchiveOutputStorage::S3 { .. } + ) { + const ERROR_MESSAGE: &str = "The compression coordinator only supports S3 archive outputs."; + tracing::error!(ERROR_MESSAGE); + anyhow::bail!(ERROR_MESSAGE); + } + + let credentials = package::credentials::Credentials { + database: package::credentials::Database { + password: secrecy::SecretString::new( + std::env::var("CLP_DB_PASS") + .inspect_err(|e| { + tracing::error!( + error = % e, + "Failed to read the database password from `CLP_DB_PASS`." + ); + })? + .into_boxed_str(), + ), + user: std::env::var("CLP_DB_USER").inspect_err(|e| { + tracing::error!( + error = % e, + "Failed to read the database user from `CLP_DB_USER`." + ); + })?, + }, + }; + + let coordinator_config = config.compression_coordinator.ok_or_else(|| { + const ERROR_MESSAGE: &str = "Compression coordinator configuration is missing."; + tracing::error!(ERROR_MESSAGE); + anyhow::anyhow!(ERROR_MESSAGE) + })?; + + let spider_config = config.spider.ok_or_else(|| { + const ERROR_MESSAGE: &str = "Spider configuration is missing."; + tracing::error!(ERROR_MESSAGE); + anyhow::anyhow!(ERROR_MESSAGE) + })?; + + let db_pool = create_clp_db_mysql_pool( + &config.database, + &credentials.database, + coordinator_config.database_connection_pool_size.get(), + ) + .await + .inspect_err(|e| tracing::error!(error = % e, "Failed to create the database pool."))?; + + let (coordinator, cancellation_token) = + compression_coordinator::coordination::Coordinator::new( + &coordinator_config, + &spider_config, + db_pool, + config.database, + ) + .await?; + + let mut coordinator_handle = tokio::spawn(coordinator.run()); + + let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("failed to listen for SIGTERM"); + + // `None` if a shutdown signal arrived while the coordinator is still running; `Some` if the + // coordinator returned on its own (an early exit, possibly on error). + let early_exit_result = tokio::select! { + _ = sigterm.recv() => { + tracing::info!("Received SIGTERM."); + None + } + result = tokio::signal::ctrl_c() => { + if let Err(e) = result { + tracing::error!(error = % e, "Failed to listen to ctrl-c."); + } + tracing::info!("Forcefully shutting down."); + None + } + join_result = &mut coordinator_handle => Some(join_result), + }; + + // Request a graceful stop. A no-op if the coordinator has already returned. + cancellation_token.cancel(); + + let join_result = if let Some(join_result) = early_exit_result { + join_result + } else { + let termination_timeout = + Duration::from_secs(coordinator_config.termination_timeout_secs.get()); + if let Ok(join_result) = + tokio::time::timeout(termination_timeout, &mut coordinator_handle).await + { + join_result + } else { + tracing::warn!( + "The compression coordinator did not stop within {termination_timeout:?}. \ + Aborting." + ); + coordinator_handle.abort(); + return Ok(()); + } + }; + + match join_result { + Ok(Ok(())) => { + tracing::info!("Compression coordinator stopped."); + Ok(()) + } + Ok(Err(e)) => { + tracing::error!(error = % e, "Compression coordinator returned on error."); + Err(anyhow::anyhow!( + "Compression coordinator returned on error." + )) + } + Err(err) => { + const ERROR_MESSAGE: &str = "Failed to join the compression coordinator."; + tracing::error!(error = % err, ERROR_MESSAGE); + Err(anyhow::anyhow!(ERROR_MESSAGE)) + } + } } diff --git a/components/compression-coordinator/src/coordination.rs b/components/compression-coordinator/src/coordination.rs new file mode 100644 index 0000000000..0873e9d42d --- /dev/null +++ b/components/compression-coordinator/src/coordination.rs @@ -0,0 +1,571 @@ +//! The coordinator poll loop that discovers pending CLP compression jobs and dispatches them to +//! Spider. + +use std::{sync::Arc, time::Duration}; + +use clp_rust_utils::{ + clp_config::package::config::{ + CompressionCoordinator as CoordinatorConfig, + Database as DatabaseConfig, + Spider as SpiderConfig, + SpiderResourceGroup, + }, + job_config::{ClpIoConfig, CompressionJobId, CompressionJobStatus}, + serde::BrotliMsgpack, +}; +use const_format::formatcp; +use spider_client::SpiderClient; +use spider_core::{ + task::{ExecutionPolicy, TimeoutPolicy}, + types::id::{JobId as SpiderJobId, ResourceGroupId}, +}; +use tokio::{select, time::Instant}; +use tokio_util::sync::CancellationToken; +use tonic::transport::Endpoint; + +use crate::{ + Error, + job_handle::{S3CompressionJobHandle, SpiderOption}, +}; + +/// Coordinator for fetching new compression jobs and submitting them to Spider. +pub struct Coordinator { + resource_group_id: ResourceGroupId, + spider_client: SpiderClient, + db_pool: sqlx::MySqlPool, + db_config: DatabaseConfig, + spider_option: Arc, + is_first_fetch: bool, + job_polling_interval: Duration, + cancellation_token: CancellationToken, +} + +impl Coordinator { + /// Factory function. + /// + /// On construction, this recovers compression jobs that a previous coordinator instance had + /// already submitted to Spider (those still [`CompressionJobStatus::Running`] with a Spider job + /// ID) by spawning a detached handle to drive each one to completion. + /// + /// # Returns + /// + /// A tuple on success, containing: + /// + /// * The constructed [`Coordinator`]. + /// * The [`CancellationToken`] the caller uses to request shutdown. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`Error::InvalidEndpoint`] if the Spider host and port do not form a valid endpoint. + /// * Forwards [`SpiderClient::builder`]'s connection return values on failure. + /// * Forwards [`get_or_create_resource_group_id`]'s return values on failure. + /// * Forwards [`Self::fetch_submitted_running_jobs`]'s return values on failure. + pub async fn new( + coordinator_config: &CoordinatorConfig, + spider_config: &SpiderConfig, + db_pool: sqlx::MySqlPool, + db_config: DatabaseConfig, + ) -> Result<(Self, CancellationToken), Error> { + let spider_host = spider_config.host.as_str(); + let spider_port = spider_config.port; + let endpoint_str = format!("http://{spider_host}:{spider_port}"); + let endpoint = Endpoint::from_shared(endpoint_str) + .inspect_err(|e| { + tracing::error!(error = % e, "Failed to create Spider endpoint."); + }) + .map_err(|e| Error::InvalidEndpoint(e.to_string()))?; + let spider_client = SpiderClient::builder(endpoint) + .connect() + .await + .inspect_err(|e| { + tracing::error!(error = % e, "Failed to connect to Spider."); + })?; + let resource_group_id = get_or_create_resource_group_id( + &coordinator_config.resource_group, + &spider_client, + &db_pool, + ) + .await + .inspect_err(|e| { + tracing::error!(error = % e, "Failed to get or create resource group."); + })?; + + let spider_option = Arc::new(SpiderOption { + compression_task_max_retry: coordinator_config.compression_task_max_retry, + commit_task_execution_policy: ExecutionPolicy { + max_num_instances: 1, + max_num_retry: coordinator_config.commit_task_max_retry, + timeout_policy: TimeoutPolicy { + soft_timeout_ms: coordinator_config.commit_task_soft_timeout_secs.get() * 1000, + hard_timeout_ms: coordinator_config.commit_task_hard_timeout_secs.get() * 1000, + }, + }, + initial_poll_backoff: Duration::from_millis( + coordinator_config + .result_polling + .init_backoff_millisecs + .get(), + ), + max_poll_backoff: Duration::from_millis( + coordinator_config + .result_polling + .max_backoff_millisecs + .get(), + ), + }); + + let cancellation_token = CancellationToken::new(); + + let coordinator = Self { + resource_group_id, + spider_client, + db_pool, + db_config, + spider_option, + is_first_fetch: true, + job_polling_interval: Duration::from_millis( + coordinator_config.job_polling_interval_millisecs.get(), + ), + cancellation_token: cancellation_token.clone(), + }; + + for (job_id, spider_job_id, clp_io_config) in + coordinator.fetch_submitted_running_jobs().await? + { + tracing::info!( + job_id = % job_id, + spider_job_id = % spider_job_id, + "Recovering a previously submitted job." + ); + let Ok(job_handle) = coordinator.create_job_handle(job_id, clp_io_config).await else { + continue; + }; + tokio::spawn(async move { + let _ = job_handle.recover(spider_job_id).await.inspect_err(|e| { + tracing::error!( + error = % e, + job_id = % job_id, + spider_job_id = % spider_job_id, + "The recovered compression job failed." + ); + }); + }); + } + + Ok((coordinator, cancellation_token)) + } + + /// Runs the coordinator's poll loop until cancelled. + /// + /// On each iteration, this method fetches the pending compression jobs, spawns a detached + /// handle to drive each one, and then sleeps until the next poll or until the cancellation + /// token is triggered. The jobs dispatched in the iteration are marked once the sleep elapses, + /// so their update does not contend with concurrent job submissions during the poll interval. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::schedule_new_jobs`]'s return values on failure. + /// * Forwards [`Self::mark_jobs_dispatched`]'s return values on failure. + pub async fn run(mut self) -> Result<(), Error> { + let cancellation_token = self.cancellation_token.clone(); + loop { + let now = Instant::now(); + + let dispatched_job_ids; + select! { + () = cancellation_token.cancelled() => { + break; + } + result = self.schedule_new_jobs() => { + dispatched_job_ids = result.inspect_err(|e| { + tracing::error!(error = % e, "Failed to schedule new jobs."); + })?; + } + } + + let elapsed = now.elapsed(); + let sleep_duration = self.job_polling_interval.saturating_sub(elapsed); + if sleep_duration.is_zero() { + tokio::task::yield_now().await; + } else if tokio::time::timeout(sleep_duration, cancellation_token.cancelled()) + .await + .is_ok() + { + break; + } + + self.mark_jobs_dispatched(&dispatched_job_ids).await?; + } + + tracing::info!("Coordinator shutting down."); + Ok(()) + } + + /// Marks the compression job identified by `job_id` as [`CompressionJobStatus::Failed`]. + /// + /// This is a best-effort update; if it fails, the error is logged and otherwise ignored. + async fn mark_job_failed(&self, job_id: CompressionJobId, status_msg: &str) { + const QUERY: &str = formatcp!( + "UPDATE `{table}` SET `status` = ?, `status_msg` = ?, `update_time` = \ + CURRENT_TIMESTAMP() WHERE `id` = ?;", + table = COMPRESSION_JOB_TABLE_NAME, + ); + tracing::info!(job_id = % job_id, "Failing the compression job."); + if let Err(e) = sqlx::query(QUERY) + .bind(CompressionJobStatus::Failed) + .bind(status_msg) + .bind(job_id) + .execute(&self.db_pool) + .await + { + tracing::error!( + error = % e, + job_id = % job_id, + "Failed to mark the compression job as failed." + ); + } + } + + /// Fetches the pending compression jobs and spawns a detached handle to drive each one. + /// + /// + /// A job whose config cannot be deserialized is marked [`CompressionJobStatus::Failed`] and + /// skipped; a job whose handle cannot be constructed is skipped as well (and marked + /// [`CompressionJobStatus::Failed`] unless its input config is unsupported, in which case it is + /// left for the legacy Celery-based compression scheduler). + /// + /// # Returns + /// + /// The IDs of the fetched jobs that were dispatched in this poll. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::fetch_new_job_rows`]'s return values on failure. + async fn schedule_new_jobs(&mut self) -> Result, Error> { + let new_job_rows = self.fetch_new_job_rows().await.inspect_err(|e| { + tracing::error!(error = % e, "Failed to fetch new jobs from database."); + })?; + let dispatched_job_ids: Vec = + new_job_rows.iter().map(|row| row.id).collect(); + for job_row in new_job_rows { + let job_id = job_row.id; + let clp_io_config: ClpIoConfig = + match BrotliMsgpack::deserialize(&job_row.serialized_clp_io_config) { + Ok(clp_io_config) => clp_io_config, + Err(e) => { + tracing::error!( + error = % e, + job_id = % job_id, + "Failed to deserialize CLP I/O config. Skipping." + ); + self.mark_job_failed( + job_id, + &format!("Failed to deserialize CLP I/O config: {e}"), + ) + .await; + continue; + } + }; + tracing::info!(job_id = % job_id, "Scheduling new job."); + let Ok(job_handle) = self.create_job_handle(job_id, clp_io_config).await else { + continue; + }; + tokio::spawn(async move { + let _ = job_handle.run().await.inspect_err(|e| { + tracing::error!( + error = % e, + job_id = % job_id, + "Failed to schedule S3 compression job." + ); + }); + }); + } + Ok(dispatched_job_ids) + } + + /// Marks the compression jobs identified by `job_ids` with the current dispatch time. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`sqlx::Pool::begin`]'s return values on failure. + /// * Forwards [`sqlx::query::Query::execute`]'s return values on failure. + /// * Forwards [`sqlx::Transaction::commit`]'s return values on failure. + async fn mark_jobs_dispatched(&self, job_ids: &[CompressionJobId]) -> Result<(), Error> { + if job_ids.is_empty() { + return Ok(()); + } + + let mut tx = self.db_pool.begin().await?; + for chunk in job_ids.chunks(1000) { + let mut query_builder = sqlx::QueryBuilder::::new(formatcp!( + "UPDATE `{table}` SET `dispatch_time` = CURRENT_TIMESTAMP() WHERE `id` IN (", + table = COMPRESSION_JOB_TABLE_NAME, + )); + let mut separated_ids = query_builder.separated(", "); + for job_id in chunk { + separated_ids.push_bind(job_id); + } + query_builder.push(");"); + query_builder.build().execute(&mut *tx).await?; + } + tx.commit().await?; + + Ok(()) + } + + /// Constructs an [`S3CompressionJobHandle`] for the given job. + /// + /// A construction failure is logged, and the job is marked [`CompressionJobStatus::Failed`] for + /// any failure other than an unsupported input config, which is only warned and left for + /// another handler. + /// + /// # Returns + /// + /// The constructed [`S3CompressionJobHandle`] on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`S3CompressionJobHandle::new`]'s return values on failure. + async fn create_job_handle( + &self, + job_id: CompressionJobId, + clp_io_config: ClpIoConfig, + ) -> Result, Error> { + let result = S3CompressionJobHandle::new( + self.db_pool.clone(), + self.db_config.clone(), + job_id, + self.spider_client.clone(), + self.resource_group_id, + clp_io_config, + self.spider_option.clone(), + ); + + if let Err(e) = &result { + if matches!(e, Error::UnsupportedInputConfig) { + tracing::warn!( + error = % e, + job_id = % job_id, + "Unsupported input config. Skipping." + ); + } else { + tracing::error!( + error = % e, + job_id = % job_id, + "Failed to create S3 job handle. Skipping." + ); + self.mark_job_failed( + job_id, + &format!("Failed to create the compression job handle: {e}"), + ) + .await; + } + } + + result + } + + /// Fetches the pending compression jobs to dispatch. + /// + /// The first fetch after startup returns every [`CompressionJobStatus::Pending`] job so that + /// jobs a previous coordinator instance had already dispatched but not started are + /// re-dispatched. Every subsequent fetch returns only [`CompressionJobStatus::Pending`] jobs + /// whose dispatch time is still not set. + /// + /// # Returns + /// + /// A vector of rows projected from the compression job table on success, each row represents a + /// pending compression job. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`sqlx::query::QueryAs::fetch_all`]'s return values on failure. + async fn fetch_new_job_rows(&mut self) -> Result, Error> { + const FIRST_FETCH_QUERY: &str = formatcp!( + "SELECT `id`, `clp_config` FROM `{table}` WHERE `status` = ? ORDER BY `id` ASC;", + table = COMPRESSION_JOB_TABLE_NAME, + ); + const SUBSEQUENT_FETCH_QUERY: &str = formatcp!( + "SELECT `id`, `clp_config` FROM `{table}` WHERE `status` = ? AND `dispatch_time` IS \ + NULL ORDER BY `id` ASC;", + table = COMPRESSION_JOB_TABLE_NAME, + ); + + let query = if self.is_first_fetch { + self.is_first_fetch = false; + FIRST_FETCH_QUERY + } else { + SUBSEQUENT_FETCH_QUERY + }; + let rows = sqlx::query_as::<_, PendingJobRowProjection>(query) + .bind(CompressionJobStatus::Pending) + .fetch_all(&self.db_pool) + .await?; + + Ok(rows) + } + + /// Fetches jobs that are still in [`CompressionJobStatus::Running`] and were previously + /// submitted by the compression coordinator. + /// + /// A running job whose config cannot be deserialized is marked [`CompressionJobStatus::Failed`] + /// and skipped. + /// + /// # Returns + /// + /// A vector of tuples on success, each tuple containing: + /// + /// * The compression job ID. + /// * The Spider job ID. + /// * The IO config of the compression job. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`sqlx::query::QueryAs::fetch_all`]'s return values on failure. + async fn fetch_submitted_running_jobs( + &self, + ) -> Result, Error> { + const QUERY: &str = formatcp!( + "SELECT `id`, `spider_id`, `clp_config` FROM `{table}` WHERE `status` = ? AND \ + `spider_id` IS NOT NULL;", + table = COMPRESSION_JOB_TABLE_NAME, + ); + + let mut recovery_context = Vec::new(); + for row in sqlx::query_as::<_, RunningJobRowProjection>(QUERY) + .bind(CompressionJobStatus::Running) + .fetch_all(&self.db_pool) + .await? + { + let clp_io_config: ClpIoConfig = + match BrotliMsgpack::deserialize(&row.serialized_clp_io_config) { + Ok(clp_io_config) => clp_io_config, + Err(e) => { + tracing::error!( + error = % e, + job_id = % row.id, + "Failed to deserialize CLP I/O config of a running job. The database \ + might be corrupted. Skipping." + ); + self.mark_job_failed( + row.id, + &format!("Failed to deserialize CLP I/O config: {e}"), + ) + .await; + continue; + } + }; + recovery_context.push((row.id, row.spider_job_id, clp_io_config)); + } + + Ok(recovery_context) + } +} + +const COMPRESSION_JOB_TABLE_NAME: &str = "compression_jobs"; + +/// A projection of the columns read from a [`CompressionJobStatus::Pending`] compression job row. +#[derive(Debug, sqlx::FromRow)] +struct PendingJobRowProjection { + id: CompressionJobId, + #[sqlx(rename = "clp_config")] + serialized_clp_io_config: Vec, +} + +/// A projection of the columns read from a [`CompressionJobStatus::Running`] compression job row. +#[derive(Debug, sqlx::FromRow)] +struct RunningJobRowProjection { + id: CompressionJobId, + #[sqlx(rename = "spider_id")] + spider_job_id: SpiderJobId, + #[sqlx(rename = "clp_config")] + serialized_clp_io_config: Vec, +} + +/// Retrieves the Spider resource group ID for the configured resource group, registering it if it +/// does not yet exist. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards [`sqlx::query::Query::execute`]'s return values on failure. +/// * Forwards [`SpiderClient::add_resource_group`]'s return values on failure. +async fn get_or_create_resource_group_id( + resource_group_config: &SpiderResourceGroup, + spider_client: &SpiderClient, + db_pool: &sqlx::MySqlPool, +) -> Result { + const SPIDER_RESOURCE_GROUP_TABLE_NAME: &str = "spider_resource_groups"; + + const CREATE_TABLE_QUERY: &str = formatcp!( + "CREATE TABLE IF NOT EXISTS `{table}` ( + `rg_name` VARCHAR(255) NOT NULL, + `rg_id` BIGINT UNSIGNED NOT NULL, + PRIMARY KEY (`rg_name`) USING BTREE + ) ROW_FORMAT=DYNAMIC", + table = SPIDER_RESOURCE_GROUP_TABLE_NAME, + ); + const SELECT_QUERY: &str = formatcp!( + "SELECT `rg_id` FROM `{table}` WHERE `rg_name` = ?;", + table = SPIDER_RESOURCE_GROUP_TABLE_NAME, + ); + const INSERT_QUERY: &str = formatcp!( + "INSERT INTO `{table}` (`rg_name`, `rg_id`) VALUES (?, ?);", + table = SPIDER_RESOURCE_GROUP_TABLE_NAME, + ); + + sqlx::query(CREATE_TABLE_QUERY).execute(db_pool).await?; + + let resource_group = resource_group_config.name.as_str(); + let existing_rg_id: Option = sqlx::query_scalar(SELECT_QUERY) + .bind(resource_group) + .fetch_optional(db_pool) + .await?; + if let Some(spider_rg_id) = existing_rg_id { + tracing::info!( + resource_group = % resource_group, + spider_rg_id = % spider_rg_id, + "Resource group already registered. Returning Spider resource group ID." + ); + return Ok(ResourceGroupId::from(spider_rg_id)); + } + + // NOTE: For now, Spider does not enforce resource group credential validation. The password is + // hardcoded to be the same as the username. + let resource_group_id = spider_client + .add_resource_group( + resource_group.to_owned(), + resource_group.as_bytes().to_vec(), + ) + .await?; + + sqlx::query(INSERT_QUERY) + .bind(resource_group) + .bind(resource_group_id.get()) + .execute(db_pool) + .await + .inspect_err(|e| { + tracing::error!( + error = % e, + "Failed to insert resource group into database. This might be a race condition. \ + Restart the service to retry." + ); + })?; + + Ok(resource_group_id) +} diff --git a/components/compression-coordinator/src/error.rs b/components/compression-coordinator/src/error.rs index f8814abcbb..b3da15ed29 100644 --- a/components/compression-coordinator/src/error.rs +++ b/components/compression-coordinator/src/error.rs @@ -22,6 +22,9 @@ pub enum Error { #[error("invalid dataset: {0}")] InvalidDataset(String), + #[error("invalid endpoint: {0}")] + InvalidEndpoint(String), + #[error("failed to create metadata table `{table}`: {source}")] MetadataTableCreation { table: String, diff --git a/components/compression-coordinator/src/lib.rs b/components/compression-coordinator/src/lib.rs index bf3c599978..8683323784 100644 --- a/components/compression-coordinator/src/lib.rs +++ b/components/compression-coordinator/src/lib.rs @@ -1,6 +1,7 @@ //! Compression-job submission and S3 input partitioning for a Spider cluster. pub mod compression_job_submitter; +pub mod coordination; mod error; pub mod job_handle; pub mod partition;