diff --git a/src/common/meta/src/key/flow.rs b/src/common/meta/src/key/flow.rs index 50c7e0ac38a..5ef0bbf39fb 100644 --- a/src/common/meta/src/key/flow.rs +++ b/src/common/meta/src/key/flow.rs @@ -19,7 +19,7 @@ pub mod flow_state; mod flownode_addr_helper; pub(crate) mod flownode_flow; pub(crate) mod table_flow; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashSet}; use std::ops::Deref; use std::sync::Arc; @@ -42,7 +42,7 @@ pub use crate::key::flow::table_flow::{TableFlowManager, TableFlowManagerRef}; use crate::key::txn_helper::TxnOpGetResponseSet; use crate::key::{DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey}; use crate::kv_backend::KvBackendRef; -use crate::kv_backend::txn::Txn; +use crate::kv_backend::txn::{Txn, TxnOp}; use crate::rpc::store::BatchDeleteRequest; pub const FLOW_KEY_PREFIX: &str = "__flow"; @@ -134,6 +134,31 @@ impl<'a, T: MetadataKey<'a, T>> MetadataKey<'a, FlowScoped> for FlowScoped } } +/// Builds txn ops that replace the `old_keys` with the `new_kvs`. +/// +/// Only deletes the old keys absent from `new_kvs`: etcd rejects a txn that +/// writes the same key more than once ("duplicate key given in txn request"). +fn build_replace_ops( + old_keys: impl IntoIterator>, + new_kvs: Vec<(Vec, Vec)>, +) -> Vec { + let new_keys = new_kvs + .iter() + .map(|(key, _)| key.as_slice()) + .collect::>(); + let mut ops = old_keys + .into_iter() + .filter(|key| !new_keys.contains(key.as_slice())) + .map(TxnOp::Delete) + .collect::>(); + ops.extend( + new_kvs + .into_iter() + .map(|(key, value)| TxnOp::Put(key, value)), + ); + ops +} + pub type FlowMetadataManagerRef = Arc; /// The manager of metadata, provides ability to: @@ -456,6 +481,7 @@ mod tests { use futures::TryStreamExt; use table::metadata::TableId; use table::table_name::TableName; + use uuid::Uuid; use super::*; use crate::FlownodeId; @@ -463,8 +489,10 @@ mod tests { use crate::key::flow::table_flow::TableFlowKey; use crate::key::node_address::{NodeAddressKey, NodeAddressValue}; use crate::key::{FlowPartitionId, MetadataValue}; - use crate::kv_backend::KvBackend; + use crate::kv_backend::chroot::ChrootKvBackend; + use crate::kv_backend::etcd::EtcdStore; use crate::kv_backend::memory::MemoryKvBackend; + use crate::kv_backend::{KvBackend, KvBackendRef}; use crate::peer::Peer; use crate::rpc::store::PutRequest; @@ -1189,4 +1217,151 @@ mod tests { err ); } + + #[test] + fn test_build_replace_ops() { + let ops = build_replace_ops( + vec![b"a".to_vec(), b"b".to_vec()], + vec![ + (b"b".to_vec(), b"1".to_vec()), + (b"c".to_vec(), b"2".to_vec()), + ], + ); + assert_eq!( + ops, + vec![ + TxnOp::Delete(b"a".to_vec()), + TxnOp::Put(b"b".to_vec(), b"1".to_vec()), + TxnOp::Put(b"c".to_vec(), b"2".to_vec()), + ] + ); + } + + /// Returns an etcd backend isolated by a chroot if `GT_ETCD_ENDPOINTS` is set, + /// because only etcd rejects txns that write the same key more than once. + async fn build_kv_backend() -> KvBackendRef { + let endpoints = std::env::var("GT_ETCD_ENDPOINTS").unwrap_or_default(); + if endpoints.is_empty() { + return Arc::new(MemoryKvBackend::default()); + } + let endpoints = endpoints.split(',').collect::>(); + let backend = EtcdStore::with_endpoints(endpoints, 128).await.unwrap(); + let chroot = format!("test_update_flow_metadata/{}/", Uuid::new_v4()); + Arc::new(ChrootKvBackend::new(chroot.into(), backend)) + } + + async fn assert_flow_metadata( + manager: &FlowMetadataManager, + flow_id: FlowId, + flow_info: &FlowInfoValue, + ) { + let got = manager.flow_info_manager().get(flow_id).await.unwrap(); + assert_eq!(got.as_ref(), Some(flow_info)); + + let routes = manager.flow_route_manager().routes(flow_id).await.unwrap(); + let expected_routes = flow_info + .flownode_ids() + .iter() + .map(|(partition_id, flownode_id)| { + ( + FlowRouteKey::new(flow_id, *partition_id), + FlowRouteValue { + peer: Peer::empty(*flownode_id), + }, + ) + }) + .collect::>(); + assert_eq!(routes, expected_routes); + + for flownode_id in 1..=4 { + let flows = manager + .flownode_flow_manager() + .flows(flownode_id) + .try_collect::>() + .await + .unwrap(); + let expected = flow_info + .flownode_ids() + .iter() + .filter(|(_, id)| **id == flownode_id) + .map(|(partition_id, _)| (flow_id, *partition_id)) + .collect::>(); + assert_eq!(flows, expected, "flownode {flownode_id}"); + } + + for table_id in 1024..=1027 { + let flows = manager + .table_flow_manager() + .flows(table_id) + .await + .unwrap() + .into_iter() + .map(|(key, _)| (key.flownode_id(), key.partition_id())) + .collect::>(); + let expected = if flow_info.source_table_ids().contains(&table_id) { + flow_info + .flownode_ids() + .iter() + .map(|(partition_id, flownode_id)| (*flownode_id, *partition_id)) + .collect() + } else { + vec![] + }; + assert_eq!(flows, expected, "table {table_id}"); + } + } + + #[tokio::test] + async fn test_update_flow_metadata_with_overlapping_keys() { + let flow_metadata_manager = FlowMetadataManager::new(build_kv_backend().await); + let flow_id = 10; + let to_routes = |flow_info: &FlowInfoValue| { + flow_info + .flownode_ids() + .iter() + .map(|(partition_id, flownode_id)| { + ( + *partition_id, + FlowRouteValue { + peer: Peer::empty(*flownode_id), + }, + ) + }) + .collect::>() + }; + + let mut current = + test_flow_info_value("flow", [(0, 1u64), (1, 2u64)].into(), vec![1024, 1025]); + flow_metadata_manager + .create_flow_metadata(flow_id, current.clone(), to_routes(¤t)) + .await + .unwrap(); + assert_flow_metadata(&flow_metadata_manager, flow_id, ¤t).await; + + let updates = [ + // Unchanged partitions, flownodes and source tables. + ([(0, 1u64), (1, 2u64)].into(), vec![1024, 1025]), + // Partially overlapping flownodes and source tables. + ([(0, 1u64), (1, 3u64)].into(), vec![1025, 1026]), + // Fully changed partitions, flownodes and source tables. + ([(2, 4u64)].into(), vec![1027]), + ]; + for (i, (flownode_ids, source_table_ids)) in updates.into_iter().enumerate() { + let mut new_flow_info = current.clone(); + new_flow_info.raw_sql = format!("update {i}"); + new_flow_info.flownode_ids = flownode_ids; + new_flow_info.source_table_ids = source_table_ids; + flow_metadata_manager + .update_flow_metadata( + flow_id, + &DeserializedValueWithBytes::from_inner(current.clone()), + &new_flow_info, + to_routes(&new_flow_info), + ) + .await + .unwrap(); + assert_flow_metadata(&flow_metadata_manager, flow_id, &new_flow_info).await; + current = new_flow_info; + } + } } diff --git a/src/common/meta/src/key/flow/flow_route.rs b/src/common/meta/src/key/flow/flow_route.rs index d402e1acef7..0cd62a3b909 100644 --- a/src/common/meta/src/key/flow/flow_route.rs +++ b/src/common/meta/src/key/flow/flow_route.rs @@ -20,7 +20,7 @@ use snafu::OptionExt; use crate::error::{self, Result}; use crate::key::flow::flow_info::FlowInfoValue; -use crate::key::flow::{FlowScoped, flownode_addr_helper}; +use crate::key::flow::{FlowScoped, build_replace_ops, flownode_addr_helper}; use crate::key::node_address::NodeAddressKey; use crate::key::{ BytesAdapter, DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey, MetadataValue, @@ -216,25 +216,26 @@ impl FlowRouteManager { /// Builds a update flow routes transaction. /// /// Puts `__flow/route/{flow_id}/{partition_id}` keys. - /// Also removes `__flow/route/{flow_id}/{old_partition_id}` keys. + /// Also removes the `__flow/route/{flow_id}/{old_partition_id}` keys that are not put. pub(crate) fn build_update_txn>( &self, flow_id: FlowId, current_flow_info: &DeserializedValueWithBytes, flow_routes: I, ) -> Result { - let del_txns = current_flow_info.flownode_ids().keys().map(|partition_id| { - let key = FlowRouteKey::new(flow_id, *partition_id).to_bytes(); - Ok(TxnOp::Delete(key)) - }); + let old_keys = current_flow_info + .flownode_ids() + .keys() + .map(|partition_id| FlowRouteKey::new(flow_id, *partition_id).to_bytes()); + let new_kvs = flow_routes + .into_iter() + .map(|(partition_id, route)| { + let key = FlowRouteKey::new(flow_id, partition_id).to_bytes(); + Ok((key, route.try_as_raw_value()?)) + }) + .collect::>>()?; - let put_txns = flow_routes.into_iter().map(|(partition_id, route)| { - let key = FlowRouteKey::new(flow_id, partition_id).to_bytes(); - - Ok(TxnOp::Put(key, route.try_as_raw_value()?)) - }); - let txns = del_txns.chain(put_txns).collect::>>()?; - Ok(Txn::new().and_then(txns)) + Ok(Txn::new().and_then(build_replace_ops(old_keys, new_kvs))) } async fn remap_flow_route_addresses( diff --git a/src/common/meta/src/key/flow/flownode_flow.rs b/src/common/meta/src/key/flow/flownode_flow.rs index 2ba690c167e..83d49229fac 100644 --- a/src/common/meta/src/key/flow/flownode_flow.rs +++ b/src/common/meta/src/key/flow/flownode_flow.rs @@ -20,8 +20,8 @@ use snafu::OptionExt; use crate::FlownodeId; use crate::error::{self, Result}; -use crate::key::flow::FlowScoped; use crate::key::flow::flow_info::FlowInfoValue; +use crate::key::flow::{FlowScoped, build_replace_ops}; use crate::key::{BytesAdapter, DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey}; use crate::kv_backend::KvBackendRef; use crate::kv_backend::txn::{Txn, TxnOp}; @@ -218,28 +218,30 @@ impl FlownodeFlowManager { /// Builds a update flownode flow transaction. /// /// Puts `__flownode_flow/{flownode_id}/{flow_id}/{partition_id}` keys. - /// Remove the old `__flownode_flow/{old_flownode_id}/{flow_id}/{old_partition_id}` keys. + /// Removes the old `__flownode_flow/{old_flownode_id}/{flow_id}/{old_partition_id}` + /// keys that are not put. pub(crate) fn build_update_txn>( &self, flow_id: FlowId, current_flow_info: &DeserializedValueWithBytes, flownode_ids: I, ) -> Txn { - let del_txns = + let old_keys = current_flow_info .flownode_ids() .iter() .map(|(partition_id, flownode_id)| { - let key = FlownodeFlowKey::new(*flownode_id, flow_id, *partition_id).to_bytes(); - TxnOp::Delete(key) + FlownodeFlowKey::new(*flownode_id, flow_id, *partition_id).to_bytes() }); - let put_txns = flownode_ids.into_iter().map(|(partition_id, flownode_id)| { - let key = FlownodeFlowKey::new(flownode_id, flow_id, partition_id).to_bytes(); - TxnOp::Put(key, vec![]) - }); - let txns = del_txns.chain(put_txns).collect::>(); + let new_kvs = flownode_ids + .into_iter() + .map(|(partition_id, flownode_id)| { + let key = FlownodeFlowKey::new(flownode_id, flow_id, partition_id).to_bytes(); + (key, vec![]) + }) + .collect(); - Txn::new().and_then(txns) + Txn::new().and_then(build_replace_ops(old_keys, new_kvs)) } } diff --git a/src/common/meta/src/key/flow/table_flow.rs b/src/common/meta/src/key/flow/table_flow.rs index d334ca02181..ae0fcb525f1 100644 --- a/src/common/meta/src/key/flow/table_flow.rs +++ b/src/common/meta/src/key/flow/table_flow.rs @@ -24,7 +24,7 @@ use table::metadata::TableId; use crate::FlownodeId; use crate::error::{self, Result}; use crate::key::flow::flow_info::FlowInfoValue; -use crate::key::flow::{FlowScoped, flownode_addr_helper}; +use crate::key::flow::{FlowScoped, build_replace_ops, flownode_addr_helper}; use crate::key::node_address::NodeAddressKey; use crate::key::{ BytesAdapter, DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey, MetadataValue, @@ -245,8 +245,9 @@ impl TableFlowManager { /// Builds a update table flow transaction. /// /// Puts `__flow/source_table/{table_id}/{node_id}/{flow_id}/{partition_id}` keys, - /// Also remove previous - /// `__flow/source_table/{table_id}/{old_node_id}/{flow_id}/{partition_id}` keys. + /// Also removes the previous + /// `__flow/source_table/{table_id}/{old_node_id}/{flow_id}/{partition_id}` keys + /// that are not put. pub fn build_update_txn( &self, flow_id: FlowId, @@ -254,22 +255,24 @@ impl TableFlowManager { table_flow_values: Vec<(FlowPartitionId, TableFlowValue)>, source_table_ids: &[TableId], ) -> Result { - let mut txns = Vec::with_capacity(2 * source_table_ids.len() * table_flow_values.len()); - - // first remove the old keys - for (part_id, node_id) in current_flow_info.flownode_ids() { - for source_table_id in current_flow_info.source_table_ids() { - txns.push(TxnOp::Delete( - TableFlowKey::new(*source_table_id, *node_id, flow_id, *part_id).to_bytes(), - )); - } - } + let old_keys = current_flow_info + .flownode_ids() + .iter() + .flat_map(|(part_id, node_id)| { + current_flow_info + .source_table_ids() + .iter() + .map(move |source_table_id| { + TableFlowKey::new(*source_table_id, *node_id, flow_id, *part_id).to_bytes() + }) + }); + let mut new_kvs = Vec::with_capacity(source_table_ids.len() * table_flow_values.len()); for (partition_id, table_flow_value) in table_flow_values { let flownode_id = table_flow_value.peer.id; let value = table_flow_value.try_as_raw_value()?; for source_table_id in source_table_ids { - txns.push(TxnOp::Put( + new_kvs.push(( TableFlowKey::new(*source_table_id, flownode_id, flow_id, partition_id) .to_bytes(), value.clone(), @@ -277,7 +280,7 @@ impl TableFlowManager { } } - Ok(Txn::new().and_then(txns)) + Ok(Txn::new().and_then(build_replace_ops(old_keys, new_kvs))) } async fn remap_table_flow_addresses(