From 8ba58cd1b5e18327ddd59a8e64c66efb138f86ee Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 8 Oct 2026 10:29:26 +0000 Subject: [PATCH] fix(metasrv): retry failed heartbeat address updates (#9437) * fix(metasrv): retry failed heartbeat address updates Signed-off-by: WenyXu * refactor(metasrv): move shared heartbeat helpers to utils Signed-off-by: WenyXu --------- Signed-off-by: WenyXu (cherry picked from commit 65b9a356921fc5157fc74ec7119b4c0988a128c7) Signed-off-by: WenyXu --- src/meta-srv/AGENTS.md | 6 +- src/meta-srv/src/handler.rs | 1 + .../handler/collect_cluster_info_handler.rs | 247 +++++++++------ .../src/handler/collect_stats_handler.rs | 209 +++++++++---- .../src/handler/remap_flow_peer_handler.rs | 136 +++++---- src/meta-srv/src/handler/test_utils.rs | 37 ++- src/meta-srv/src/handler/utils.rs | 281 ++++++++++++++++++ 7 files changed, 714 insertions(+), 203 deletions(-) create mode 100644 src/meta-srv/src/handler/utils.rs diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index e1e23581525..5f9f573978c 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -38,7 +38,11 @@ abstraction, election, key encoding, and the DDL manager live in `common-meta`; - **Heartbeat**: `service/heartbeat.rs` receives datanode/frontend heartbeats and runs the handler chain in `handler/`; responses carry mailbox messages (DDL - results, cache invalidation). + results, cache invalidation). Node information is refreshed in memory on each + heartbeat. Datanode statistics collection and flownode remapping share + `handler/utils.rs` for per-node address persistence. Datanode address + updates retain the statistics handler's per-node ordering. Address epochs are + confirmed only after a successful save; cache invalidation remains best effort. - **DDL**: `service/procedure.rs` (`ddl`) is leader-only and hands off to `common-meta`'s `DdlManager`, executed via `common-procedure`. - **Region migration**: `procedure/region_migration/manager.rs` plus per-step diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 320f26fae7c..6bf71674a82 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -81,6 +81,7 @@ pub mod publish_heartbeat_handler; pub mod region_lease_handler; pub mod remap_flow_peer_handler; pub mod response_header_handler; +mod utils; #[cfg(test)] pub mod test_utils; diff --git a/src/meta-srv/src/handler/collect_cluster_info_handler.rs b/src/meta-srv/src/handler/collect_cluster_info_handler.rs index 6bee86ad711..12fee395136 100644 --- a/src/meta-srv/src/handler/collect_cluster_info_handler.rs +++ b/src/meta-srv/src/handler/collect_cluster_info_handler.rs @@ -46,30 +46,21 @@ impl HeartbeatHandler for CollectFrontendClusterInfoHandler { ctx: &mut Context, _acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; let frontend_workloads = get_frontend_workloads(req.node_workloads.as_ref()); - let value = NodeInfo { - peer, - last_activity_ts: common_time::util::current_time_millis(), - status: NodeStatus::Frontend(FrontendStatus { - workloads: frontend_workloads, - }), - version: info.version, - git_commit: info.git_commit, - start_time_ms: info.start_time_ms, - total_cpu_millicores: info.total_cpu_millicores, - total_memory_bytes: info.total_memory_bytes, - cpu_usage_millicores: info.cpu_usage_millicores, - memory_usage_bytes: info.memory_usage_bytes, - hostname: info.hostname, - env_vars, - }; - - put_into_memory_store(ctx, key, value).await?; + update + .save_to_memory( + ctx, + common_time::util::current_time_millis(), + NodeStatus::Frontend(FrontendStatus { + workloads: frontend_workloads, + }), + ) + .await?; Ok(HandleControl::Continue) } @@ -89,29 +80,20 @@ impl HeartbeatHandler for CollectFlownodeClusterInfoHandler { ctx: &mut Context, _acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; let flownode_workloads = get_flownode_workloads(req.node_workloads.as_ref()); - let value = NodeInfo { - peer, - last_activity_ts: common_time::util::current_time_millis(), - status: NodeStatus::Flownode(FlownodeStatus { - workloads: flownode_workloads, - }), - version: info.version, - git_commit: info.git_commit, - start_time_ms: info.start_time_ms, - total_cpu_millicores: info.total_cpu_millicores, - total_memory_bytes: info.total_memory_bytes, - cpu_usage_millicores: info.cpu_usage_millicores, - memory_usage_bytes: info.memory_usage_bytes, - hostname: info.hostname, - env_vars, - }; - - put_into_memory_store(ctx, key, value).await?; + update + .save_to_memory( + ctx, + common_time::util::current_time_millis(), + NodeStatus::Flownode(FlownodeStatus { + workloads: flownode_workloads, + }), + ) + .await?; Ok(HandleControl::Continue) } @@ -132,7 +114,7 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { ctx: &mut Context, acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; @@ -147,16 +129,64 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { .count(); let follower_regions = stat.region_stats.len() - leader_regions; - let value = NodeInfo { + update + .save_to_memory( + ctx, + stat.timestamp_millis, + NodeStatus::Datanode(DatanodeStatus { + rcus: stat.rcus, + wcus: stat.wcus, + leader_regions, + follower_regions, + workloads: stat.datanode_workloads.clone(), + }), + ) + .await?; + + Ok(HandleControl::Continue) + } +} + +/// Common fields collected for one heartbeat's in-memory node information +/// update. Persistent routing addresses have a separate lifecycle. +struct NodeInfoUpdate { + key: NodeInfoKey, + peer: Peer, + info: PbNodeInfo, + env_vars: HashMap, +} + +impl NodeInfoUpdate { + fn from_request(request: &HeartbeatRequest) -> Option { + let key = NodeInfoKey::new(request)?; + let peer = request.peer.clone()?; + let info = request.info.clone()?; + let env_vars = EnvVars::from_extensions(&request.extensions) + .inspect_err(|e| { + warn!(e; "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer); + }) + .unwrap_or_default() + .map(|e| e.vars) + .unwrap_or_default(); + Some(Self { + key, peer, - last_activity_ts: stat.timestamp_millis, - status: NodeStatus::Datanode(DatanodeStatus { - rcus: stat.rcus, - wcus: stat.wcus, - leader_regions, - follower_regions, - workloads: stat.datanode_workloads.clone(), - }), + info, + env_vars, + }) + } + + async fn save_to_memory( + self, + ctx: &Context, + last_activity_ts: i64, + status: NodeStatus, + ) -> Result<()> { + let info = self.info; + let value = NodeInfo { + peer: self.peer, + last_activity_ts, + status, version: info.version, git_commit: info.git_commit, start_time_ms: info.start_time_ms, @@ -165,53 +195,88 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { cpu_usage_millicores: info.cpu_usage_millicores, memory_usage_bytes: info.memory_usage_bytes, hostname: info.hostname, - env_vars, + env_vars: self.env_vars, }; - - put_into_memory_store(ctx, key, value).await?; - - Ok(HandleControl::Continue) + let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; + ctx.in_memory + .put(PutRequest { + key: (&self.key).into(), + value, + ..Default::default() + }) + .await + .context(SaveClusterInfoSnafu)?; + Ok(()) } } -fn extract_base_info( - request: &HeartbeatRequest, -) -> Option<(NodeInfoKey, Peer, PbNodeInfo, HashMap)> { - let HeartbeatRequest { peer, info, .. } = request; - let key = NodeInfoKey::new(request)?; - let Some(peer) = &peer else { - return None; - }; - let Some(info) = &info else { - return None; - }; +#[cfg(test)] +mod tests { + use api::v1::meta::RequestHeader; + use common_meta::datanode::Stat; - let env_vars = EnvVars::from_extensions(&request.extensions) - .inspect_err(|e| { - warn!(e; - "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer + use super::*; + use crate::handler::test_utils::TestEnv; + + #[tokio::test] + async fn test_node_info_refreshes_without_epoch_change() { + let mut ctx = TestEnv::new().ctx(); + let handlers: Vec<(Role, Box)> = vec![ + (Role::Frontend, Box::new(CollectFrontendClusterInfoHandler)), + (Role::Datanode, Box::new(CollectDatanodeClusterInfoHandler)), + (Role::Flownode, Box::new(CollectFlownodeClusterInfoHandler)), + ]; + for (role, handler) in handlers { + let mut req = HeartbeatRequest { + header: Some(RequestHeader { + role: role as i32, + ..Default::default() + }), + peer: Some(Peer { + id: 7, + addr: "node:3001".into(), + }), + info: Some(PbNodeInfo { + version: "v1".into(), + hostname: "host".into(), + ..Default::default() + }), + node_epoch: 10, + ..Default::default() + }; + let mut acc = HeartbeatAccumulator { + stat: Some(Stat { + timestamp_millis: 123, + rcus: 5, + wcus: 6, + ..Default::default() + }), + ..Default::default() + }; + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + let key: Vec = (&NodeInfoKey::new(&req).unwrap()).into(); + let first = + NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap(); + assert_eq!(first.peer, req.peer.clone().unwrap()); + assert_eq!(first.version, "v1"); + assert_eq!(first.hostname, "host"); + req.info.as_mut().unwrap().version = "v2".into(); + req.info.as_mut().unwrap().memory_usage_bytes = 42; + acc.stat.as_mut().unwrap().timestamp_millis = 456; + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + let second = + NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap(); + assert_eq!(second.version, "v2"); + assert_eq!(second.memory_usage_bytes, 42); + assert_eq!( + serde_json::to_value(&second.status).unwrap(), + serde_json::to_value(&first.status).unwrap() ); - }) - .unwrap_or_default() - .map(|e| e.vars) - .unwrap_or_default(); - - Some((key, peer.clone(), info.clone(), env_vars)) -} - -async fn put_into_memory_store(ctx: &mut Context, key: NodeInfoKey, value: NodeInfo) -> Result<()> { - let key = (&key).into(); - let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; - let put_req = PutRequest { - key, - value, - ..Default::default() - }; - - ctx.in_memory - .put(put_req) - .await - .context(SaveClusterInfoSnafu)?; - - Ok(()) + if role == Role::Datanode { + assert_eq!(second.last_activity_ts, 456); + } + // Node information must not be persisted to the configured backend. + assert!(ctx.kv_backend.get(&key).await.unwrap().is_none()); + } + } } diff --git a/src/meta-srv/src/handler/collect_stats_handler.rs b/src/meta-srv/src/handler/collect_stats_handler.rs index cc653f9982c..4b58beaf0fb 100644 --- a/src/meta-srv/src/handler/collect_stats_handler.rs +++ b/src/meta-srv/src/handler/collect_stats_handler.rs @@ -15,19 +15,19 @@ use std::cmp::Ordering; use std::sync::Arc; -use api::v1::meta::{HeartbeatRequest, Role}; +use api::v1::meta::{HeartbeatRequest, Peer, Role}; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; use common_meta::instruction::CacheIdent; -use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue}; -use common_meta::key::{MetadataKey, MetadataValue}; -use common_meta::peer::Peer; +use common_meta::key::MetadataKey; +use common_meta::key::node_address::NodeAddressKey; use common_meta::rpc::store::PutRequest; -use common_telemetry::{error, info, warn}; +use common_telemetry::{error, warn}; use dashmap::DashMap; use snafu::ResultExt; use tokio::sync::Mutex; use crate::error::{self, Result}; +use crate::handler::utils::{self, NodeAddressUpdater}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; @@ -74,6 +74,7 @@ const DEFAULT_FLUSH_STATS_FACTOR: usize = 3; pub struct CollectStatsHandler { stats_cache: DashMap>>, flush_stats_factor: usize, + address_updater: NodeAddressUpdater, } impl Default for CollectStatsHandler { @@ -87,6 +88,7 @@ impl CollectStatsHandler { Self { flush_stats_factor: flush_stats_factor.unwrap_or(DEFAULT_FLUSH_STATS_FACTOR), stats_cache: DashMap::default(), + address_updater: NodeAddressUpdater::default(), } } } @@ -119,6 +121,7 @@ impl HeartbeatHandler for CollectStatsHandler { let key: Vec = key.into(); + let current_epoch = current_stat.node_epoch; let refresh = if let Some(epoch) = epoch_stats.epoch() { match current_stat.node_epoch.cmp(&epoch) { Ordering::Greater => { @@ -146,11 +149,13 @@ impl HeartbeatHandler for CollectStatsHandler { true }; - // Need to refresh the [datanode -> address] mapping - if refresh { - // Safety: `epoch_stats.stats` is not empty - let last = epoch_stats.stats.last().unwrap(); - rewrite_node_address(ctx, last).await; + // Keep address updates and statistics in the original per-node order. + // Address acknowledgement is independent of the statistics epoch so a + // failed save can be retried by the next heartbeat of the same epoch. + if epoch_stats.epoch() == Some(current_epoch) + && let Some(stat) = epoch_stats.stats.last() + { + update_datanode_address(&self.address_updater, ctx, stat).await; } if !refresh && epoch_stats.len() < self.flush_stats_factor { @@ -178,48 +183,26 @@ impl HeartbeatHandler for CollectStatsHandler { } } -async fn rewrite_node_address(ctx: &Context, stat: &Stat) { - let peer = Peer { - id: stat.id, - addr: stat.addr.clone(), - }; - let key = NodeAddressKey::with_datanode(peer.id).to_bytes(); - if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() { - let put = PutRequest { - key, - value, - prev_kv: false, - }; - - match ctx.leader_cached_kv_backend.put(put).await { - Ok(_) => { - info!( - "Successfully updated datanode `NodeAddressValue`: {:?}", - peer - ); - // broadcast invalidating cache - let cache_idents = stat - .table_ids() - .into_iter() - .map(CacheIdent::TableId) - .collect::>(); - if let Err(e) = ctx - .cache_invalidator - .invalidate(&Default::default(), &cache_idents) - .await - { - error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer); - } - } - Err(e) => { - error!(e; "Failed to update datanode `NodeAddressValue`: {:?}", peer); - } - } - } else { - warn!( - "Failed to serialize datanode `NodeAddressValue`: {:?}", - peer - ); +/// Applies datanode-specific address and table-cache updates. The caller keeps +/// its statistics lock so notifications and statistics retain their ordering. +async fn update_datanode_address(updater: &NodeAddressUpdater, ctx: &Context, stat: &Stat) { + let updated = updater + .update_if_needed((Role::Datanode, stat.id), stat.node_epoch, || async { + let peer = Peer { id: stat.id, addr: stat.addr.clone() }; + let key = NodeAddressKey::with_datanode(stat.id).to_bytes(); + utils::save_node_address(ctx, key, peer).await + }) + .await + .inspect_err(|e| { + error!(e; "Failed to update datanode address, node: {}, address: {}", stat.id, stat.addr); + }); + if let Ok(true) = updated { + let cache_idents = stat + .table_ids() + .into_iter() + .map(CacheIdent::TableId) + .collect::>(); + utils::invalidate_address_caches(ctx, stat.id, &cache_idents).await; } } @@ -553,4 +536,126 @@ mod tests { handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); } } + + #[tokio::test] + async fn test_persistence_failure_retries_but_notification_failure_does_not() { + use crate::handler::collect_stats_handler::CollectStatsHandler; + use crate::handler::test_utils::FailingCacheInvalidator; + use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; + use crate::service::store::cached_kv::LeaderCachedKvBackend; + use api::v1::meta::HeartbeatRequest; + use common_meta::key::MetadataValue; + use common_meta::key::node_address::NodeAddressValue; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::rpc::store::PutResponse; + let mut ctx = TestEnv::new().ctx(); + let puts = Arc::new(std::sync::Mutex::new(Vec::::new())); + let backend = MockKvBackendBuilder::default() + .put_fn(Arc::new({ + let puts = puts.clone(); + move |req: PutRequest| { + let mut puts = puts.lock().unwrap(); + puts.push(req); + if puts.len() == 1 { + return common_meta::error::UnexpectedSnafu { + err_msg: "injected persistence failure", + } + .fail(); + } + Ok(PutResponse::default()) + } + }) as _) + .build() + .unwrap(); + ctx.leader_cached_kv_backend = + Arc::new(LeaderCachedKvBackend::with_always_leader(Arc::new(backend))); + let invalidator = Arc::new(FailingCacheInvalidator::default()); + ctx.cache_invalidator = invalidator.clone(); + let handler = CollectStatsHandler::default(); + let mut stat = Stat { + id: 101, + addr: "dn-original".into(), + node_epoch: 7, + ..Default::default() + }; + let req = HeartbeatRequest::default(); + + // A backend error must not interrupt subsequent heartbeat handlers. + assert!(matches!( + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + } + ) + .await + .unwrap(), + HandleControl::Continue + )); + assert_eq!(puts.lock().unwrap().len(), 1); + assert_eq!(invalidator.attempts(), 0); + + let stats_key: Vec = common_meta::datanode::DatanodeStatKey { node_id: 101 }.into(); + assert!(ctx.in_memory.get(&stats_key).await.unwrap().is_some()); + + // Same epoch retries persistence using the current heartbeat's address. + stat.addr = "dn-retry".into(); + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + { + let puts = puts.lock().unwrap(); + assert_eq!(puts.len(), 2); + assert_eq!(puts[1].key, NodeAddressKey::with_datanode(101).to_bytes()); + assert_eq!( + NodeAddressValue::try_from_raw_value(&puts[1].value) + .unwrap() + .peer + .addr, + "dn-retry" + ); + } + assert_eq!(invalidator.attempts(), 1); + + // Notification remains best-effort: its failure neither rewrites nor retries. + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(puts.lock().unwrap().len(), 2); + assert_eq!(invalidator.attempts(), 1); + + stat.node_epoch = 8; + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(puts.lock().unwrap().len(), 3); + assert_eq!(invalidator.attempts(), 2); + } } diff --git a/src/meta-srv/src/handler/remap_flow_peer_handler.rs b/src/meta-srv/src/handler/remap_flow_peer_handler.rs index f6b3c67a9e5..186cefa19e8 100644 --- a/src/meta-srv/src/handler/remap_flow_peer_handler.rs +++ b/src/meta-srv/src/handler/remap_flow_peer_handler.rs @@ -12,22 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{HeartbeatRequest, Peer, Role}; +use api::v1::meta::{HeartbeatRequest, Role}; use common_meta::instruction::CacheIdent; -use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue}; -use common_meta::key::{MetadataKey, MetadataValue}; -use common_meta::rpc::store::PutRequest; -use common_telemetry::{error, info, warn}; -use dashmap::DashMap; +use common_meta::key::MetadataKey; +use common_meta::key::node_address::NodeAddressKey; +use common_telemetry::error; use crate::Result; +use crate::handler::utils::{self, NodeAddressUpdater}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; #[derive(Debug, Default)] pub struct RemapFlowPeerHandler { - /// flow_node_id -> epoch - epoch_cache: DashMap, + address_updater: NodeAddressUpdater, } #[async_trait::async_trait] @@ -45,61 +43,83 @@ impl HeartbeatHandler for RemapFlowPeerHandler { let Some(peer) = req.peer.as_ref() else { return Ok(HandleControl::Continue); }; - - let current_epoch = req.node_epoch; - let flow_node_id = peer.id; - - let refresh = if let Some(mut epoch) = self.epoch_cache.get_mut(&flow_node_id) { - if current_epoch > *epoch.value() { - *epoch.value_mut() = current_epoch; - true - } else { - false - } - } else { - self.epoch_cache.insert(flow_node_id, current_epoch); - true - }; - - if refresh { - rewrite_node_address(ctx, peer).await; + let updated = self + .address_updater + .update_if_needed((Role::Flownode, peer.id), req.node_epoch, || { + let key = NodeAddressKey::with_flownode(peer.id).to_bytes(); + utils::save_node_address(ctx, key, peer.clone()) + }) + .await + .inspect_err(|e| { + // Address errors must not prevent later heartbeat handlers from running. + error!(e; "Failed to update flownode address, peer: {:?}", peer); + }); + if let Ok(true) = updated { + let cache_idents = [CacheIdent::FlowNodeAddressChange(peer.id)]; + utils::invalidate_address_caches(ctx, peer.id, &cache_idents).await; } - Ok(HandleControl::Continue) } } -async fn rewrite_node_address(ctx: &mut Context, peer: &Peer) { - let key = NodeAddressKey::with_flownode(peer.id).to_bytes(); - if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() { - let put = PutRequest { - key, - value, - prev_kv: false, - }; +#[cfg(test)] +mod tests { + use super::*; + use crate::handler::test_utils::TestEnv; + use api::v1::meta::Peer; + use std::sync::Arc; + #[tokio::test] + async fn test_notification_failure_keeps_persisted_epoch() { + use crate::handler::{HeartbeatAccumulator, HeartbeatHandler}; + use api::v1::meta::HeartbeatRequest; - match ctx.leader_cached_kv_backend.put(put).await { - Ok(_) => { - info!("Successfully updated flow `NodeAddressValue`: {:?}", peer); - // broadcast invalidating cache to all frontends - let cache_idents = vec![CacheIdent::FlowNodeAddressChange(peer.id)]; - info!( - "Invalidate flow node cache for new address with cache idents: {:?}", - cache_idents - ); - if let Err(e) = ctx - .cache_invalidator - .invalidate(&Default::default(), &cache_idents) - .await - { - error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer); - } - } - Err(e) => { - error!(e; "Failed to update flow `NodeAddressValue`: {:?}", peer); - } - } - } else { - warn!("Failed to serialize flow `NodeAddressValue`: {:?}", peer); + use crate::handler::test_utils::FailingCacheInvalidator; + use common_meta::key::MetadataValue; + use common_meta::key::node_address::NodeAddressValue; + let mut ctx = TestEnv::new().ctx(); + let invalidator = Arc::new(FailingCacheInvalidator::default()); + ctx.cache_invalidator = invalidator.clone(); + let handler = RemapFlowPeerHandler::default(); + let mut req = HeartbeatRequest { + peer: Some(Peer { + id: 1, + addr: "first".into(), + }), + node_epoch: 1, + ..Default::default() + }; + let key = NodeAddressKey::with_flownode(1).to_bytes(); + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + req.peer.as_mut().unwrap().addr = "same-epoch".into(); + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + assert_eq!(invalidator.attempts(), 1); + let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!( + NodeAddressValue::try_from_raw_value(&kv.value) + .unwrap() + .peer + .addr, + "first" + ); + req.node_epoch = 2; + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + assert_eq!(invalidator.attempts(), 2); + let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!( + NodeAddressValue::try_from_raw_value(&kv.value) + .unwrap() + .peer + .addr, + "same-epoch" + ); } } diff --git a/src/meta-srv/src/handler/test_utils.rs b/src/meta-srv/src/handler/test_utils.rs index d8241c785ea..c2afb3cb14a 100644 --- a/src/meta-srv/src/handler/test_utils.rs +++ b/src/meta-srv/src/handler/test_utils.rs @@ -13,9 +13,13 @@ // limitations under the License. use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; -use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; +use common_meta::cache_invalidator::{ + CacheInvalidator, CacheInvalidatorRef, Context as InvalidationContext, DummyCacheInvalidator, +}; use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL; +use common_meta::instruction::CacheIdent; use common_meta::key::{TableMetadataManager, TableMetadataManagerRef}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; @@ -97,3 +101,34 @@ impl TestEnv { } } } + +/// Counts failed notification attempts without requiring heartbeat connections. +#[derive(Default)] +pub(crate) struct FailingCacheInvalidator { + attempts: AtomicUsize, +} + +impl FailingCacheInvalidator { + pub(crate) fn attempts(&self) -> usize { + self.attempts.load(Ordering::Relaxed) + } +} + +#[async_trait::async_trait] +impl CacheInvalidator for FailingCacheInvalidator { + async fn invalidate( + &self, + _: &InvalidationContext, + _: &[CacheIdent], + ) -> common_meta::error::Result<()> { + self.attempts.fetch_add(1, Ordering::Relaxed); + common_meta::error::UnexpectedSnafu { + err_msg: "injected notification failure", + } + .fail() + } + + fn invalidate_all(&self) -> common_meta::error::Result<()> { + Ok(()) + } +} diff --git a/src/meta-srv/src/handler/utils.rs b/src/meta-srv/src/handler/utils.rs new file mode 100644 index 00000000000..96624d87306 --- /dev/null +++ b/src/meta-srv/src/handler/utils.rs @@ -0,0 +1,281 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::future::Future; +use std::sync::Arc; + +use api::v1::meta::{Peer, Role}; +use common_meta::instruction::CacheIdent; +use common_meta::key::MetadataValue; +use common_meta::key::node_address::NodeAddressValue; +use common_meta::rpc::store::PutRequest; +use common_telemetry::{error, info}; +use dashmap::DashMap; +use snafu::ResultExt; +use tokio::sync::Mutex; + +use crate::Result; +use crate::error::{InvalidNodeInfoFormatSnafu, KvBackendSnafu}; +use crate::metasrv::Context; + +/// Serializes address changes per node and acknowledges only successful updates. +/// Node information refreshed on every heartbeat has a separate lifecycle. +#[derive(Debug, Default)] +pub(crate) struct NodeAddressUpdater { + states: DashMap<(Role, u64), Arc>>, +} + +#[derive(Debug, Default)] +struct NodeAddressEpochs { + /// Rejects stale heartbeats even after a newer epoch failed to save. + seen_epoch: Option, + /// Suppresses retries only after persistence succeeds. + saved_epoch: Option, +} + +impl NodeAddressUpdater { + /// Saves a new or previously failed epoch. Returns true only when this call + /// saved the address, or false when the heartbeat is stale or already saved. + /// Cache invalidation is separate and does not affect acknowledgement. + pub(crate) async fn update_if_needed( + &self, + node: (Role, u64), + epoch: u64, + save: F, + ) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { + // Never retain a DashMap guard across await, even while waiting for a + // different node: its entry may live in the same shard. + let state = { + let entry = self.states.entry(node).or_default(); + Arc::clone(entry.value()) + }; + let mut state = state.lock().await; + if state.seen_epoch.is_some_and(|seen| epoch < seen) || state.saved_epoch == Some(epoch) { + return Ok(false); + } + + // Seeing an epoch rejects stale heartbeats, but only a successful save + // suppresses retries of that same epoch. Cancellation also leaves it retryable. + state.seen_epoch = Some(epoch); + save().await?; + state.saved_epoch = Some(epoch); + Ok(true) + } +} + +/// Persists an address through the leader cache. +pub(crate) async fn save_node_address(ctx: &Context, key: Vec, peer: Peer) -> Result<()> { + let address = NodeAddressValue::new(peer); + let value = address + .try_as_raw_value() + .context(InvalidNodeInfoFormatSnafu)?; + ctx.leader_cached_kv_backend + .put(PutRequest { + key, + value, + prev_kv: false, + }) + .await + .context(KvBackendSnafu)?; + info!("Successfully updated node address: {:?}", address.peer); + Ok(()) +} + +/// Preserves the existing best-effort notification behavior. Notification +/// failures do not undo a persisted address or cause same-epoch rewrites. +pub(crate) async fn invalidate_address_caches( + ctx: &Context, + node_id: u64, + cache_idents: &[CacheIdent], +) { + if let Err(e) = ctx + .cache_invalidator + .invalidate(&Default::default(), cache_idents) + .await + { + error!(e; "Failed to invalidate address caches for node {}: {:?}", node_id, cache_idents); + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use tokio::sync::Notify; + use tokio::time::timeout; + + use super::*; + + #[tokio::test] + async fn test_failed_save_retries_same_epoch_and_rejects_stale_heartbeats() { + let updater = NodeAddressUpdater::default(); + let key = (Role::Datanode, 1); + assert!( + updater + .update_if_needed(key, 2, || async { + crate::error::UnexpectedSnafu { + violated: "injected save failure", + } + .fail() + }) + .await + .is_err() + ); + assert!( + !updater + .update_if_needed(key, 1, || async { + panic!("stale heartbeat must not save"); + }) + .await + .unwrap() + ); + assert!( + updater + .update_if_needed(key, 2, || async { Ok(()) }) + .await + .unwrap() + ); + assert!( + !updater + .update_if_needed(key, 2, || async { + panic!("successful epoch must not save again"); + }) + .await + .unwrap() + ); + assert!( + updater + .update_if_needed(key, 3, || async { Ok(()) }) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_cancelled_save_can_be_retried() { + let updater = NodeAddressUpdater::default(); + let key = (Role::Datanode, 1); + let entered = Notify::new(); + { + let first = updater.update_if_needed(key, 0, || async { + entered.notify_one(); + std::future::pending::>().await + }); + tokio::pin!(first); + tokio::select! { + _ = entered.notified() => {}, + _ = &mut first => panic!("save must remain pending"), + } + } + assert!( + timeout( + Duration::from_secs(2), + updater.update_if_needed(key, 0, || async { Ok(()) }) + ) + .await + .unwrap() + .unwrap() + ); + } + + #[test] + fn test_same_node_serialized_other_nodes_progress() { + // A watchdog outside the runtime detects synchronous DashMap deadlocks; + // a Tokio timeout alone cannot fire if the only worker is blocked. + let (progress_tx, progress_rx) = std::sync::mpsc::sync_channel(1); + let release = Arc::new(Notify::new()); + let finish = Arc::new(Notify::new()); + let worker = std::thread::spawn({ + let release = release.clone(); + let finish = finish.clone(); + move || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap() + .block_on(async move { + let updater = Arc::new(NodeAddressUpdater::default()); + let entered = Arc::new(Notify::new()); + + let first = tokio::spawn({ + let updater = updater.clone(); + let entered = entered.clone(); + let release = release.clone(); + async move { + updater + .update_if_needed((Role::Datanode, 1), 1, || async { + entered.notify_one(); + release.notified().await; + Ok(()) + }) + .await + } + }); + entered.notified().await; + let second = tokio::spawn({ + let updater = updater.clone(); + async move { + updater + .update_if_needed((Role::Datanode, 1), 2, || async { Ok(()) }) + .await + } + }); + tokio::task::yield_now().await; + assert!(!second.is_finished()); + // Exercise shared DashMap shards on a single-thread runtime while the + // first node is suspended in external I/O. + timeout(Duration::from_secs(2), async { + for id in 0_u64..128 { + assert!( + updater + .update_if_needed((Role::Datanode, id + 2), 1, || async { + Ok(()) + }) + .await + .unwrap() + ); + } + }) + .await + .unwrap(); + progress_tx.send(()).unwrap(); + finish.notified().await; + timeout(Duration::from_secs(2), async { + assert!(first.await.unwrap().unwrap()); + assert!(second.await.unwrap().unwrap()); + }) + .await + .unwrap(); + assert!( + !updater + .update_if_needed((Role::Datanode, 1), 1, || async { + panic!("old epoch must not overwrite a newer one"); + }) + .await + .unwrap() + ); + }); + } + }); + let progress = progress_rx.recv_timeout(Duration::from_secs(5)); + release.notify_one(); + finish.notify_one(); + progress.expect("other node updates must progress while one node save is blocked"); + worker.join().unwrap(); + } +}