diff --git a/Cargo.lock b/Cargo.lock index 3c490b08db..f90af1ffdb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3271,6 +3271,27 @@ dependencies = [ "pulldown-cmark", ] +[[package]] +name = "query-coordinator" +version = "0.13.1-dev" +dependencies = [ + "async-trait", + "clp-rust-utils", + "const_format", + "mongodb", + "non-empty-string", + "rmp-serde", + "serde", + "spider-client", + "spider-core", + "sqlx", + "thiserror", + "tokio", + "tokio-util", + "tonic", + "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..b8982cd8e2 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,19 @@ 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, + `dispatch_time` DATETIME 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/clp_config/package/config.rs b/components/clp-rust-utils/src/clp_config/package/config.rs index 1c9fa72d71..e751424257 100644 --- a/components/clp-rust-utils/src/clp_config/package/config.rs +++ b/components/clp-rust-utils/src/clp_config/package/config.rs @@ -10,6 +10,7 @@ use serde::Deserialize; use crate::clp_config::AwsAuthentication; use crate::clp_config::S3Config; use crate::dataset::resolve_dataset_name; +use crate::types::non_empty_string::ExpectedNonEmpty; /// Mirror of `clp_py_utils.clp_config.ClpConfig`. /// @@ -280,6 +281,17 @@ pub struct ResultsCache { pub db_name: String, } +impl ResultsCache { + /// Returns the MongoDB URI for the results cache database. + #[must_use] + pub fn uri(&self) -> NonEmptyString { + NonEmptyString::from_string(format!( + "mongodb://{}:{}/{}", + self.host, self.port, self.db_name + )) + } +} + impl Default for ResultsCache { fn default() -> Self { Self { @@ -381,6 +393,20 @@ impl ArchiveOutput { .to_string_lossy() .into_owned() } + + /// Derives the S3 object key of an archive in a dataset. + /// + /// # Returns + /// + /// The dataset's archive storage directory joined with `archive_id`, where a `None` dataset + /// resolves to `default`. + #[must_use] + pub fn dataset_archive_object_key(&self, dataset: Option<&str>, archive_id: &str) -> String { + format!( + "{}/{archive_id}", + self.dataset_archive_storage_directory(dataset) + ) + } } impl Default for ArchiveOutput { @@ -477,6 +503,37 @@ impl Default for Telemetry { } } +/// Query coordinator configuration. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq)] +#[serde(default)] +pub struct QueryCoordinator { + pub resource_group: SpiderResourceGroup, + pub job_polling_interval_millisecs: NonZeroU64, + pub max_concurrent_jobs: NonZeroUsize, + pub result_polling: PollingBackoff, +} + +impl Default for QueryCoordinator { + fn default() -> Self { + Self { + resource_group: SpiderResourceGroup { + name: NonEmptyString::new("query-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"), + max_concurrent_jobs: NonZeroUsize::new(1000) + .expect("default maximum number of concurrent jobs 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 coordinator configuration. #[derive(Clone, Debug, Deserialize, Eq, PartialEq)] #[serde(default)] @@ -699,6 +756,38 @@ mod tests { ); } + #[test] + fn dataset_archive_object_key_joins_prefix_dataset_and_id() { + use non_empty_string::NonEmptyString; + + use crate::clp_config::AwsAuthentication; + use crate::clp_config::S3Config; + use crate::types::non_empty_string::ExpectedNonEmpty; + + let archive_output = ArchiveOutput { + storage: ArchiveOutputStorage::S3 { + staging_directory: "var/data/staged-archives".to_owned(), + s3_config: S3Config { + bucket: NonEmptyString::from_static_str("bucket"), + region_code: None, + key_prefix: NonEmptyString::from_static_str("LIB1/"), + endpoint_url: None, + aws_authentication: AwsAuthentication::Default, + }, + }, + ..ArchiveOutput::default() + }; + + assert_eq!( + archive_output.dataset_archive_object_key(None, "abc"), + "LIB1/default/abc" + ); + assert_eq!( + archive_output.dataset_archive_object_key(Some("mydataset"), "abc"), + "LIB1/mydataset/abc" + ); + } + #[test] fn deserialize_database_ignores_provided_table_prefix() { let database_json = serde_json::json!({ diff --git a/components/clp-rust-utils/src/job_config/search.rs b/components/clp-rust-utils/src/job_config/search.rs index dded4e7c39..47cc069eb5 100644 --- a/components/clp-rust-utils/src/job_config/search.rs +++ b/components/clp-rust-utils/src/job_config/search.rs @@ -1,10 +1,17 @@ +use non_empty_string::NonEmptyString; use num_enum::IntoPrimitive; use num_enum::TryFromPrimitive; use serde::Deserialize; use serde::Serialize; +use strum::EnumString; +use utoipa::ToSchema; 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)] @@ -34,8 +41,22 @@ pub struct SearchJobConfig { } /// Mirror of `job_orchestration.scheduler.constants.QueryJobStatus`. Must be kept in sync. -#[derive(Clone, Debug, Deserialize, Eq, IntoPrimitive, PartialEq, Serialize, TryFromPrimitive)] +#[derive( + Clone, + Copy, + Debug, + Deserialize, + EnumString, + Eq, + IntoPrimitive, + PartialEq, + Serialize, + ToSchema, + TryFromPrimitive, + sqlx::Type, +)] #[repr(i32)] +#[strum(ascii_case_insensitive)] pub enum QueryJobStatus { Pending = 0, Running = 1, 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..cbeb8f48c2 --- /dev/null +++ b/components/clp-rust-utils/src/task_io/query.rs @@ -0,0 +1,128 @@ +//! 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)] +#[serde(deny_unknown_fields)] +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)] +#[serde(deny_unknown_fields, tag = "type")] +pub enum OutputHandle { + /// The results cache, addressed by a MongoDB URI whose path names the database. The collection + /// is the query job's ID. + #[serde(rename = "results_cache")] + ResultsCache { uri: NonEmptyString }, + + /// A file per archive. Not yet supported by the Spider query flow. + #[serde(rename = "file")] + File, +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU32; + + use non_empty_string::NonEmptyString; + + use super::ClpSQueryOption; + use super::OutputHandle; + use crate::types::non_empty_string::ExpectedNonEmpty; + + #[test] + fn clp_s_query_option_with_timestamp_bounds_round_trips_through_msgpack() { + let expected = ClpSQueryOption { + query_string: NonEmptyString::from_static_str("level:error"), + max_num_results: Some(NonZeroU32::new(1_000).expect("1,000 is nonzero")), + begin_timestamp_millisecs: Some(1_700_000_000_001), + end_timestamp_millisecs: Some(1_700_000_000_999), + ignore_case: true, + }; + + let serialized = rmp_serde::to_vec(&expected).expect("query options should serialize"); + let actual: ClpSQueryOption = + rmp_serde::from_slice(&serialized).expect("query options should deserialize"); + + assert_eq!(expected, actual); + } + + #[test] + fn clp_s_query_option_without_timestamp_bounds_round_trips_through_msgpack() { + let expected = ClpSQueryOption { + query_string: NonEmptyString::from_static_str("*"), + max_num_results: Some(NonZeroU32::new(1).expect("1 is nonzero")), + begin_timestamp_millisecs: None, + end_timestamp_millisecs: None, + ignore_case: false, + }; + + let serialized = rmp_serde::to_vec(&expected).expect("query options should serialize"); + let actual: ClpSQueryOption = + rmp_serde::from_slice(&serialized).expect("query options should deserialize"); + + assert_eq!(expected, actual); + } + + #[test] + fn clp_s_query_option_without_max_num_results_round_trips_through_msgpack() { + let expected = ClpSQueryOption { + query_string: NonEmptyString::from_static_str("*"), + max_num_results: None, + begin_timestamp_millisecs: None, + end_timestamp_millisecs: None, + ignore_case: false, + }; + + let serialized = rmp_serde::to_vec(&expected).expect("query options should serialize"); + let actual: ClpSQueryOption = + rmp_serde::from_slice(&serialized).expect("query options should deserialize"); + + assert_eq!(expected, actual); + } + + #[test] + fn output_handle_results_cache_round_trips_through_msgpack() { + let expected = OutputHandle::ResultsCache { + uri: NonEmptyString::from_static_str("mongodb://results-cache:27017/clp-query-results"), + }; + + let serialized = rmp_serde::to_vec(&expected).expect("output handle should serialize"); + let actual: OutputHandle = + rmp_serde::from_slice(&serialized).expect("output handle should deserialize"); + + assert_eq!(expected, actual); + } + + #[test] + fn output_handle_file_round_trips_through_msgpack() { + let expected = OutputHandle::File; + + let serialized = rmp_serde::to_vec(&expected).expect("output handle should serialize"); + let actual: OutputHandle = + rmp_serde::from_slice(&serialized).expect("output handle should deserialize"); + + assert_eq!(expected, actual); + } +} diff --git a/components/clp-tdl-package/README.md b/components/clp-tdl-package/README.md index b22ac07af7..0ab8164359 100644 --- a/components/clp-tdl-package/README.md +++ b/components/clp-tdl-package/README.md @@ -12,3 +12,7 @@ documented below. * `compression::clp_s_s3_compress`: Compress inputs from S3 using `clp-s`. * `compression::commit`: Commit compression task outcomes to the CLP metadata database. + +### Query + +* `query::clp_s_search`: Search a single archive using the `clp-s` engine. 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..fcec9fbce8 100644 --- a/components/clp-tdl-package/src/task/compression/compress.rs +++ b/components/clp-tdl-package/src/task/compression/compress.rs @@ -10,12 +10,8 @@ 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; use clp_rust_utils::clp_config::package::config::Database; use clp_rust_utils::clp_config::package::config::SpiderTaskExecutorConfig; @@ -30,6 +26,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. @@ -130,7 +128,9 @@ pub(super) fn compress( ArchiveFinisher { client: client.clone(), bucket: bucket.clone(), - key: create_archive_s3_key(&config.archive_output, dataset.as_deref(), &archive.id), + key: config + .archive_output + .dataset_archive_object_key(dataset.as_deref(), &archive.id), indexer_bin: indexer_bin.clone(), database: config.database.clone(), dataset: dataset.clone(), @@ -345,74 +345,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 +517,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 @@ -614,23 +537,6 @@ fn extract_s3_output_config(config: &SpiderTaskExecutorConfig) -> anyhow::Result } } -/// Builds the S3 object key for an archive by appending `archive_id` to -/// [`ArchiveOutput::dataset_archive_storage_directory`]. -/// -/// # Returns -/// -/// The archive's S3 object key. -fn create_archive_s3_key( - archive_output: &ArchiveOutput, - dataset: Option<&str>, - archive_id: &str, -) -> String { - format!( - "{}/{archive_id}", - archive_output.dataset_archive_storage_directory(dataset) - ) -} - /// Uploads a local file to S3 through `PutObject`. /// /// # Errors @@ -905,10 +811,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; use clp_rust_utils::clp_config::package::config::ClpDbNames; use clp_rust_utils::clp_config::package::config::Database; use clp_rust_utils::task_io::compression::ArchiveMetadata; @@ -921,9 +823,7 @@ mod tests { use super::build_indexer_args; use super::build_log_converter_args; 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 +847,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!( @@ -1105,30 +983,6 @@ mod tests { ); } - #[test] - fn archive_s3_key_joins_prefix_dataset_and_id() { - let archive_output = ArchiveOutput { - storage: ArchiveOutputStorage::S3 { - staging_directory: "var/data/staged-archives".to_owned(), - s3_config: S3Config { - bucket: NonEmptyString::try_from("bucket".to_string()) - .expect("bucket is non-empty"), - region_code: None, - key_prefix: NonEmptyString::try_from("LIB1/".to_string()) - .expect("key prefix is non-empty"), - endpoint_url: None, - aws_authentication: AwsAuthentication::Default, - }, - }, - ..ArchiveOutput::default() - }; - - assert_eq!( - create_archive_s3_key(&archive_output, None, "abc"), - "LIB1/default/abc" - ); - } - #[test] fn build_indexer_args_uses_mysql_and_expected_order() { let database = Database { 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..562692501a --- /dev/null +++ b/components/clp-tdl-package/src/task/query/mod.rs @@ -0,0 +1,32 @@ +//! The query tasks: the `#[task]` wrappers Spider invokes and their implementations. + +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::TdlError; +use spider_tdl::task; + +mod search; + +#[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: NonEmptyString, + output_handle: OutputHandle, +) -> Result<(), TdlError> { + search::search( + &ctx, + crate::common::spider_task_executor_config(), + query_job_id, + &clp_s_query_option, + archive_id.into_inner(), + dataset.as_ref().map(NonEmptyString::as_str), + &output_handle, + ) + .map_err(|e| TdlError::ExecutionError(format!("{e:#}"))) +} diff --git a/components/clp-tdl-package/src/task/query/search.rs b/components/clp-tdl-package/src/task/query/search.rs new file mode 100644 index 0000000000..571175ba82 --- /dev/null +++ b/components/clp-tdl-package/src/task/query/search.rs @@ -0,0 +1,704 @@ +//! The `clp-s` search worker that queries a single archive. + +use std::ffi::OsString; +use std::io::Read; +use std::path::Path; +use std::path::PathBuf; +use std::process::Command; +use std::process::Stdio; + +use anyhow::Context; +use clp_rust_utils::aws::AWS_DEFAULT_REGION; +use clp_rust_utils::clp_config::package::config::ArchiveOutputStorage; +use clp_rust_utils::clp_config::package::config::SpiderTaskExecutorConfig; +use clp_rust_utils::clp_config::package::config::StorageEngine; +use clp_rust_utils::dataset::resolve_dataset_name; +use clp_rust_utils::job_config::QueryJobId; +use clp_rust_utils::s3::generate_s3_url; +use clp_rust_utils::task_io::query::ClpSQueryOption; +use clp_rust_utils::task_io::query::OutputHandle; +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; + +/// Searches one archive with clp-s, handles the search results according to the given +/// `output_handle`. +/// +/// A pure worker function called by a spider-tdl task wrapper, which formats any returned +/// `anyhow::Error` into a user-space TDL error. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * The configured storage engine is not [`StorageEngine::ClpS`]. +/// * `output_handle` is not [`OutputHandle::ResultsCache`]. The current implementation only +/// supports result cache output streaming. +/// * Forwards [`resolve_archive_input`]'s return values on failure. +/// * Forwards [`run_clp_s_search`]'s return values on failure. +pub(super) fn search( + ctx: &spider_tdl::TaskContext, + config: &SpiderTaskExecutorConfig, + query_job_id: QueryJobId, + clp_s_query_option: &ClpSQueryOption, + archive_id: String, + dataset: Option<&str>, + output_handle: &OutputHandle, +) -> anyhow::Result<()> { + if StorageEngine::ClpS != config.package.storage_engine { + anyhow::bail!("the clp-s query task requires the `clp-s` storage engine"); + } + let OutputHandle::ResultsCache { uri } = output_handle else { + anyhow::bail!("unsupported query output handler"); + }; + + let dataset = resolve_dataset_name(dataset); + + tracing::info!( + job_id = % ctx.job_id, + task_id = % ctx.task_id, + task_instance_id = % ctx.task_instance_id, + query_job_id = % query_job_id, + dataset = % dataset, + archive_id = % archive_id, + "clp-s query task started.", + ); + + let clp_home = clp_home(); + let (archive_selector, credential_env) = + resolve_archive_input(&runtime(), clp_home, config, dataset, archive_id).inspect_err( + |e| { + tracing::error!( + job_id = % ctx.job_id, + task_id = % ctx.task_id, + task_instance_id = % ctx.task_instance_id, + query_job_id = % query_job_id, + error = % e, + "Failed to resolve the archive input." + ); + }, + )?; + let args = build_clp_s_search_args_for_result_cache( + &archive_selector, + clp_s_query_option, + uri.as_str(), + query_job_id, + dataset, + ); + run_clp_s_search(&clp_binary_path(clp_home, "clp-s"), args, &credential_env)?; + + tracing::info!( + job_id = % ctx.job_id, + task_id = % ctx.task_id, + task_instance_id = % ctx.task_instance_id, + query_job_id = % query_job_id, + "clp-s query task completed successfully.", + ); + Ok(()) +} + +/// Selector for clp-s to address the archive to search. +enum ArchiveSelector { + /// A local dataset archives directory plus the `--archive-id` selecting one archive in it. + Directory { path: PathBuf, archive_id: String }, + + /// The URL of an S3-hosted archive, read with `--auth s3`. + ObjectUrl(String), +} + +/// Resolves how clp-s addresses the archive, and the credential env vars it needs to read it. +/// +/// # Returns +/// +/// A tuple containing: +/// +/// * The archive selector. +/// * The credential env vars clp-s should run with, which are empty for filesystem-backed archive +/// output. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * The archive's object key is empty. +/// * Forwards [`generate_s3_url`]'s return values on failure. +/// * Forwards [`s3_credential_env`]'s return values on failure. +fn resolve_archive_input( + runtime: &tokio::runtime::Handle, + clp_home: &Path, + config: &SpiderTaskExecutorConfig, + dataset: &str, + archive_id: String, +) -> anyhow::Result<(ArchiveSelector, Vec<(&'static str, String)>)> { + let s3_config = match &config.archive_output.storage { + ArchiveOutputStorage::Fs { .. } => { + return Ok(( + ArchiveSelector::Directory { + path: config.abs_archive_output_staging(clp_home).join(dataset), + archive_id, + }, + Vec::new(), + )); + } + ArchiveOutputStorage::S3 { s3_config, .. } => s3_config, + }; + + let object_key = config + .archive_output + .dataset_archive_object_key(Some(dataset), &archive_id); + let object_key = NonEmptyString::try_from(object_key) + .map_err(|_| anyhow::anyhow!("archive object key must not be empty"))?; + let url = generate_s3_url( + s3_config.endpoint_url.as_ref().map(NonEmptyString::as_str), + s3_config.region_code.as_ref().map(NonEmptyString::as_str), + &s3_config.bucket, + &object_key, + )?; + let region = s3_config + .region_code + .as_ref() + .map_or(AWS_DEFAULT_REGION, NonEmptyString::as_str); + let credential_env = s3_credential_env(runtime, region, &s3_config.aws_authentication)?; + Ok((ArchiveSelector::ObjectUrl(url), credential_env)) +} + +/// Builds the clp-s command-line arguments for a single-archive search writing to the result cache. +/// +/// # Returns +/// +/// The ordered clp-s arguments. +fn build_clp_s_search_args_for_result_cache( + archive_selector: &ArchiveSelector, + clp_s_query_option: &ClpSQueryOption, + result_cache_uri: &str, + query_job_id: QueryJobId, + dataset: &str, +) -> Vec { + let mut args = vec![OsString::from("s")]; + match archive_selector { + ArchiveSelector::Directory { path, archive_id } => { + args.push(path.as_os_str().to_os_string()); + args.push(OsString::from("--archive-id")); + args.push(OsString::from(archive_id)); + } + ArchiveSelector::ObjectUrl(url) => { + args.push(OsString::from(url)); + args.push(OsString::from("--auth")); + args.push(OsString::from("s3")); + } + } + + args.push(OsString::from(clp_s_query_option.query_string.as_str())); + if let Some(begin_timestamp_millisecs) = clp_s_query_option.begin_timestamp_millisecs { + args.push(OsString::from("--tge")); + args.push(OsString::from(begin_timestamp_millisecs.to_string())); + } + if let Some(end_timestamp_millisecs) = clp_s_query_option.end_timestamp_millisecs { + args.push(OsString::from("--tle")); + args.push(OsString::from(end_timestamp_millisecs.to_string())); + } + if clp_s_query_option.ignore_case { + args.push(OsString::from("--ignore-case")); + } + + args.extend([ + OsString::from("results-cache"), + OsString::from("--uri"), + OsString::from(result_cache_uri), + OsString::from("--collection"), + OsString::from(query_job_id.to_string()), + ]); + if let Some(max_num_results) = clp_s_query_option.max_num_results { + args.push(OsString::from("--max-num-results")); + args.push(OsString::from(max_num_results.to_string())); + } + args.extend([OsString::from("--dataset"), OsString::from(dataset)]); + args +} + +/// Runs clp-s with the given search arguments, blocking until it exits. +/// +/// # Observability +/// +/// This method logs errors on failure before returning to the caller. `clp-s`' stderr is logged if +/// successfully captured. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * clp-s exits with a non-zero status. +/// * Forwards [`Command::spawn`]'s return values on failure. +/// * Forwards [`std::process::Child::wait`]'s return values on failure. +/// +/// # Panics +/// +/// Panics if clp-s's piped stderr is unexpectedly absent (should be unreachable). +fn run_clp_s_search( + clp_s_bin: &Path, + args: Vec, + credential_env: &[(&'static str, String)], +) -> anyhow::Result<()> { + let mut child = Command::new(clp_s_bin) + .args(args) + .envs(credential_env.iter().cloned()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn() + .with_context(|| format!("failed to spawn clp-s at {}", clp_s_bin.display())) + .inspect_err(|e| { + tracing::error!( + error = % e, + clp_s_bin = % clp_s_bin.display(), + "Failed to spawn clp-s.", + ); + })?; + + let mut stderr = child + .stderr + .take() + .expect("piped stderr should always be present"); + + let mut captured_stderr = String::new(); + if let Err(e) = stderr.read_to_string(&mut captured_stderr) { + captured_stderr = format!("failed to read clp-s stderr: {e}"); + } + + let status = child + .wait() + .context("failed to wait for clp-s to exit") + .inspect_err(|e| { + tracing::error!( + error = % e, + stderr = % captured_stderr, + "Failed to wait for clp-s to exit." + ); + })?; + if !status.success() { + tracing::error!( + status = status.code(), + stderr = % captured_stderr, + "clp-s exited on failure." + ); + anyhow::bail!("clp-s exited on error with status={status}"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::ffi::OsString; + use std::num::NonZeroU32; + use std::path::Path; + 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; + use clp_rust_utils::clp_config::package::config::Package; + use clp_rust_utils::clp_config::package::config::SpiderTaskExecutorConfig; + use clp_rust_utils::clp_config::package::config::StorageEngine; + use clp_rust_utils::task_io::query::ClpSQueryOption; + use clp_rust_utils::task_io::query::OutputHandle; + use clp_rust_utils::types::non_empty_string::ExpectedNonEmpty; + use non_empty_string::NonEmptyString; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::id::TaskId; + use spider_tdl::TaskContext; + + use super::ArchiveSelector; + use super::build_clp_s_search_args_for_result_cache; + use super::resolve_archive_input; + use super::search; + + /// # Returns + /// + /// A query option with no timestamp bounds, no result limit, and case-sensitive matching. + fn unbounded_query_option() -> ClpSQueryOption { + ClpSQueryOption { + query_string: NonEmptyString::from_static_str("level: \"ERROR\""), + max_num_results: None, + begin_timestamp_millisecs: None, + end_timestamp_millisecs: None, + ignore_case: false, + } + } + + /// # Returns + /// + /// The AWS credentials the S3-backed test configs authenticate with. + fn test_aws_authentication() -> AwsAuthentication { + AwsAuthentication::Credentials { + credentials: AwsCredentials { + access_key_id: "the-access-key".to_string(), + secret_access_key: "the-secret-key".to_string(), + session_token: None, + }, + } + } + + /// # Returns + /// + /// An [`ArchiveSelector`] addressing `archive-id` under `/archives/ds1`. + fn directory_selector() -> ArchiveSelector { + ArchiveSelector::Directory { + path: PathBuf::from("/archives/ds1"), + archive_id: "archive-id".to_string(), + } + } + + /// # Returns + /// + /// A [`SpiderTaskExecutorConfig`] whose archive output is S3-backed with the given staging + /// directory and endpoint URL. + /// + /// # Panics + /// + /// Panics if any of the static S3 config strings is empty. + fn s3_backed_config( + staging_directory: &str, + endpoint_url: Option<&'static str>, + ) -> SpiderTaskExecutorConfig { + SpiderTaskExecutorConfig { + package: Package { + storage_engine: StorageEngine::ClpS, + }, + archive_output: ArchiveOutput { + storage: ArchiveOutputStorage::S3 { + staging_directory: staging_directory.to_owned(), + s3_config: S3Config { + bucket: NonEmptyString::from_static_str("bucket"), + region_code: None, + key_prefix: NonEmptyString::from_static_str("LIB1/"), + endpoint_url: endpoint_url.map(NonEmptyString::from_static_str), + aws_authentication: test_aws_authentication(), + }, + }, + ..ArchiveOutput::default() + }, + ..SpiderTaskExecutorConfig::default() + } + } + + /// # Returns + /// + /// A [`TaskContext`] for a non-commit task. + /// + /// # Panics + /// + /// Panics if [`TaskContext::new`] returns an error. + fn task_context() -> TaskContext { + TaskContext::new( + JobId::random(), + TaskId::Index(0), + 1, + ResourceGroupId::random(), + None, + ) + .expect("a non-commit task context without graph outputs is valid") + } + + #[test] + fn build_clp_s_search_args_for_result_cache_fs_with_timestamps_and_ignore_case() { + let clp_s_query_option = ClpSQueryOption { + query_string: NonEmptyString::from_static_str("level: \"ERROR\""), + max_num_results: Some(NonZeroU32::new(7).expect("7 is nonzero")), + begin_timestamp_millisecs: Some(1_310_138_944_000), + end_timestamp_millisecs: Some(1_311_208_074_120), + ignore_case: true, + }; + + assert_eq!( + build_clp_s_search_args_for_result_cache( + &directory_selector(), + &clp_s_query_option, + "mongodb://results-cache:27017/clp-query-results", + 42, + "ds1", + ), + vec![ + OsString::from("s"), + OsString::from("/archives/ds1"), + OsString::from("--archive-id"), + OsString::from("archive-id"), + OsString::from("level: \"ERROR\""), + OsString::from("--tge"), + OsString::from("1310138944000"), + OsString::from("--tle"), + OsString::from("1311208074120"), + OsString::from("--ignore-case"), + OsString::from("results-cache"), + OsString::from("--uri"), + OsString::from("mongodb://results-cache:27017/clp-query-results"), + OsString::from("--collection"), + OsString::from("42"), + OsString::from("--max-num-results"), + OsString::from("7"), + OsString::from("--dataset"), + OsString::from("ds1"), + ] + ); + } + + #[test] + fn build_clp_s_search_args_for_result_cache_omits_max_num_results_when_unset() { + let clp_s_query_option = ClpSQueryOption { + query_string: NonEmptyString::from_static_str("level: \"ERROR\""), + max_num_results: None, + begin_timestamp_millisecs: Some(1_310_138_944_000), + end_timestamp_millisecs: Some(1_311_208_074_120), + ignore_case: true, + }; + + assert_eq!( + build_clp_s_search_args_for_result_cache( + &directory_selector(), + &clp_s_query_option, + "mongodb://results-cache:27017/clp-query-results", + 42, + "ds1", + ), + vec![ + OsString::from("s"), + OsString::from("/archives/ds1"), + OsString::from("--archive-id"), + OsString::from("archive-id"), + OsString::from("level: \"ERROR\""), + OsString::from("--tge"), + OsString::from("1310138944000"), + OsString::from("--tle"), + OsString::from("1311208074120"), + OsString::from("--ignore-case"), + OsString::from("results-cache"), + OsString::from("--uri"), + OsString::from("mongodb://results-cache:27017/clp-query-results"), + OsString::from("--collection"), + OsString::from("42"), + OsString::from("--dataset"), + OsString::from("ds1"), + ] + ); + } + + #[test] + fn build_clp_s_search_args_for_result_cache_fs_without_timestamps_or_ignore_case() { + assert_eq!( + build_clp_s_search_args_for_result_cache( + &directory_selector(), + &unbounded_query_option(), + "mongodb://results-cache:27017/clp-query-results", + 42, + "default", + ), + vec![ + OsString::from("s"), + OsString::from("/archives/ds1"), + OsString::from("--archive-id"), + OsString::from("archive-id"), + OsString::from("level: \"ERROR\""), + OsString::from("results-cache"), + OsString::from("--uri"), + OsString::from("mongodb://results-cache:27017/clp-query-results"), + OsString::from("--collection"), + OsString::from("42"), + OsString::from("--dataset"), + OsString::from("default"), + ] + ); + } + + #[test] + fn build_clp_s_search_args_for_result_cache_s3_uses_object_url_and_no_archive_id() { + assert_eq!( + build_clp_s_search_args_for_result_cache( + &ArchiveSelector::ObjectUrl( + "https://bucket.s3.amazonaws.com/LIB1/ds1/archive-id".to_string() + ), + &unbounded_query_option(), + "mongodb://results-cache:27017/clp-query-results", + 42, + "ds1", + ), + vec![ + OsString::from("s"), + OsString::from("https://bucket.s3.amazonaws.com/LIB1/ds1/archive-id"), + OsString::from("--auth"), + OsString::from("s3"), + OsString::from("level: \"ERROR\""), + OsString::from("results-cache"), + OsString::from("--uri"), + OsString::from("mongodb://results-cache:27017/clp-query-results"), + OsString::from("--collection"), + OsString::from("42"), + OsString::from("--dataset"), + OsString::from("ds1"), + ] + ); + } + + #[test] + fn resolve_archive_input_fs_joins_dataset_and_returns_no_credentials() -> anyhow::Result<()> { + let runtime = tokio::runtime::Runtime::new().expect("failed to create Tokio runtime"); + let config = SpiderTaskExecutorConfig { + package: Package { + storage_engine: StorageEngine::ClpS, + }, + archive_output: ArchiveOutput { + storage: ArchiveOutputStorage::Fs { + directory: "var/data/archives".to_owned(), + }, + ..ArchiveOutput::default() + }, + ..SpiderTaskExecutorConfig::default() + }; + + let (selector, credential_env) = resolve_archive_input( + runtime.handle(), + Path::new("/clp"), + &config, + "ds1", + "archive-id".to_string(), + )?; + + let ArchiveSelector::Directory { path, archive_id } = selector else { + panic!("expected a directory selector"); + }; + assert_eq!(path, PathBuf::from("/clp/var/data/archives/ds1")); + assert_eq!(archive_id, "archive-id"); + assert_eq!(credential_env, &[]); + + Ok(()) + } + + #[test] + fn resolve_archive_input_s3_builds_object_url_and_credentials() -> anyhow::Result<()> { + let runtime = tokio::runtime::Runtime::new().expect("failed to create Tokio runtime"); + let config = s3_backed_config("var/data/staged-archives", None); + + let (selector, credential_env) = resolve_archive_input( + runtime.handle(), + Path::new("/clp"), + &config, + "ds1", + "archive-id".to_string(), + )?; + + let ArchiveSelector::ObjectUrl(url) = selector else { + panic!("expected an object-URL selector"); + }; + assert_eq!(url, "https://bucket.s3.amazonaws.com/LIB1/ds1/archive-id"); + assert_eq!( + credential_env, + vec![ + ("AWS_ACCESS_KEY_ID", "the-access-key".to_string()), + ("AWS_SECRET_ACCESS_KEY", "the-secret-key".to_string()), + ] + ); + + Ok(()) + } + + #[test] + fn resolve_archive_input_s3_ignores_staging_directory() -> anyhow::Result<()> { + const STAGING_DIRECTORY: &str = "/wrong-staging-directory"; + + let runtime = tokio::runtime::Runtime::new().expect("failed to create Tokio runtime"); + let config = s3_backed_config(STAGING_DIRECTORY, None); + + let (selector, _) = resolve_archive_input( + runtime.handle(), + Path::new("/clp"), + &config, + "ds1", + "archive-id".to_string(), + )?; + let args = build_clp_s_search_args_for_result_cache( + &selector, + &unbounded_query_option(), + "mongodb://results-cache:27017/clp-query-results", + 42, + "ds1", + ); + + assert!( + !args + .iter() + .any(|arg| arg.to_string_lossy().contains(STAGING_DIRECTORY)) + ); + + Ok(()) + } + + #[test] + fn resolve_archive_input_s3_custom_endpoint_uses_path_style_url() -> anyhow::Result<()> { + let runtime = tokio::runtime::Runtime::new().expect("failed to create Tokio runtime"); + let config = s3_backed_config("var/data/staged-archives", Some("http://minio:9000")); + + let (selector, _) = resolve_archive_input( + runtime.handle(), + Path::new("/clp"), + &config, + "ds1", + "archive-id".to_string(), + )?; + + let ArchiveSelector::ObjectUrl(url) = selector else { + panic!("expected an object-URL selector"); + }; + assert_eq!(url, "http://minio:9000/bucket/LIB1/ds1/archive-id"); + + Ok(()) + } + + #[test] + fn search_rejects_file_output_handle() { + let config = SpiderTaskExecutorConfig { + package: Package { + storage_engine: StorageEngine::ClpS, + }, + ..SpiderTaskExecutorConfig::default() + }; + + let error = search( + &task_context(), + &config, + 42, + &unbounded_query_option(), + "archive-id".to_string(), + None, + &OutputHandle::File, + ) + .expect_err("the file output handler is unsupported"); + + assert!(error.to_string().contains("unsupported")); + } + + #[test] + fn search_rejects_non_clp_s_storage_engine() { + let config = SpiderTaskExecutorConfig::default(); + assert_eq!(config.package.storage_engine, StorageEngine::Clp); + + let error = search( + &task_context(), + &config, + 42, + &unbounded_query_option(), + "archive-id".to_string(), + None, + &OutputHandle::ResultsCache { + uri: NonEmptyString::from_static_str( + "mongodb://results-cache:27017/clp-query-results", + ), + }, + ) + .expect_err("the clp storage engine is unsupported"); + + assert!(error.to_string().contains("clp-s")); + } +} 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..425c048676 --- /dev/null +++ b/components/query-coordinator/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "query-coordinator" +version = { workspace = true } +edition = { workspace = true } + +[dependencies] +async-trait = { workspace = true } +clp-rust-utils = { workspace = true } +const_format = { workspace = true } +mongodb = { workspace = true } +non-empty-string = { workspace = true } +rmp-serde = { workspace = true } +serde = { workspace = true } +spider-client = { workspace = true } +spider-core = { workspace = true } +sqlx = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true } +tonic = { workspace = true } +tracing = { workspace = true } diff --git a/components/query-coordinator/src/coordination.rs b/components/query-coordinator/src/coordination.rs new file mode 100644 index 0000000000..3127445ca5 --- /dev/null +++ b/components/query-coordinator/src/coordination.rs @@ -0,0 +1,671 @@ +//! The coordinator poll loop that discovers pending CLP query jobs and dispatches them to +//! Spider. +//! +//! The coordinator is responsible for the query jobs in the `query_jobs` table that +//! are in one of the following states: +//! +//! | `status` | `spider_id` | `dispatch_time` | Description | +//! |----------|-------------|-----------------|--------------------------------------------------| +//! | PENDING | NULL | NULL | New jobs awaiting dispatch. | +//! | PENDING | NULL | NOT NULL | Jobs dispatched but not yet submitted to Spider. | +//! | RUNNING | NOT NULL | NOT NULL | Jobs submitted to Spider. | +//! +//! NOTE: +//! +//! * These are the only legal states for a job that hasn't terminated. +//! * A non-NULL `dispatch_time` indicates that the coordinator has picked up the job and granted it +//! permission to run under the concurrency limit. + +use std::sync::Arc; +use std::time::Duration; + +use clp_rust_utils::clp_config::package::config::Database as DatabaseConfig; +use clp_rust_utils::clp_config::package::config::QueryCoordinator as CoordinatorConfig; +use clp_rust_utils::clp_config::package::config::ResultsCache as ResultsCacheConfig; +use clp_rust_utils::clp_config::package::config::Spider as SpiderConfig; +use clp_rust_utils::clp_config::package::config::SpiderResourceGroup; +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::job_config::QueryJobType; +use clp_rust_utils::job_config::SearchJobConfig; +use clp_rust_utils::task_io::query::OutputHandle; +use const_format::formatcp; +use mongodb::options::ClientOptions; +use spider_client::SpiderClient; +use spider_core::types::id::JobId as SpiderJobId; +use spider_core::types::id::ResourceGroupId; +use tokio::select; +use tokio::sync::Semaphore; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; +use tonic::transport::Endpoint; + +use crate::Error; +use crate::job_handle::QueryJobHandle; +use crate::job_handle::QueryJobHandleContext; +use crate::job_handle::SpiderOption; + +/// Coordinator for fetching new query jobs and submitting them to Spider. +pub struct Coordinator { + resource_group_id: ResourceGroupId, + spider_client: SpiderClient, + db_pool: sqlx::MySqlPool, + job_handle_context: Arc, + is_first_fetch: bool, + job_polling_interval: Duration, + cancellation_token: CancellationToken, + job_handler_sem: Arc, +} + +impl Coordinator { + /// Factory function. + /// + /// On construction, this recovers query jobs that a previous coordinator instance had + /// already submitted to Spider (those still [`QueryJobStatus::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::InvalidConfiguration`] if the query coordinator configuration is invalid. + /// * [`Error::InvalidEndpoint`] if the Spider host and port do not form a valid endpoint. + /// * Forwards [`create_results_cache_database`]'s return values on failure. + /// * 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, + results_cache_config: &ResultsCacheConfig, + ) -> Result<(Self, CancellationToken), Error> { + let max_concurrent_jobs = coordinator_config.max_concurrent_jobs.get(); + if max_concurrent_jobs > Semaphore::MAX_PERMITS { + return Err(Error::InvalidConfiguration(format!( + "`max_concurrent_jobs` must not exceed {}, got {max_concurrent_jobs}", + Semaphore::MAX_PERMITS, + ))); + } + + let results_cache_uri = results_cache_config.uri(); + let results_cache = create_results_cache_database( + results_cache_uri.as_str(), + &results_cache_config.db_name, + ) + .await + .inspect_err(|e| { + tracing::error!(error = % e, "Failed to create the results cache client."); + })?; + + 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 job_handle_context = Arc::new(QueryJobHandleContext { + db_pool: db_pool.clone(), + db_config, + results_cache, + output_handle: OutputHandle::ResultsCache { + uri: results_cache_uri, + }, + spider_option: SpiderOption { + // Use the initial polling delay as the fixed interval for query jobs. + poll_interval: Duration::from_millis( + coordinator_config + .result_polling + .init_backoff_millisecs + .get(), + ), + }, + }); + + let cancellation_token = CancellationToken::new(); + + let coordinator = Self { + resource_group_id, + spider_client, + db_pool, + job_handle_context, + is_first_fetch: true, + job_polling_interval: Duration::from_millis( + coordinator_config.job_polling_interval_millisecs.get(), + ), + cancellation_token: cancellation_token.clone(), + job_handler_sem: Arc::new(Semaphore::new(max_concurrent_jobs)), + }; + + // NOTE: The current implementation does not enforce concurrency limits for recovered jobs + // since they were already submitted to Spider. See #2472. + for (job_id, spider_job_id, search_job_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, search_job_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 query job failed." + ); + }); + }); + } + + Ok((coordinator, cancellation_token)) + } + + /// Runs the coordinator's poll loop until cancelled. + /// + /// On each iteration, this method fetches the pending query 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 query job identified by `job_id` as [`QueryJobStatus::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: QueryJobId, status_msg: &str) { + const QUERY: &str = formatcp!( + "UPDATE `{table}` SET `status` = ?, `status_msg` = LEFT(?, 512) WHERE `id` = ?;", + table = QUERY_JOBS_TABLE_NAME, + ); + tracing::info!(job_id = % job_id, "Failing the query job."); + if let Err(e) = sqlx::query(QUERY) + .bind(QueryJobStatus::Failed) + .bind(status_msg) + .bind(job_id) + .execute(&self.db_pool) + .await + { + tracing::error!( + error = % e, + job_id = % job_id, + "Failed to mark the query job as failed." + ); + } + } + + /// Fetches pending query jobs and spawns a detached handle to drive each one as permitted + /// by the job-handler semaphore. + /// + /// A job whose config cannot be deserialized is marked [`QueryJobStatus::Failed`] and + /// skipped; a job whose handle cannot be constructed is marked failed and skipped as well. + /// Aggregation jobs are left undispatched for another scheduler. + /// + /// # Returns + /// + /// The IDs of the fetched jobs that were dispatched in this poll. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`Error::Semaphore`] if acquiring a job handler permit from `job_handler_sem` fails. + /// * Forwards [`Self::fetch_new_job_rows`]'s return values on failure. + async fn schedule_new_jobs(&mut self) -> Result, Error> { + if self.job_handler_sem.available_permits() == 0 { + return Ok(Vec::new()); + } + + 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 mut dispatched_job_ids = Vec::new(); + for job_row in new_job_rows { + let job_id = job_row.id; + let search_job_config: SearchJobConfig = + match rmp_serde::from_slice(&job_row.serialized_search_job_config) { + Ok(search_job_config) => search_job_config, + Err(e) => { + tracing::error!( + error = % e, + job_id = % job_id, + "Failed to deserialize search job config. Skipping." + ); + self.mark_job_failed( + job_id, + &format!("Failed to deserialize search job config: {e}"), + ) + .await; + continue; + } + }; + if search_job_config.aggregation_config.is_some() { + continue; + } + tracing::info!(job_id = % job_id, "Scheduling new job."); + let Ok(job_handle) = self.create_job_handle(job_id, search_job_config).await else { + continue; + }; + + let permit = self + .job_handler_sem + .clone() + .acquire_owned() + .await + .map_err(|e| { + Error::Semaphore(format!("failed to acquire a job handler permit: {e}")) + })?; + + dispatched_job_ids.push(job_id); + tokio::spawn(async move { + let _permit = permit; + let _ = job_handle.run().await.inspect_err(|e| { + tracing::error!( + error = % e, + job_id = % job_id, + "Failed to schedule query job." + ); + }); + }); + } + Ok(dispatched_job_ids) + } + + /// Marks the query 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: &[QueryJobId]) -> 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` = COALESCE(`dispatch_time`, \ + CURRENT_TIMESTAMP()) WHERE `id` IN (", + table = QUERY_JOBS_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 a [`QueryJobHandle`] for the given job. + /// + /// A construction failure is logged, and the job is marked [`QueryJobStatus::Failed`]. + /// + /// # Returns + /// + /// The constructed [`QueryJobHandle`] on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`QueryJobHandle::new`]'s return values on failure. + async fn create_job_handle( + &self, + job_id: QueryJobId, + search_job_config: SearchJobConfig, + ) -> Result, Error> { + let result = QueryJobHandle::new( + self.job_handle_context.clone(), + job_id, + self.spider_client.clone(), + self.resource_group_id, + search_job_config, + ); + + if let Err(e) = &result { + tracing::error!( + error = % e, + job_id = % job_id, + "Failed to create query job handle. Skipping." + ); + self.mark_job_failed( + job_id, + &format!("Failed to create the query job handle: {e}"), + ) + .await; + } + + result + } + + /// Fetches pending query jobs eligible for dispatch. + /// + /// The first fetch after startup returns every [`QueryJobStatus::Pending`] job whose + /// `dispatch_time` is set, so that jobs dispatched but not started by the previous coordinator + /// instance can be re-dispatched. No explicit limit is imposed because: + /// + /// * This query runs only once, so limiting it could leave previously dispatched jobs + /// unfetched. + /// * The recovery set is bounded by the previous coordinator's concurrency limit. + /// + /// Every subsequent fetch returns only [`QueryJobStatus::Pending`] jobs whose dispatch + /// time is not set. Pages are scanned in ID order past unsupported aggregation jobs until + /// enough eligible rows fill the available permits or no rows remain. + /// + /// # Returns + /// + /// A vector of rows projected from the query job table on success, each row represents a + /// pending query 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`, `job_config` FROM `{table}` WHERE `type` = ? AND `status` = ? AND \ + `dispatch_time` IS NOT NULL ORDER BY `id` ASC;", + table = QUERY_JOBS_TABLE_NAME, + ); + const SUBSEQUENT_FETCH_QUERY: &str = formatcp!( + "SELECT `id`, `job_config` FROM `{table}` WHERE `type` = ? AND `status` = ? AND \ + `dispatch_time` IS NULL AND `id` > ? ORDER BY `id` ASC LIMIT ?;", + table = QUERY_JOBS_TABLE_NAME, + ); + + if self.is_first_fetch { + self.is_first_fetch = false; + return sqlx::query_as::<_, PendingJobRowProjection>(FIRST_FETCH_QUERY) + .bind(i32::from(QueryJobType::SearchOrAggregation)) + .bind(QueryJobStatus::Pending) + .fetch_all(&self.db_pool) + .await + .map_err(Into::into); + } + + let limit = self.job_handler_sem.available_permits(); + let mut rows = Vec::new(); + let mut last_id = i32::MIN; + while rows.len() < limit { + let remaining = limit - rows.len(); + let batch = sqlx::query_as::<_, PendingJobRowProjection>(SUBSEQUENT_FETCH_QUERY) + .bind(i32::from(QueryJobType::SearchOrAggregation)) + .bind(QueryJobStatus::Pending) + .bind(last_id) + .bind( + i64::try_from(remaining) + .expect("limit is bounded by Semaphore::MAX_PERMITS, which fits in i64"), + ) + .fetch_all(&self.db_pool) + .await?; + let exhausted = batch.len() < remaining; + for row in batch { + last_id = row.id; + // Keep malformed rows so scheduling can report their configuration errors. + if rmp_serde::from_slice::(&row.serialized_search_job_config) + .is_ok_and(|config| config.aggregation_config.is_some()) + { + continue; + } + rows.push(row); + } + if exhausted { + break; + } + } + + Ok(rows) + } + + /// Fetches jobs that are still in [`QueryJobStatus::Running`] and were previously + /// submitted by the query coordinator. + /// + /// A running job whose config cannot be deserialized is marked [`QueryJobStatus::Failed`] + /// and skipped. + /// + /// # Returns + /// + /// A vector of tuples on success, each tuple containing: + /// + /// * The query job ID. + /// * The Spider job ID. + /// * The search config of the query 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`, `job_config` FROM `{table}` WHERE `type` = ? AND `status` \ + = ? AND `spider_id` IS NOT NULL;", + table = QUERY_JOBS_TABLE_NAME, + ); + + let mut recovery_context = Vec::new(); + for row in sqlx::query_as::<_, RunningJobRowProjection>(QUERY) + .bind(i32::from(QueryJobType::SearchOrAggregation)) + .bind(QueryJobStatus::Running) + .fetch_all(&self.db_pool) + .await? + { + let search_job_config: SearchJobConfig = match rmp_serde::from_slice( + &row.serialized_search_job_config, + ) { + Ok(search_job_config) => search_job_config, + Err(e) => { + tracing::error!( + error = % e, + job_id = % row.id, + "Failed to deserialize search job config of a running job. The database \ + might be corrupted. Skipping." + ); + self.mark_job_failed( + row.id, + &format!("Failed to deserialize search job config: {e}"), + ) + .await; + continue; + } + }; + if search_job_config.aggregation_config.is_some() { + continue; + } + recovery_context.push((row.id, row.spider_job_id, search_job_config)); + } + + Ok(recovery_context) + } +} + +/// A projection of the columns read from a [`QueryJobStatus::Pending`] query job row. +#[derive(Debug, sqlx::FromRow)] +struct PendingJobRowProjection { + id: QueryJobId, + #[sqlx(rename = "job_config")] + serialized_search_job_config: Vec, +} + +/// A projection of the columns read from a [`QueryJobStatus::Running`] query job row. +#[derive(Debug, sqlx::FromRow)] +struct RunningJobRowProjection { + id: QueryJobId, + #[sqlx(rename = "spider_id")] + spider_job_id: SpiderJobId, + #[sqlx(rename = "job_config")] + serialized_search_job_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) +} + +/// Creates a client for the results cache database at `uri`. +/// +/// # Errors +/// +/// Forwards errors from MongoDB client-option parsing and client construction. +async fn create_results_cache_database( + uri: &str, + db_name: &str, +) -> Result { + let mut client_options = ClientOptions::parse(uri).await?; + client_options.direct_connection = Some(true); + Ok(mongodb::Client::with_options(client_options)?.database(db_name)) +} diff --git a/components/query-coordinator/src/error.rs b/components/query-coordinator/src/error.rs new file mode 100644 index 0000000000..84455b706a --- /dev/null +++ b/components/query-coordinator/src/error.rs @@ -0,0 +1,38 @@ +//! The crate-level error type for the query coordinator. + +/// Errors returned by the query coordinator. +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid coordinator configuration: {0}")] + InvalidConfiguration(String), + + #[error("invalid Spider endpoint: {0}")] + InvalidEndpoint(String), + + #[error("semaphore error: {0}")] + Semaphore(String), + + #[error("invalid query job configuration: {0}")] + InvalidQueryJobConfig(String), + + #[error("failed to update SQL database: {0}")] + SqlxNoRowsAffected(String), + + #[error("mongodb error: {0}")] + Mongo(#[from] mongodb::error::Error), + + #[error("spider request failure: {0}")] + SpiderClient(#[from] spider_client::error::ClientError), + + #[error("sqlx error: {0}")] + Sqlx(#[from] sqlx::Error), + + #[error("number of query tasks {0} exceeds `i32::MAX`")] + TooManyQueryTasks(usize), + + #[error("failed to build the query task graph: {0}")] + TaskGraph(#[from] spider_core::task::Error), + + #[error("failed to serialize a task input: {0}")] + TaskInputSerialization(#[from] rmp_serde::encode::Error), +} diff --git a/components/query-coordinator/src/job_handle.rs b/components/query-coordinator/src/job_handle.rs new file mode 100644 index 0000000000..4d122041c9 --- /dev/null +++ b/components/query-coordinator/src/job_handle.rs @@ -0,0 +1,366 @@ +//! Lifecycle management for one coordinator-planned query job. + +use std::num::NonZeroU32; +use std::sync::Arc; +use std::time::Duration; + +use clp_rust_utils::clp_config::package::config::Database; +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::job_config::SearchJobConfig; +use clp_rust_utils::task_io::query::ClpSQueryOption; +use clp_rust_utils::task_io::query::OutputHandle; +use const_format::formatcp; +use non_empty_string::NonEmptyString; +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; + +/// Options for a query job running in Spider. +pub struct SpiderOption { + pub poll_interval: Duration, +} + +/// Resources shared by query job handles created by the coordinator. +pub struct QueryJobHandleContext { + pub db_pool: MySqlPool, + pub db_config: Database, + pub results_cache: mongodb::Database, + pub output_handle: OutputHandle, + pub spider_option: SpiderOption, +} + +/// Handles the asynchronous submission of a query job and the retrieval of its result. +/// +/// # Type Parameters +/// +/// * `SubmitterType` - The type of the job submitter for Spider job submission. +pub struct QueryJobHandle { + context: Arc, + query_job_id: QueryJobId, + job_submitter: SubmitterType, + resource_group_id: ResourceGroupId, + _search_job_config: SearchJobConfig, + clp_s_query_option: ClpSQueryOption, +} + +impl QueryJobHandle { + /// Factory function. + /// + /// # Returns + /// + /// A newly created [`QueryJobHandle`] for the given query job configuration. + /// + /// # Errors + /// + /// Returns an error if the query string is empty. + pub fn new( + context: Arc, + query_job_id: QueryJobId, + job_submitter: SubmitterType, + resource_group_id: ResourceGroupId, + search_job_config: SearchJobConfig, + ) -> Result { + let query_string = NonEmptyString::try_from(search_job_config.query_string.clone()) + .map_err(|_| { + Error::InvalidQueryJobConfig("query string must not be empty".to_owned()) + })?; + let clp_s_query_option = ClpSQueryOption { + query_string, + max_num_results: NonZeroU32::new(search_job_config.max_num_results), + begin_timestamp_millisecs: search_job_config.begin_timestamp, + end_timestamp_millisecs: search_job_config.end_timestamp, + ignore_case: search_job_config.ignore_case, + }; + + Ok(Self { + context, + query_job_id, + job_submitter, + resource_group_id, + _search_job_config: search_job_config, + clp_s_query_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. + /// + /// If no matching row is found when persisting the Spider ID, the job may have been cancelled, + /// deleted, or claimed by another coordinator job handler. Anyhow, this handle no longer owns + /// it, so it skips trying to report a job failure. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::plan_and_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."); + + match self.plan_and_submit().await { + Ok(Some(spider_job_id)) => self.to_completion(spider_job_id).await, + Ok(None) => Ok(()), + Err(error) => { + if !matches!(error, Error::SqlxNoRowsAffected(_)) { + self.report_failure(&error).await; + } + Err(error) + } + } + } + + /// Plans the query inputs and submits the query job, or marks it as succeeded if no archives + /// are selected. + /// + /// # Returns + /// + /// On success, the submitted Spider job ID, or `None` if no archives are selected. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::prepare_task_inputs`]'s return values on failure. + /// * Forwards [`Self::update_job_status`]'s return values on failure when no archives are + /// selected. + /// * Forwards [`Self::submit`]'s return values on failure. + async fn plan_and_submit(&self) -> Result, Error> { + let archives_to_search = self.prepare_task_inputs().await?; + if archives_to_search.is_empty() { + if !self + .update_job_status(QueryJobStatus::Pending, QueryJobStatus::Succeeded, None) + .await? + { + return Err(Error::SqlxNoRowsAffected(format!( + "no pending query job row found for query job {}", + self.query_job_id + ))); + } + return Ok(None); + } + self.submit(archives_to_search).await.map(Some) + } + + /// 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 [`QueryJobSubmitter::submit_query_job`]'s return values on failure. + /// * Forwards [`Self::persist_spider_job_id`]'s return values on failure. + async fn submit( + &self, + archives_to_search: Vec<(ArchiveMetadata, ExecutionPolicy)>, + ) -> Result { + let num_tasks = archives_to_search.len(); + let persisted_num_tasks = + i32::try_from(num_tasks).map_err(|_| Error::TooManyQueryTasks(num_tasks))?; + let spider_job_id = self + .job_submitter + .submit_query_job( + self.query_job_id, + self.resource_group_id, + self.clp_s_query_option.clone(), + self.context.output_handle.clone(), + archives_to_search, + ) + .await?; + + tracing::info!( + query_job_id = % self.query_job_id, + spider_job_id = % spider_job_id, + num_tasks, + "Query job submitted.", + ); + + self.persist_spider_job_id(spider_job_id, persisted_num_tasks) + .await?; + Ok(spider_job_id) + } + + /// Prepares the archive inputs and execution policies for the query tasks. + /// + /// # Returns + /// + /// The archives to search and their execution policies on success. + /// + /// # Errors + /// + /// Returns an error if archive input preparation fails. + async fn prepare_task_inputs(&self) -> Result, Error> { + todo!("prepare query task inputs") + } + + /// Persists the Spider job ID and marks the query job as running. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`Error::SqlxNoRowsAffected`] if no pending query job row was updated. + /// * Forwards [`sqlx::query::Query::execute`]'s return values on failure. + async fn persist_spider_job_id( + &self, + spider_job_id: SpiderJobId, + num_tasks: i32, + ) -> Result<(), Error> { + let query = formatcp!( + "UPDATE `{QUERY_JOBS_TABLE_NAME}` SET `spider_id` = ?, `status` = ?, `num_tasks` = ?, \ + `start_time` = CURRENT_TIMESTAMP(3) WHERE `id` = ? AND `status` = ?" + ); + let query = sqlx::query(query) + .bind(spider_job_id.get()) + .bind(QueryJobStatus::Running) + .bind(num_tasks) + .bind(self.query_job_id) + .bind(QueryJobStatus::Pending); + if !execute_update(query, &self.context.db_pool).await? { + return Err(Error::SqlxNoRowsAffected(format!( + "no pending query job row found for query job {} (Spider job ID {})", + self.query_job_id, spider_job_id + ))); + } + Ok(()) + } + + /// Starts the Spider job if needed, waits for completion, and finalizes the query job. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::update_job_status`]'s return values on failure. + /// * 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.context.spider_option.poll_interval) + .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, None), + QueryJobOutcome::Failed { error_message } => ( + QueryJobStatus::Failed, + Some(format!("The Spider query job failed: {error_message}")), + ), + QueryJobOutcome::Cancelled => ( + QueryJobStatus::Cancelled, + Some("The Spider query job was cancelled.".to_owned()), + ), + }; + if !self + .update_job_status(QueryJobStatus::Running, status, status_message.as_deref()) + .await? + { + return Err(Error::SqlxNoRowsAffected(format!( + "no running query job row found for query job {}", + self.query_job_id + ))); + } + Ok(()) + } + + /// 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.", + ); + + let _ = self + .update_job_status( + QueryJobStatus::Pending, + QueryJobStatus::Failed, + Some(&format!("Query job orchestration failed: {error}")), + ) + .await + .inspect_err(|status_error| { + tracing::error!( + query_job_id = % self.query_job_id, + error = % status_error, + "Failed to persist the query job failure.", + ); + }); + } + + /// Updates the query job status in the CLP database. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`sqlx::query::Query::execute`]'s return values on failure. + async fn update_job_status( + &self, + from: QueryJobStatus, + to: QueryJobStatus, + msg: Option<&str>, + ) -> Result { + let query = sqlx::query(formatcp!( + "UPDATE `{QUERY_JOBS_TABLE_NAME}` SET `status` = ?, `status_msg` = ? WHERE `id` = ? \ + AND `status` = ?" + )) + .bind(to) + .bind(msg.unwrap_or_default()) + .bind(self.query_job_id) + .bind(from); + execute_update(query, &self.context.db_pool).await + } +} + +/// Executes an SQL update and reports whether any row was affected. +async fn execute_update( + query: sqlx::query::Query<'_, sqlx::MySql, sqlx::mysql::MySqlArguments>, + db_pool: &MySqlPool, +) -> Result { + let result = query.execute(db_pool).await?; + Ok(result.rows_affected() > 0) +} diff --git a/components/query-coordinator/src/lib.rs b/components/query-coordinator/src/lib.rs new file mode 100644 index 0000000000..e7ef8ef995 --- /dev/null +++ b/components/query-coordinator/src/lib.rs @@ -0,0 +1,9 @@ +//! Coordination for CLP query jobs. + +pub mod coordination; +mod error; +pub mod job_handle; +pub mod plan; +pub mod query_job_submitter; + +pub use error::Error; diff --git a/components/query-coordinator/src/plan.rs b/components/query-coordinator/src/plan.rs new file mode 100644 index 0000000000..72bacab1e8 --- /dev/null +++ b/components/query-coordinator/src/plan.rs @@ -0,0 +1,4 @@ +//! Planning options for query jobs. + +/// Placeholder for query-job planning options. +pub struct PlanningOption {} 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..db4a6931b8 --- /dev/null +++ b/components/query-coordinator/src/query_job_submitter/mod.rs @@ -0,0 +1,97 @@ +//! 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 serde::Deserialize; +use serde::Serialize; +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(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub enum QueryJobOutcome { + /// The job completed successfully. + Succeeded, + + /// The job failed with the given error. + Failed { error_message: String }, + + /// The job was cancelled before reaching completion. + Cancelled, +} + +/// 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. + /// * `poll_interval` - The delay after each non-terminal job-state poll. + /// + /// # 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, + poll_interval: 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..1489410799 --- /dev/null +++ b/components/query-coordinator/src/query_job_submitter/spider.rs @@ -0,0 +1,175 @@ +//! [`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::DataTypeDescriptor; +use spider_core::task::ExecutionPolicy; +use spider_core::task::TaskDescriptor; +use spider_core::task::TaskGraph; +use spider_core::task::TdlContext; +use spider_core::task::ValueTypeDescriptor; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::io::TaskInput; + +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 { + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`build_query_task_graph`]'s return values on failure. + /// * Forwards [`SpiderClient::submit_job`]'s return values on failure. + 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 { + let (graph, inputs) = build_query_task_graph( + query_job_id, + &clp_s_query_option, + &output_handle, + archives_to_search, + )?; + let spider_job_id = self.submit_job(resource_group_id, &graph, inputs).await?; + + tracing::info!( + query_job_id = % query_job_id, + spider_job_id = % spider_job_id, + num_tasks = graph.get_num_tasks(), + "Submitted query job to Spider.", + ); + + Ok(spider_job_id) + } + + /// # 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, + poll_interval: Duration, + ) -> Result { + match self.start_job(spider_job_id).await { + Ok(_) | Err(ClientError::InvalidJobState(_)) => {} + Err(error) => return Err(error.into()), + } + + let terminal_state = loop { + let state = self.get_job_state(spider_job_id).await?; + if state.is_terminal() { + break state; + } + tokio::time::sleep(poll_interval).await; + }; + + 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::Cancelled, + _ => unreachable!("a terminal Spider state must have a terminal outcome"), + }) + } +} + +/// Builds independent query tasks and their positionally ordered external inputs. +/// +/// # Returns +/// +/// A tuple on success, containing: +/// +/// * The constructed task graph. +/// * The positionally ordered external inputs. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards [`TaskGraph::new`]'s return values on failure. +/// * Forwards [`ValueTypeDescriptor::struct_from_name`]'s return values on failure. +/// * Forwards [`TaskGraph::insert_task`]'s return values on failure. +/// * Forwards [`rmp_serde::to_vec`]'s return values on failure. +fn build_query_task_graph( + query_job_id: QueryJobId, + clp_s_query_option: &ClpSQueryOption, + output_handle: &OutputHandle, + archives_to_search: Vec<(ArchiveMetadata, ExecutionPolicy)>, +) -> Result<(TaskGraph, Vec), Error> { + // NOTE: Keep these names and the input order in sync with the TDL package definitions. + const CLP_TDL_PACKAGE_NAME: &str = "clp"; + const QUERY_TASK_FUNC: &str = "query::clp_s_search"; + + let mut graph = TaskGraph::new(None, None)?; + + let mut inputs = Vec::new(); + for (archive, execution_policy) in archives_to_search { + graph.insert_task(TaskDescriptor { + tdl_context: TdlContext { + package: CLP_TDL_PACKAGE_NAME.to_owned(), + task_func: QUERY_TASK_FUNC.to_owned(), + }, + execution_policy: Some(execution_policy), + inputs: vec![ + DataTypeDescriptor::Value(ValueTypeDescriptor::int32()), + DataTypeDescriptor::Value(ValueTypeDescriptor::struct_from_name( + "ClpSQueryOption", + )?), + DataTypeDescriptor::Value(ValueTypeDescriptor::struct_from_name( + "Option", + )?), + DataTypeDescriptor::Value(ValueTypeDescriptor::struct_from_name("NonEmptyString")?), + DataTypeDescriptor::Value(ValueTypeDescriptor::struct_from_name("OutputHandle")?), + ], + outputs: vec![], + input_sources: None, + })?; + inputs.push(TaskInput::ValuePayload(rmp_serde::to_vec(&query_job_id)?)); + inputs.push(TaskInput::ValuePayload(rmp_serde::to_vec( + clp_s_query_option, + )?)); + inputs.push(TaskInput::ValuePayload(rmp_serde::to_vec( + &archive.dataset, + )?)); + inputs.push(TaskInput::ValuePayload(rmp_serde::to_vec(&archive.id)?)); + inputs.push(TaskInput::ValuePayload(rmp_serde::to_vec(output_handle)?)); + } + + Ok((graph, inputs)) +}