datanode/heartbeat/handler/
remap_manifest.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::{
16    InstructionError, InstructionReply, RemapManifest, RemapManifestReply,
17};
18use common_telemetry::{error, info, warn};
19use store_api::region_engine::RemapManifestsRequest;
20
21use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
22
23pub struct RemapManifestHandler;
24
25#[async_trait::async_trait]
26impl InstructionHandler for RemapManifestHandler {
27    type Instruction = RemapManifest;
28    async fn handle(
29        &self,
30        ctx: &HandlerContext,
31        request: Self::Instruction,
32    ) -> Option<InstructionReply> {
33        let RemapManifest {
34            region_id,
35            input_regions,
36            region_mapping,
37            new_partition_exprs,
38        } = request;
39        info!(
40            "Datanode received remap manifest request, region_id: {}, input_regions: {}, target_regions: {}",
41            region_id,
42            input_regions.len(),
43            new_partition_exprs.len()
44        );
45        let Some(leader) = ctx.region_server.is_region_leader(region_id) else {
46            warn!("Region: {} is not found", region_id);
47            return Some(InstructionReply::RemapManifest(RemapManifestReply {
48                exists: false,
49                manifest_paths: Default::default(),
50                error: None,
51            }));
52        };
53
54        if !leader {
55            warn!("Region: {} is not leader", region_id);
56            return Some(InstructionReply::RemapManifest(RemapManifestReply {
57                exists: true,
58                manifest_paths: Default::default(),
59                error: Some(InstructionError::legacy_internal_retryable(
60                    "Region is not leader",
61                )),
62            }));
63        }
64
65        let reply = match ctx
66            .region_server
67            .remap_manifests(RemapManifestsRequest {
68                region_id,
69                input_regions,
70                region_mapping,
71                new_partition_exprs,
72            })
73            .await
74        {
75            Ok(result) => InstructionReply::RemapManifest(RemapManifestReply {
76                exists: true,
77                manifest_paths: result.manifest_paths,
78                error: None,
79            }),
80            Err(e) => {
81                error!(
82                    e;
83                    "Remap manifests failed on datanode, region_id: {}",
84                    region_id
85                );
86                InstructionReply::RemapManifest(RemapManifestReply {
87                    exists: true,
88                    manifest_paths: Default::default(),
89                    error: Some(InstructionError::from_error(&e)),
90                })
91            }
92        };
93
94        Some(reply)
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use std::collections::HashMap;
101    use std::sync::Arc;
102
103    use common_meta::instruction::RemapManifest;
104    use common_meta::kv_backend::memory::MemoryKvBackend;
105    use datatypes::value::Value;
106    use mito2::config::MitoConfig;
107    use mito2::engine::MITO_ENGINE_NAME;
108    use mito2::test_util::{CreateRequestBuilder, TestEnv};
109    use partition::expr::{PartitionExpr, col};
110    use store_api::path_utils::table_dir;
111    use store_api::region_engine::RegionRole;
112    use store_api::region_request::{
113        EnterStagingRequest, RegionRequest, StagingPartitionDirective,
114    };
115    use store_api::storage::RegionId;
116
117    use crate::heartbeat::handler::remap_manifest::RemapManifestHandler;
118    use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
119    use crate::region_server::RegionServer;
120    use crate::tests::{MockRegionEngine, mock_region_server};
121
122    #[tokio::test]
123    async fn test_region_not_exist() {
124        let mut mock_region_server = mock_region_server();
125        let (mock_engine, _) = MockRegionEngine::new(MITO_ENGINE_NAME);
126        mock_region_server.register_engine(mock_engine);
127        let kv_backend = Arc::new(MemoryKvBackend::new());
128        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
129        let region_id = RegionId::new(1024, 1);
130        let reply = RemapManifestHandler
131            .handle(
132                &handler_context,
133                RemapManifest {
134                    region_id,
135                    input_regions: vec![],
136                    region_mapping: HashMap::new(),
137                    new_partition_exprs: HashMap::new(),
138                },
139            )
140            .await
141            .unwrap();
142        let reply = &reply.expect_remap_manifest_reply();
143        assert!(!reply.exists);
144        assert!(reply.error.is_none());
145        assert!(reply.manifest_paths.is_empty());
146    }
147
148    #[tokio::test]
149    async fn test_region_not_leader() {
150        let mock_region_server = mock_region_server();
151        let region_id = RegionId::new(1024, 1);
152        let (mock_engine, _) =
153            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
154                region_engine.mock_role = Some(Some(RegionRole::Follower));
155                region_engine.handle_request_mock_fn = Some(Box::new(|_, _| Ok(0)));
156            });
157        mock_region_server.register_test_region(region_id, mock_engine);
158        let kv_backend = Arc::new(MemoryKvBackend::new());
159        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
160        let reply = RemapManifestHandler
161            .handle(
162                &handler_context,
163                RemapManifest {
164                    region_id,
165                    input_regions: vec![],
166                    region_mapping: HashMap::new(),
167                    new_partition_exprs: HashMap::new(),
168                },
169            )
170            .await
171            .unwrap();
172        let reply = reply.expect_remap_manifest_reply();
173        assert!(reply.exists);
174        assert!(reply.error.is_some());
175    }
176
177    fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr {
178        col(col_name)
179            .gt_eq(Value::Int64(start))
180            .and(col(col_name).lt(Value::Int64(end)))
181    }
182
183    async fn prepare_region(region_server: &RegionServer) {
184        let region_specs = [
185            (RegionId::new(1024, 1), range_expr("x", 0, 50)),
186            (RegionId::new(1024, 2), range_expr("x", 50, 100)),
187        ];
188
189        for (region_id, partition_expr) in region_specs {
190            let builder = CreateRequestBuilder::new();
191            let mut create_req = builder.build();
192            create_req.table_dir = table_dir("test", 1024);
193            region_server
194                .handle_request(region_id, RegionRequest::Create(create_req))
195                .await
196                .unwrap();
197            region_server
198                .handle_request(
199                    region_id,
200                    RegionRequest::EnterStaging(EnterStagingRequest {
201                        partition_directive: StagingPartitionDirective::UpdatePartitionExpr(
202                            partition_expr.as_json_str().unwrap(),
203                        ),
204                    }),
205                )
206                .await
207                .unwrap();
208        }
209    }
210
211    #[tokio::test]
212    async fn test_remap_manifest() {
213        common_telemetry::init_default_ut_logging();
214        let mut region_server = mock_region_server();
215        let region_id = RegionId::new(1024, 1);
216        let mut engine_env = TestEnv::new().await;
217        let engine = engine_env.create_engine(MitoConfig::default()).await;
218        region_server.register_engine(Arc::new(engine.clone()));
219        prepare_region(&region_server).await;
220
221        let kv_backend = Arc::new(MemoryKvBackend::new());
222        let handler_context = HandlerContext::new_for_test(region_server, kv_backend);
223        let region_id2 = RegionId::new(1024, 2);
224        let reply = RemapManifestHandler
225            .handle(
226                &handler_context,
227                RemapManifest {
228                    region_id,
229                    input_regions: vec![region_id, region_id2],
230                    region_mapping: HashMap::from([
231                        (region_id, vec![region_id]),
232                        (region_id2, vec![region_id]),
233                    ]),
234                    new_partition_exprs: HashMap::from([(
235                        region_id,
236                        range_expr("x", 0, 100).as_json_str().unwrap(),
237                    )]),
238                },
239            )
240            .await
241            .unwrap();
242        let reply = reply.expect_remap_manifest_reply();
243        assert!(reply.exists);
244        assert!(reply.error.is_none(), "{}", reply.error.unwrap());
245        assert_eq!(reply.manifest_paths.len(), 1);
246
247        // Remap failed
248        let reply = RemapManifestHandler
249            .handle(
250                &handler_context,
251                RemapManifest {
252                    region_id,
253                    input_regions: vec![region_id],
254                    region_mapping: HashMap::from([
255                        (region_id, vec![region_id]),
256                        (region_id2, vec![region_id]),
257                    ]),
258                    new_partition_exprs: HashMap::from([(
259                        region_id,
260                        range_expr("x", 0, 100).as_json_str().unwrap(),
261                    )]),
262                },
263            )
264            .await
265            .unwrap();
266        let reply = reply.expect_remap_manifest_reply();
267        assert!(reply.exists);
268        assert!(reply.error.is_some());
269        assert!(reply.manifest_paths.is_empty());
270    }
271}
datanode/heartbeat/handler/remap_manifest.rs

datanode/heartbeat/handler/
remap_manifest.rs
