datanode/heartbeat/handler/
close_region.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 common_meta::RegionIdent;
16use common_meta::instruction::{InstructionError, InstructionReply, SimpleReply};
17use common_telemetry::warn;
18use futures::future::join_all;
19use store_api::region_request::{RegionCloseRequest, RegionRequest};
20use store_api::storage::RegionId;
21
22use crate::error;
23use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
24
25#[derive(Debug, Clone, Copy, Default)]
26pub struct CloseRegionsHandler;
27
28#[async_trait::async_trait]
29impl InstructionHandler for CloseRegionsHandler {
30    type Instruction = Vec<RegionIdent>;
31
32    async fn handle(
33        &self,
34        ctx: &HandlerContext,
35        region_idents: Self::Instruction,
36    ) -> Option<InstructionReply> {
37        let region_ids = region_idents
38            .into_iter()
39            .map(|region_ident| RegionId::new(region_ident.table_id, region_ident.region_number))
40            .collect::<Vec<_>>();
41
42        let futs = region_ids.iter().map(|region_id| {
43            ctx.region_server.handle_request(
44                *region_id,
45                RegionRequest::Close(RegionCloseRequest::default()),
46            )
47        });
48
49        let results = join_all(futs).await;
50
51        let mut errors = vec![];
52        for (region_id, result) in region_ids.into_iter().zip(results) {
53            match result {
54                Ok(_) => (),
55                Err(error::Error::RegionNotFound { .. }) => {
56                    warn!(
57                        "Received a close regions instruction from meta, but target region:{} is not found.",
58                        region_id
59                    );
60                }
61                Err(err) => errors.push(format!("region:{region_id}: {err:?}")),
62            }
63        }
64
65        if errors.is_empty() {
66            return Some(InstructionReply::CloseRegions(SimpleReply {
67                result: true,
68                error: None,
69            }));
70        }
71
72        Some(InstructionReply::CloseRegions(SimpleReply {
73            result: false,
74            error: Some(InstructionError::legacy_internal_retryable(
75                errors.join("; "),
76            )),
77        }))
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use std::assert_matches;
84    use std::sync::Arc;
85
86    use common_meta::RegionIdent;
87    use common_meta::heartbeat::handler::{HandleControl, HeartbeatResponseHandler};
88    use common_meta::heartbeat::mailbox::MessageMeta;
89    use common_meta::instruction::Instruction;
90    use common_meta::kv_backend::memory::MemoryKvBackend;
91    use mito2::config::MitoConfig;
92    use mito2::engine::MITO_ENGINE_NAME;
93    use mito2::test_util::{CreateRequestBuilder, TestEnv};
94    use store_api::region_request::RegionRequest;
95    use store_api::storage::RegionId;
96
97    use crate::heartbeat::handler::RegionHeartbeatResponseHandler;
98    use crate::heartbeat::handler::tests::HeartbeatResponseTestEnv;
99    use crate::tests::mock_region_server;
100
101    fn close_regions_instruction(region_ids: impl IntoIterator<Item = RegionId>) -> Instruction {
102        let region_idents = region_ids
103            .into_iter()
104            .map(|region_id| RegionIdent {
105                table_id: region_id.table_id(),
106                region_number: region_id.region_number(),
107                datanode_id: 2,
108                engine: MITO_ENGINE_NAME.to_string(),
109            })
110            .collect();
111
112        Instruction::CloseRegions(region_idents)
113    }
114
115    #[tokio::test]
116    async fn test_close_regions() {
117        common_telemetry::init_default_ut_logging();
118
119        let mut region_server = mock_region_server();
120        let kv_backend = Arc::new(MemoryKvBackend::new());
121        let heartbeat_handler =
122            RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend);
123        let mut engine_env = TestEnv::with_prefix("close-regions").await;
124        let engine = engine_env.create_engine(MitoConfig::default()).await;
125        region_server.register_engine(Arc::new(engine.clone()));
126        let region_id = RegionId::new(1024, 1);
127        let region_id1 = RegionId::new(1024, 2);
128
129        let builder = CreateRequestBuilder::new();
130        let create_req = builder.build();
131        region_server
132            .handle_request(region_id, RegionRequest::Create(create_req))
133            .await
134            .unwrap();
135
136        let create_req1 = builder.build();
137        region_server
138            .handle_request(region_id1, RegionRequest::Create(create_req1))
139            .await
140            .unwrap();
141        let meta = MessageMeta::new_test(1, "test", "dn-1", "meta-0");
142        let instruction =
143            close_regions_instruction([region_id, region_id1, RegionId::new(1024, 3)]);
144        let mut heartbeat_env = HeartbeatResponseTestEnv::new();
145        let mut ctx = heartbeat_env.create_handler_ctx((meta, Default::default(), instruction));
146        let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
147        assert_matches!(control, HandleControl::Continue);
148
149        let (_, reply) = heartbeat_env.receiver.recv().await.unwrap();
150        let reply = reply.expect_close_regions_reply();
151        assert!(reply.result);
152        assert!(reply.error.is_none());
153        assert!(!engine.is_region_exists(region_id));
154        assert!(!engine.is_region_exists(region_id1));
155        assert!(!engine.is_region_exists(RegionId::new(1024, 3)));
156    }
157}
datanode/heartbeat/handler/close_region.rs

datanode/heartbeat/handler/
close_region.rs
