Skip to main content

datanode/heartbeat/handler/
file_ref.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_error::ext::ErrorExt;
16use common_meta::instruction::{
17    GetFileRefs, GetFileRefsReply, GetPackedFileRefs, GetPackedFileRefsReply, InstructionError,
18    InstructionReply, PackedFileRefsManifest,
19};
20use store_api::storage::FileRefsManifest;
21
22use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
23
24pub struct GetFileRefsHandler;
25
26#[async_trait::async_trait]
27impl InstructionHandler for GetFileRefsHandler {
28    type Instruction = GetFileRefs;
29
30    async fn handle(
31        &self,
32        ctx: &HandlerContext,
33        get_file_refs: Self::Instruction,
34    ) -> Option<InstructionReply> {
35        let region_server = &ctx.region_server;
36
37        // Get the MitoEngine
38        let Some(mito_engine) = region_server.mito_engine() else {
39            return Some(InstructionReply::GetFileRefs(GetFileRefsReply {
40                file_refs_manifest: FileRefsManifest::default(),
41                success: false,
42                error: Some(InstructionError::legacy_internal_retryable(
43                    "MitoEngine not found",
44                )),
45            }));
46        };
47        match mito_engine
48            .get_snapshot_of_file_refs(get_file_refs.query_regions, get_file_refs.related_regions)
49            .await
50        {
51            Ok(all_file_refs) => {
52                // Return the file references
53                Some(InstructionReply::GetFileRefs(GetFileRefsReply {
54                    file_refs_manifest: all_file_refs,
55                    success: true,
56                    error: None,
57                }))
58            }
59            Err(e) => Some(InstructionReply::GetFileRefs(GetFileRefsReply {
60                file_refs_manifest: FileRefsManifest::default(),
61                success: false,
62                error: Some(InstructionError::new(
63                    e.status_code(),
64                    format!("Failed to get file refs: {}", e.output_msg()),
65                    e.retry_hint(),
66                )),
67            })),
68        }
69    }
70}
71
72pub struct GetPackedFileRefsHandler;
73
74#[async_trait::async_trait]
75impl InstructionHandler for GetPackedFileRefsHandler {
76    type Instruction = GetPackedFileRefs;
77
78    async fn handle(
79        &self,
80        ctx: &HandlerContext,
81        get_file_refs: Self::Instruction,
82    ) -> Option<InstructionReply> {
83        let region_server = &ctx.region_server;
84        let Some(mito_engine) = region_server.mito_engine() else {
85            return Some(InstructionReply::GetPackedFileRefs(
86                GetPackedFileRefsReply {
87                    packed_file_refs_manifest: PackedFileRefsManifest::default(),
88                    success: false,
89                    error: Some(InstructionError::legacy_internal_retryable(
90                        "MitoEngine not found",
91                    )),
92                },
93            ));
94        };
95        match mito_engine
96            .get_snapshot_of_file_refs(get_file_refs.query_regions, get_file_refs.related_regions)
97            .await
98        {
99            Ok(all_file_refs) => Some(InstructionReply::GetPackedFileRefs(
100                GetPackedFileRefsReply {
101                    packed_file_refs_manifest: PackedFileRefsManifest::from_manifest(
102                        &all_file_refs,
103                    ),
104                    success: true,
105                    error: None,
106                },
107            )),
108            Err(e) => Some(InstructionReply::GetPackedFileRefs(
109                GetPackedFileRefsReply {
110                    packed_file_refs_manifest: PackedFileRefsManifest::default(),
111                    success: false,
112                    error: Some(InstructionError::new(
113                        e.status_code(),
114                        format!("Failed to get file refs: {}", e.output_msg()),
115                        e.retry_hint(),
116                    )),
117                },
118            )),
119        }
120    }
121}