netdata_plugin/

collector.rs

// SPDX-License-Identifier: AGPL-3.0-or-later
// 
use super::{Begin, Chart, Dimension, Instruction, Set};

use std::collections::HashMap;
use std::time::Instant;
use thiserror::Error;
use validator::{Validate, ValidationErrors};

use std::io::Write;

/// Error codes used by [Collector] and its methods.
#[derive(Error, Debug)]
pub enum CollectorError {
    /// The referred chart is not yet registert in the collector.
    #[error("unknown Chart: type.id = {0}")]
    UnkownChart(String),
    /// The referred chart is not yet registert in the collector.
    #[error("unknown Dimension: id = {0}")]
    UnkownDimension(String),
    /// One of the field entries violates the formal requirements.
    #[error(transparent)]
    ValidationErrors(#[from] ValidationErrors),
    /// Can't Write to the provided writer.
    #[error(transparent)]
    IOError(#[from] std::io::Error),
}

/// Internal List of active Dimensions.
#[derive(Default)]
struct CollectedDimensionInfo {
    /// Prepared value.
    value: i64,
    /// Will be set to `true` if already commited.
    commited: bool,
}

/// Internal List of active Charts.
#[derive(Default)]
struct CollectedChartInfo {
    dimensions: HashMap<String, CollectedDimensionInfo>,
    /// Timestamp of last commit.
    last_commit: Option<Instant>,
}

/// A High-Level interface to run the data collecting `BEGIN`-`SET`-`END` loop
/// and setup [Chart] and [Dimension] info in an efficient manner.
///
/// It's used roughly like this:
///
/// Example:
///
/// ```rust
/// use std::error;
/// use netdata_plugin::{Chart, Dimension};
/// use netdata_plugin::collector::Collector;
///
/// fn main() -> Result<(), Box<dyn error::Error>> {
///     // prepare collector
///     let mut writer = std::io::stdout();
///     let mut c = Collector::new(&mut writer);
///
///     // add charts and their associated dimensions
///     c.add_chart(&mut Chart { type_id: "mytype.something" /*...*/, ..Default::default()})?;
///     // writes CHART ...
///     c.add_dimension( "mytype.something", &Dimension {id: "d1", ..Default::default()})?;
///     c.add_dimension( "mytype.something", &Dimension {id: "d2", ..Default::default()})?;
///     // writes DIMENSION ...
///
///     // data collecting loop
///     loop {
///         c.prepare_value("mytype.something", "d1", 4242)?;
///         c.prepare_value("mytype.something", "d2", 4242)?;
///         c.commit_chart("mytype.something").unwrap();
///         // The BEGIN ... - SET ... - END block will be written to stdout
///
///         std::thread::sleep(std::time::Duration::from_secs(1));
///         break; // stop the demonstration ;)
///     }
///     Ok(())
/// }
/// ```
#[derive()]
pub struct Collector<'a, W: Write> {
    charts: HashMap<String, CollectedChartInfo>,
    writer: &'a mut W,
}

impl<'a, W: Write> Collector<'a, W> {
    /// Initilaize the data store to manage the list of active [Chart]s.
    ///
    /// In most cases you can use a mutable `stdout` reference as writer.
    pub fn new(writer: &'a mut W) -> Self {
        Collector {
            charts: HashMap::new(),
            writer: writer,
        }
    }

    /// Registers a new [Chart] in this [Collector] and write
    /// a `CHART...`-command to the plugins `stdout` to add the Chart properties
    /// on the host side, too.
    ///
    /// The `type_id` of the [Chart] entry will be checkt for
    /// [formal validity](super::validate_type_id()) in
    /// advance and may raise [ValidationErrors](CollectorError::ValidationErrors).
    pub fn add_chart(&mut self, chart: &Chart) -> Result<(), CollectorError> {
        chart.validate()?;
        let cci = CollectedChartInfo {
            ..Default::default()
        };
        self.charts.insert(format!("{}", chart.type_id), cci);
        writeln!(self.writer, "{}", chart)?;
        Ok(())
    }

