datanode/heartbeat/handler/
open_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::instruction::{InstructionReply, OpenRegion, SimpleReply};
16use common_meta::wal_provider::serialize_wal_options;
17use common_telemetry::info;
18use store_api::path_utils::table_dir;
19use store_api::region_request::{PathType, RegionOpenRequest};
20use store_api::storage::RegionId;
21
22use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
23
24pub struct OpenRegionsHandler {
25    pub open_region_parallelism: usize,
26}
27
28#[async_trait::async_trait]
29impl InstructionHandler for OpenRegionsHandler {
30    type Instruction = Vec<OpenRegion>;
31    async fn handle(
32        &self,
33        ctx: &HandlerContext,
34        open_regions: Self::Instruction,
35    ) -> Option<InstructionReply> {
36        let requests = open_regions
37            .into_iter()
38            .map(|open_region| {
39                let OpenRegion {
40                    region_ident,
41                    region_storage_path,
42                    mut region_options,
43                    region_wal_options,
44                    skip_wal_replay,
45                    reason,
46                    requirements,
47                } = open_region;
48                let region_id = RegionId::new(region_ident.table_id, region_ident.region_number);
49                info!(
50                    "Received open region instruction, region_id: {region_id}, reason: {reason:?}"
51                );
52                if let Err(err) =
53                    serialize_wal_options(&mut region_options, region_id, &region_wal_options)
54                {
55                    return Err(format!(
56                        "Failed to serialize WAL options for region {region_id}: {err:?}"
57                    ));
58                }
59                let request = RegionOpenRequest {
60                    engine: region_ident.engine,
61                    table_dir: table_dir(&region_storage_path, region_id.table_id()),
62                    path_type: PathType::Bare,
63                    options: region_options,
64                    skip_wal_replay,
65                    checkpoint: None,
66                    requirements,
67                };
68                Ok((region_id, request))
69            })
70            .collect::<Result<Vec<_>, _>>();
71        let requests = match requests {
72            Ok(requests) => requests,
73            Err(error) => {
74                return Some(InstructionReply::OpenRegions(SimpleReply {
75                    result: false,
76                    error: Some(error),
77                }));
78            }
79        };
80
81        let result = ctx
82            .region_server
83            .handle_batch_open_requests(self.open_region_parallelism, requests, false)
84            .await;
85        let success = result.is_ok();
86        let error = result.as_ref().map_err(|e| format!("{e:?}")).err();
87
88        Some(InstructionReply::OpenRegions(SimpleReply {
89            result: success,
90            error,
91        }))
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use std::assert_matches;
98    use std::collections::HashMap;
99    use std::sync::Arc;
100
101    use common_meta::RegionIdent;
102    use common_meta::heartbeat::handler::{HandleControl, HeartbeatResponseHandler};
103    use common_meta::heartbeat::mailbox::MessageMeta;
104    use common_meta::instruction::{Instruction, OpenRegion};
105    use common_meta::kv_backend::memory::MemoryKvBackend;
106    use mito2::config::MitoConfig;
107    use mito2::engine::MITO_ENGINE_NAME;
108    use mito2::test_util::{CreateRequestBuilder, TestEnv};
109    use store_api::path_utils::table_dir;
110    use store_api::region_request::{RegionCloseRequest, RegionRequest, RegionRequirements};
111    use store_api::storage::RegionId;
112
113    use crate::heartbeat::handler::RegionHeartbeatResponseHandler;
114    use crate::heartbeat::handler::tests::HeartbeatResponseTestEnv;
115    use crate::tests::mock_region_server;
116
117    fn open_regions_instruction(
118        region_ids: impl IntoIterator<Item = RegionId>,
119        storage_path: &str,
120    ) -> Instruction {
121        let region_idents = region_ids
122            .into_iter()
123            .map(|region_id| {
124                OpenRegion::new(
125                    RegionIdent {
126                        datanode_id: 0,
127                        table_id: region_id.table_id(),
128                        region_number: region_id.region_number(),
129                        engine: MITO_ENGINE_NAME.to_string(),
130                    },
131                    storage_path,
132                    HashMap::new(),
133                    HashMap::new(),
134                    false,
135                    None,
136                    RegionRequirements::empty(),
137                )
138            })
139            .collect();
140
141        Instruction::OpenRegions(region_idents)
142    }
143
144    #[tokio::test]
145    async fn test_open_regions() {
146        common_telemetry::init_default_ut_logging();
147
148        let mut region_server = mock_region_server();
149        let kv_backend = Arc::new(MemoryKvBackend::new());
150        let heartbeat_handler =
151            RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend);
152        let mut engine_env = TestEnv::with_prefix("open-regions").await;
153        let engine = engine_env.create_engine(MitoConfig::default()).await;
154        region_server.register_engine(Arc::new(engine.clone()));
155        let region_id = RegionId::new(1024, 1);
156        let region_id1 = RegionId::new(1024, 2);
157        let storage_path = "test";
158        let builder = CreateRequestBuilder::new();
159        let mut create_req = builder.build();
160        create_req.table_dir = table_dir(storage_path, region_id.table_id());
161        region_server
162            .handle_request(region_id, RegionRequest::Create(create_req))
163            .await
164            .unwrap();
165        let mut create_req1 = builder.build();
166        create_req1.table_dir = table_dir(storage_path, region_id1.table_id());
167        region_server
168            .handle_request(region_id1, RegionRequest::Create(create_req1))
169            .await
170            .unwrap();
171        region_server
172            .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {}))
173            .await
174            .unwrap();
175        region_server
176            .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {}))
177            .await
178            .unwrap();
179
180        let meta = MessageMeta::new_test(1, "test", "dn-1", "me-0");
181        let instruction = open_regions_instruction([region_id, region_id1], storage_path);
182        let mut heartbeat_env = HeartbeatResponseTestEnv::new();
183        let mut ctx = heartbeat_env.create_handler_ctx((meta, Default::default(), instruction));
184        let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
185        assert_matches!(control, HandleControl::Continue);
186        let (_, reply) = heartbeat_env.receiver.recv().await.unwrap();
187
188        let reply = reply.expect_open_regions_reply();
189        assert!(reply.result);
190        assert!(reply.error.is_none());
191
192        assert!(engine.is_region_exists(region_id));
193        assert!(engine.is_region_exists(region_id1));
194    }
195}
datanode/heartbeat/handler/open_region.rs

datanode/heartbeat/handler/
open_region.rs
