mito2/worker/
handle_remap.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;
16use std::time::Instant;
17
18use common_error::ext::BoxedError;
19use common_telemetry::{debug, info};
20use futures::future::try_join_all;
21use partition::expr::PartitionExpr;
22use snafu::{OptionExt, ResultExt};
23use store_api::storage::RegionId;
24
25use crate::error::{self, FetchManifestsSnafu, InvalidRequestSnafu, MissingManifestSnafu, Result};
26use crate::region::{MitoRegionRef, RegionMetadataLoader};
27use crate::remap_manifest::RemapManifest;
28use crate::request::RemapManifestsRequest;
29use crate::sst::location::region_dir_from_table_dir;
30use crate::worker::RegionWorkerLoop;
31
32impl<S> RegionWorkerLoop<S> {
33    pub(crate) fn handle_remap_manifests_request(&mut self, request: RemapManifestsRequest) {
34        let region_id = request.region_id;
35        let sender = request.sender;
36        let region = match self.regions.staging_region(region_id) {
37            Ok(region) => region,
38            Err(e) => {
39                let _ = sender.send(Err(e));
40                return;
41            }
42        };
43
44        let same_table = request
45            .input_regions
46            .iter()
47            .map(|r| r.table_id())
48            .all(|t| t == region_id.table_id());
49
50        if !same_table {
51            let _ = sender.send(
52                InvalidRequestSnafu {
53                    region_id,
54                    reason: "Input regions must be from the same table",
55                }
56                .fail(),
57            );
58            return;
59        }
60
61        let region_metadata_loader = RegionMetadataLoader::new(
62            self.config.clone(),
63            self.object_store_manager.clone(),
64            self.node_lease.is_enabled(),
65        );
66        common_runtime::spawn_global(async move {
67            let result = Self::fetch_and_remap_manifests(
68                region,
69                region_metadata_loader,
70                request.input_regions,
71                request.new_partition_exprs,
72                request.region_mapping,
73            )
74            .await;
75
76            let _ = sender.send(result);
77        });
78    }
79
80    // Fetches manifests for input regions, remaps them according to the provided
81    // mapping and partition expressions.
82    //
83    // Returns a map from each new region to its relative staging manifest path.
84    async fn fetch_and_remap_manifests(
85        region: MitoRegionRef,
86        region_metadata_loader: RegionMetadataLoader,
87        input_regions: Vec<RegionId>,
88        new_partition_exprs: HashMap<RegionId, PartitionExpr>,
89        region_mapping: HashMap<RegionId, Vec<RegionId>>,
90    ) -> Result<HashMap<RegionId, String>> {
91        let mut tasks = Vec::with_capacity(input_regions.len());
92        let region_options = region.version().options.clone();
93        let table_dir = region.table_dir();
94        let path_type = region.path_type();
95        let now = Instant::now();
96        for input_region in &input_regions {
97            let region_dir = region_dir_from_table_dir(table_dir, *input_region, path_type);
98            let region_options = region_options.clone();
99            let moved_region_metadata_loader = region_metadata_loader.clone();
100            tasks.push(async move {
101                moved_region_metadata_loader
102                    .load_manifest(&region_dir, &region_options)
103                    .await
104            });
105        }
106        let results = try_join_all(tasks)
107            .await
108            .map_err(BoxedError::new)
109            .context(FetchManifestsSnafu)?;
110        let manifests = results
111            .into_iter()
112            .zip(input_regions)
113            .map(|(manifest_res, region_id)| {
114                let manifest = manifest_res.context(MissingManifestSnafu { region_id })?;
115                Ok((region_id, (*manifest).clone()))
116            })
117            .collect::<Result<HashMap<_, _>>>()?;
118        let mut mapper = RemapManifest::new(manifests, new_partition_exprs, region_mapping);
119        let remap_result = mapper.remap_manifests()?;
120
121        // Write new manifests to staging blob storage.
122        let manifest_manager = region.manifest_ctx.manifest_manager.write().await;
123        let manifest_storage = manifest_manager.store();
124        let staging_blob_storage = manifest_storage.staging_storage().blob_storage().clone();
125        let mut tasks = Vec::with_capacity(remap_result.new_manifests.len());
126
127        for (remap_region_id, manifest) in &remap_result.new_manifests {
128            let bytes = serde_json::to_vec(&manifest).context(error::SerializeManifestSnafu {
129                region_id: *remap_region_id,
130            })?;
131            let key = remap_manifest_key(remap_region_id);
132            tasks.push(async {
133                debug!(
134                    "Putting manifest to staging blob storage, region_id: {}, key: {}",
135                    *remap_region_id, key
136                );
137                staging_blob_storage.put(&key, bytes).await?;
138                Ok((*remap_region_id, key))
139            });
140        }
141        let r = try_join_all(tasks).await?;
142        info!(
143            "Remap manifests cost: {:?}, region: {}",
144            now.elapsed(),
145            region.region_id
146        );
147
148        Ok(r.into_iter().collect::<HashMap<_, _>>())
149    }
150}
151
152fn remap_manifest_key(region_id: &RegionId) -> String {
153    format!("remap_manifest_{}", region_id.as_u64())
154}
mito2/worker/handle_remap.rs

mito2/worker/
handle_remap.rs
