1use std::collections::HashMap;
16use std::sync::Arc;
17
18use api::v1::meta::{HeartbeatRequest, RegionLease, Role};
19use async_trait::async_trait;
20use common_meta::key::TableMetadataManagerRef;
21use common_meta::region_keeper::MemoryRegionKeeperRef;
22use common_telemetry::error;
23use store_api::region_engine::GrantedRegion;
24use store_api::storage::RegionId;
25
26use crate::error::Result;
27use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
28use crate::metasrv::Context;
29use crate::region::RegionLeaseKeeper;
30use crate::region::lease_keeper::{
31 RegionLeaseInfo, RegionLeaseKeeperRef, RenewRegionLeasesResponse,
32};
33
34pub struct RegionLeaseHandler {
35 region_lease_seconds: u64,
36 region_lease_keeper: RegionLeaseKeeperRef,
37 customized_region_lease_renewer: Option<CustomizedRegionLeaseRenewerRef>,
38}
39
40pub type CustomizedRegionLeaseRenewerRef = Arc<dyn CustomizedRegionLeaseRenewer>;
41
42pub trait CustomizedRegionLeaseRenewer: Send + Sync {
43 fn renew(
44 &self,
45 ctx: &mut Context,
46 regions: HashMap<RegionId, RegionLeaseInfo>,
47 ) -> Vec<GrantedRegion>;
48}
49
50impl RegionLeaseHandler {
51 pub fn new(
52 region_lease_seconds: u64,
53 table_metadata_manager: TableMetadataManagerRef,
54 memory_region_keeper: MemoryRegionKeeperRef,
55 customized_region_lease_renewer: Option<CustomizedRegionLeaseRenewerRef>,
56 ) -> Self {
57 let region_lease_keeper =
58 RegionLeaseKeeper::new(table_metadata_manager, memory_region_keeper.clone());
59
60 Self {
61 region_lease_seconds,
62 region_lease_keeper: Arc::new(region_lease_keeper),
63 customized_region_lease_renewer,
64 }
65 }
66}
67
68#[async_trait]
69impl HeartbeatHandler for RegionLeaseHandler {
70 fn is_acceptable(&self, role: Role) -> bool {
71 role == Role::Datanode
72 }
73
74 async fn handle(
75 &self,
76 req: &HeartbeatRequest,
77 ctx: &mut Context,
78 acc: &mut HeartbeatAccumulator,
79 ) -> Result<HandleControl> {
80 let Some(stat) = acc.stat.as_ref() else {
81 return Ok(HandleControl::Continue);
82 };
83
84 let regions = stat.regions();
85 let datanode_id = stat.id;
86
87 match self
88 .region_lease_keeper
89 .renew_region_leases(datanode_id, ®ions)
90 .await
91 {
92 Ok(RenewRegionLeasesResponse {
93 non_exists,
94 renewed,
95 }) => {
96 let renewed = if let Some(renewer) = &self.customized_region_lease_renewer {
97 renewer
98 .renew(ctx, renewed)
99 .into_iter()
100 .map(|region| region.into())
101 .collect()
102 } else {
103 renewed
104 .into_iter()
105 .map(|(region_id, region_lease_info)| {
106 GrantedRegion::new(region_id, region_lease_info.role).into()
107 })
108 .collect::<Vec<_>>()
109 };
110
111 acc.region_lease = Some(RegionLease {
112 regions: renewed,
113 duration_since_epoch: req.duration_since_epoch,
114 lease_seconds: self.region_lease_seconds,
115 closeable_region_ids: non_exists.iter().map(|region| region.as_u64()).collect(),
116 });
117 acc.inactive_region_ids = non_exists;
118 }
119 Err(e) => {
120 error!(e; "Failed to renew region leases for datanode: {datanode_id:?}, regions: {:?}", regions);
121 }
124 }
125
126 Ok(HandleControl::Continue)
127 }
128}
129
130#[cfg(test)]
131mod test {
132
133 use std::collections::{HashMap, HashSet};
134 use std::sync::Arc;
135
136 use common_meta::datanode::{RegionManifestInfo, RegionStat, Stat};
137 use common_meta::distributed_time_constants::default_distributed_time_constants;
138 use common_meta::key::TableMetadataManager;
139 use common_meta::key::table_route::TableRouteValue;
140 use common_meta::key::test_utils::new_test_table_info;
141 use common_meta::kv_backend::memory::MemoryKvBackend;
142 use common_meta::kv_backend::test_util::MockKvBackendBuilder;
143 use common_meta::peer::Peer;
144 use common_meta::region_keeper::MemoryRegionKeeper;
145 use common_meta::rpc::router::{LeaderState, Region, RegionRoute};
146 use store_api::region_engine::RegionRole;
147 use store_api::storage::RegionId;
148
149 use super::*;
150 use crate::metasrv::builder::MetasrvBuilder;
151
152 fn new_test_keeper() -> RegionLeaseKeeper {
153 let store = Arc::new(MemoryKvBackend::new());
154
155 let table_metadata_manager = Arc::new(TableMetadataManager::new(store));
156
157 let memory_region_keeper = Arc::new(MemoryRegionKeeper::default());
158 RegionLeaseKeeper::new(table_metadata_manager, memory_region_keeper)
159 }
160
161 fn new_empty_region_stat(region_id: RegionId, role: RegionRole) -> RegionStat {
162 RegionStat {
163 id: region_id,
164 role,
165 rcus: 0,
166 wcus: 0,
167 approximate_bytes: 0,
168 engine: String::new(),
169 num_rows: 0,
170 memtable_size: 0,
171 manifest_size: 0,
172 sst_size: 0,
173 sst_num: 0,
174 index_size: 0,
175 region_manifest: RegionManifestInfo::Mito {
176 manifest_version: 0,
177 flushed_entry_id: 0,
178 file_removed_cnt: 0,
179 },
180 data_topic_latest_entry_id: 0,
181 metadata_topic_latest_entry_id: 0,
182 written_bytes: 0,
183 query_cpu_time: 0,
184 query_scanned_bytes: 0,
185 }
186 }
187
188 #[tokio::test]
189 async fn test_handle_upgradable_follower() {
190 let datanode_id = 1;
191 let region_number = 1u32;
192 let table_id = 10;
193 let region_id = RegionId::new(table_id, region_number);
194 let another_region_id = RegionId::new(table_id, region_number + 1);
195 let peer = Peer::empty(datanode_id);
196 let follower_peer = Peer::empty(datanode_id + 1);
197 let table_info = new_test_table_info(table_id);
198
199 let region_routes = vec![RegionRoute {
200 region: Region::new_test(region_id),
201 leader_peer: Some(peer.clone()),
202 follower_peers: vec![follower_peer.clone()],
203 ..Default::default()
204 }];
205
206 let keeper = new_test_keeper();
207 let table_metadata_manager = keeper.table_metadata_manager();
208
209 table_metadata_manager
210 .create_table_metadata(
211 table_info,
212 TableRouteValue::physical(region_routes),
213 HashMap::default(),
214 )
215 .await
216 .unwrap();
217
218 let builder = MetasrvBuilder::new();
219 let metasrv = builder.build().await.unwrap();
220 let ctx = &mut metasrv.new_ctx();
221
222 let acc = &mut HeartbeatAccumulator::default();
223
224 acc.stat = Some(Stat {
225 id: peer.id,
226 region_stats: vec![
227 new_empty_region_stat(region_id, RegionRole::Follower),
228 new_empty_region_stat(another_region_id, RegionRole::Follower),
229 ],
230 ..Default::default()
231 });
232
233 let req = HeartbeatRequest {
234 duration_since_epoch: 1234,
235 ..Default::default()
236 };
237
238 let opening_region_keeper = Arc::new(MemoryRegionKeeper::default());
239
240 let handler = RegionLeaseHandler::new(
241 default_distributed_time_constants().region_lease.as_secs(),
242 table_metadata_manager.clone(),
243 opening_region_keeper.clone(),
244 None,
245 );
246
247 handler.handle(&req, ctx, acc).await.unwrap();
248
249 assert_region_lease(acc, vec![GrantedRegion::new(region_id, RegionRole::Leader)]);
250 assert_eq!(acc.inactive_region_ids, HashSet::from([another_region_id]));
251 assert_eq!(
252 acc.region_lease.as_ref().unwrap().closeable_region_ids,
253 vec![another_region_id]
254 );
255
256 let acc = &mut HeartbeatAccumulator::default();
257
258 acc.stat = Some(Stat {
259 id: follower_peer.id,
260 region_stats: vec![
261 new_empty_region_stat(region_id, RegionRole::Follower),
262 new_empty_region_stat(another_region_id, RegionRole::Follower),
263 ],
264 ..Default::default()
265 });
266
267 handler.handle(&req, ctx, acc).await.unwrap();
268
269 assert_eq!(
270 acc.region_lease.as_ref().unwrap().lease_seconds,
271 default_distributed_time_constants().region_lease.as_secs()
272 );
273
274 assert_region_lease(
275 acc,
276 vec![GrantedRegion::new(region_id, RegionRole::Follower)],
277 );
278 assert_eq!(acc.inactive_region_ids, HashSet::from([another_region_id]));
279 assert_eq!(
280 acc.region_lease.as_ref().unwrap().closeable_region_ids,
281 vec![another_region_id]
282 );
283
284 let opening_region_id = RegionId::new(table_id, region_number + 2);
285 let _guard = opening_region_keeper
286 .register_with_role(follower_peer.id, opening_region_id, RegionRole::Follower)
287 .unwrap();
288
289 let acc = &mut HeartbeatAccumulator::default();
290
291 acc.stat = Some(Stat {
292 id: follower_peer.id,
293 region_stats: vec![
294 new_empty_region_stat(region_id, RegionRole::Follower),
295 new_empty_region_stat(another_region_id, RegionRole::Follower),
296 new_empty_region_stat(opening_region_id, RegionRole::Follower),
297 ],
298 ..Default::default()
299 });
300
301 handler.handle(&req, ctx, acc).await.unwrap();
302
303 assert_eq!(
304 acc.region_lease.as_ref().unwrap().lease_seconds,
305 default_distributed_time_constants().region_lease.as_secs()
306 );
307
308 assert_region_lease(
309 acc,
310 vec![
311 GrantedRegion::new(region_id, RegionRole::Follower),
312 GrantedRegion::new(opening_region_id, RegionRole::Follower),
313 ],
314 );
315 assert_eq!(acc.inactive_region_ids, HashSet::from([another_region_id]));
316 assert_eq!(
317 acc.region_lease.as_ref().unwrap().closeable_region_ids,
318 vec![another_region_id]
319 );
320 }
321
322 #[tokio::test]
323
324 async fn test_handle_downgradable_leader() {
325 let datanode_id = 1;
326 let region_number = 1u32;
327 let table_id = 10;
328 let region_id = RegionId::new(table_id, region_number);
329 let another_region_id = RegionId::new(table_id, region_number + 1);
330 let no_exist_region_id = RegionId::new(table_id, region_number + 2);
331 let peer = Peer::empty(datanode_id);
332 let follower_peer = Peer::empty(datanode_id + 1);
333 let table_info = new_test_table_info(table_id);
334
335 let region_routes = vec![
336 RegionRoute {
337 region: Region::new_test(region_id),
338 leader_peer: Some(peer.clone()),
339 follower_peers: vec![follower_peer.clone()],
340 leader_state: Some(LeaderState::Downgrading),
341 leader_down_since: Some(1),
342 write_route_policy: None,
343 },
344 RegionRoute {
345 region: Region::new_test(another_region_id),
346 leader_peer: Some(peer.clone()),
347 ..Default::default()
348 },
349 ];
350
351 let keeper = new_test_keeper();
352 let table_metadata_manager = keeper.table_metadata_manager();
353
354 table_metadata_manager
355 .create_table_metadata(
356 table_info,
357 TableRouteValue::physical(region_routes),
358 HashMap::default(),
359 )
360 .await
361 .unwrap();
362
363 let builder = MetasrvBuilder::new();
364 let metasrv = builder.build().await.unwrap();
365 let ctx = &mut metasrv.new_ctx();
366
367 let req = HeartbeatRequest {
368 duration_since_epoch: 1234,
369 ..Default::default()
370 };
371
372 let acc = &mut HeartbeatAccumulator::default();
373
374 acc.stat = Some(Stat {
375 id: peer.id,
376 region_stats: vec![
377 new_empty_region_stat(region_id, RegionRole::Leader),
378 new_empty_region_stat(another_region_id, RegionRole::Leader),
379 new_empty_region_stat(no_exist_region_id, RegionRole::Leader),
380 ],
381 ..Default::default()
382 });
383
384 let handler = RegionLeaseHandler::new(
385 default_distributed_time_constants().region_lease.as_secs(),
386 table_metadata_manager.clone(),
387 Default::default(),
388 None,
389 );
390
391 handler.handle(&req, ctx, acc).await.unwrap();
392
393 assert_region_lease(
394 acc,
395 vec![
396 GrantedRegion::new(region_id, RegionRole::DowngradingLeader),
397 GrantedRegion::new(another_region_id, RegionRole::Leader),
398 ],
399 );
400 assert_eq!(acc.inactive_region_ids, HashSet::from([no_exist_region_id]));
401 }
402
403 #[tokio::test]
404 async fn test_handle_staging_leader() {
405 let datanode_id = 1;
406 let region_number = 1u32;
407 let table_id = 10;
408 let region_id = RegionId::new(table_id, region_number);
409 let peer = Peer::empty(datanode_id);
410 let table_info = new_test_table_info(table_id);
411
412 let region_routes = vec![RegionRoute {
413 region: Region::new_test(region_id),
414 leader_peer: Some(peer.clone()),
415 leader_state: Some(LeaderState::Staging),
416 ..Default::default()
417 }];
418
419 let keeper = new_test_keeper();
420 let table_metadata_manager = keeper.table_metadata_manager();
421
422 table_metadata_manager
423 .create_table_metadata(
424 table_info,
425 TableRouteValue::physical(region_routes),
426 HashMap::default(),
427 )
428 .await
429 .unwrap();
430
431 let builder = MetasrvBuilder::new();
432 let metasrv = builder.build().await.unwrap();
433 let ctx = &mut metasrv.new_ctx();
434
435 let req = HeartbeatRequest {
436 duration_since_epoch: 1234,
437 ..Default::default()
438 };
439
440 let acc = &mut HeartbeatAccumulator::default();
441 acc.stat = Some(Stat {
442 id: peer.id,
443 region_stats: vec![new_empty_region_stat(region_id, RegionRole::StagingLeader)],
444 ..Default::default()
445 });
446
447 let handler = RegionLeaseHandler::new(
448 default_distributed_time_constants().region_lease.as_secs(),
449 table_metadata_manager.clone(),
450 Default::default(),
451 None,
452 );
453
454 handler.handle(&req, ctx, acc).await.unwrap();
455
456 assert_region_lease(
457 acc,
458 vec![GrantedRegion::new(region_id, RegionRole::StagingLeader)],
459 );
460 }
461
462 fn assert_region_lease(acc: &HeartbeatAccumulator, expected: Vec<GrantedRegion>) {
463 let region_lease = acc.region_lease.as_ref().unwrap().clone();
464 let granted: Vec<GrantedRegion> = region_lease
465 .regions
466 .into_iter()
467 .map(Into::into)
468 .collect::<Vec<_>>();
469
470 let granted = granted
471 .into_iter()
472 .map(|region| (region.region_id, region))
473 .collect::<HashMap<_, _>>();
474
475 let expected = expected
476 .into_iter()
477 .map(|region| (region.region_id, region))
478 .collect::<HashMap<_, _>>();
479
480 assert_eq!(granted, expected);
481 }
482
483 #[tokio::test]
484 async fn test_handle_renew_region_lease_failure() {
485 common_telemetry::init_default_ut_logging();
486 let kv = MockKvBackendBuilder::default()
487 .batch_get_fn(Arc::new(|_| {
488 common_meta::error::UnexpectedSnafu {
489 err_msg: "mock err",
490 }
491 .fail()
492 }) as _)
493 .build()
494 .unwrap();
495 let kvbackend = Arc::new(kv);
496 let table_metadata_manager = Arc::new(TableMetadataManager::new(kvbackend));
497
498 let datanode_id = 1;
499 let region_number = 1u32;
500 let table_id = 10;
501 let region_id = RegionId::new(table_id, region_number);
502 let another_region_id = RegionId::new(table_id, region_number + 1);
503 let no_exist_region_id = RegionId::new(table_id, region_number + 2);
504 let peer = Peer::empty(datanode_id);
505
506 let builder = MetasrvBuilder::new();
507 let metasrv = builder.build().await.unwrap();
508 let ctx = &mut metasrv.new_ctx();
509
510 let req = HeartbeatRequest {
511 duration_since_epoch: 1234,
512 ..Default::default()
513 };
514
515 let acc = &mut HeartbeatAccumulator::default();
516 acc.stat = Some(Stat {
517 id: peer.id,
518 region_stats: vec![
519 new_empty_region_stat(region_id, RegionRole::Leader),
520 new_empty_region_stat(another_region_id, RegionRole::Leader),
521 new_empty_region_stat(no_exist_region_id, RegionRole::Leader),
522 ],
523 ..Default::default()
524 });
525 let handler = RegionLeaseHandler::new(
526 default_distributed_time_constants().region_lease.as_secs(),
527 table_metadata_manager.clone(),
528 Default::default(),
529 None,
530 );
531 handler.handle(&req, ctx, acc).await.unwrap();
532
533 assert!(acc.region_lease.is_none());
534 assert!(acc.inactive_region_ids.is_empty());
535 }
536}