meta_srv/handler/
response_header_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, ResponseHeader, Role, PROTOCOL_VERSION};
16
17use crate::error::Result;
18use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
19use crate::metasrv::Context;
20
21pub struct ResponseHeaderHandler;
22
23#[async_trait::async_trait]
24impl HeartbeatHandler for ResponseHeaderHandler {
25    fn is_acceptable(&self, _: Role) -> bool {
26        true
27    }
28
29    async fn handle(
30        &self,
31        _req: &HeartbeatRequest,
32        _ctx: &mut Context,
33        acc: &mut HeartbeatAccumulator,
34    ) -> Result<HandleControl> {
35        let res_header = ResponseHeader {
36            protocol_version: PROTOCOL_VERSION,
37            ..Default::default()
38        };
39        acc.header = Some(res_header);
40        Ok(HandleControl::Continue)
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use api::v1::meta::RequestHeader;
47    use common_telemetry::tracing_context::W3cTrace;
48
49    use super::*;
50    use crate::handler::test_utils::TestEnv;
51
52    #[tokio::test]
53    async fn test_handle_heartbeat_resp_header() {
54        let env = TestEnv::new();
55        let mut ctx = env.ctx();
56
57        let req = HeartbeatRequest {
58            header: Some(RequestHeader::new(2, Role::Datanode, W3cTrace::new())),
59            ..Default::default()
60        };
61        let mut acc = HeartbeatAccumulator::default();
62
63        let response_handler = ResponseHeaderHandler {};
64        response_handler
65            .handle(&req, &mut ctx, &mut acc)
66            .await
67            .unwrap();
68    }
69}