diff --git a/Cargo.lock b/Cargo.lock index 3c490b08db..85032a8a0c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3271,6 +3271,21 @@ dependencies = [ "pulldown-cmark", ] +[[package]] +name = "query-coordinator" +version = "0.13.1-dev" +dependencies = [ + "async-trait", + "clp-rust-utils", + "non-empty-string", + "spider-client", + "spider-core", + "sqlx", + "thiserror", + "tokio", + "tracing", +] + [[package]] name = "quote" version = "1.0.47" diff --git a/Cargo.toml b/Cargo.toml index d493d3e5fd..5b1de8aa1c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,8 @@ members = [ "components/clp-rust-utils", "components/clp-tdl-package", "components/compression-coordinator", - "components/log-ingestor" + "components/log-ingestor", + "components/query-coordinator", ] resolver = "3" diff --git a/components/clp-py-utils/clp_py_utils/initialize-orchestration-db.py b/components/clp-py-utils/clp_py_utils/initialize-orchestration-db.py index 5af48908a6..1dd800cd99 100644 --- a/components/clp-py-utils/clp_py_utils/initialize-orchestration-db.py +++ b/components/clp-py-utils/clp_py_utils/initialize-orchestration-db.py @@ -134,15 +134,18 @@ def main(argv): `id` INT NOT NULL AUTO_INCREMENT, `type` INT NOT NULL, `status` INT NOT NULL DEFAULT '{QueryJobStatus.PENDING}', + `status_msg` VARCHAR(512) NOT NULL DEFAULT '', `creation_time` DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3), `num_tasks` INT NOT NULL DEFAULT '0', `num_tasks_completed` INT NOT NULL DEFAULT '0', `start_time` DATETIME(3) NULL DEFAULT NULL, `duration` FLOAT NULL DEFAULT NULL, `job_config` MEDIUMBLOB NOT NULL, + `spider_id` BIGINT UNSIGNED NULL DEFAULT NULL, PRIMARY KEY (`id`) USING BTREE, INDEX `CREATION_TIME` (`creation_time`) USING BTREE, - INDEX `JOB_STATUS` (`status`) USING BTREE + INDEX `JOB_STATUS` (`status`) USING BTREE, + INDEX `JOB_SPIDER_ID` (`spider_id`) USING BTREE ) ROW_FORMAT=DYNAMIC """ ) diff --git a/components/clp-rust-utils/src/job_config/search.rs b/components/clp-rust-utils/src/job_config/search.rs index dded4e7c39..ea2d6651d9 100644 --- a/components/clp-rust-utils/src/job_config/search.rs +++ b/components/clp-rust-utils/src/job_config/search.rs @@ -1,3 +1,4 @@ +use non_empty_string::NonEmptyString; use num_enum::IntoPrimitive; use num_enum::TryFromPrimitive; use serde::Deserialize; @@ -5,6 +6,10 @@ use serde::Serialize; pub const QUERY_JOBS_TABLE_NAME: &str = "query_jobs"; +pub type ArchiveId = NonEmptyString; + +pub type QueryJobId = i32; + /// Mirror of `job_orchestration.scheduler.job_config.AggregationConfig`. Must be kept in sync. #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] #[serde(default)] diff --git a/components/clp-rust-utils/src/task_io.rs b/components/clp-rust-utils/src/task_io.rs index 376ef623d5..09f7d4812e 100644 --- a/components/clp-rust-utils/src/task_io.rs +++ b/components/clp-rust-utils/src/task_io.rs @@ -1 +1,2 @@ pub mod compression; +pub mod query; diff --git a/components/clp-rust-utils/src/task_io/query.rs b/components/clp-rust-utils/src/task_io/query.rs new file mode 100644 index 0000000000..b5f53a9c4b --- /dev/null +++ b/components/clp-rust-utils/src/task_io/query.rs @@ -0,0 +1,31 @@ +//! Protocol types exchanged with the Spider (Huntsman) tasks that run CLP query jobs. + +use std::num::NonZeroU32; + +use non_empty_string::NonEmptyString; +use serde::Deserialize; +use serde::Serialize; + +/// `clp-s` options for a query job. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct ClpSQueryOption { + /// The query string passed positionally to `clp-s`. + pub query_string: NonEmptyString, + + /// The per-archive result limit. When absent, the task omits `--max-num-results` and uses the + /// `clp-s` default. + pub max_num_results: Option, + + /// Inclusive `--tge` bound in Unix epoch milliseconds. + pub begin_timestamp_millisecs: Option, + + /// Inclusive `--tle` bound in Unix epoch milliseconds. + pub end_timestamp_millisecs: Option, + + /// Whether `clp-s` performs a case-insensitive search. + pub ignore_case: bool, +} + +/// The output handler that `clp-s` writes a query task's results to. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub enum OutputHandle {} diff --git a/components/clp-tdl-package/src/lib.rs b/components/clp-tdl-package/src/lib.rs index 42aa104fb8..750047df6a 100644 --- a/components/clp-tdl-package/src/lib.rs +++ b/components/clp-tdl-package/src/lib.rs @@ -1,4 +1,4 @@ -//! Spider TDL task package `clp`: the CLP compression tasks the Spider task executor loads. +//! Spider TDL package `clp`, providing CLP compression and query tasks for Spider task executors. pub mod common; mod task; @@ -28,5 +28,9 @@ fn package_init() -> Result<(), TdlError> { spider_tdl::register_tdl_package! { package_name: "clp", init: package_init, - tasks: [task::compression::s3_compress_task, task::compression::commit_task], + tasks: [ + task::compression::s3_compress_task, + task::compression::commit_task, + task::query::clp_s_search_task, + ], } diff --git a/components/clp-tdl-package/src/task/compression/compress.rs b/components/clp-tdl-package/src/task/compression/compress.rs index a1ccc35b72..0cdf035603 100644 --- a/components/clp-tdl-package/src/task/compression/compress.rs +++ b/components/clp-tdl-package/src/task/compression/compress.rs @@ -10,10 +10,7 @@ use std::process::Command; use std::process::Stdio; use anyhow::Context; -use aws_config::BehaviorVersion; -use aws_sdk_s3::config::ProvideCredentials; use clp_rust_utils::aws::AWS_DEFAULT_REGION; -use clp_rust_utils::clp_config::AwsAuthentication; use clp_rust_utils::clp_config::S3Config; use clp_rust_utils::clp_config::package::config::ArchiveOutput; use clp_rust_utils::clp_config::package::config::ArchiveOutputStorage; @@ -30,6 +27,8 @@ use non_empty_string::NonEmptyString; use crate::common::clp_home; use crate::common::runtime; +use crate::task::utils::clp_binary_path; +use crate::task::utils::s3_credential_env; /// Compresses the given S3 objects into archives, uploads them to S3, and returns their metadata /// for the commit task. @@ -345,74 +344,6 @@ fn build_s3_logs_list(input_source: &S3InputSource) -> anyhow::Result { Ok(list) } -/// Resolves the AWS credential env vars clp-s needs to access the S3 objects. -/// -/// # Returns -/// -/// The env-var name-value pairs with the following environment variables set: -/// -/// * `AWS_ACCESS_KEY_ID` -/// * `AWS_SECRET_ACCESS_KEY` -/// * `AWS_SESSION_TOKEN` (if any) -/// -/// # Errors -/// -/// Returns an error if: -/// -/// * The default AWS SDK credential provider chain has no provider. -/// * Forwards [`ProvideCredentials::provide_credentials`]'s return values on failure. -fn s3_credential_env( - runtime: &tokio::runtime::Handle, - region: &str, - auth: &AwsAuthentication, -) -> anyhow::Result> { - /// The env var holding the AWS access key ID. - const AWS_ACCESS_KEY_ID_ENV_VAR: &str = "AWS_ACCESS_KEY_ID"; - - /// The env var holding the AWS secret access key. - const AWS_SECRET_ACCESS_KEY_ENV_VAR: &str = "AWS_SECRET_ACCESS_KEY"; - - /// The env var holding the AWS session token. - const AWS_SESSION_TOKEN_ENV_VAR: &str = "AWS_SESSION_TOKEN"; - - let (access_key_id, secret_access_key, session_token) = match auth { - AwsAuthentication::Credentials { credentials } => ( - credentials.access_key_id.clone(), - credentials.secret_access_key.clone(), - credentials.session_token.clone(), - ), - AwsAuthentication::Default => { - let sdk_config = runtime.block_on( - aws_config::defaults(BehaviorVersion::latest()) - .region(aws_sdk_s3::config::Region::new(region.to_string())) - .load(), - ); - let provider = sdk_config - .credentials_provider() - .context("default AWS SDK credential provider is unavailable")?; - let credentials = runtime - .block_on(provider.provide_credentials()) - .context("failed to resolve credentials from the default AWS SDK provider chain")?; - ( - credentials.access_key_id().to_string(), - credentials.secret_access_key().to_string(), - credentials - .session_token() - .map(std::string::ToString::to_string), - ) - } - }; - - let mut env = vec![ - (AWS_ACCESS_KEY_ID_ENV_VAR, access_key_id), - (AWS_SECRET_ACCESS_KEY_ENV_VAR, secret_access_key), - ]; - if let Some(session_token) = session_token { - env.push((AWS_SESSION_TOKEN_ENV_VAR, session_token)); - } - Ok(env) -} - /// Parses a single clp-s `--print-archive-stats` stdout line into an [`ArchiveMetadata`]. /// /// NOTE: clp-s emits a superset of [`ArchiveMetadata`]'s fields per line; unknown fields are @@ -585,15 +516,6 @@ fn build_log_converter_args(output_dir: &Path, inputs_from_path: &Path) -> Vec PathBuf { - clp_home.join("bin").join(binary) -} - /// Resolves the S3 config the archives are uploaded to from `config`. /// /// # Returns @@ -905,7 +827,6 @@ mod tests { use std::path::PathBuf; use clp_rust_utils::clp_config::AwsAuthentication; - use clp_rust_utils::clp_config::AwsCredentials; use clp_rust_utils::clp_config::S3Config; use clp_rust_utils::clp_config::package::config::ArchiveOutput; use clp_rust_utils::clp_config::package::config::ArchiveOutputStorage; @@ -923,7 +844,6 @@ mod tests { use super::build_s3_logs_list; use super::create_archive_s3_key; use super::parse_archive_stats; - use super::s3_credential_env; #[test] fn build_s3_logs_list_default_endpoint() -> anyhow::Result<()> { @@ -947,28 +867,6 @@ mod tests { Ok(()) } - #[test] - fn s3_credential_env_credentials() { - let runtime = tokio::runtime::Runtime::new().expect("failed to create Tokio runtime"); - let auth = AwsAuthentication::Credentials { - credentials: AwsCredentials { - access_key_id: "the-access-key".to_string(), - secret_access_key: "the-secret-key".to_string(), - session_token: Some("the-session-token".to_string()), - }, - }; - - assert_eq!( - s3_credential_env(runtime.handle(), "us-east-1", &auth) - .expect("failed to resolve credentials"), - vec![ - ("AWS_ACCESS_KEY_ID", "the-access-key".to_string()), - ("AWS_SECRET_ACCESS_KEY", "the-secret-key".to_string()), - ("AWS_SESSION_TOKEN", "the-session-token".to_string()), - ] - ); - } - #[test] fn parse_archive_stats_ignores_extra_keys() { let line = concat!( diff --git a/components/clp-tdl-package/src/task/mod.rs b/components/clp-tdl-package/src/task/mod.rs index f672b3ee36..06418e0022 100644 --- a/components/clp-tdl-package/src/task/mod.rs +++ b/components/clp-tdl-package/src/task/mod.rs @@ -1,3 +1,5 @@ //! The task implementations this package registers with Spider. pub mod compression; +pub mod query; +pub mod utils; diff --git a/components/clp-tdl-package/src/task/query/mod.rs b/components/clp-tdl-package/src/task/query/mod.rs new file mode 100644 index 0000000000..76697069a5 --- /dev/null +++ b/components/clp-tdl-package/src/task/query/mod.rs @@ -0,0 +1,21 @@ +//! The query-task signatures registered with Spider. + +use clp_rust_utils::job_config::ArchiveId; +use clp_rust_utils::job_config::QueryJobId; +use clp_rust_utils::task_io::query::ClpSQueryOption; +use clp_rust_utils::task_io::query::OutputHandle; +use non_empty_string::NonEmptyString; +use spider_tdl::TaskContext; +use spider_tdl::task; + +#[task(name = "query::clp_s_search")] +pub(crate) fn clp_s_search_task( + _ctx: TaskContext, + _query_job_id: QueryJobId, + _clp_s_query_option: ClpSQueryOption, + _dataset: Option, + _archive_id: ArchiveId, + _output_handle: OutputHandle, +) -> Result<(), spider_tdl::TdlError> { + todo!("clp-s search task is not implemented") +} diff --git a/components/clp-tdl-package/src/task/utils.rs b/components/clp-tdl-package/src/task/utils.rs new file mode 100644 index 0000000000..197eed0c2e --- /dev/null +++ b/components/clp-tdl-package/src/task/utils.rs @@ -0,0 +1,116 @@ +//! Helpers shared by the tasks that invoke CLP's core binaries. + +use std::path::Path; +use std::path::PathBuf; + +use anyhow::Context; +use aws_config::BehaviorVersion; +use aws_sdk_s3::config::ProvideCredentials; +use clp_rust_utils::clp_config::AwsAuthentication; + +/// Resolves the path of a CLP binary under `clp_home`, joining `bin/{binary}`. +/// +/// # Returns +/// +/// The path to the named binary under the CLP installation. +pub(super) fn clp_binary_path(clp_home: &Path, binary: &str) -> PathBuf { + clp_home.join("bin").join(binary) +} + +/// Resolves the AWS credential env vars clp-s needs to access the S3 objects. +/// +/// # Returns +/// +/// The env-var name-value pairs with the following environment variables set: +/// +/// * `AWS_ACCESS_KEY_ID` +/// * `AWS_SECRET_ACCESS_KEY` +/// * `AWS_SESSION_TOKEN` (if any) +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * The default AWS SDK credential provider chain has no provider. +/// * Forwards [`ProvideCredentials::provide_credentials`]'s return values on failure. +pub(super) fn s3_credential_env( + runtime: &tokio::runtime::Handle, + region: &str, + auth: &AwsAuthentication, +) -> anyhow::Result> { + /// The env var holding the AWS access key ID. + const AWS_ACCESS_KEY_ID_ENV_VAR: &str = "AWS_ACCESS_KEY_ID"; + + /// The env var holding the AWS secret access key. + const AWS_SECRET_ACCESS_KEY_ENV_VAR: &str = "AWS_SECRET_ACCESS_KEY"; + + /// The env var holding the AWS session token. + const AWS_SESSION_TOKEN_ENV_VAR: &str = "AWS_SESSION_TOKEN"; + + let (access_key_id, secret_access_key, session_token) = match auth { + AwsAuthentication::Credentials { credentials } => ( + credentials.access_key_id.clone(), + credentials.secret_access_key.clone(), + credentials.session_token.clone(), + ), + AwsAuthentication::Default => { + let sdk_config = runtime.block_on( + aws_config::defaults(BehaviorVersion::latest()) + .region(aws_sdk_s3::config::Region::new(region.to_string())) + .load(), + ); + let provider = sdk_config + .credentials_provider() + .context("default AWS SDK credential provider is unavailable")?; + let credentials = runtime + .block_on(provider.provide_credentials()) + .context("failed to resolve credentials from the default AWS SDK provider chain")?; + ( + credentials.access_key_id().to_string(), + credentials.secret_access_key().to_string(), + credentials + .session_token() + .map(std::string::ToString::to_string), + ) + } + }; + + let mut env = vec![ + (AWS_ACCESS_KEY_ID_ENV_VAR, access_key_id), + (AWS_SECRET_ACCESS_KEY_ENV_VAR, secret_access_key), + ]; + if let Some(session_token) = session_token { + env.push((AWS_SESSION_TOKEN_ENV_VAR, session_token)); + } + Ok(env) +} + +#[cfg(test)] +mod tests { + use clp_rust_utils::clp_config::AwsAuthentication; + use clp_rust_utils::clp_config::AwsCredentials; + + use super::s3_credential_env; + + #[test] + fn s3_credential_env_credentials() { + let runtime = tokio::runtime::Runtime::new().expect("failed to create Tokio runtime"); + let auth = AwsAuthentication::Credentials { + credentials: AwsCredentials { + access_key_id: "the-access-key".to_string(), + secret_access_key: "the-secret-key".to_string(), + session_token: Some("the-session-token".to_string()), + }, + }; + + assert_eq!( + s3_credential_env(runtime.handle(), "us-east-1", &auth) + .expect("failed to resolve credentials"), + vec![ + ("AWS_ACCESS_KEY_ID", "the-access-key".to_string()), + ("AWS_SECRET_ACCESS_KEY", "the-secret-key".to_string()), + ("AWS_SESSION_TOKEN", "the-session-token".to_string()), + ] + ); + } +} diff --git a/components/query-coordinator/Cargo.toml b/components/query-coordinator/Cargo.toml new file mode 100644 index 0000000000..c5bb002a41 --- /dev/null +++ b/components/query-coordinator/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "query-coordinator" +version = { workspace = true } +edition = { workspace = true } + +[dependencies] +async-trait = { workspace = true } +clp-rust-utils = { workspace = true } +non-empty-string = { workspace = true } +spider-client = { workspace = true } +spider-core = { workspace = true } +sqlx = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true } +tracing = { workspace = true } diff --git a/components/query-coordinator/src/error.rs b/components/query-coordinator/src/error.rs new file mode 100644 index 0000000000..af25857ab9 --- /dev/null +++ b/components/query-coordinator/src/error.rs @@ -0,0 +1,30 @@ +//! The crate-level error type for the query coordinator. + +/// Errors returned by the query coordinator. +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("query job {0} is no longer pending")] + JobNotPending(clp_rust_utils::job_config::QueryJobId), + + #[error("spider request failure: {0}")] + SpiderClient(#[from] spider_client::error::ClientError), + + #[error("sqlx error: {0}")] + Sqlx(#[from] sqlx::Error), + + #[error("failed to persist terminal status for query job {query_job_id}: {source}")] + TerminalStatusPersistence { + /// The query job whose terminal state could not be persisted. + query_job_id: clp_rust_utils::job_config::QueryJobId, + + /// The persistence failure. + #[source] + source: sqlx::Error, + }, + + #[error("number of query tasks {0} exceeds `i32::MAX`")] + TooManyQueryTasks(usize), + + #[error("no archives were selected for the query job")] + NoArchivesToSearch, +} diff --git a/components/query-coordinator/src/job_handle.rs b/components/query-coordinator/src/job_handle.rs new file mode 100644 index 0000000000..02a8dd6965 --- /dev/null +++ b/components/query-coordinator/src/job_handle.rs @@ -0,0 +1,308 @@ +//! Lifecycle management for one coordinator-planned query job. + +use std::sync::Arc; +use std::time::Duration; + +use clp_rust_utils::job_config::QUERY_JOBS_TABLE_NAME; +use clp_rust_utils::job_config::QueryJobId; +use clp_rust_utils::job_config::QueryJobStatus; +use clp_rust_utils::task_io::query::ClpSQueryOption; +use clp_rust_utils::task_io::query::OutputHandle; +use spider_core::task::ExecutionPolicy; +use spider_core::types::id::JobId as SpiderJobId; +use spider_core::types::id::ResourceGroupId; +use sqlx::MySqlPool; + +use crate::Error; +use crate::query_job_submitter::ArchiveMetadata; +use crate::query_job_submitter::QueryJobOutcome; +use crate::query_job_submitter::QueryJobSubmitter; + +/// Spider polling options shared by query-job handles. +pub struct SpiderOption { + /// Initial delay after a non-terminal Spider job-state poll. + pub initial_poll_backoff: Duration, + + /// Maximum delay between Spider job-state polls. + pub max_poll_backoff: Duration, +} + +/// Drives one already-planned query job through submission and terminal persistence. +/// +/// # Type Parameters +/// +/// * `SubmitterType` - The type of the job submitter for Spider job submission. +pub struct QueryJobHandle { + db_pool: MySqlPool, + query_job_id: QueryJobId, + job_submitter: SubmitterType, + resource_group_id: ResourceGroupId, + clp_s_query_option: ClpSQueryOption, + output_handle: OutputHandle, + archives_to_search: Vec<(ArchiveMetadata, ExecutionPolicy)>, + spider_option: Arc, +} + +impl QueryJobHandle { + /// Factory function. + /// + /// # Returns + /// + /// A newly created [`QueryJobHandle`] for the given already-planned query job. + pub const fn new( + db_pool: MySqlPool, + query_job_id: QueryJobId, + job_submitter: SubmitterType, + resource_group_id: ResourceGroupId, + clp_s_query_option: ClpSQueryOption, + output_handle: OutputHandle, + archives_to_search: Vec<(ArchiveMetadata, ExecutionPolicy)>, + spider_option: Arc, + ) -> Self { + Self { + db_pool, + query_job_id, + job_submitter, + resource_group_id, + clp_s_query_option, + output_handle, + archives_to_search, + spider_option, + } + } + + /// Submits the prepared graph and drives the query job to a terminal state. + /// + /// On a submission failure, this method makes a best-effort attempt to mark the CLP query job + /// as failed before returning the original error. After the job is durably running, monitoring + /// and terminal-persistence failures leave it running so recovery can reattach to Spider. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::submit`]'s return values on failure. + /// * Forwards [`Self::to_completion`]'s return values on failure. + pub async fn run(self) -> Result<(), Error> { + tracing::info!(query_job_id = % self.query_job_id, "Starting query job."); + + let spider_job_id = match self.submit().await { + Ok(spider_job_id) => spider_job_id, + Err(error) => { + if !matches!(error, Error::JobNotPending(_)) { + self.report_failure(&error).await; + } + return Err(error); + } + }; + self.to_completion(spider_job_id).await + } + + /// Resumes a query job that was already submitted to Spider. + /// + /// The caller must ensure `spider_job_id` belongs to this CLP query job. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::to_completion`]'s return values on failure. + pub async fn recover(self, spider_job_id: SpiderJobId) -> Result<(), Error> { + tracing::info!( + query_job_id = % self.query_job_id, + spider_job_id = % spider_job_id, + "Recovering query job.", + ); + + self.to_completion(spider_job_id).await + } + + /// Submits the query job to Spider and persists its running state. + /// + /// # Returns + /// + /// The submitted Spider job ID on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`Error::TooManyQueryTasks`] if the number of query tasks exceeds `i32`'s range. + /// * Forwards [`Self::submit_to_spider`]'s return values on failure. + /// * Forwards [`Self::persist_submission`]'s return values on failure. + async fn submit(&self) -> Result { + let num_tasks = self.archives_to_search.len(); + if num_tasks == 0 { + return Err(Error::NoArchivesToSearch); + } + let persisted_num_tasks = + i32::try_from(num_tasks).map_err(|_| Error::TooManyQueryTasks(num_tasks))?; + let spider_job_id = self.submit_to_spider().await?; + + tracing::info!( + query_job_id = % self.query_job_id, + spider_job_id = % spider_job_id, + num_tasks, + "Query job submitted.", + ); + + self.persist_submission(spider_job_id, persisted_num_tasks) + .await?; + Ok(spider_job_id) + } + + /// Submits the prepared query graph to Spider. + /// + /// # Returns + /// + /// The submitted Spider job ID on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`QueryJobSubmitter::submit_query_job`]'s return values on failure. + async fn submit_to_spider(&self) -> Result { + self.job_submitter + .submit_query_job( + self.query_job_id, + self.resource_group_id, + self.clp_s_query_option.clone(), + self.output_handle.clone(), + self.archives_to_search.clone(), + ) + .await + } + + /// Persists the Spider job ID and marks the query job as running. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`Error::JobNotPending`] if the query job is no longer pending. + /// * Forwards [`sqlx::query::Query::execute`]'s return values on failure. + async fn persist_submission( + &self, + spider_job_id: SpiderJobId, + num_tasks: i32, + ) -> Result<(), Error> { + let query = format!( + "UPDATE `{QUERY_JOBS_TABLE_NAME}` SET `spider_id` = ?, `status` = ?, `num_tasks` = ?, \ + `start_time` = CURRENT_TIMESTAMP(3) WHERE `id` = ? AND `status` = ?" + ); + let result = sqlx::query(&query) + .bind(spider_job_id.get()) + .bind(i32::from(QueryJobStatus::Running)) + .bind(num_tasks) + .bind(self.query_job_id) + .bind(i32::from(QueryJobStatus::Pending)) + .execute(&self.db_pool) + .await?; + + if 1 != result.rows_affected() { + return Err(Error::JobNotPending(self.query_job_id)); + } + Ok(()) + } + + /// Waits for the associated Spider job to complete and finalizes the query job. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`Error::TerminalStatusPersistence`] if the terminal query-job status cannot be persisted. + /// * Forwards [`QueryJobSubmitter::run_query_job_to_completion`]'s return values on failure. + async fn to_completion(&self, spider_job_id: SpiderJobId) -> Result<(), Error> { + let outcome = self + .job_submitter + .run_query_job_to_completion( + spider_job_id, + self.spider_option.initial_poll_backoff, + self.spider_option.max_poll_backoff, + ) + .await?; + + tracing::info!( + query_job_id = % self.query_job_id, + spider_job_id = % spider_job_id, + outcome = ? outcome, + "Query job reached a terminal Spider state.", + ); + + let (status, status_message) = match outcome { + QueryJobOutcome::Succeeded => (QueryJobStatus::Succeeded, String::new()), + QueryJobOutcome::Failed { error_message } => ( + QueryJobStatus::Failed, + format!("The Spider query job failed: {error_message}"), + ), + QueryJobOutcome::UnexpectedlyCancelled => ( + QueryJobStatus::Failed, + "The Spider query job was unexpectedly cancelled.".to_string(), + ), + }; + self.update_terminal_status(status, &status_message, QueryJobStatus::Running) + .await + .map_err(|source| Error::TerminalStatusPersistence { + query_job_id: self.query_job_id, + source, + }) + } + + /// Reports a query-job orchestration failure. + /// + /// Logs the original error and makes a best-effort attempt to mark the query job as failed. If + /// terminal-status persistence fails, the status-update error is logged and otherwise ignored. + async fn report_failure(&self, error: &Error) { + tracing::error!( + query_job_id = % self.query_job_id, + error = % error, + "Query-job orchestration failed.", + ); + + if let Err(status_error) = self + .update_terminal_status( + QueryJobStatus::Failed, + &format!("Query-job orchestration failed: {error}"), + QueryJobStatus::Pending, + ) + .await + { + tracing::error!( + query_job_id = % self.query_job_id, + error = % status_error, + "Failed to persist the query-job failure.", + ); + } + } + + /// Updates a query job only when it has the expected non-terminal status. + /// A zero-row update is treated as success so an ineligible or missing job row is left + /// unchanged. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`sqlx::query::Query::execute`]'s return values on failure. + async fn update_terminal_status( + &self, + status: QueryJobStatus, + status_message: &str, + expected_status: QueryJobStatus, + ) -> Result<(), sqlx::Error> { + let query = format!( + "UPDATE `{QUERY_JOBS_TABLE_NAME}` SET `status` = ?, `status_msg` = LEFT(?, 512), \ + `duration` = CASE WHEN `start_time` IS NULL THEN 0 ELSE TIMESTAMPDIFF(MICROSECOND, \ + `start_time`, CURRENT_TIMESTAMP(3)) / 1000000.0 END WHERE `id` = ? AND `status` = ?" + ); + let query = sqlx::query(&query) + .bind(i32::from(status)) + .bind(status_message) + .bind(self.query_job_id) + .bind(i32::from(expected_status)); + query.execute(&self.db_pool).await?; + Ok(()) + } +} diff --git a/components/query-coordinator/src/lib.rs b/components/query-coordinator/src/lib.rs new file mode 100644 index 0000000000..32869c80d1 --- /dev/null +++ b/components/query-coordinator/src/lib.rs @@ -0,0 +1,7 @@ +//! Coordination for CLP query jobs. + +mod error; +pub mod job_handle; +pub mod query_job_submitter; + +pub use error::Error; diff --git a/components/query-coordinator/src/query_job_submitter/mod.rs b/components/query-coordinator/src/query_job_submitter/mod.rs new file mode 100644 index 0000000000..8b6679ca2c --- /dev/null +++ b/components/query-coordinator/src/query_job_submitter/mod.rs @@ -0,0 +1,100 @@ +//! The query-job submission interface. + +mod spider; + +use std::time::Duration; + +use async_trait::async_trait; +use clp_rust_utils::job_config::ArchiveId; +use clp_rust_utils::job_config::QueryJobId; +use clp_rust_utils::task_io::query::ClpSQueryOption; +use clp_rust_utils::task_io::query::OutputHandle; +use non_empty_string::NonEmptyString; +use spider_core::task::ExecutionPolicy; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; + +use crate::Error; + +/// Identifies an archive handled by query tasks. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ArchiveMetadata { + /// The archive's ID. + pub id: ArchiveId, + + /// The archive's dataset, or `None` for the default dataset. + pub dataset: Option, + + /// The archive's compressed size in bytes. + pub size: u64, +} + +/// The terminal outcome of a query job. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum QueryJobOutcome { + /// Every archive query completed successfully. + Succeeded, + + /// At least one archive query failed. + Failed { + /// The error reported by Spider. + error_message: String, + }, + + /// Spider cancelled the job unexpectedly. User-requested cancellation is outside the MVP. + UnexpectedlyCancelled, +} + +/// Drives CLP query jobs on a Spider (Huntsman) cluster. +#[async_trait] +pub trait QueryJobSubmitter: Clone + Send + Sync { + /// Builds the query task graph for the given archives and registers it with Spider, without + /// starting it. + /// + /// # Parameters + /// + /// * `query_job_id` - The unique ID of the CLP query job. + /// * `resource_group_id` - The Spider resource group to register the job under. + /// * `clp_s_query_option` - `clp-s` query options shared by every task in the job. + /// * `output_handle` - The output handle selecting how the query outputs are returned. + /// * `archives_to_search` - The archives to search, each represents a query task paired with + /// the task execution policy. + /// + /// # Returns + /// + /// The job ID issued by Spider on success. + /// + /// # Errors + /// + /// Implementations must document their error conditions. + async fn submit_query_job( + &self, + query_job_id: QueryJobId, + resource_group_id: ResourceGroupId, + clp_s_query_option: ClpSQueryOption, + output_handle: OutputHandle, + archives_to_search: Vec<(ArchiveMetadata, ExecutionPolicy)>, + ) -> Result; + + /// Idempotently starts `spider_job_id` and waits for it to reach a terminal state. + /// + /// # Parameters + /// + /// * `spider_job_id` - The ID of the Spider job to start and monitor. + /// * `initial_poll_backoff` - The initial delay after a non-terminal job-state poll. + /// * `max_poll_backoff` - The maximum delay between job-state polls. + /// + /// # Returns + /// + /// The terminal query-job outcome on success. + /// + /// # Errors + /// + /// Implementations must document their error conditions. + async fn run_query_job_to_completion( + &self, + spider_job_id: JobId, + initial_poll_backoff: Duration, + max_poll_backoff: Duration, + ) -> Result; +} diff --git a/components/query-coordinator/src/query_job_submitter/spider.rs b/components/query-coordinator/src/query_job_submitter/spider.rs new file mode 100644 index 0000000000..40563c15dd --- /dev/null +++ b/components/query-coordinator/src/query_job_submitter/spider.rs @@ -0,0 +1,90 @@ +//! [`QueryJobSubmitter`] implementation for [`spider_client::SpiderClient`]. + +use std::time::Duration; + +use async_trait::async_trait; +use clp_rust_utils::job_config::QueryJobId; +use clp_rust_utils::task_io::query::ClpSQueryOption; +use clp_rust_utils::task_io::query::OutputHandle; +use spider_client::SpiderClient; +use spider_client::error::ClientError; +use spider_core::job::JobState; +use spider_core::task::ExecutionPolicy; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; + +use crate::Error; +use crate::query_job_submitter::ArchiveMetadata; +use crate::query_job_submitter::QueryJobOutcome; +use crate::query_job_submitter::QueryJobSubmitter; + +#[async_trait] +impl QueryJobSubmitter for SpiderClient { + async fn submit_query_job( + &self, + _query_job_id: QueryJobId, + _resource_group_id: ResourceGroupId, + _clp_s_query_option: ClpSQueryOption, + _output_handle: OutputHandle, + _archives_to_search: Vec<(ArchiveMetadata, ExecutionPolicy)>, + ) -> Result { + todo!("construct and submit the clp-s query task graph") + } + + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`SpiderClient::start_job`]'s return values on failure, except + /// [`ClientError::InvalidJobState`]. + /// * Forwards [`SpiderClient::get_job_state`]'s return values on failure. + /// + /// # Panics + /// + /// Panics if Spider returns a terminal state without a corresponding [`QueryJobOutcome`]. + async fn run_query_job_to_completion( + &self, + spider_job_id: JobId, + initial_poll_backoff: Duration, + max_poll_backoff: Duration, + ) -> Result { + const POLL_BACKOFF_FACTOR: u32 = 2; + + match self.start_job(spider_job_id).await { + Ok(_) | Err(ClientError::InvalidJobState(_)) => {} + Err(error) => return Err(error.into()), + } + + let mut backoff = initial_poll_backoff.min(max_poll_backoff); + let terminal_state = loop { + let state = self.get_job_state(spider_job_id).await?; + if state.is_terminal() { + break state; + } + tokio::time::sleep(backoff).await; + backoff = backoff + .saturating_mul(POLL_BACKOFF_FACTOR) + .min(max_poll_backoff); + }; + + Ok(match terminal_state { + JobState::Succeeded => QueryJobOutcome::Succeeded, + JobState::Failed => { + let error_message = match self.get_job_error(spider_job_id).await { + Ok(error_message) => error_message, + Err(error) => { + tracing::warn!( + spider_job_id = % spider_job_id, + error = % error, + "Failed to fetch the Spider job error.", + ); + format!("") + } + }; + QueryJobOutcome::Failed { error_message } + } + JobState::Cancelled => QueryJobOutcome::UnexpectedlyCancelled, + _ => unreachable!("a terminal Spider state must have a terminal outcome"), + }) + } +}