use std::borrow::Cow;
use std::cmp::Ordering;
use std::hash::{Hash, Hasher};
pub type Label = (Cow<'static, str>, Cow<'static, str>);
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct LabelSet {
pairs: Vec<Label>,
}
impl LabelSet {
pub const EMPTY: LabelSet = LabelSet { pairs: Vec::new() };
#[inline]
pub const fn new() -> Self {
Self::EMPTY
}
#[inline]
pub fn len(&self) -> usize {
self.pairs.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.pairs.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> + '_ {
self.pairs.iter().map(|(k, v)| (k.as_ref(), v.as_ref()))
}
pub fn add(
&mut self,
key: impl Into<Cow<'static, str>>,
value: impl Into<Cow<'static, str>>,
) -> &mut Self {
let key = key.into();
let value = value.into();
match self
.pairs
.binary_search_by(|(k, _)| (k.as_ref()).cmp(key.as_ref()))
{
Ok(idx) => self.pairs[idx].1 = value,
Err(idx) => self.pairs.insert(idx, (key, value)),
}
self
}
#[must_use]
pub fn with(
mut self,
key: impl Into<Cow<'static, str>>,
value: impl Into<Cow<'static, str>>,
) -> Self {
self.add(key, value);
self
}
pub fn get(&self, key: &str) -> Option<&str> {
self.pairs
.binary_search_by(|(k, _)| (k.as_ref()).cmp(key))
.ok()
.map(|idx| self.pairs[idx].1.as_ref())
}
pub fn remove(&mut self, key: &str) -> bool {
if let Ok(idx) = self.pairs.binary_search_by(|(k, _)| (k.as_ref()).cmp(key)) {
self.pairs.remove(idx);
true
} else {
false
}
}
pub fn to_prometheus(&self) -> String {
if self.pairs.is_empty() {
return String::new();
}
let mut out = String::with_capacity(2 + self.pairs.len() * 16);
out.push('{');
for (i, (k, v)) in self.pairs.iter().enumerate() {
if i > 0 {
out.push(',');
}
out.push_str(k);
out.push_str("=\"");
escape_prometheus_value(&mut out, v);
out.push('"');
}
out.push('}');
out
}
pub fn to_statsd(&self) -> String {
if self.pairs.is_empty() {
return String::new();
}
let mut out = String::with_capacity(2 + self.pairs.len() * 16);
out.push_str("|#");
for (i, (k, v)) in self.pairs.iter().enumerate() {
if i > 0 {
out.push(',');
}
out.push_str(k);
out.push(':');
for c in v.chars() {
if matches!(c, '|' | ',' | '\n' | ':') {
out.push('_');
} else {
out.push(c);
}
}
}
out
}
}
impl Hash for LabelSet {
fn hash<H: Hasher>(&self, state: &mut H) {
self.pairs.len().hash(state);
for (k, v) in &self.pairs {
k.as_ref().hash(state);
v.as_ref().hash(state);
}
}
}
impl PartialOrd for LabelSet {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for LabelSet {
fn cmp(&self, other: &Self) -> Ordering {
let a = self.pairs.iter().map(|(k, v)| (k.as_ref(), v.as_ref()));
let b = other.pairs.iter().map(|(k, v)| (k.as_ref(), v.as_ref()));
a.cmp(b)
}
}
impl<K, V> FromIterator<(K, V)> for LabelSet
where
K: Into<Cow<'static, str>>,
V: Into<Cow<'static, str>>,
{
fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
let mut s = Self::new();
for (k, v) in iter {
s.add(k, v);
}
s
}
}
impl<K, V, const N: usize> From<[(K, V); N]> for LabelSet
where
K: Into<Cow<'static, str>>,
V: Into<Cow<'static, str>>,
{
fn from(arr: [(K, V); N]) -> Self {
arr.into_iter().collect()
}
}
#[cfg(feature = "serde")]
impl serde::Serialize for LabelSet {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
use serde::ser::SerializeMap;
let mut map = serializer.serialize_map(Some(self.pairs.len()))?;
for (k, v) in &self.pairs {
map.serialize_entry(k.as_ref(), v.as_ref())?;
}
map.end()
}
}
fn escape_prometheus_value(out: &mut String, v: &str) {
for c in v.chars() {
match c {
'\\' => out.push_str("\\\\"),
'"' => out.push_str("\\\""),
'\n' => out.push_str("\\n"),
c => out.push(c),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty_set_is_empty_and_renders_empty() {
let l = LabelSet::EMPTY;
assert!(l.is_empty());
assert_eq!(l.len(), 0);
assert_eq!(l.to_prometheus(), "");
assert_eq!(l.to_statsd(), "");
}
#[test]
fn add_keeps_sorted_and_deduplicates() {
let mut l = LabelSet::new();
l.add("status", "200");
l.add("method", "GET");
l.add("region", "us");
let pairs: Vec<_> = l.iter().collect();
assert_eq!(
pairs,
vec![("method", "GET"), ("region", "us"), ("status", "200")]
);
l.add("status", "500");
assert_eq!(l.get("status"), Some("500"));
assert_eq!(l.len(), 3);
}
#[test]
fn equal_sets_hash_equal_regardless_of_insertion_order() {
use std::collections::hash_map::DefaultHasher;
fn hash(l: &LabelSet) -> u64 {
let mut h = DefaultHasher::new();
l.hash(&mut h);
h.finish()
}
let a = LabelSet::from([("a", "1"), ("b", "2"), ("c", "3")]);
let b = LabelSet::from([("c", "3"), ("a", "1"), ("b", "2")]);
assert_eq!(a, b);
assert_eq!(hash(&a), hash(&b));
}
#[test]
fn remove_keeps_invariants() {
let mut l = LabelSet::from([("a", "1"), ("b", "2"), ("c", "3")]);
assert!(l.remove("b"));
assert!(!l.remove("b"));
assert_eq!(l.iter().collect::<Vec<_>>(), vec![("a", "1"), ("c", "3")]);
}
#[test]
fn prometheus_rendering_escapes_correctly() {
let l = LabelSet::from([("path", r#"/foo "bar"\baz"#), ("note", "line1\nline2")]);
let s = l.to_prometheus();
assert_eq!(s, r#"{note="line1\nline2",path="/foo \"bar\"\\baz"}"#);
}
#[test]
fn statsd_rendering_sanitises_specials() {
let l = LabelSet::from([("k1", "with|pipe"), ("k2", "with,comma:colon")]);
let s = l.to_statsd();
assert_eq!(s, "|#k1:with_pipe,k2:with_comma_colon");
}
#[test]
fn ordering_is_lexicographic_over_pairs() {
let a = LabelSet::from([("a", "1")]);
let b = LabelSet::from([("a", "2")]);
let c = LabelSet::from([("b", "0")]);
let mut v = vec![c.clone(), b.clone(), a.clone()];
v.sort();
assert_eq!(v, vec![a, b, c]);
}
#[cfg(feature = "serde")]
#[test]
fn serde_serializes_as_map() {
let l = LabelSet::from([("method", "GET"), ("status", "200")]);
let j = serde_json::to_string(&l).unwrap();
assert_eq!(j, r#"{"method":"GET","status":"200"}"#);
}
}