diff --git a/src/flow/mod.rs b/src/flow/mod.rs index f6eb49f..91d716f 100644 --- a/src/flow/mod.rs +++ b/src/flow/mod.rs @@ -8,3 +8,27 @@ pub fn get_flow_identifier(flow: &ValidationFlow) -> String { flow.r#type, flow.project_slug, flow.project_id, flow.flow_id ) } + +pub fn key_has_flow_id(key: &str, flow_id: i64) -> bool { + key.rsplit_once('.') + .and_then(|(_, id)| id.parse::().ok()) + == Some(flow_id) +} + +#[cfg(test)] +mod tests { + use super::key_has_flow_id; + + #[test] + fn matches_flow_id_in_final_key_segment() { + assert!(key_has_flow_id("CRON.test.1.1", 1)); + assert!(key_has_flow_id("REST.project.42.123", 123)); + } + + #[test] + fn rejects_partial_or_non_final_flow_id_matches() { + assert!(!key_has_flow_id("CRON.test.1.11", 1)); + assert!(!key_has_flow_id("CRON.test.1.1.extra", 1)); + assert!(!key_has_flow_id("CRON.test.1.invalid", 1)); + } +} diff --git a/src/sagittarius/flow_service_client_impl.rs b/src/sagittarius/flow_service_client_impl.rs index a788656..e277d90 100644 --- a/src/sagittarius/flow_service_client_impl.rs +++ b/src/sagittarius/flow_service_client_impl.rs @@ -1,5 +1,6 @@ use crate::{ - authorization::authorization::get_authorization_metadata, flow::get_flow_identifier, + authorization::authorization::get_authorization_metadata, + flow::{get_flow_identifier, key_has_flow_id}, telemetry::metrics, }; use futures::{StreamExt, TryStreamExt}; @@ -28,12 +29,6 @@ fn module_config_stats(configs: &tucana::shared::ModuleConfigurations) -> (usize (project_count, config_count) } -fn key_has_flow_id(key: &str, flow_id: i64) -> bool { - key.rsplit_once('.') - .and_then(|(_, id)| id.parse::().ok()) - == Some(flow_id) -} - #[derive(Clone)] pub struct SagittariusFlowClient { store: Arc, diff --git a/src/sagittarius/test_execution_client_impl.rs b/src/sagittarius/test_execution_client_impl.rs index d84b0d7..c204e68 100644 --- a/src/sagittarius/test_execution_client_impl.rs +++ b/src/sagittarius/test_execution_client_impl.rs @@ -4,7 +4,7 @@ In some conditions Sagittarius can't connect to Aquila Thus Aquila sends a `Logon` request to connect to Sagittarius establishing the connection */ -use futures::StreamExt; +use futures::{StreamExt, TryStreamExt}; use prost::Message; use std::{ collections::HashMap, @@ -20,7 +20,7 @@ use tucana::sagittarius::execution_service_client::ExecutionServiceClient; use tucana::sagittarius::{ExecutionLogonRequest, Logon}; use tucana::shared::{ExecutionFlow, ExecutionResult, ValidationFlow}; -use crate::authorization::authorization::get_authorization_metadata; +use crate::{authorization::authorization::get_authorization_metadata, flow::key_has_flow_id}; const EXECUTION_FLOW_ID_TTL: Duration = Duration::from_secs(30 * 60); const MAX_EXECUTION_FLOW_IDS: usize = 10_000; @@ -276,7 +276,44 @@ impl SagittariusTestExecutionServiceClient { } async fn load_validation_flow(&self, flow_id: i64) -> Option { - match self.store.get(format!("*.*.*.{}", flow_id)).await { + let mut keys = match self.store.keys().await { + Ok(keys) => keys, + Err(err) => { + log::error!( + "Failed to list validation flow keys flow_id={} error={:?}", + flow_id, + err + ); + return None; + } + }; + + let key = loop { + match keys.try_next().await { + Ok(Some(key)) if key_has_flow_id(&key, flow_id) => break key, + Ok(Some(_)) => {} + Ok(None) => { + log::error!("Validation flow was not found flow_id={}", flow_id); + return None; + } + Err(err) => { + log::error!( + "Failed while scanning validation flow keys flow_id={} error={:?}", + flow_id, + err + ); + return None; + } + } + }; + + log::debug!( + "Resolved validation flow key flow_id={} key={}", + flow_id, + key + ); + + match self.store.get(&key).await { Ok(Some(bytes)) => match ValidationFlow::decode(bytes) { Ok(flow) => { log::debug!( @@ -298,13 +335,18 @@ impl SagittariusTestExecutionServiceClient { } }, Ok(None) => { - log::error!("Validation flow was not found flow_id={}", flow_id); + log::error!( + "Validation flow disappeared after key resolution flow_id={} key={}", + flow_id, + key + ); None } Err(err) => { log::error!( - "Failed to fetch validation flow flow_id={} error={:?}", + "Failed to fetch validation flow flow_id={} key={} error={:?}", flow_id, + key, err ); None