datanode/heartbeat/handler/
sync_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::{
16    InstructionError, InstructionReply, SyncRegion, SyncRegionReply, SyncRegionsReply,
17};
18use common_telemetry::{debug, error, info, warn};
19use futures::future::join_all;
20
21use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
22
23/// Handler for [SyncRegion] instruction.
24/// It syncs the region from a manifest or another region.
25#[derive(Debug, Clone, Copy, Default)]
26pub struct SyncRegionHandler;
27
28#[async_trait::async_trait]
29impl InstructionHandler for SyncRegionHandler {
30    type Instruction = Vec<SyncRegion>;
31
32    /// Handles a batch of [SyncRegion] instructions.
33    async fn handle(
34        &self,
35        ctx: &HandlerContext,
36        regions: Self::Instruction,
37    ) -> Option<InstructionReply> {
38        info!("Received sync region instructions: {:?}", regions);
39        let futures = regions
40            .into_iter()
41            .map(|sync_region| Self::handle_sync_region(ctx, sync_region));
42        let results = join_all(futures).await;
43        debug!("Sync region results: {:?}", results);
44
45        Some(InstructionReply::SyncRegions(SyncRegionsReply::new(
46            results,
47        )))
48    }
49}
50
51impl SyncRegionHandler {
52    /// Handles a single [SyncRegion] instruction.
53    async fn handle_sync_region(
54        ctx: &HandlerContext,
55        SyncRegion { region_id, request }: SyncRegion,
56    ) -> SyncRegionReply {
57        let Some(writable) = ctx.region_server.is_region_leader(region_id) else {
58            warn!("Region: {} is not found", region_id);
59            return SyncRegionReply {
60                region_id,
61                ready: false,
62                exists: false,
63                error: None,
64            };
65        };
66
67        if !writable {
68            warn!("Region: {} is not writable", region_id);
69            return SyncRegionReply {
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        match ctx.region_server.sync_region(region_id, request).await {
80            Ok(_) => {
81                info!("Successfully synced region: {}", region_id);
82                SyncRegionReply {
83                    region_id,
84                    ready: true,
85                    exists: true,
86                    error: None,
87                }
88            }
89            Err(e) => {
90                error!(e; "Failed to sync region: {}", region_id);
91                SyncRegionReply {
92                    region_id,
93                    ready: false,
94                    exists: true,
95                    error: Some(InstructionError::from_error(&e)),
96                }
97            }
98        }
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use std::sync::Arc;
105
106    use common_meta::kv_backend::memory::MemoryKvBackend;
107    use store_api::metric_engine_consts::METRIC_ENGINE_NAME;
108    use store_api::region_engine::{RegionRole, SyncRegionFromRequest};
109    use store_api::storage::RegionId;
110
111    use crate::heartbeat::handler::sync_region::SyncRegionHandler;
112    use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
113    use crate::tests::{MockRegionEngine, mock_region_server};
114
115    #[tokio::test]
116    async fn test_handle_sync_region_not_found() {
117        let mut mock_region_server = mock_region_server();
118        let (mock_engine, _) = MockRegionEngine::new(METRIC_ENGINE_NAME);
119        mock_region_server.register_engine(mock_engine);
120
121        let kv_backend = Arc::new(MemoryKvBackend::new());
122        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
123        let handler = SyncRegionHandler;
124
125        let region_id = RegionId::new(1024, 1);
126        let sync_region = common_meta::instruction::SyncRegion {
127            region_id,
128            request: SyncRegionFromRequest::from_manifest(Default::default()),
129        };
130
131        let reply = handler
132            .handle(&handler_context, vec![sync_region])
133            .await
134            .unwrap()
135            .expect_sync_regions_reply();
136
137        assert_eq!(reply.len(), 1);
138        assert_eq!(reply[0].region_id, region_id);
139        assert!(!reply[0].exists);
140        assert!(!reply[0].ready);
141    }
142
143    #[tokio::test]
144    async fn test_handle_sync_region_not_writable() {
145        let mock_region_server = mock_region_server();
146        let region_id = RegionId::new(1024, 1);
147        let (mock_engine, _) = MockRegionEngine::with_custom_apply_fn(METRIC_ENGINE_NAME, |r| {
148            r.mock_role = Some(Some(RegionRole::Follower));
149        });
150        mock_region_server.register_test_region(region_id, mock_engine);
151
152        let kv_backend = Arc::new(MemoryKvBackend::new());
153        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
154        let handler = SyncRegionHandler;
155
156        let sync_region = common_meta::instruction::SyncRegion {
157            region_id,
158            request: SyncRegionFromRequest::from_manifest(Default::default()),
159        };
160
161        let reply = handler
162            .handle(&handler_context, vec![sync_region])
163            .await
164            .unwrap()
165            .expect_sync_regions_reply();
166
167        assert_eq!(reply.len(), 1);
168        assert_eq!(reply[0].region_id, region_id);
169        assert!(reply[0].exists);
170        assert!(!reply[0].ready);
171        assert!(reply[0].error.is_some());
172    }
173
174    #[tokio::test]
175    async fn test_handle_sync_region_success() {
176        let mock_region_server = mock_region_server();
177        let region_id = RegionId::new(1024, 1);
178        let (mock_engine, _) = MockRegionEngine::with_custom_apply_fn(METRIC_ENGINE_NAME, |r| {
179            r.mock_role = Some(Some(RegionRole::Leader));
180        });
181        mock_region_server.register_test_region(region_id, mock_engine);
182
183        let kv_backend = Arc::new(MemoryKvBackend::new());
184        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
185        let handler = SyncRegionHandler;
186
187        let sync_region = common_meta::instruction::SyncRegion {
188            region_id,
189            request: SyncRegionFromRequest::from_manifest(Default::default()),
190        };
191
192        let reply = handler
193            .handle(&handler_context, vec![sync_region])
194            .await
195            .unwrap()
196            .expect_sync_regions_reply();
197
198        assert_eq!(reply.len(), 1);
199        assert_eq!(reply[0].region_id, region_id);
200        assert!(reply[0].exists);
201        assert!(reply[0].ready);
202        assert!(reply[0].error.is_none());
203    }
204}
datanode/heartbeat/handler/sync_region.rs

datanode/heartbeat/handler/
sync_region.rs
