meta_srv/handler/
remap_flow_peer_handler.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use api::v1::meta::{HeartbeatRequest, Role};
16use common_meta::instruction::CacheIdent;
17use common_meta::key::MetadataKey;
18use common_meta::key::node_address::NodeAddressKey;
19use common_telemetry::error;
20
21use crate::Result;
22use crate::handler::utils::{self, NodeAddressUpdater};
23use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
24use crate::metasrv::Context;
25
26#[derive(Debug, Default)]
27pub struct RemapFlowPeerHandler {
28    address_updater: NodeAddressUpdater,
29}
30
31#[async_trait::async_trait]
32impl HeartbeatHandler for RemapFlowPeerHandler {
33    fn is_acceptable(&self, role: Role) -> bool {
34        role == Role::Flownode
35    }
36
37    async fn handle(
38        &self,
39        req: &HeartbeatRequest,
40        ctx: &mut Context,
41        _acc: &mut HeartbeatAccumulator,
42    ) -> Result<HandleControl> {
43        let Some(peer) = req.peer.as_ref() else {
44            return Ok(HandleControl::Continue);
45        };
46        let updated = self
47            .address_updater
48            .update_if_needed((Role::Flownode, peer.id), req.node_epoch, || {
49                let key = NodeAddressKey::with_flownode(peer.id).to_bytes();
50                utils::save_node_address(ctx, key, peer.clone())
51            })
52            .await
53            .inspect_err(|e| {
54                // Address errors must not prevent later heartbeat handlers from running.
55                error!(e; "Failed to update flownode address, peer: {:?}", peer);
56            });
57        if let Ok(true) = updated {
58            let cache_idents = [CacheIdent::FlowNodeAddressChange(peer.id)];
59            utils::invalidate_address_caches(ctx, peer.id, &cache_idents).await;
60        }
61        Ok(HandleControl::Continue)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use crate::handler::test_utils::TestEnv;
69    use api::v1::meta::Peer;
70    use std::sync::Arc;
71    #[tokio::test]
72    async fn test_notification_failure_keeps_persisted_epoch() {
73        use crate::handler::{HeartbeatAccumulator, HeartbeatHandler};
74        use api::v1::meta::HeartbeatRequest;
75
76        use crate::handler::test_utils::FailingCacheInvalidator;
77        use common_meta::key::MetadataValue;
78        use common_meta::key::node_address::NodeAddressValue;
79        let mut ctx = TestEnv::new().ctx();
80        let invalidator = Arc::new(FailingCacheInvalidator::default());
81        ctx.cache_invalidator = invalidator.clone();
82        let handler = RemapFlowPeerHandler::default();
83        let mut req = HeartbeatRequest {
84            peer: Some(Peer {
85                id: 1,
86                addr: "first".into(),
87            }),
88            node_epoch: 1,
89            ..Default::default()
90        };
91        let key = NodeAddressKey::with_flownode(1).to_bytes();
92        handler
93            .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default())
94            .await
95            .unwrap();
96        req.peer.as_mut().unwrap().addr = "same-epoch".into();
97        handler
98            .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default())
99            .await
100            .unwrap();
101        assert_eq!(invalidator.attempts(), 1);
102        let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap();
103        assert_eq!(
104            NodeAddressValue::try_from_raw_value(&kv.value)
105                .unwrap()
106                .peer
107                .addr,
108            "first"
109        );
110        req.node_epoch = 2;
111        handler
112            .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default())
113            .await
114            .unwrap();
115        assert_eq!(invalidator.attempts(), 2);
116        let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap();
117        assert_eq!(
118            NodeAddressValue::try_from_raw_value(&kv.value)
119                .unwrap()
120                .peer
121                .addr,
122            "same-epoch"
123        );
124    }
125}
meta_srv/handler/remap_flow_peer_handler.rs

meta_srv/handler/
remap_flow_peer_handler.rs
