use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
pub const CHANNEL_VERSION: &str = "1";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum MetricType {
Gauge,
Counter,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "kebab-case")]
pub enum ChildMessage {
Ready,
Metric {
name: String,
value: f64,
#[serde(rename = "type", skip_serializing_if = "Option::is_none", default)]
metric_type: Option<MetricType>,
#[serde(skip_serializing_if = "BTreeMap::is_empty", default)]
labels: BTreeMap<String, String>,
#[serde(skip_serializing_if = "Option::is_none", default)]
help: Option<String>,
},
ActionReply {
action: String,
body: String,
#[serde(skip_serializing_if = "Option::is_none", default)]
id: Option<u64>,
},
LambLabel {
pid: u32,
label: LambLabel,
},
Note {
name: String,
value: f64,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(try_from = "String", into = "String")]
pub struct LambLabel(String);
impl LambLabel {
pub const MAX_CHARS: usize = 64;
pub fn new(label: impl Into<String>) -> Result<Self, LambLabelError> {
let label = label.into();
let chars = label.chars().count();
if chars > Self::MAX_CHARS {
return Err(LambLabelError::TooLong { chars });
}
if label.chars().any(char::is_control) {
return Err(LambLabelError::ControlCharacter);
}
Ok(Self(label))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
}
impl TryFrom<String> for LambLabel {
type Error = LambLabelError;
fn try_from(label: String) -> Result<Self, Self::Error> {
Self::new(label)
}
}
impl From<LambLabel> for String {
fn from(label: LambLabel) -> Self {
label.0
}
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LambLabelError {
TooLong {
chars: usize,
},
ControlCharacter,
}
impl core::fmt::Display for LambLabelError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::TooLong { chars } => write!(
f,
"a lamb label holds at most {} characters, and this one has {chars}",
LambLabel::MAX_CHARS
),
Self::ControlCharacter => f.write_str("a lamb label cannot hold a control character"),
}
}
}
impl core::error::Error for LambLabelError {}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "kebab-case")]
pub enum ShepherdMessage {
Shutdown,
Action {
name: String,
#[serde(skip_serializing_if = "Option::is_none", default)]
params: Option<String>,
id: u64,
},
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn ready_wire_fixture_round_trips() {
let fixture = r#"{"kind":"ready"}"#;
assert_eq!(
serde_json::from_str::<ChildMessage>(fixture).unwrap(),
ChildMessage::Ready
);
assert_eq!(
serde_json::to_string(&ChildMessage::Ready).unwrap(),
fixture
);
}
#[test]
fn metric_wire_fixture_round_trips() {
let fixture = r#"{"kind":"metric","name":"rps","value":42.0}"#;
let msg = ChildMessage::Metric {
name: "rps".to_string(),
value: 42.0,
metric_type: None,
labels: BTreeMap::new(),
help: None,
};
assert_eq!(serde_json::from_str::<ChildMessage>(fixture).unwrap(), msg);
assert_eq!(serde_json::to_string(&msg).unwrap(), fixture);
}
#[test]
fn a_metric_with_type_labels_and_help_round_trips() {
let fixture = concat!(
r#"{"kind":"metric","name":"queue_depth","value":3.0,"type":"counter","#,
r#""labels":{"model":"qwen"},"help":"Requests waiting for the lease."}"#,
);
let msg = ChildMessage::Metric {
name: "queue_depth".to_string(),
value: 3.0,
metric_type: Some(MetricType::Counter),
labels: [("model".to_string(), "qwen".to_string())].into(),
help: Some("Requests waiting for the lease.".to_string()),
};
assert_eq!(serde_json::from_str::<ChildMessage>(fixture).unwrap(), msg);
assert_eq!(serde_json::to_string(&msg).unwrap(), fixture);
}
#[test]
fn a_gauge_type_round_trips_as_its_own_spelling() {
let fixture = r#"{"kind":"metric","name":"rps","value":1.0,"type":"gauge"}"#;
let msg = ChildMessage::Metric {
name: "rps".to_string(),
value: 1.0,
metric_type: Some(MetricType::Gauge),
labels: BTreeMap::new(),
help: None,
};
assert_eq!(serde_json::from_str::<ChildMessage>(fixture).unwrap(), msg);
assert_eq!(serde_json::to_string(&msg).unwrap(), fixture);
}
#[test]
fn an_unknown_metric_type_fails_to_decode() {
for frame in [
r#"{"kind":"metric","name":"rps","value":1.0,"type":"histogram"}"#,
r#"{"kind":"metric","name":"rps","value":1.0,"type":"Counter"}"#,
] {
assert!(
serde_json::from_str::<ChildMessage>(frame).is_err(),
"{frame} decoded"
);
}
}
#[test]
fn an_action_reply_without_an_id_round_trips() {
let fixture = r#"{"kind":"action-reply","action":"gc","body":"ok"}"#;
let msg = ChildMessage::ActionReply {
action: "gc".to_string(),
body: "ok".to_string(),
id: None,
};
assert_eq!(serde_json::from_str::<ChildMessage>(fixture).unwrap(), msg);
assert_eq!(serde_json::to_string(&msg).unwrap(), fixture);
}
#[test]
fn an_action_reply_with_an_echoed_id_round_trips() {
let fixture = r#"{"kind":"action-reply","action":"gc","body":"ok","id":7}"#;
let msg = ChildMessage::ActionReply {
action: "gc".to_string(),
body: "ok".to_string(),
id: Some(7),
};
assert_eq!(serde_json::from_str::<ChildMessage>(fixture).unwrap(), msg);
assert_eq!(serde_json::to_string(&msg).unwrap(), fixture);
}
#[test]
fn shutdown_wire_fixture_round_trips() {
let fixture = r#"{"kind":"shutdown"}"#;
assert_eq!(
serde_json::from_str::<ShepherdMessage>(fixture).unwrap(),
ShepherdMessage::Shutdown
);
assert_eq!(
serde_json::to_string(&ShepherdMessage::Shutdown).unwrap(),
fixture
);
}
#[test]
fn an_action_carries_its_id_with_or_without_params() {
let bare = r#"{"kind":"action","name":"gc","id":7}"#;
let bare_msg = ShepherdMessage::Action {
name: "gc".to_string(),
params: None,
id: 7,
};
assert_eq!(serde_json::to_string(&bare_msg).unwrap(), bare);
assert_eq!(
serde_json::from_str::<ShepherdMessage>(bare).unwrap(),
bare_msg
);
let with_params = r#"{"kind":"action","name":"set-log-level","params":"debug","id":8}"#;
let with_params_msg = ShepherdMessage::Action {
name: "set-log-level".to_string(),
params: Some("debug".to_string()),
id: 8,
};
assert_eq!(
serde_json::to_string(&with_params_msg).unwrap(),
with_params
);
assert_eq!(
serde_json::from_str::<ShepherdMessage>(with_params).unwrap(),
with_params_msg
);
}
#[test]
fn a_lamb_label_round_trips() {
let fixture = r#"{"kind":"lamb-label","pid":4312,"label":"worker 1"}"#;
let msg = ChildMessage::LambLabel {
pid: 4312,
label: LambLabel::new("worker 1").unwrap(),
};
assert_eq!(serde_json::from_str::<ChildMessage>(fixture).unwrap(), msg);
assert_eq!(serde_json::to_string(&msg).unwrap(), fixture);
}
#[test]
fn an_empty_lamb_label_is_on_the_wire_and_clears() {
let fixture = r#"{"kind":"lamb-label","pid":4312,"label":""}"#;
let ChildMessage::LambLabel { label, .. } =
serde_json::from_str::<ChildMessage>(fixture).unwrap()
else {
panic!("{fixture} decoded as another kind");
};
assert!(label.is_empty());
assert!(
serde_json::from_str::<ChildMessage>(r#"{"kind":"lamb-label","pid":4312}"#).is_err()
);
}
#[test]
fn a_note_round_trips_with_any_name() {
let fixture = r#"{"kind":"note","name":"lease.web.want.1760500000","value":1.0}"#;
let msg = ChildMessage::Note {
name: "lease.web.want.1760500000".to_string(),
value: 1.0,
};
assert_eq!(serde_json::from_str::<ChildMessage>(fixture).unwrap(), msg);
assert_eq!(serde_json::to_string(&msg).unwrap(), fixture);
}
#[test]
fn a_note_without_a_value_fails_to_decode() {
assert!(serde_json::from_str::<ChildMessage>(r#"{"kind":"note","name":"x"}"#).is_err());
}
#[test]
fn a_label_is_counted_in_characters_up_to_the_limit() {
let at_limit = "é".repeat(LambLabel::MAX_CHARS);
assert_eq!(LambLabel::new(at_limit.clone()).unwrap().as_str(), at_limit);
assert_eq!(
LambLabel::new(format!("{at_limit}e")),
Err(LambLabelError::TooLong {
chars: LambLabel::MAX_CHARS + 1
})
);
}
#[test]
fn a_label_with_a_control_character_is_refused_on_the_wire_too() {
for label in ["a\nb", "\u{1b}[2J", "tab\there", "\u{9b}31m"] {
assert_eq!(
LambLabel::new(label),
Err(LambLabelError::ControlCharacter),
"{label:?}"
);
let frame = serde_json::json!({"kind": "lamb-label", "pid": 1, "label": label});
assert!(
serde_json::from_value::<ChildMessage>(frame).is_err(),
"{label:?} decoded"
);
}
}
#[test]
fn a_label_error_says_the_limit_and_the_length() {
assert_eq!(
LambLabelError::TooLong { chars: 70 }.to_string(),
"a lamb label holds at most 64 characters, and this one has 70"
);
}
}