1use bytesize::ByteSize;
2use clap::ValueEnum;
3use derive_more::{Debug, Display, From, Into};
4use maelstrom_macro::pocket_definition;
5use serde::{
6 de::{self, Deserializer, Visitor},
7 Deserialize, Serialize,
8};
9use slog::Level;
10use std::{
11 error, fmt, io,
12 net::{SocketAddr, ToSocketAddrs},
13 num::ParseIntError,
14 result,
15 str::FromStr,
16};
17use strum::EnumString;
18
19#[pocket_definition(export)]
20#[derive(Clone, Copy, Debug, Display, PartialEq, Serialize)]
21#[serde(transparent)]
22#[debug("{_0:?}")]
23#[display("{_0}")]
24pub struct BrokerAddr(SocketAddr);
25
26impl BrokerAddr {
27 pub fn new(inner: SocketAddr) -> Self {
28 BrokerAddr(inner)
29 }
30
31 pub fn inner(&self) -> &SocketAddr {
32 &self.0
33 }
34
35 pub fn into_inner(self) -> SocketAddr {
36 self.0
37 }
38}
39
40impl FromStr for BrokerAddr {
41 type Err = io::Error;
42
43 fn from_str(value: &str) -> Result<Self, Self::Err> {
44 let addrs: Vec<SocketAddr> = value.to_socket_addrs()?.collect();
45 Ok(BrokerAddr(*addrs.first().unwrap()))
48 }
49}
50
51impl TryFrom<String> for BrokerAddr {
52 type Error = io::Error;
53
54 fn try_from(value: String) -> Result<Self, Self::Error> {
55 value.parse()
56 }
57}
58
59impl From<BrokerAddr> for String {
60 fn from(b: BrokerAddr) -> Self {
61 b.to_string()
62 }
63}
64
65impl<'de> Deserialize<'de> for BrokerAddr {
66 fn deserialize<D>(deserializer: D) -> Result<BrokerAddr, D::Error>
67 where
68 D: Deserializer<'de>,
69 {
70 struct BrokerAddrVisitor;
71 impl Visitor<'_> for BrokerAddrVisitor {
72 type Value = BrokerAddr;
73
74 fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
75 formatter.write_str("a socket address")
76 }
77
78 fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
79 where
80 E: de::Error,
81 {
82 BrokerAddr::from_str(v).map_err(de::Error::custom)
83 }
84 }
85 if deserializer.is_human_readable() {
86 deserializer.deserialize_str(BrokerAddrVisitor)
87 } else {
88 SocketAddr::deserialize(deserializer).map(BrokerAddr::new)
89 }
90 }
91}
92
93#[derive(Debug, Display)]
94#[display("{_0}")]
95pub struct StringError(pub String);
96
97impl StringError {
98 pub fn new<T: ToString>(inner: T) -> Self {
99 Self(inner.to_string())
100 }
101}
102
103macro_rules! byte_size_u64_from_impls {
104 ($name:ident) => {
105 impl From<$name> for u64 {
106 fn from(l: $name) -> u64 {
107 (l.0).0
108 }
109 }
110
111 impl From<u64> for $name {
112 fn from(v: u64) -> Self {
113 ByteSize(v).into()
114 }
115 }
116 };
117}
118
119impl error::Error for StringError {}
120
121#[pocket_definition(export)]
122#[derive(Clone, Copy, Debug, Display, Deserialize, Eq, PartialEq, From, Into)]
123#[serde(transparent)]
124#[debug("{_0:?}")]
125#[display("{_0}")]
126pub struct CacheSize(#[serde(with = "bytesize_serde")] ByteSize);
127
128byte_size_u64_from_impls!(CacheSize);
129
130impl Default for CacheSize {
131 fn default() -> Self {
132 Self(ByteSize::gb(1))
133 }
134}
135
136impl FromStr for CacheSize {
137 type Err = StringError;
138 fn from_str(s: &str) -> Result<Self, Self::Err> {
139 Ok(Self(
140 <ByteSize as FromStr>::from_str(s).map_err(StringError)?,
141 ))
142 }
143}
144
145#[derive(Clone, Copy, Debug, Deserialize, EnumString, Serialize, ValueEnum)]
146#[clap(rename_all = "kebab_case")]
147#[serde(rename_all = "kebab-case")]
148#[strum(serialize_all = "kebab-case")]
149pub enum LogLevel {
150 Error,
151 Warning,
152 Info,
153 Debug,
154}
155
156impl LogLevel {
157 pub fn as_slog_level(&self) -> Level {
158 match self {
159 LogLevel::Error => slog::Level::Error,
160 LogLevel::Warning => slog::Level::Warning,
161 LogLevel::Info => slog::Level::Info,
162 LogLevel::Debug => slog::Level::Debug,
163 }
164 }
165}
166
167#[pocket_definition(export)]
168#[derive(Clone, Copy, Debug, Display, Deserialize, Eq, PartialEq, From, Into)]
169#[debug("{_0:?}")]
170#[display("{_0}")]
171#[serde(transparent)]
172pub struct InlineLimit(#[serde(with = "bytesize_serde")] ByteSize);
173
174byte_size_u64_from_impls!(InlineLimit);
175
176impl Default for InlineLimit {
177 fn default() -> Self {
178 Self(ByteSize::mb(1))
179 }
180}
181
182impl FromStr for InlineLimit {
183 type Err = StringError;
184 fn from_str(s: &str) -> Result<Self, Self::Err> {
185 Ok(Self(
186 <ByteSize as FromStr>::from_str(s).map_err(StringError)?,
187 ))
188 }
189}
190
191#[pocket_definition(export)]
192#[derive(Clone, Copy, Debug, Display, Deserialize, Into)]
193#[serde(try_from = "u16")]
194#[into(usize, u16, u32)]
195#[debug("{_0:?}")]
196#[display("{_0}")]
197pub struct Slots(u16);
198
199impl Slots {
200 pub fn inner(&self) -> &u16 {
201 &self.0
202 }
203
204 pub fn into_inner(self) -> u16 {
205 self.0
206 }
207}
208
209impl TryFrom<u16> for Slots {
210 type Error = String;
211
212 fn try_from(value: u16) -> Result<Self, Self::Error> {
213 if value < 1 {
214 Err("value must be at least 1".to_string())
215 } else if value > 1000 {
216 Err("value must be less than 1000".to_string())
217 } else {
218 Ok(Slots(value))
219 }
220 }
221}
222
223macro_rules! slots_try_from_impl {
224 ($($int_type:ty),*) => {$(
225 impl TryFrom<$int_type> for Slots {
226 type Error = String;
227
228 fn try_from(value: $int_type) -> Result<Self, Self::Error> {
229 let value: u16 = value
230 .try_into()
231 .map_err(|_| "value must be less than 1000".to_string())?;
232 value.try_into()
233 }
234 }
235 )*}
236}
237
238slots_try_from_impl!(usize, u32);
239
240impl Default for Slots {
241 fn default() -> Self {
242 Self::try_from(num_cpus::get()).unwrap()
243 }
244}
245
246impl FromStr for Slots {
247 type Err = SlotsFromStrError;
248 fn from_str(slots: &str) -> result::Result<Self, Self::Err> {
249 let slots = u16::from_str(slots).map_err(SlotsFromStrError::Parse)?;
250 Self::try_from(slots).map_err(SlotsFromStrError::Bounds)
251 }
252}
253
254#[derive(Debug, Display)]
255pub enum SlotsFromStrError {
256 #[display("{_0}")]
257 Parse(ParseIntError),
258 #[display("{_0}")]
259 Bounds(String),
260}
261
262impl error::Error for SlotsFromStrError {}
263
264#[pocket_definition(export)]
265#[derive(Copy, Clone, Debug, EnumString, Deserialize)]
266#[serde(rename_all = "kebab-case")]
267#[strum(serialize_all = "kebab-case")]
268pub enum ClusterCommunicationStrategy {
269 Tcp = 0,
270 #[serde(rename = "github")]
271 #[strum(serialize = "github")]
272 GitHub = 1,
273}
274
275impl ValueEnum for ClusterCommunicationStrategy {
276 fn value_variants<'a>() -> &'a [Self] {
277 &[Self::Tcp, Self::GitHub][..]
278 }
279
280 fn to_possible_value(&self) -> Option<clap::builder::PossibleValue> {
281 match self {
282 Self::Tcp => Some(clap::builder::PossibleValue::new("tcp")),
283 Self::GitHub => Some(clap::builder::PossibleValue::new("github")),
284 }
285 }
286}
287
288#[cfg(test)]
289mod tests {
290 use super::{Debug, *};
291 use serde_test::{assert_de_tokens, assert_tokens, Configure, Token};
292 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
293
294 const LOCALHOST4: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1234);
295 const LOCALHOST6: SocketAddr = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 1234);
296
297 #[test]
298 fn broker_addr_from_str() {
299 assert_eq!(
300 BrokerAddr::from_str("127.0.0.1:1234").unwrap(),
301 BrokerAddr::new(LOCALHOST4),
302 );
303 assert_eq!(
304 BrokerAddr::from_str("localhost:1234").unwrap(),
305 BrokerAddr::new(LOCALHOST4),
306 );
307 assert_eq!(
308 BrokerAddr::from_str("[::1]:1234").unwrap(),
309 BrokerAddr::new(LOCALHOST6),
310 );
311 assert_eq!(
312 BrokerAddr::from_str("ip6-localhost:1234").unwrap(),
313 BrokerAddr::new(LOCALHOST6),
314 );
315 }
316
317 #[test]
318 fn broker_addr_into_string() {
319 assert_eq!(String::from(BrokerAddr::new(LOCALHOST4)), "127.0.0.1:1234");
320 assert_eq!(String::from(BrokerAddr::new(LOCALHOST6)), "[::1]:1234");
321 }
322
323 #[test]
324 fn broker_addr_display() {
325 assert_eq!(
326 format!("{}", BrokerAddr::new(LOCALHOST4)),
327 format!("{LOCALHOST4}")
328 );
329 assert_eq!(
330 format!("{}", BrokerAddr::new(LOCALHOST6)),
331 format!("{LOCALHOST6}")
332 );
333 }
334
335 #[test]
336 fn broker_addr_debug() {
337 assert_eq!(
338 format!("{:?}", BrokerAddr::new(LOCALHOST4)),
339 format!("{LOCALHOST4:?}")
340 );
341 assert_eq!(
342 format!("{:?}", BrokerAddr::new(LOCALHOST6)),
343 format!("{LOCALHOST6:?}")
344 );
345 }
346
347 #[test]
348 fn broker_addr_serialize_and_deserialize_human_readable() {
349 assert_tokens(
350 &BrokerAddr::new(LOCALHOST4).readable(),
351 &[Token::String("127.0.0.1:1234")],
352 );
353 assert_tokens(
354 &BrokerAddr::new(LOCALHOST6).readable(),
355 &[Token::String("[::1]:1234")],
356 );
357 }
358
359 #[test]
360 fn broker_addr_serialize_and_deserialize_compact() {
361 assert_tokens(
362 &BrokerAddr::new(LOCALHOST4).compact(),
363 &[
364 Token::NewtypeVariant {
365 name: "SocketAddr",
366 variant: "V4",
367 },
368 Token::Tuple { len: 2 },
369 Token::Tuple { len: 4 },
370 Token::U8(127),
371 Token::U8(0),
372 Token::U8(0),
373 Token::U8(1),
374 Token::TupleEnd,
375 Token::U16(1234),
376 Token::TupleEnd,
377 ],
378 );
379 assert_tokens(
380 &BrokerAddr::new(LOCALHOST6).compact(),
381 &[
382 Token::NewtypeVariant {
383 name: "SocketAddr",
384 variant: "V6",
385 },
386 Token::Tuple { len: 2 },
387 Token::Tuple { len: 16 },
388 Token::U8(0),
389 Token::U8(0),
390 Token::U8(0),
391 Token::U8(0),
392 Token::U8(0),
393 Token::U8(0),
394 Token::U8(0),
395 Token::U8(0),
396 Token::U8(0),
397 Token::U8(0),
398 Token::U8(0),
399 Token::U8(0),
400 Token::U8(0),
401 Token::U8(0),
402 Token::U8(0),
403 Token::U8(1),
404 Token::TupleEnd,
405 Token::U16(1234),
406 Token::TupleEnd,
407 ],
408 );
409 }
410
411 #[test]
412 fn broker_addr_deserialize_from_string() {
413 assert_de_tokens(
414 &BrokerAddr::new(LOCALHOST4).readable(),
415 &[Token::String("127.0.0.1:1234")],
416 );
417 assert_de_tokens(
418 &BrokerAddr::new(LOCALHOST6).readable(),
419 &[Token::String("[::1]:1234")],
420 );
421 }
422
423 #[test]
424 fn broker_addr_deserialize_from_hostname_string() {
425 assert_de_tokens(
426 &BrokerAddr::new(LOCALHOST4).readable(),
427 &[Token::String("localhost:1234")],
428 );
429 assert_de_tokens(
430 &BrokerAddr::new(LOCALHOST6).readable(),
431 &[Token::String("ip6-localhost:1234")],
432 );
433 }
434
435 #[derive(Debug, Deserialize, Eq, PartialEq)]
436 struct TwoCacheSizes {
437 one: CacheSize,
438 two: CacheSize,
439 }
440
441 #[test]
442 fn cache_size() {
443 assert_eq!(
444 CacheSize::from(ByteSize::b(1234)),
445 CacheSize::from(ByteSize::b(1234))
446 );
447 assert_eq!(
448 format!("{:?}", CacheSize::from(ByteSize::b(1234))),
449 "1.2 KB"
450 );
451 assert_eq!(format!("{}", CacheSize::from(ByteSize::b(1234))), "1.2 KB");
452 assert_eq!(
453 CacheSize::from_str("10 MB").unwrap(),
454 CacheSize::from(ByteSize::mb(10))
455 );
456 let two_cache_sizes: TwoCacheSizes = toml::from_str(
457 r#"
458 one = 1000
459 two = "1.23 MB"
460 "#,
461 )
462 .unwrap();
463 assert_eq!(
464 two_cache_sizes,
465 TwoCacheSizes {
466 one: "1 kB".parse().unwrap(),
467 two: "1.23 MB".parse().unwrap()
468 }
469 );
470 }
471
472 #[derive(Debug, Deserialize, Eq, PartialEq)]
473 struct TwoInlineLimits {
474 one: InlineLimit,
475 two: InlineLimit,
476 }
477
478 #[test]
479 fn inline_limit() {
480 assert_eq!(
481 InlineLimit::from(ByteSize::b(1234)),
482 InlineLimit::from(ByteSize::b(1234))
483 );
484 assert_eq!(
485 format!("{:?}", InlineLimit::from(ByteSize::b(1234))),
486 "1.2 KB"
487 );
488 assert_eq!(
489 format!("{}", InlineLimit::from(ByteSize::b(1234))),
490 "1.2 KB"
491 );
492 assert_eq!(
493 InlineLimit::from_str("10 MB").unwrap(),
494 InlineLimit::from(ByteSize::mb(10))
495 );
496 let two_inline_limits: TwoInlineLimits = toml::from_str(
497 r#"
498 one = 1000
499 two = "1.23 MB"
500 "#,
501 )
502 .unwrap();
503 assert_eq!(
504 two_inline_limits,
505 TwoInlineLimits {
506 one: "1 kB".parse().unwrap(),
507 two: "1.23 MB".parse().unwrap()
508 }
509 );
510 }
511}