fix(meta): backport flow metadata txn dedup to release/5e02efa9 (#9816) (#9823)

CREATE OR REPLACE FLOW deleted and then put the same flow route, flownode
flow and table flow keys in a single txn when the replacement kept the
partitions, flownodes or source tables. etcd rejects such a txn with
"duplicate key given in txn request", and the error surfaced to clients as
a masked "Internal error: 1003".

Only delete the old keys that are not put again, via a shared
build_replace_ops helper, so a replacement keeps its atomic txn while never
writing one key twice.

Adds regression coverage for unchanged, partially overlapping and fully
changed metadata keys. The tests use an isolated etcd chroot when
GT_ETCD_ENDPOINTS is set (only etcd rejects duplicate writes) and an
in-memory backend otherwise.

Closes #9810.

(cherry picked from commit da714ae72e)

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

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<T>> for FlowScoped<T>
     }
 }
 
+/// 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<Item = Vec<u8>>,
+    new_kvs: Vec<(Vec<u8>, Vec<u8>)>,
+) -> Vec<TxnOp> {
+    let new_keys = new_kvs
+        .iter()
+        .map(|(key, _)| key.as_slice())
+        .collect::<HashSet<_>>();
+    let mut ops = old_keys
+        .into_iter()
+        .filter(|key| !new_keys.contains(key.as_slice()))
+        .map(TxnOp::Delete)
+        .collect::<Vec<_>>();
+    ops.extend(
+        new_kvs
+            .into_iter()
+            .map(|(key, value)| TxnOp::Put(key, value)),
+    );
+    ops
+}
+
 pub type FlowMetadataManagerRef = Arc<FlowMetadataManager>;
 
 /// 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::<Vec<_>>();
+        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::<Vec<_>>();
+        assert_eq!(routes, expected_routes);
+
+        for flownode_id in 1..=4 {
+            let flows = manager
+                .flownode_flow_manager()
+                .flows(flownode_id)
+                .try_collect::<Vec<_>>()
+                .await
+                .unwrap();
+            let expected = flow_info
+                .flownode_ids()
+                .iter()
+                .filter(|(_, id)| **id == flownode_id)
+                .map(|(partition_id, _)| (flow_id, *partition_id))
+                .collect::<Vec<_>>();
+            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::<Vec<_>>();
+            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::<Vec<_>>()
+        };
+
+        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(&current))
+            .await
+            .unwrap();
+        assert_flow_metadata(&flow_metadata_manager, flow_id, &current).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;
+        }
+    }
 }

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<I: IntoIterator<Item = (FlowPartitionId, FlowRouteValue)>>(
         &self,
         flow_id: FlowId,
         current_flow_info: &DeserializedValueWithBytes<FlowInfoValue>,
         flow_routes: I,
     ) -> Result<Txn> {
-        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::<Result<Vec<_>>>()?;
 
-        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::<Result<Vec<_>>>()?;
-        Ok(Txn::new().and_then(txns))
+        Ok(Txn::new().and_then(build_replace_ops(old_keys, new_kvs)))
     }
 
     async fn remap_flow_route_addresses(

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<I: IntoIterator<Item = (FlowPartitionId, FlownodeId)>>(
         &self,
         flow_id: FlowId,
         current_flow_info: &DeserializedValueWithBytes<FlowInfoValue>,
         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::<Vec<_>>();
+        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))
     }
 }
 

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