meta_srv/handler/
flow_state_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::{FlowStat, HeartbeatRequest, Role};
16use common_meta::key::flow::flow_state::{FlowStateManager, FlowStateValue};
17use snafu::ResultExt;
18
19use crate::error::{FlowStateHandlerSnafu, Result};
20use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
21use crate::metasrv::Context;
22
23pub struct FlowStateHandler {
24    flow_state_manager: FlowStateManager,
25}
26
27impl FlowStateHandler {
28    pub fn new(flow_state_manager: FlowStateManager) -> Self {
29        Self { flow_state_manager }
30    }
31}
32
33#[async_trait::async_trait]
34impl HeartbeatHandler for FlowStateHandler {
35    fn is_acceptable(&self, role: Role) -> bool {
36        role == Role::Flownode
37    }
38
39    async fn handle(
40        &self,
41        req: &HeartbeatRequest,
42        _ctx: &mut Context,
43        _acc: &mut HeartbeatAccumulator,
44    ) -> Result<HandleControl> {
45        if let Some(FlowStat {
46            flow_stat_size,
47            flow_last_exec_time_map,
48        }) = &req.flow_stat
49        {
50            let state_size = flow_stat_size
51                .iter()
52                .map(|(k, v)| (*k, *v as usize))
53                .collect();
54            let last_exec_time_map = flow_last_exec_time_map
55                .iter()
56                .map(|(k, v)| (*k, *v))
57                .collect();
58            let value: FlowStateValue = FlowStateValue::new(state_size, last_exec_time_map);
59            self.flow_state_manager
60                .put(value)
61                .await
62                .context(FlowStateHandlerSnafu)?;
63        }
64        Ok(HandleControl::Continue)
65    }
66}