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, Peer, Role};
16use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue};
17use common_meta::key::{MetadataKey, MetadataValue};
18use common_meta::rpc::store::PutRequest;
19use common_telemetry::{error, info, warn};
20use dashmap::DashMap;
21
22use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
23use crate::metasrv::Context;
24use crate::Result;
25
26#[derive(Debug, Default)]
27pub struct RemapFlowPeerHandler {
28    /// flow_node_id -> epoch
29    epoch_cache: DashMap<u64, u64>,
30}
31
32#[async_trait::async_trait]
33impl HeartbeatHandler for RemapFlowPeerHandler {
34    fn is_acceptable(&self, role: Role) -> bool {
35        role == Role::Flownode
36    }
37
38    async fn handle(
39        &self,
40        req: &HeartbeatRequest,
41        ctx: &mut Context,
42        _acc: &mut HeartbeatAccumulator,
43    ) -> Result<HandleControl> {
44        let Some(peer) = req.peer.as_ref() else {
45            return Ok(HandleControl::Continue);
46        };
47
48        let current_epoch = req.node_epoch;
49        let flow_node_id = peer.id;
50
51        let refresh = if let Some(mut epoch) = self.epoch_cache.get_mut(&flow_node_id) {
52            if current_epoch > *epoch.value() {
53                *epoch.value_mut() = current_epoch;
54                true
55            } else {
56                false
57            }
58        } else {
59            self.epoch_cache.insert(flow_node_id, current_epoch);
60            true
61        };
62
63        if refresh {
64            rewrite_node_address(ctx, peer).await;
65        }
66
67        Ok(HandleControl::Continue)
68    }
69}
70
71async fn rewrite_node_address(ctx: &mut Context, peer: &Peer) {
72    let key = NodeAddressKey::with_flownode(peer.id).to_bytes();
73    if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() {
74        let put = PutRequest {
75            key,
76            value,
77            prev_kv: false,
78        };
79
80        match ctx.leader_cached_kv_backend.put(put).await {
81            Ok(_) => {
82                info!("Successfully updated flow `NodeAddressValue`: {:?}", peer);
83                // TODO(discord): broadcast invalidating cache to all frontends
84            }
85            Err(e) => {
86                error!(e; "Failed to update flow `NodeAddressValue`: {:?}", peer);
87            }
88        }
89    } else {
90        warn!("Failed to serialize flow `NodeAddressValue`: {:?}", peer);
91    }
92}
meta_srv/handler/remap_flow_peer_handler.rs

meta_srv/handler/
remap_flow_peer_handler.rs
