meta_srv/
key.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
15mod datanode;
16mod flownode;
17mod node_scheduling;
18
19use std::str::FromStr;
20
21use api::v1::meta::heartbeat_request::NodeWorkloads;
22pub use datanode::*;
23pub use flownode::*;
24pub use node_scheduling::*;
25use serde::{Deserialize, Serialize};
26use snafu::{OptionExt, ResultExt, ensure};
27
28use crate::error;
29
30macro_rules! impl_from_str_lease_key {
31    ($key_type:ty, $pattern:expr) => {
32        impl FromStr for $key_type {
33            type Err = error::Error;
34
35            fn from_str(key: &str) -> error::Result<Self> {
36                let caps = $pattern
37                    .captures(key)
38                    .context(error::InvalidLeaseKeySnafu { key })?;
39
40                ensure!(caps.len() == 3, error::InvalidLeaseKeySnafu { key });
41                let node_id = caps[2].to_string();
42                let node_id: u64 = node_id.parse().context(error::ParseNumSnafu {
43                    err_msg: format!("invalid node_id: {node_id}"),
44                })?;
45
46                Ok(Self { node_id })
47            }
48        }
49    };
50}
51
52impl_from_str_lease_key!(FlownodeLeaseKey, FLOWNODE_LEASE_KEY_PATTERN);
53impl_from_str_lease_key!(DatanodeLeaseKey, DATANODE_LEASE_KEY_PATTERN);
54
55macro_rules! impl_try_from_lease_key {
56    ($key_type:ty, $prefix:expr) => {
57        impl TryFrom<Vec<u8>> for $key_type {
58            type Error = error::Error;
59
60            fn try_from(bytes: Vec<u8>) -> error::Result<Self> {
61                String::from_utf8(bytes)
62                    .context(error::LeaseKeyFromUtf8Snafu {})
63                    .map(|x| x.parse())?
64            }
65        }
66
67        impl TryFrom<$key_type> for Vec<u8> {
68            type Error = error::Error;
69
70            fn try_from(key: $key_type) -> error::Result<Self> {
71                Ok(format!("{}-0-{}", $prefix, key.node_id).into_bytes())
72            }
73        }
74    };
75}
76
77impl_try_from_lease_key!(FlownodeLeaseKey, FLOWNODE_LEASE_PREFIX);
78impl_try_from_lease_key!(DatanodeLeaseKey, DATANODE_LEASE_PREFIX);
79
80#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
81pub struct LeaseValue {
82    // last activity
83    pub timestamp_millis: i64,
84    pub node_addr: String,
85    pub workloads: NodeWorkloads,
86}
87
88impl FromStr for LeaseValue {
89    type Err = error::Error;
90
91    fn from_str(value: &str) -> crate::Result<Self> {
92        serde_json::from_str(value).context(error::DeserializeFromJsonSnafu {
93            input_len: value.len(),
94        })
95    }
96}
97
98impl TryFrom<Vec<u8>> for LeaseValue {
99    type Error = error::Error;
100
101    fn try_from(bytes: Vec<u8>) -> crate::Result<Self> {
102        String::from_utf8(bytes)
103            .context(error::LeaseValueFromUtf8Snafu {})
104            .map(|x| x.parse())?
105    }
106}
107
108impl TryFrom<LeaseValue> for Vec<u8> {
109    type Error = error::Error;
110
111    fn try_from(value: LeaseValue) -> crate::Result<Self> {
112        Ok(serde_json::to_string(&value)
113            .context(error::SerializeToJsonSnafu {
114                input: format!("{value:?}"),
115            })?
116            .into_bytes())
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use api::v1::meta::DatanodeWorkloads;
123
124    use super::*;
125
126    #[test]
127    fn test_lease_value_round_trip() {
128        let value = LeaseValue {
129            timestamp_millis: 111,
130            node_addr: "127.0.0.1:3002".to_string(),
131            workloads: NodeWorkloads::Datanode(DatanodeWorkloads { types: vec![] }),
132        };
133
134        let value_bytes: Vec<u8> = value.clone().try_into().unwrap();
135        let new_value: LeaseValue = value_bytes.try_into().unwrap();
136
137        assert_eq!(new_value, value);
138    }
139}
meta_srv/key.rs

meta_srv/
key.rs
