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

datanode/heartbeat/handler/
apply_staging_manifest.rs
