datanode/heartbeat/
handler.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 async_trait::async_trait;
16use common_meta::error::{InvalidHeartbeatResponseSnafu, Result as MetaResult};
17use common_meta::heartbeat::handler::{
18    HandleControl, HeartbeatResponseHandler, HeartbeatResponseHandlerContext,
19};
20use common_meta::instruction::{Instruction, InstructionReply};
21use common_meta::kv_backend::KvBackendRef;
22use common_telemetry::error;
23use common_telemetry::tracing_context::FutureExt;
24use snafu::OptionExt;
25use store_api::storage::GcReport;
26use strum::AsRefStr;
27
28mod apply_staging_manifest;
29mod close_region;
30mod downgrade_region;
31mod enter_staging;
32mod file_ref;
33mod flush_region;
34mod gc_worker;
35mod open_region;
36mod remap_manifest;
37mod sync_region;
38mod upgrade_region;
39
40use crate::heartbeat::handler::apply_staging_manifest::ApplyStagingManifestsHandler;
41use crate::heartbeat::handler::close_region::CloseRegionsHandler;
42use crate::heartbeat::handler::downgrade_region::DowngradeRegionsHandler;
43use crate::heartbeat::handler::enter_staging::EnterStagingRegionsHandler;
44use crate::heartbeat::handler::file_ref::GetFileRefsHandler;
45use crate::heartbeat::handler::flush_region::FlushRegionsHandler;
46use crate::heartbeat::handler::gc_worker::GcRegionsHandler;
47use crate::heartbeat::handler::open_region::OpenRegionsHandler;
48use crate::heartbeat::handler::remap_manifest::RemapManifestHandler;
49use crate::heartbeat::handler::sync_region::SyncRegionHandler;
50use crate::heartbeat::handler::upgrade_region::UpgradeRegionsHandler;
51use crate::heartbeat::task_tracker::TaskTracker;
52use crate::region_server::RegionServer;
53
54/// The handler for [`Instruction`]s.
55#[derive(Clone)]
56pub struct RegionHeartbeatResponseHandler {
57    region_server: RegionServer,
58    downgrade_tasks: TaskTracker<()>,
59    flush_tasks: TaskTracker<()>,
60    open_region_parallelism: usize,
61    gc_tasks: TaskTracker<GcReport>,
62    kv_backend: KvBackendRef,
63}
64
65#[async_trait::async_trait]
66pub trait InstructionHandler: Send + Sync {
67    type Instruction;
68    async fn handle(
69        &self,
70        ctx: &HandlerContext,
71        instruction: Self::Instruction,
72    ) -> Option<InstructionReply>;
73}
74
75#[derive(Clone)]
76pub struct HandlerContext {
77    pub region_server: RegionServer,
78    pub downgrade_tasks: TaskTracker<()>,
79    pub flush_tasks: TaskTracker<()>,
80    pub gc_tasks: TaskTracker<GcReport>,
81    pub kv_backend: KvBackendRef,
82}
83
84impl HandlerContext {
85    #[cfg(test)]
86    pub fn new_for_test(region_server: RegionServer, kv_backend: KvBackendRef) -> Self {
87        Self {
88            region_server,
89            downgrade_tasks: TaskTracker::new(),
90            flush_tasks: TaskTracker::new(),
91            gc_tasks: TaskTracker::new(),
92            kv_backend,
93        }
94    }
95}
96
97impl RegionHeartbeatResponseHandler {
98    /// Returns the [RegionHeartbeatResponseHandler].
99    pub fn new(region_server: RegionServer, kv_backend: KvBackendRef) -> Self {
100        Self {
101            region_server,
102            downgrade_tasks: TaskTracker::new(),
103            flush_tasks: TaskTracker::new(),
104            // Default to half of the number of CPUs.
105            open_region_parallelism: (num_cpus::get() / 2).max(1),
106            gc_tasks: TaskTracker::new(),
107            kv_backend,
108        }
109    }
110
111    /// Sets the parallelism for opening regions.
112    pub fn with_open_region_parallelism(mut self, parallelism: usize) -> Self {
113        self.open_region_parallelism = parallelism;
114        self
115    }
116
117    fn build_handler(
118        &self,
119        instruction: &Instruction,
120    ) -> MetaResult<Option<Box<InstructionHandlers>>> {
121        match instruction {
122            Instruction::CloseRegions(_) => Ok(Some(Box::new(CloseRegionsHandler.into()))),
123            Instruction::OpenRegions(_) => Ok(Some(Box::new(
124                OpenRegionsHandler {
125                    open_region_parallelism: self.open_region_parallelism,
126                }
127                .into(),
128            ))),
129            Instruction::FlushRegions(_) => Ok(Some(Box::new(FlushRegionsHandler.into()))),
130            Instruction::DowngradeRegions(_) => Ok(Some(Box::new(DowngradeRegionsHandler.into()))),
131            Instruction::UpgradeRegions(_) => Ok(Some(Box::new(
132                UpgradeRegionsHandler {
133                    upgrade_region_parallelism: self.open_region_parallelism,
134                }
135                .into(),
136            ))),
137            Instruction::GetFileRefs(_) => Ok(Some(Box::new(GetFileRefsHandler.into()))),
138            Instruction::GcRegions(_) => Ok(Some(Box::new(GcRegionsHandler.into()))),
139            Instruction::InvalidateCaches(_) => InvalidHeartbeatResponseSnafu.fail(),
140            Instruction::Suspend => Ok(None),
141            Instruction::EnterStagingRegions(_) => {
142                Ok(Some(Box::new(EnterStagingRegionsHandler.into())))
143            }
144            Instruction::SyncRegions(_) => Ok(Some(Box::new(SyncRegionHandler.into()))),
145            Instruction::RemapManifest(_) => Ok(Some(Box::new(RemapManifestHandler.into()))),
146            Instruction::ApplyStagingManifests(_) => {
147                Ok(Some(Box::new(ApplyStagingManifestsHandler.into())))
148            }
149        }
150    }
151}
152
153#[allow(clippy::enum_variant_names)]
154#[derive(AsRefStr)]
155pub enum InstructionHandlers {
156    CloseRegions(CloseRegionsHandler),
157    OpenRegions(OpenRegionsHandler),
158    FlushRegions(FlushRegionsHandler),
159    DowngradeRegions(DowngradeRegionsHandler),
160    UpgradeRegions(UpgradeRegionsHandler),
161    GetFileRefs(GetFileRefsHandler),
162    GcRegions(GcRegionsHandler),
163    EnterStagingRegions(EnterStagingRegionsHandler),
164    SyncRegions(SyncRegionHandler),
165    RemapManifest(RemapManifestHandler),
166    ApplyStagingManifests(ApplyStagingManifestsHandler),
167}
168
169impl InstructionHandlers {
170    // Returns the string representation of the instruction handler.
171    pub fn as_ref_str(&self) -> &str {
172        self.as_ref()
173    }
174}
175
176macro_rules! impl_from_handler {
177    ($($handler:ident => $variant:ident),*) => {
178        $(
179            impl From<$handler> for InstructionHandlers {
180                fn from(handler: $handler) -> Self {
181                    InstructionHandlers::$variant(handler)
182                }
183            }
184        )*
185    };
186}
187
188impl_from_handler!(
189    CloseRegionsHandler => CloseRegions,
190    OpenRegionsHandler => OpenRegions,
191    FlushRegionsHandler => FlushRegions,
192    DowngradeRegionsHandler => DowngradeRegions,
193    UpgradeRegionsHandler => UpgradeRegions,
194    GetFileRefsHandler => GetFileRefs,
195    GcRegionsHandler => GcRegions,
196    EnterStagingRegionsHandler => EnterStagingRegions,
197    SyncRegionHandler => SyncRegions,
198    RemapManifestHandler => RemapManifest,
199    ApplyStagingManifestsHandler => ApplyStagingManifests
200);
201
202macro_rules! dispatch_instr {
203    (
204        $( $instr_variant:ident => $handler_variant:ident ),* $(,)?
205    ) => {
206        impl InstructionHandlers {
207            pub async fn handle(
208                &self,
209                ctx: &HandlerContext,
210                instruction: Instruction,
211            ) -> Option<InstructionReply> {
212                match (self, instruction) {
213                    $(
214                        (
215                            InstructionHandlers::$handler_variant(handler),
216                            Instruction::$instr_variant(instr),
217                        ) => handler.handle(ctx, instr).await,
218                    )*
219                    // Safety: must be used in pairs with `build_handler`.
220                    _ => unreachable!(),
221                }
222            }
223            /// Check whether this instruction is acceptable by any handler.
224            pub fn is_acceptable(instruction: &Instruction) -> bool {
225                matches!(
226                    instruction,
227                    $(
228                        Instruction::$instr_variant { .. }
229                    )|*
230                )
231            }
232        }
233    };
234}
235
236dispatch_instr!(
237    CloseRegions => CloseRegions,
238    OpenRegions => OpenRegions,
239    FlushRegions => FlushRegions,
240    DowngradeRegions => DowngradeRegions,
241    UpgradeRegions => UpgradeRegions,
242    GetFileRefs => GetFileRefs,
243    GcRegions => GcRegions,
244    EnterStagingRegions => EnterStagingRegions,
245    SyncRegions => SyncRegions,
246    RemapManifest => RemapManifest,
247    ApplyStagingManifests => ApplyStagingManifests,
248);
249
250#[async_trait]
251impl HeartbeatResponseHandler for RegionHeartbeatResponseHandler {
252    fn is_acceptable(&self, ctx: &HeartbeatResponseHandlerContext) -> bool {
253        if let Some((_, _, instruction)) = ctx.incoming_message.as_ref() {
254            return InstructionHandlers::is_acceptable(instruction);
255        }
256        false
257    }
258
259    async fn handle(&self, ctx: &mut HeartbeatResponseHandlerContext) -> MetaResult<HandleControl> {
260        let (meta, tracing_ctx, instruction) = ctx
261            .incoming_message
262            .take()
263            .context(InvalidHeartbeatResponseSnafu)?;
264
265        let mailbox = ctx.mailbox.clone();
266        if let Some(handler) = self.build_handler(&instruction)? {
267            let context = HandlerContext {
268                region_server: self.region_server.clone(),
269                downgrade_tasks: self.downgrade_tasks.clone(),
270                flush_tasks: self.flush_tasks.clone(),
271                gc_tasks: self.gc_tasks.clone(),
272                kv_backend: self.kv_backend.clone(),
273            };
274            let span = tracing_ctx.attach(tracing::info_span!(
275                "RegionHeartbeatResponseHandler::handle",
276                from = %meta.from,
277                to = %meta.to,
278                handler = %handler.as_ref_str(),
279            ));
280            let _handle = common_runtime::spawn_global(async move {
281                let reply = handler.handle(&context, instruction).trace(span).await;
282                if let Some(reply) = reply
283                    && let Err(e) = mailbox.send((meta, reply)).await
284                {
285                    let error = e.to_string();
286                    let (meta, reply) = e.0;
287                    error!("Failed to send reply {reply} to {meta:?}: {error}");
288                }
289            });
290        }
291
292        Ok(HandleControl::Continue)
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use std::assert_matches;
299    use std::collections::HashMap;
300    use std::sync::Arc;
301    use std::time::Duration;
302
303    use common_meta::RegionIdent;
304    use common_meta::heartbeat::mailbox::{
305        HeartbeatMailbox, IncomingMessage, MailboxRef, MessageMeta,
306    };
307    use common_meta::instruction::{
308        DowngradeRegion, EnterStagingRegion, OpenRegion, StagingPartitionDirective, UpgradeRegion,
309    };
310    use common_meta::kv_backend::memory::MemoryKvBackend;
311    use mito2::config::MitoConfig;
312    use mito2::engine::MITO_ENGINE_NAME;
313    use mito2::test_util::{CreateRequestBuilder, TestEnv};
314    use store_api::path_utils::table_dir;
315    use store_api::region_engine::RegionRole;
316    use store_api::region_request::{RegionCloseRequest, RegionRequest, RegionRequirements};
317    use store_api::storage::RegionId;
318    use tokio::sync::mpsc::{self, Receiver};
319
320    use super::*;
321    use crate::error;
322    use crate::tests::mock_region_server;
323
324    pub struct HeartbeatResponseTestEnv {
325        pub(crate) mailbox: MailboxRef,
326        pub(crate) receiver: Receiver<(MessageMeta, InstructionReply)>,
327    }
328
329    impl HeartbeatResponseTestEnv {
330        pub fn new() -> Self {
331            let (tx, rx) = mpsc::channel(8);
332            let mailbox = Arc::new(HeartbeatMailbox::new(tx));
333
334            HeartbeatResponseTestEnv {
335                mailbox,
336                receiver: rx,
337            }
338        }
339
340        pub fn create_handler_ctx(
341            &self,
342            incoming_message: IncomingMessage,
343        ) -> HeartbeatResponseHandlerContext {
344            HeartbeatResponseHandlerContext {
345                mailbox: self.mailbox.clone(),
346                response: Default::default(),
347                incoming_message: Some(incoming_message),
348            }
349        }
350    }
351
352    #[test]
353    fn test_is_acceptable() {
354        common_telemetry::init_default_ut_logging();
355        let region_server = mock_region_server();
356        let kv_backend = Arc::new(MemoryKvBackend::new());
357        let heartbeat_handler =
358            RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend);
359        let heartbeat_env = HeartbeatResponseTestEnv::new();
360        let meta = MessageMeta::new_test(1, "test", "dn-1", "me-0");
361
362        // Open region
363        let region_id = RegionId::new(1024, 1);
364        let storage_path = "test";
365        let instruction = open_region_instruction(region_id, storage_path);
366        assert!(
367            heartbeat_handler.is_acceptable(&heartbeat_env.create_handler_ctx((
368                meta.clone(),
369                Default::default(),
370                instruction
371            )))
372        );
373
374        // Close region
375        let instruction = close_region_instruction(region_id);
376        assert!(
377            heartbeat_handler.is_acceptable(&heartbeat_env.create_handler_ctx((
378                meta.clone(),
379                Default::default(),
380                instruction
381            )))
382        );
383
384        // Downgrade region
385        let instruction = Instruction::DowngradeRegions(vec![DowngradeRegion {
386            region_id: RegionId::new(2048, 1),
387            flush_timeout: Some(Duration::from_secs(1)),
388        }]);
389        assert!(
390            heartbeat_handler.is_acceptable(&heartbeat_env.create_handler_ctx((
391                meta.clone(),
392                Default::default(),
393                instruction
394            )))
395        );
396
397        // Upgrade region
398        let instruction = Instruction::UpgradeRegions(vec![UpgradeRegion {
399            region_id,
400            ..Default::default()
401        }]);
402        assert!(
403            heartbeat_handler.is_acceptable(&heartbeat_env.create_handler_ctx((
404                meta.clone(),
405                Default::default(),
406                instruction
407            )))
408        );
409
410        // Enter staging region
411        let instruction = Instruction::EnterStagingRegions(vec![EnterStagingRegion {
412            region_id,
413            partition_directive: StagingPartitionDirective::UpdatePartitionExpr("".to_string()),
414        }]);
415        assert!(
416            heartbeat_handler.is_acceptable(&heartbeat_env.create_handler_ctx((
417                meta,
418                Default::default(),
419                instruction
420            )))
421        );
422    }
423
424    fn close_region_instruction(region_id: RegionId) -> Instruction {
425        Instruction::CloseRegions(vec![RegionIdent {
426            table_id: region_id.table_id(),
427            region_number: region_id.region_number(),
428            datanode_id: 2,
429            engine: MITO_ENGINE_NAME.to_string(),
430        }])
431    }
432
433    fn open_region_instruction(region_id: RegionId, path: &str) -> Instruction {
434        Instruction::OpenRegions(vec![OpenRegion::new(
435            RegionIdent {
436                table_id: region_id.table_id(),
437                region_number: region_id.region_number(),
438                datanode_id: 2,
439                engine: MITO_ENGINE_NAME.to_string(),
440            },
441            path,
442            HashMap::new(),
443            HashMap::new(),
444            false,
445            None,
446            RegionRequirements::empty(),
447        )])
448    }
449
450    #[tokio::test]
451    async fn test_close_region() {
452        common_telemetry::init_default_ut_logging();
453
454        let mut region_server = mock_region_server();
455        let kv_backend = Arc::new(MemoryKvBackend::new());
456        let heartbeat_handler =
457            RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend);
458
459        let mut engine_env = TestEnv::with_prefix("close-region").await;
460        let engine = engine_env.create_engine(MitoConfig::default()).await;
461        region_server.register_engine(Arc::new(engine));
462        let region_id = RegionId::new(1024, 1);
463
464        let builder = CreateRequestBuilder::new();
465        let create_req = builder.build();
466        region_server
467            .handle_request(region_id, RegionRequest::Create(create_req))
468            .await
469            .unwrap();
470
471        let mut heartbeat_env = HeartbeatResponseTestEnv::new();
472
473        // Should be ok, if we try to close it twice.
474        for _ in 0..2 {
475            let meta = MessageMeta::new_test(1, "test", "dn-1", "me-0");
476            let instruction = close_region_instruction(region_id);
477
478            let mut ctx = heartbeat_env.create_handler_ctx((meta, Default::default(), instruction));
479            let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
480            assert_matches!(control, HandleControl::Continue);
481
482            let (_, reply) = heartbeat_env.receiver.recv().await.unwrap();
483
484            if let InstructionReply::CloseRegions(reply) = reply {
485                assert!(reply.result);
486                assert!(reply.error.is_none());
487            } else {
488                unreachable!()
489            }
490
491            assert_matches!(
492                region_server
493                    .set_region_role(region_id, RegionRole::Leader)
494                    .unwrap_err(),
495                error::Error::RegionNotFound { .. }
496            );
497        }
498    }
499
500    #[tokio::test]
501    async fn test_open_region_ok() {
502        common_telemetry::init_default_ut_logging();
503
504        let mut region_server = mock_region_server();
505        let kv_backend = Arc::new(MemoryKvBackend::new());
506        let heartbeat_handler =
507            RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend);
508
509        let mut engine_env = TestEnv::with_prefix("open-region").await;
510        let engine = engine_env.create_engine(MitoConfig::default()).await;
511        region_server.register_engine(Arc::new(engine));
512        let region_id = RegionId::new(1024, 1);
513
514        let builder = CreateRequestBuilder::new();
515        let mut create_req = builder.build();
516        let storage_path = "test";
517        create_req.table_dir = table_dir(storage_path, region_id.table_id());
518
519        region_server
520            .handle_request(region_id, RegionRequest::Create(create_req))
521            .await
522            .unwrap();
523
524        region_server
525            .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {}))
526            .await
527            .unwrap();
528        let mut heartbeat_env = HeartbeatResponseTestEnv::new();
529
530        // Should be ok, if we try to open it twice.
531        for _ in 0..2 {
532            let meta = MessageMeta::new_test(1, "test", "dn-1", "me-0");
533            let instruction = open_region_instruction(region_id, storage_path);
534
535            let mut ctx = heartbeat_env.create_handler_ctx((meta, Default::default(), instruction));
536            let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
537            assert_matches!(control, HandleControl::Continue);
538
539            let (_, reply) = heartbeat_env.receiver.recv().await.unwrap();
540
541            if let InstructionReply::OpenRegions(reply) = reply {
542                assert!(reply.result);
543                assert!(reply.error.is_none());
544            } else {
545                unreachable!()
546            }
547        }
548    }
549
550    #[tokio::test]
551    async fn test_open_not_exists_region() {
552        common_telemetry::init_default_ut_logging();
553
554        let mut region_server = mock_region_server();
555        let kv_backend = Arc::new(MemoryKvBackend::new());
556        let heartbeat_handler =
557            RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend);
558
559        let mut engine_env = TestEnv::with_prefix("open-not-exists-region").await;
560        let engine = engine_env.create_engine(MitoConfig::default()).await;
561        region_server.register_engine(Arc::new(engine));
562        let region_id = RegionId::new(1024, 1);
563        let storage_path = "test";
564
565        let mut heartbeat_env = HeartbeatResponseTestEnv::new();
566
567        let meta = MessageMeta::new_test(1, "test", "dn-1", "me-0");
568        let instruction = open_region_instruction(region_id, storage_path);
569
570        let mut ctx = heartbeat_env.create_handler_ctx((meta, Default::default(), instruction));
571        let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
572        assert_matches!(control, HandleControl::Continue);
573
574        let (_, reply) = heartbeat_env.receiver.recv().await.unwrap();
575
576        if let InstructionReply::OpenRegions(reply) = reply {
577            assert!(!reply.result);
578            assert!(reply.error.is_some());
579        } else {
580            unreachable!()
581        }
582    }
583
584    #[tokio::test]
585    async fn test_downgrade_region() {
586        common_telemetry::init_default_ut_logging();
587
588        let mut region_server = mock_region_server();
589        let kv_backend = Arc::new(MemoryKvBackend::new());
590        let heartbeat_handler =
591            RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend);
592
593        let mut engine_env = TestEnv::with_prefix("downgrade-region").await;
594        let engine = engine_env.create_engine(MitoConfig::default()).await;
595        region_server.register_engine(Arc::new(engine));
596        let region_id = RegionId::new(1024, 1);
597
598        let builder = CreateRequestBuilder::new();
599        let mut create_req = builder.build();
600        let storage_path = "test";
601        create_req.table_dir = table_dir(storage_path, region_id.table_id());
602
603        region_server
604            .handle_request(region_id, RegionRequest::Create(create_req))
605            .await
606            .unwrap();
607
608        let mut heartbeat_env = HeartbeatResponseTestEnv::new();
609
610        // Should be ok, if we try to downgrade it twice.
611        for _ in 0..2 {
612            let meta = MessageMeta::new_test(1, "test", "dn-1", "me-0");
613            let instruction = Instruction::DowngradeRegions(vec![DowngradeRegion {
614                region_id,
615                flush_timeout: Some(Duration::from_secs(1)),
616            }]);
617
618            let mut ctx = heartbeat_env.create_handler_ctx((meta, Default::default(), instruction));
619            let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
620            assert_matches!(control, HandleControl::Continue);
621
622            let (_, reply) = heartbeat_env.receiver.recv().await.unwrap();
623
624            let reply = &reply.expect_downgrade_regions_reply()[0];
625            assert!(reply.exists);
626            assert!(reply.error.is_none());
627            assert_eq!(reply.last_entry_id.unwrap(), 0);
628        }
629
630        // Downgrades a not exists region.
631        let meta = MessageMeta::new_test(1, "test", "dn-1", "me-0");
632        let instruction = Instruction::DowngradeRegions(vec![DowngradeRegion {
633            region_id: RegionId::new(2048, 1),
634            flush_timeout: Some(Duration::from_secs(1)),
635        }]);
636        let mut ctx = heartbeat_env.create_handler_ctx((meta, Default::default(), instruction));
637        let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
638        assert_matches!(control, HandleControl::Continue);
639
640        let (_, reply) = heartbeat_env.receiver.recv().await.unwrap();
641
642        let reply = reply.expect_downgrade_regions_reply();
643        assert!(!reply[0].exists);
644        assert!(reply[0].error.is_none());
645        assert!(reply[0].last_entry_id.is_none());
646    }
647}
datanode/heartbeat/handler.rs

datanode/heartbeat/
handler.rs
