feat(query-coordinator): Add query job coordination loop.

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"
 

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
                 """
             )

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!({

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,

components/clp-rust-utils/src/task_io.rs

@@ -1 +1,2 @@
 pub mod compression;
+pub mod query;

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<NonZeroU32>,
+
+    /// Inclusive `--tge` bound in Unix epoch milliseconds.
+    pub begin_timestamp_millisecs: Option<i64>,
+
+    /// Inclusive `--tle` bound in Unix epoch milliseconds.
+    pub end_timestamp_millisecs: Option<i64>,
+
+    /// 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);
+    }
+}

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.

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,
+    ],
 }