meta_srv/procedure/repartition/
allocate_region.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::any::Any;
16
17use common_procedure::{Context as ProcedureContext, Status};
18use serde::{Deserialize, Serialize};
19
20use crate::error::Result;
21use crate::procedure::repartition::dispatch::Dispatch;
22use crate::procedure::repartition::plan::{AllocationPlanEntry, RepartitionPlanEntry};
23use crate::procedure::repartition::{Context, State};
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct AllocateRegion {
27    plan_entries: Vec<AllocationPlanEntry>,
28}
29
30impl AllocateRegion {
31    pub fn new(plan_entries: Vec<AllocationPlanEntry>) -> Self {
32        Self { plan_entries }
33    }
34}
35
36#[async_trait::async_trait]
37#[typetag::serde]
38impl State for AllocateRegion {
39    async fn next(
40        &mut self,
41        ctx: &mut Context,
42        _procedure_ctx: &ProcedureContext,
43    ) -> Result<(Box<dyn State>, Status)> {
44        let region_to_allocate = self
45            .plan_entries
46            .iter()
47            .map(|p| p.regions_to_allocate)
48            .sum::<usize>();
49
50        if region_to_allocate == 0 {
51            let repartition_plan_entries = self
52                .plan_entries
53                .iter()
54                .map(RepartitionPlanEntry::from_allocation_plan_entry)
55                .collect::<Vec<_>>();
56            ctx.persistent_ctx.plans = repartition_plan_entries;
57            return Ok((Box::new(Dispatch), Status::executing(true)));
58        }
59
60        // TODO(weny): allocate regions.
61        todo!()
62    }
63
64    fn as_any(&self) -> &dyn Any {
65        self
66    }
67}