diff --git a/src/common/meta/src/key/flow/flow_state.rs b/src/common/meta/src/key/flow/flow_state.rs index 77ccd6206a..5423cef37b 100644 --- a/src/common/meta/src/key/flow/flow_state.rs +++ b/src/common/meta/src/key/flow/flow_state.rs @@ -16,16 +16,20 @@ use std::collections::BTreeMap; use std::sync::Arc; use serde::{Deserialize, Serialize}; +use tokio::sync::Mutex; use crate::error::{self, Result}; use crate::key::flow::FlowScoped; use crate::key::{FlowId, MetadataKey, MetadataValue}; use crate::kv_backend::KvBackendRef; -use crate::rpc::store::PutRequest; +use crate::rpc::store::{PutRequest, RangeRequest}; /// The entire FlowId to Flow Size's Map is stored directly in the value part of the key. pub const FLOW_STATE_KEY: &str = "state"; +/// The inner prefix (under `state/`) of the per-flownode flow state keys. +pub const FLOW_STATE_NODE_KEY_PREFIX: &str = "node"; + /// The key of flow state. #[derive(Debug, Clone, Copy, PartialEq)] struct FlowStateKeyInner; @@ -92,6 +96,99 @@ impl<'a> MetadataKey<'a, FlowStateKey> for FlowStateKey { } } +/// The inner key of a per-flownode flow state entry: `state/node/{node_id}`. +/// +/// `node_id` is the operator-configured flownode id (unique within the +/// cluster; a flownode requires `node_id` in its config, see +/// `src/cmd/src/flownode.rs`). It is the same value reported as +/// `HeartbeatRequest.header.member_id` — the canonical identity metasrv uses +/// for flownodes, see `get_node_id` in `src/meta-srv/src/service/heartbeat.rs` +/// — and as `HeartbeatRequest.peer.id`. +#[derive(Debug, Clone, PartialEq)] +struct FlowStateNodeKeyInner { + node_id: u64, +} + +impl FlowStateNodeKeyInner { + pub fn new(node_id: u64) -> Self { + Self { node_id } + } +} + +impl<'a> MetadataKey<'a, FlowStateNodeKeyInner> for FlowStateNodeKeyInner { + fn to_bytes(&self) -> Vec { + format!( + "{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/{}", + self.node_id + ) + .into_bytes() + } + + fn from_bytes(bytes: &'a [u8]) -> Result { + let key = std::str::from_utf8(bytes).map_err(|e| { + error::InvalidMetadataSnafu { + err_msg: format!( + "FlowStateNodeKeyInner '{}' is not a valid UTF8 string: {e}", + String::from_utf8_lossy(bytes) + ), + } + .build() + })?; + let prefix = format!("{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/"); + let Some(node_id) = key.strip_prefix(&prefix) else { + return Err(error::InvalidMetadataSnafu { + err_msg: format!("Invalid FlowStateNodeKeyInner '{key}'"), + } + .build()); + }; + let node_id = node_id.parse::().map_err(|_| { + error::InvalidMetadataSnafu { + err_msg: format!("Invalid node id '{node_id}' in FlowStateNodeKeyInner '{key}'"), + } + .build() + })?; + Ok(FlowStateNodeKeyInner::new(node_id)) + } +} + +/// The key stores the per-flownode flow state report. +/// +/// The layout: `__flow/state/node/{node_id}`. +/// +/// Per-node keys live in the in-memory KV (same as the global `__flow/state` +/// key), so they are automatically cleared when metasrv resets the in-memory +/// KV on leader change; no separate cleanup is needed. +pub struct FlowStateNodeKey(FlowScoped); + +impl FlowStateNodeKey { + /// Returns the [FlowStateNodeKey] of the given node. + pub fn new(node_id: u64) -> FlowStateNodeKey { + FlowStateNodeKey(FlowScoped::new(FlowStateNodeKeyInner::new(node_id))) + } + + /// Returns the full key prefix of all per-node flow state keys: + /// `__flow/state/node/`. + pub fn prefix() -> Vec { + format!( + "{}{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/", + FlowScoped::::PREFIX + ) + .into_bytes() + } +} + +impl<'a> MetadataKey<'a, FlowStateNodeKey> for FlowStateNodeKey { + fn to_bytes(&self) -> Vec { + self.0.to_bytes() + } + + fn from_bytes(bytes: &'a [u8]) -> Result { + Ok(FlowStateNodeKey( + FlowScoped::::from_bytes(bytes)?, + )) + } +} + /// The value of flow state size #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] pub struct FlowStateValue { @@ -125,13 +222,25 @@ pub type FlowStateManagerRef = Arc; /// /// This is only used in distributed mode. When meta-srv use heartbeat to update the flow stat report /// and frontned use get to get the latest flow stat report. +/// +/// Per-flownode reports are stored under `__flow/state/node/{node_id}` keys in +/// the in-memory KV (not in a separate in-process map), so a metasrv leader +/// change — which resets the in-memory KV — automatically clears all per-node +/// state without leaving stale entries behind. pub struct FlowStateManager { in_memory: KvBackendRef, + /// Serializes the critical section of [`FlowStateManager::merge`] + /// (write per-node key -> scan -> aggregate -> write global key). It holds + /// no long-lived state; per-node reports live in the in-memory KV. + merge_lock: Mutex<()>, } impl FlowStateManager { pub fn new(in_memory: KvBackendRef) -> Self { - Self { in_memory } + Self { + in_memory, + merge_lock: Mutex::new(()), + } } pub async fn get(&self) -> Result> { @@ -150,6 +259,78 @@ impl FlowStateManager { self.in_memory.put(req).await?; Ok(()) } + + /// Merges a flow state report from a single flownode into the global view. + /// + /// `node_id` is the operator-configured flownode id (unique within the + /// cluster; reported as `HeartbeatRequest.header.member_id`, the canonical + /// metasrv identity of a flownode, and equal to `peer.id`). + /// + /// Reports are tracked per node under `__flow/state/node/{node_id}` in the + /// in-memory KV. A new report from the same node unconditionally replaces + /// that node's previous entry: reports are processed strictly in arrival + /// order, so no epoch comparison is made and a wall-clock rollback after a + /// node restart cannot permanently drop the node's later reports. The + /// global `FlowStateValue` is then aggregated over all per-node entries: + /// for the same flow, `last_exec_time_map` takes the max reported + /// timestamp and `state_size` takes the max reported size across nodes (a + /// flow normally runs on a single active flownode, so max is a safe + /// approximation). The aggregated value is written into the in-memory KV + /// under the global key `__flow/state` via the same path as `put`, keeping + /// `get()` behavior unchanged. + /// + /// Per-node keys are cleared automatically when the in-memory KV is reset + /// on a leader change. Known limitation: entries of dropped flows are not + /// proactively removed here. Once a flow is dropped its metadata is gone, + /// so the flows table join simply can't see the stale entry; it becomes + /// user-invisible until the owning flownode reports again (or stops + /// heartbeating forever, in which case the stale flow id lingers in the + /// aggregate but is never joined against any flow metadata). + pub async fn merge(&self, node_id: u64, incoming: FlowStateValue) -> Result<()> { + let _guard = self.merge_lock.lock().await; + + // 1. Store this node's latest report under its per-node key. + let node_key = FlowStateNodeKey::new(node_id).to_bytes(); + let value = incoming.try_as_raw_value()?; + let req = PutRequest::new().with_key(node_key).with_value(value); + self.in_memory.put(req).await?; + + // 2. Read back every per-node report and aggregate them. + let req = RangeRequest::new().with_prefix(FlowStateNodeKey::prefix()); + let resp = self.in_memory.range(req).await?; + let mut state_size = BTreeMap::new(); + let mut last_exec_time_map = BTreeMap::new(); + let mut start_time_map = BTreeMap::new(); + for kv in resp.kvs { + let state = FlowStateValue::try_from_raw_value(&kv.value)?; + for (flow_id, size) in state.state_size { + state_size + .entry(flow_id) + .and_modify(|v: &mut usize| *v = (*v).max(size)) + .or_insert(size); + } + for (flow_id, ts) in state.last_exec_time_map { + last_exec_time_map + .entry(flow_id) + .and_modify(|v: &mut i64| *v = (*v).max(ts)) + .or_insert(ts); + } + for (flow_id, ts) in state.start_time_map { + start_time_map + .entry(flow_id) + .and_modify(|v: &mut i64| *v = (*v).max(ts)) + .or_insert(ts); + } + } + + // 3. Write the aggregated value to the global key. + let aggregated = FlowStateValue::new(state_size, last_exec_time_map, start_time_map); + let key = FlowStateKey::new().to_bytes(); + let value = aggregated.try_as_raw_value()?; + let req = PutRequest::new().with_key(key).with_value(value); + self.in_memory.put(req).await?; + Ok(()) + } } /// Flow's state report, send regularly through heartbeat message @@ -234,4 +415,210 @@ mod tests { let decoded: FlowStateValue = serde_json::from_str(&json).unwrap(); assert_eq!(decoded, value); } + + use std::sync::Arc; + + use super::*; + use crate::kv_backend::memory::MemoryKvBackend; + + fn state(last_exec_time_map: BTreeMap) -> FlowStateValue { + FlowStateValue::new(BTreeMap::new(), last_exec_time_map, BTreeMap::new()) + } + + #[tokio::test] + async fn test_merge_keeps_reports_from_different_nodes() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + // Node A reports flow 1 executed at t1. + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + // Node B reports flow 2 executed at t2. Before the per-node merge this + // would have wiped out node A's flow 1 entry. + manager + .merge(2, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + // Node A reports flow 1 executed at t3. + manager + .merge(1, state(BTreeMap::from([(1, 300)]))) + .await + .unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + // flow 1 and flow 2 are both present, and flow 1 takes max(t1, t3). + assert_eq!(value.last_exec_time_map.get(&1), Some(&300)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_replaces_same_node_state() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + // A new report from the same node replaces the previous one. There is + // no epoch comparison: arrival order alone decides, so this is also + // what a restarted node (new epoch) hits. + manager + .merge(1, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert!(!value.last_exec_time_map.contains_key(&1)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_accepts_clock_rollback_from_same_node() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + // Simulates a flownode restart whose wall clock rolled back: the node + // first reports flow 1 at t1, then (after restart) reports a *smaller* + // timestamp t0. Because reports are processed strictly in arrival + // order (no epoch comparison), the later report must win instead of + // being permanently rejected. + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + manager + .merge(1, state(BTreeMap::from([(1, 50)]))) + .await + .unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert_eq!(value.last_exec_time_map.get(&1), Some(&50)); + } + + #[tokio::test] + async fn test_merge_aggregates_state_size_and_last_exec_time_by_max() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + // Both nodes report the same flow; the aggregate must take the max of + // state_size and last_exec_time_map across nodes. + manager + .merge( + 1, + FlowStateValue::new( + BTreeMap::from([(1, 1024)]), + BTreeMap::from([(1, 100)]), + BTreeMap::new(), + ), + ) + .await + .unwrap(); + manager + .merge( + 2, + FlowStateValue::new( + BTreeMap::from([(1, 2048)]), + BTreeMap::from([(1, 50)]), + BTreeMap::new(), + ), + ) + .await + .unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert_eq!(value.state_size.get(&1), Some(&2048)); + assert_eq!(value.last_exec_time_map.get(&1), Some(&100)); + } + + #[tokio::test] + async fn test_merge_concurrent_reports_no_lost_update() { + let manager = Arc::new(FlowStateManager::new(Arc::new(MemoryKvBackend::default()))); + + // Two nodes report concurrently; the merge lock must serialize the + // read-modify-write so neither node's report is lost. + let m1 = manager.clone(); + let h1 = tokio::spawn(async move { + m1.merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + }); + let m2 = manager.clone(); + let h2 = tokio::spawn(async move { + m2.merge(2, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + }); + h1.await.unwrap(); + h2.await.unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert_eq!(value.last_exec_time_map.get(&1), Some(&100)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_empty_report_removes_own_flows_keeps_others() { + let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default())); + + // Node A reports flow 1, node B reports flow 2. + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + manager + .merge(2, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + + // Node A reports an empty map: its own flow 1 disappears from the + // aggregate while node B's flow 2 is retained. + manager.merge(1, state(BTreeMap::new())).await.unwrap(); + + let value = manager.get().await.unwrap().unwrap(); + assert!(!value.last_exec_time_map.contains_key(&1)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_state_cleared_on_in_memory_kv_reset() { + let backend = Arc::new(MemoryKvBackend::default()); + let manager = FlowStateManager::new(backend.clone()); + + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + assert!(manager.get().await.unwrap().is_some()); + + // Simulate a metasrv leader change, which clears the in-memory KV + // (including the per-node keys, since they live in the same KV). + backend.clear(); + assert!(manager.get().await.unwrap().is_none()); + + // A fresh report works again after the reset. + manager + .merge(2, state(BTreeMap::from([(2, 200)]))) + .await + .unwrap(); + let value = manager.get().await.unwrap().unwrap(); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_merge_writes_global_key() { + let backend = Arc::new(MemoryKvBackend::default()); + let manager = FlowStateManager::new(backend.clone()); + + manager + .merge(1, state(BTreeMap::from([(1, 100)]))) + .await + .unwrap(); + + // The global key is still `__flow/state` and holds the serialized + // FlowStateValue, so existing get()/client readers are unchanged. + let dump = backend.dump(); + let global_key = "__flow/state".as_bytes().to_vec(); + assert!(dump.contains_key(&global_key)); + let value = FlowStateValue::try_from_raw_value(dump.get(&global_key).unwrap()).unwrap(); + assert_eq!(value.last_exec_time_map.get(&1), Some(&100)); + } } diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index 3a4c86b30a..a707c2d9eb 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -969,15 +969,27 @@ pub(crate) async fn gen_plan_with_matching_schema( } pub fn df_plan_to_sql(plan: &LogicalPlan) -> Result { - /// A dialect that forces identifiers to be quoted when have uppercase + /// A dialect that forces identifiers to be quoted when they contain + /// anything other than lowercase alphanumerics and underscores, or start + /// with a digit. + /// + /// Unquoted identifiers are normalized to lowercase by the SQL parser, so + /// uppercase letters need quoting to preserve case. Special characters + /// (e.g. ':' in Prometheus-style table names like + /// `kube_pod_cpu_cores:sum`, '.', '-', spaces) and digit-leading names + /// (e.g. `123metrics`) would produce invalid SQL if left unquoted, so they + /// are quoted as well. SQL keywords are intentionally not checked here: + /// quoting every ALL_KEYWORDS member would also quote common column names + /// like `number`; the unparse failure path has an InsertIntoPlan fallback. struct ForceQuoteIdentifiers; impl datafusion::sql::unparser::dialect::Dialect for ForceQuoteIdentifiers { fn identifier_quote_style(&self, identifier: &str) -> Option { - if identifier.to_lowercase() != identifier { - Some('`') - } else { - None - } + let is_plain = !identifier.is_empty() + && !identifier.starts_with(|c: char| c.is_ascii_digit()) + && identifier + .chars() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_'); + if is_plain { None } else { Some('"') } } } let unparser = Unparser::new(&ForceQuoteIdentifiers); diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index 7fbd3300c4..1356f87e8f 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/src/flow/src/batching_mode/utils/test.rs @@ -268,7 +268,7 @@ async fn test_sql_plan_convert() { let new_sql = df_plan_to_sql(&new).unwrap(); assert_eq!( - r#"SELECT `UPPERCASE_NUMBERS_WITH_TS`.`NUMBER` FROM `UPPERCASE_NUMBERS_WITH_TS`"#, + r#"SELECT "UPPERCASE_NUMBERS_WITH_TS"."NUMBER" FROM "UPPERCASE_NUMBERS_WITH_TS""#, new_sql ); } @@ -1985,3 +1985,78 @@ async fn test_gen_plan_with_matching_schema_last_non_null_rejects_extra_flow_col "{err}" ); } + +#[test] +fn test_df_plan_to_sql_quotes_colon_table_name() { + // Prometheus-style table names contain ':' (e.g. + // `kube_pod_cpu_cores:sum`). The unparser dialect must quote them, + // otherwise the re-parsed SQL is invalid (`keyword: :`). + let table = single_row_u32_table("kube_pod_cpu_cores:sum", vec!["value"]); + let provider = Arc::new(DfTableProviderAdapter::new(table)); + let table_source = Arc::new(DefaultTableSource::new(provider)); + let table_ref = TableReference::full("catalog", "schema", "kube_pod_cpu_cores:sum"); + let plan = LogicalPlanBuilder::scan(table_ref, table_source, None) + .unwrap() + .build() + .unwrap(); + + let sql = df_plan_to_sql(&plan).unwrap(); + assert!( + sql.contains("\"kube_pod_cpu_cores:sum\""), + "expected quoted table name in {sql}" + ); + // The only occurrence of `cores:sum` must be inside the quoted identifier. + assert_eq!( + sql.matches("cores:sum").count(), + 1, + "colon should only appear inside quotes in {sql}" + ); +} + +#[test] +fn test_df_plan_to_sql_does_not_quote_plain_lowercase() { + let table = single_row_u32_table("plain_table", vec!["value"]); + let provider = Arc::new(DfTableProviderAdapter::new(table)); + let table_source = Arc::new(DefaultTableSource::new(provider)); + let table_ref = TableReference::full("catalog", "schema", "plain_table"); + let plan = LogicalPlanBuilder::scan(table_ref, table_source, None) + .unwrap() + .project(vec![datafusion_expr::col("value")]) + .unwrap() + .build() + .unwrap(); + + let sql = df_plan_to_sql(&plan).unwrap(); + assert!( + sql.contains("plain_table") && !sql.contains("\"plain_table\""), + "plain lowercase table should stay unquoted in {sql}" + ); + // `value` is not a reserved word, so the column stays unquoted. + assert!( + sql.contains("plain_table.value"), + "column unquoted in {sql}" + ); + assert!(!sql.contains('`'), "no backtick quoting in {sql}"); +} + +#[test] +fn test_df_plan_to_sql_quotes_digit_leading_table_name() { + // A table literally named `123metrics` starts with a digit and must be + // quoted, otherwise the re-parsed SQL is invalid. + let table = single_row_u32_table("123metrics", vec!["value"]); + let provider = Arc::new(DfTableProviderAdapter::new(table)); + let table_source = Arc::new(DefaultTableSource::new(provider)); + let table_ref = TableReference::full("catalog", "schema", "123metrics"); + let plan = LogicalPlanBuilder::scan(table_ref, table_source, None) + .unwrap() + .project(vec![datafusion_expr::col("value")]) + .unwrap() + .build() + .unwrap(); + + let sql = df_plan_to_sql(&plan).unwrap(); + assert!( + sql.contains("\"123metrics\""), + "expected digit-leading table name quoted in {sql}" + ); +} diff --git a/src/meta-srv/src/handler/flow_state_handler.rs b/src/meta-srv/src/handler/flow_state_handler.rs index a13c18a9e0..b4abbd7e37 100644 --- a/src/meta-srv/src/handler/flow_state_handler.rs +++ b/src/meta-srv/src/handler/flow_state_handler.rs @@ -14,12 +14,27 @@ use api::v1::meta::{FlowStat, HeartbeatRequest, Role}; use common_meta::key::flow::flow_state::{FlowStateManager, FlowStateValue}; +use common_telemetry::debug; use snafu::ResultExt; use crate::error::{FlowStateHandlerSnafu, Result}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; +/// Extracts the flownode identity from a heartbeat request. +/// +/// Prefers `header.member_id` — the canonical identity metasrv uses for +/// flownodes (see `get_node_id` in `src/meta-srv/src/service/heartbeat.rs`). +/// Falls back to `peer.id`, the operator-configured node id, which is also +/// unique within the cluster (a flownode requires `node_id` in its config, see +/// `src/cmd/src/flownode.rs`). Returns `None` when neither is present. +fn node_identity(req: &HeartbeatRequest) -> Option { + req.header + .as_ref() + .map(|header| header.member_id) + .or_else(|| req.peer.as_ref().map(|peer| peer.id)) +} + pub struct FlowStateHandler { flow_state_manager: FlowStateManager, } @@ -60,11 +75,141 @@ impl HeartbeatHandler for FlowStateHandler { // mode until a follow-up PR adds heartbeat propagation. let value: FlowStateValue = FlowStateValue::new(state_size, last_exec_time_map, Default::default()); - self.flow_state_manager - .put(value) - .await - .context(FlowStateHandlerSnafu)?; + match node_identity(req) { + Some(node_id) => { + // Merge by node so that reports from different flownodes + // don't overwrite each other in the global state. + self.flow_state_manager + .merge(node_id, value) + .await + .context(FlowStateHandlerSnafu)?; + } + // No usable identity in the request: ignore the report instead + // of falling back to a whole-map replace, which would clobber + // other nodes' reports. + // Normal flownodes always carry header.member_id/peer.id; a + // report without either indicates an old or malformed client. + // Log at debug to avoid an anomalous sender spamming warn. + None => { + debug!( + "Ignore flow state report without node identity (no header.member_id and no peer.id): {value:?}" + ); + } + } } Ok(HandleControl::Continue) } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use api::v1::meta::{HeartbeatRequest, Peer, RequestHeader, Role}; + + use super::*; + use crate::handler::test_utils::TestEnv; + + #[test] + fn test_node_identity_prefers_header_member_id() { + let req = HeartbeatRequest { + header: Some(RequestHeader::new(42, Role::Flownode, HashMap::new())), + peer: Some(Peer { + id: 99, + addr: "127.0.0.1:4001".to_string(), + }), + ..Default::default() + }; + assert_eq!(node_identity(&req), Some(42)); + } + + #[test] + fn test_node_identity_falls_back_to_peer_id() { + let req = HeartbeatRequest { + header: None, + peer: Some(Peer { + id: 99, + addr: "127.0.0.1:4001".to_string(), + }), + ..Default::default() + }; + assert_eq!(node_identity(&req), Some(99)); + } + + #[test] + fn test_node_identity_none_without_header_and_peer() { + let req = HeartbeatRequest::default(); + assert_eq!(node_identity(&req), None); + } + + fn flow_stat() -> api::v1::meta::FlowStat { + api::v1::meta::FlowStat { + flow_stat_size: HashMap::from([(1, 1024)]), + flow_last_exec_time_map: HashMap::from([(1, 100)]), + } + } + + #[tokio::test] + async fn test_handle_merges_reports_from_different_nodes() { + let env = TestEnv::new(); + let ctx = env.ctx(); + let flow_state_manager = FlowStateManager::new(ctx.in_memory.clone().as_kv_backend_ref()); + let handler = FlowStateHandler::new(flow_state_manager); + + // Node 42 (header.member_id) reports flow 1. + let req_a = HeartbeatRequest { + header: Some(RequestHeader::new(42, Role::Flownode, HashMap::new())), + peer: Some(Peer { + id: 42, + addr: "127.0.0.1:4001".to_string(), + }), + flow_stat: Some(flow_stat()), + ..Default::default() + }; + let mut ctx = env.ctx(); + let mut acc = HeartbeatAccumulator::default(); + handler.handle(&req_a, &mut ctx, &mut acc).await.unwrap(); + + // Node 7 (only peer.id present) reports flow 2. + let req_b = HeartbeatRequest { + header: None, + peer: Some(Peer { + id: 7, + addr: "127.0.0.1:4007".to_string(), + }), + flow_stat: Some(api::v1::meta::FlowStat { + flow_stat_size: HashMap::from([(2, 2048)]), + flow_last_exec_time_map: HashMap::from([(2, 200)]), + }), + ..Default::default() + }; + let mut ctx = env.ctx(); + let mut acc = HeartbeatAccumulator::default(); + handler.handle(&req_b, &mut ctx, &mut acc).await.unwrap(); + + let value = handler.flow_state_manager.get().await.unwrap().unwrap(); + assert_eq!(value.last_exec_time_map.get(&1), Some(&100)); + assert_eq!(value.last_exec_time_map.get(&2), Some(&200)); + } + + #[tokio::test] + async fn test_handle_ignores_report_without_identity() { + let env = TestEnv::new(); + let ctx = env.ctx(); + let flow_state_manager = FlowStateManager::new(ctx.in_memory.clone().as_kv_backend_ref()); + let handler = FlowStateHandler::new(flow_state_manager); + + // No header and no peer: the report must be ignored and no KV written. + let req = HeartbeatRequest { + header: None, + peer: None, + flow_stat: Some(flow_stat()), + ..Default::default() + }; + let mut ctx = env.ctx(); + let mut acc = HeartbeatAccumulator::default(); + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + + assert!(handler.flow_state_manager.get().await.unwrap().is_none()); + } +}