Skip to main content

datanode/
utils.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 std::collections::HashMap;
16
17use common_meta::DatanodeId;
18use common_meta::key::datanode_table::DatanodeTableManager;
19use common_meta::key::topic_region::{TopicRegionKey, TopicRegionManager, TopicRegionValue};
20use common_meta::kv_backend::KvBackendRef;
21use common_meta::wal_provider::{extract_topic_from_wal_options, prepare_wal_options};
22use futures::TryStreamExt;
23use snafu::ResultExt;
24use store_api::path_utils::table_dir;
25use store_api::region_request::{PathType, RegionOpenRequest, ReplayCheckpoint};
26use store_api::storage::{RegionId, RegionNumber};
27use tracing::info;
28
29use crate::error::{GetMetadataSnafu, Result};
30
31/// The requests to open regions.
32pub struct RegionOpenRequests {
33    pub(crate) leader_regions: Vec<(RegionId, RegionOpenRequest)>,
34    #[cfg(feature = "enterprise")]
35    pub(crate) follower_regions: Vec<(RegionId, RegionOpenRequest)>,
36}
37
38impl RegionOpenRequests {
39    /// Splits the request set into leader and follower regions.
40    #[allow(clippy::type_complexity)]
41    pub fn into_parts(
42        self,
43    ) -> (
44        Vec<(RegionId, RegionOpenRequest)>,
45        Vec<(RegionId, RegionOpenRequest)>,
46    ) {
47        let leader_regions = self.leader_regions;
48        #[cfg(feature = "enterprise")]
49        let follower_regions = self.follower_regions;
50        #[cfg(not(feature = "enterprise"))]
51        let follower_regions = Vec::new();
52        (leader_regions, follower_regions)
53    }
54}
55
56fn group_region_by_topic(
57    region_id: RegionId,
58    region_options: &HashMap<RegionNumber, String>,
59    topic_regions: &mut HashMap<String, Vec<RegionId>>,
60) {
61    if let Some(topic) = extract_topic_from_wal_options(region_id, region_options) {
62        topic_regions.entry(topic).or_default().push(region_id);
63    }
64}
65
66fn get_replay_checkpoint(
67    region_id: RegionId,
68    topic_region_values: &Option<HashMap<RegionId, TopicRegionValue>>,
69) -> Option<ReplayCheckpoint> {
70    let topic_region_values = topic_region_values.as_ref()?;
71    let topic_region_value = topic_region_values.get(&region_id);
72    let replay_checkpoint = topic_region_value.and_then(|value| value.checkpoint);
73    replay_checkpoint.map(|checkpoint| ReplayCheckpoint {
74        entry_id: checkpoint.entry_id,
75        metadata_entry_id: checkpoint.metadata_entry_id,
76    })
77}
78
79/// Builds region-open requests from persisted metadata.
80pub async fn build_region_open_requests(
81    node_id: DatanodeId,
82    kv_backend: KvBackendRef,
83) -> Result<RegionOpenRequests> {
84    let datanode_table_manager = DatanodeTableManager::new(kv_backend.clone());
85    let table_values = datanode_table_manager
86        .tables(node_id)
87        .try_collect::<Vec<_>>()
88        .await
89        .context(GetMetadataSnafu)?;
90
91    let topic_region_manager = TopicRegionManager::new(kv_backend);
92    let mut topic_regions = HashMap::<String, Vec<RegionId>>::new();
93    let mut regions = vec![];
94    #[cfg(feature = "enterprise")]
95    let mut follower_regions = vec![];
96
97    for table_value in table_values {
98        for region_number in table_value.regions {
99            let region_id = RegionId::new(table_value.table_id, region_number);
100            // Augments region options with wal options if a wal options is provided.
101            let mut region_options = table_value.region_info.region_options.clone();
102            prepare_wal_options(
103                &mut region_options,
104                region_id,
105                &table_value.region_info.region_wal_options,
106            );
107            group_region_by_topic(
108                region_id,
109                &table_value.region_info.region_wal_options,
110                &mut topic_regions,
111            );
112
113            regions.push((
114                region_id,
115                table_value.region_info.engine.clone(),
116                table_value.region_info.region_storage_path.clone(),
117                region_options,
118            ));
119        }
120
121        #[cfg(feature = "enterprise")]
122        for region_number in table_value.follower_regions {
123            let region_id = RegionId::new(table_value.table_id, region_number);
124            // Augments region options with wal options if a wal options is provided.
125            let mut region_options = table_value.region_info.region_options.clone();
126            prepare_wal_options(
127                &mut region_options,
128                RegionId::new(table_value.table_id, region_number),
129                &table_value.region_info.region_wal_options,
130            );
131            group_region_by_topic(
132                region_id,
133                &table_value.region_info.region_wal_options,
134                &mut topic_regions,
135            );
136
137            follower_regions.push((
138                RegionId::new(table_value.table_id, region_number),
139                table_value.region_info.engine.clone(),
140                table_value.region_info.region_storage_path.clone(),
141                region_options,
142            ));
143        }
144    }
145
146    let topic_region_values = if !topic_regions.is_empty() {
147        let keys = topic_regions
148            .iter()
149            .flat_map(|(topic, regions)| {
150                regions
151                    .iter()
152                    .map(|region_id| TopicRegionKey::new(*region_id, topic))
153            })
154            .collect::<Vec<_>>();
155        let topic_region_manager = topic_region_manager
156            .batch_get(keys)
157            .await
158            .context(GetMetadataSnafu)?;
159        Some(topic_region_manager)
160    } else {
161        None
162    };
163
164    let mut leader_region_requests = Vec::with_capacity(regions.len());
165    for (region_id, engine, store_path, options) in regions {
166        let table_dir = table_dir(&store_path, region_id.table_id());
167        let checkpoint = get_replay_checkpoint(region_id, &topic_region_values);
168        info!("region_id: {}, checkpoint: {:?}", region_id, checkpoint);
169        leader_region_requests.push((
170            region_id,
171            RegionOpenRequest {
172                engine,
173                table_dir,
174                path_type: PathType::Bare,
175                options,
176                skip_wal_replay: false,
177                checkpoint,
178            },
179        ));
180    }
181
182    #[cfg(feature = "enterprise")]
183    let follower_region_requests = {
184        let mut follower_region_requests = Vec::with_capacity(follower_regions.len());
185        for (region_id, engine, store_path, options) in follower_regions {
186            let table_dir = table_dir(&store_path, region_id.table_id());
187            follower_region_requests.push((
188                region_id,
189                RegionOpenRequest {
190                    engine,
191                    table_dir,
192                    path_type: PathType::Bare,
193                    options,
194                    skip_wal_replay: true,
195                    checkpoint: None,
196                },
197            ));
198        }
199        follower_region_requests
200    };
201
202    Ok(RegionOpenRequests {
203        leader_regions: leader_region_requests,
204        #[cfg(feature = "enterprise")]
205        follower_regions: follower_region_requests,
206    })
207}