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

datanode/heartbeat/handler/
enter_staging.rs
