risingwave_common/util/
cluster_limit.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
// Copyright 2024 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::fmt::{self, Display, Formatter};

use risingwave_pb::meta::actor_count_per_parallelism::PbWorkerActorCount;
use risingwave_pb::meta::cluster_limit::PbLimit;
use risingwave_pb::meta::{PbActorCountPerParallelism, PbClusterLimit};
pub enum ClusterLimit {
    ActorCount(ActorCountPerParallelism),
}

impl From<ClusterLimit> for PbClusterLimit {
    fn from(limit: ClusterLimit) -> Self {
        match limit {
            ClusterLimit::ActorCount(actor_count_per_parallelism) => PbClusterLimit {
                limit: Some(PbLimit::ActorCount(actor_count_per_parallelism.into())),
            },
        }
    }
}

impl From<PbClusterLimit> for ClusterLimit {
    fn from(pb_limit: PbClusterLimit) -> Self {
        match pb_limit.limit.unwrap() {
            PbLimit::ActorCount(actor_count_per_parallelism) => {
                ClusterLimit::ActorCount(actor_count_per_parallelism.into())
            }
        }
    }
}

#[derive(Debug)]
pub struct WorkerActorCount {
    pub actor_count: usize,
    pub parallelism: usize,
}

impl From<WorkerActorCount> for PbWorkerActorCount {
    fn from(worker_actor_count: WorkerActorCount) -> Self {
        PbWorkerActorCount {
            actor_count: worker_actor_count.actor_count as u64,
            parallelism: worker_actor_count.parallelism as u64,
        }
    }
}

impl From<PbWorkerActorCount> for WorkerActorCount {
    fn from(pb_worker_actor_count: PbWorkerActorCount) -> Self {
        WorkerActorCount {
            actor_count: pb_worker_actor_count.actor_count as usize,
            parallelism: pb_worker_actor_count.parallelism as usize,
        }
    }
}

pub struct ActorCountPerParallelism {
    pub worker_id_to_actor_count: HashMap<u32, WorkerActorCount>,
    pub hard_limit: usize,
    pub soft_limit: usize,
}

impl From<ActorCountPerParallelism> for PbActorCountPerParallelism {
    fn from(actor_count_per_parallelism: ActorCountPerParallelism) -> Self {
        PbActorCountPerParallelism {
            worker_id_to_actor_count: actor_count_per_parallelism
                .worker_id_to_actor_count
                .into_iter()
                .map(|(k, v)| (k, v.into()))
                .collect(),
            hard_limit: actor_count_per_parallelism.hard_limit as u64,
            soft_limit: actor_count_per_parallelism.soft_limit as u64,
        }
    }
}

impl From<PbActorCountPerParallelism> for ActorCountPerParallelism {
    fn from(pb_actor_count_per_parallelism: PbActorCountPerParallelism) -> Self {
        ActorCountPerParallelism {
            worker_id_to_actor_count: pb_actor_count_per_parallelism
                .worker_id_to_actor_count
                .into_iter()
                .map(|(k, v)| (k, v.into()))
                .collect(),
            hard_limit: pb_actor_count_per_parallelism.hard_limit as usize,
            soft_limit: pb_actor_count_per_parallelism.soft_limit as usize,
        }
    }
}

impl ActorCountPerParallelism {
    pub fn exceed_hard_limit(&self) -> bool {
        self.worker_id_to_actor_count
            .values()
            .any(|v| v.actor_count > self.hard_limit.saturating_mul(v.parallelism))
    }

    pub fn exceed_soft_limit(&self) -> bool {
        self.worker_id_to_actor_count
            .values()
            .any(|v| v.actor_count > self.soft_limit.saturating_mul(v.parallelism))
    }

    pub fn exceed_limit(&self) -> bool {
        self.exceed_soft_limit() || self.exceed_hard_limit()
    }
}

impl Display for ActorCountPerParallelism {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let worker_id_to_actor_count_str: Vec<_> = self
            .worker_id_to_actor_count
            .iter()
            .map(|(k, v)| format!("{} -> {:?}", k, v))
            .collect();
        write!(
            f,
            "ActorCountPerParallelism {{ critical limit: {:?}, recommended limit: {:?}. worker_id_to_actor_count: {:?} }}",
            self.hard_limit, self.soft_limit, worker_id_to_actor_count_str
        )
    }
}