common_workload/
lib.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_telemetry::info;
16use serde::{Deserialize, Serialize};
17
18/// The workload type of the datanode.
19#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq, Hash)]
20#[serde(rename_all = "snake_case")]
21pub enum DatanodeWorkloadType {
22    /// The datanode can handle all workloads.
23    Hybrid = 0,
24}
25
26impl DatanodeWorkloadType {
27    /// Convert from `i32` to `DatanodeWorkloadType`.
28    pub fn from_i32(value: i32) -> Option<Self> {
29        match value {
30            v if v == Self::Hybrid as i32 => Some(Self::Hybrid),
31            _ => None,
32        }
33    }
34
35    /// Convert from `DatanodeWorkloadType` to `i32`.
36    pub fn to_i32(self) -> i32 {
37        self as i32
38    }
39
40    pub fn accept_ingest(&self) -> bool {
41        matches!(self, Self::Hybrid)
42    }
43}
44
45/// Sanitize the workload types.
46pub fn sanitize_workload_types(workload_types: &mut Vec<DatanodeWorkloadType>) {
47    if workload_types.is_empty() {
48        info!("The workload types is empty, use Hybrid workload type");
49        workload_types.push(DatanodeWorkloadType::Hybrid);
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[test]
58    fn test_sanitize_workload_types() {
59        let hybrid = DatanodeWorkloadType::Hybrid;
60        assert_eq!(hybrid as i32, 0);
61        let hybrid_i32 = hybrid.to_i32();
62        assert_eq!(hybrid_i32, 0);
63        assert_eq!(DatanodeWorkloadType::from_i32(hybrid_i32), Some(hybrid));
64
65        let unexpected_i32 = 100;
66        assert_eq!(DatanodeWorkloadType::from_i32(unexpected_i32), None);
67    }
68}