1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use serde::ser::{Serialize, Serializer};
use std::collections::HashMap;

/// An integer metric value.
pub enum Integer {
    /// A signed value.
    Signed(i64),

    /// An unsigned value.
    Unsigned(u64),
}

impl From<i64> for Integer {
    fn from(i: i64) -> Integer {
        Integer::Signed(i)
    }
}

impl From<u64> for Integer {
    fn from(i: u64) -> Integer {
        Integer::Unsigned(i)
    }
}

enum TreeEntry {
    Value(Integer),
    Nested(MetricsTree),
}

impl Serialize for TreeEntry {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match self {
            TreeEntry::Value(value) => match value {
                Integer::Signed(i) => serializer.serialize_i64(*i),
                Integer::Unsigned(i) => serializer.serialize_u64(*i),
            },
            TreeEntry::Nested(tree) => tree.serialize(serializer),
        }
    }
}

/// A tree-structured metrics container.
///
/// Used for building a tree structure out of scoped metrics, where each level in the tree
/// represents a nested scope.
#[derive(Default)]
pub struct MetricsTree {
    contents: HashMap<String, TreeEntry>,
}

impl MetricsTree {
    /// Inserts a single value into the tree.
    pub fn insert_value<V: Into<Integer>>(
        &mut self,
        mut levels: Vec<String>,
        key: String,
        value: V,
    ) {
        match levels.len() {
            0 => {
                self.contents.insert(key, TreeEntry::Value(value.into()));
            }
            _ => {
                let name = levels.remove(0);
                let inner = self
                    .contents
                    .entry(name)
                    .or_insert_with(|| TreeEntry::Nested(MetricsTree::default()));

                if let TreeEntry::Nested(tree) = inner {
                    tree.insert_value(levels, key, value);
                }
            }
        }
    }

    /// Inserts multiple values into the tree.
    pub fn insert_values<V: Into<Integer>>(
        &mut self,
        mut levels: Vec<String>,
        values: Vec<(String, V)>,
    ) {
        match levels.len() {
            0 => {
                for v in values.into_iter() {
                    self.contents.insert(v.0, TreeEntry::Value(v.1.into()));
                }
            }
            _ => {
                let name = levels.remove(0);
                let inner = self
                    .contents
                    .entry(name)
                    .or_insert_with(|| TreeEntry::Nested(MetricsTree::default()));

                if let TreeEntry::Nested(tree) = inner {
                    tree.insert_values(levels, values);
                }
            }
        }
    }

    /// Clears all entries in the tree.
    pub fn clear(&mut self) {
        self.contents.clear();
    }
}

impl Serialize for MetricsTree {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut sorted = self.contents.iter().collect::<Vec<_>>();
        sorted.sort_by_key(|p| p.0);

        serializer.collect_map(sorted)
    }
}