    /// Registers a new [Dimension] for a Chart of a given `type.id` in this
    /// [Collector] and write a `DIMENSION...`-command to the plugins `stdout` to add it
    /// on the host side as well.
    ///
    /// The `id` field of the [Dimension] entry will be checkt for
    /// [formal validity](super::validate_id()) in
    /// advance and may raise [ValidationErrors](CollectorError::ValidationErrors).
    pub fn add_dimension(
        &mut self,
        chart_id: &str,
        dimension: &Dimension,
    ) -> Result<(), CollectorError> {
        dimension.validate()?;
        let cci = self
            .charts
            .get_mut(chart_id)
            .ok_or(CollectorError::UnkownChart(chart_id.to_owned()))?;
        let cdi = CollectedDimensionInfo {
            ..Default::default()
        };
        cci.dimensions.insert(dimension.id.to_owned(), cdi);
        writeln!(self.writer, "{}", dimension)?;
        Ok(())
    }

    /// Update the value for one [Dimension] to be later commited
    /// with all other updated values of this [Chart]
    ///
    /// If the reffered [Dimension] or [Chart] isn't already registered in the
    /// [Collector], it will will raise [UnkownChart](CollectorError::UnkownChart) or
    /// [UnkownDimension](CollectorError::UnkownDimension). This may be used to
    /// dynamically setup the needed categories on demand.
    pub fn prepare_value(
        &mut self,
        chart_id: &str,
        dimension_id: &str,
        value: i64,
    ) -> Result<(), CollectorError> {
        let cci = self
            .charts
            .get_mut(chart_id)
            .ok_or(CollectorError::UnkownChart(chart_id.to_owned()))?;
        let mut cdi = cci
            .dimensions
            .get_mut(dimension_id)
            .ok_or(CollectorError::UnkownDimension(dimension_id.to_owned()))?;
        cdi.value = value;
        cdi.commited = false;
        Ok(())
    }

    /// Send a block containing all updated values to `stdout`.
    ///
    /// In multi threaded scenarios, this could require additional lock preventions.
    pub fn commit_chart(&mut self, chart_id: &str) -> Result<(), CollectorError> {
        let mut cci = self
            .charts
            .get_mut(chart_id)
            .ok_or(CollectorError::UnkownChart(chart_id.to_owned()))?;
        let now = Instant::now();

        let begin = Begin {
            type_id: chart_id,
            microseconds: match cci.last_commit {
                Some(t) => Some(now.duration_since(t).as_micros()),
                _ => None,
            },
        };
        cci.last_commit = Some(now);
        writeln!(self.writer, "{}", begin).unwrap();

        for (id, cdi) in cci.dimensions.iter_mut() {
            if !cdi.commited {
                writeln!(
                    self.writer,
                    "{}",
                    Set {
                        id: id,
                        value: Some(cdi.value)
                    }
                )
                .unwrap();
                cdi.commited = true;
            }
        }

        writeln!(self.writer, "{}", Instruction::END)?;
        Ok(())
    }
}

#[cfg(test)]
mod collector_tests {
    use super::{Chart, Collector, Dimension};
    use pretty_assertions::assert_eq;

    #[test]
    fn collector_test() {
        let mut redirect_buf = Vec::new();

        let mut collector = Collector::new(&mut redirect_buf);
        collector
            .add_chart(&mut Chart {
                type_id: "olsr.test_id",
                name: "test_name",
                title: "captions",
                units: "ms",
                ..Default::default()
            })
            .unwrap();
        collector
            .add_dimension(
                "olsr.test_id",
                &Dimension {
                    id: "test_dim_id",
                    name: "test_dim_name",
                    ..Default::default()
                },
            )
            .unwrap();
        collector
            .prepare_value("olsr.test_id", "test_dim_id", 4242)
            .unwrap();
        collector.commit_chart("olsr.test_id").unwrap();
        collector
            .prepare_value("olsr.test_id", "test_dim_id", 4343)
            .unwrap();
        collector.commit_chart("olsr.test_id").unwrap();

        let should_be = r#"CHART "olsr.test_id" "test_name" "captions" "ms"
DIMENSION "test_dim_id" "test_dim_name"
BEGIN "olsr.test_id"
SET "test_dim_id" = 00
END
BEGIN "olsr.test_id" 0
SET "test_dim_id" = 00
END
"#;
        let mut output = String::from_utf8(redirect_buf).unwrap();
        output = output
            .chars()
            .map(|x| if x.is_numeric() { '0' } else { x })
            .collect::<String>()
            .replace("00", "0");
        assert_eq!(output, should_be);
    }
}