use std::collections::HashMap;
use std::sync::Arc;
use crate::error::TelemetryError;
use crate::event::Event;
use crate::machine::MachineId;
use crate::sink::{NoopSink, TelemetrySink};
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum CollectionPolicy {
#[default]
Disabled,
Enabled,
}
#[derive(Clone)]
pub struct TelemetryContext {
tool: Arc<String>,
tool_version: Arc<String>,
machine_id: Arc<String>,
sink: Arc<dyn TelemetrySink>,
policy: CollectionPolicy,
}
impl std::fmt::Debug for TelemetryContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TelemetryContext")
.field("tool", &self.tool)
.field("tool_version", &self.tool_version)
.field("machine_id_len", &self.machine_id.len())
.field("policy", &self.policy)
.field("sink", &"<dyn TelemetrySink>")
.finish()
}
}
impl TelemetryContext {
pub fn builder() -> TelemetryContextBuilder {
TelemetryContextBuilder::default()
}
#[must_use]
pub const fn policy(&self) -> CollectionPolicy {
self.policy
}
pub async fn record(&self, event_name: &str) -> Result<(), TelemetryError> {
if self.policy == CollectionPolicy::Disabled {
return Ok(());
}
let event = Event::now(event_name, &*self.tool, &*self.tool_version, &*self.machine_id);
self.sink.emit(&event).await
}
pub async fn record_with_attrs(
&self,
event_name: &str,
attrs: HashMap<String, String>,
) -> Result<(), TelemetryError> {
if self.policy == CollectionPolicy::Disabled {
return Ok(());
}
let mut event = Event::now(event_name, &*self.tool, &*self.tool_version, &*self.machine_id);
event.attrs = attrs;
self.sink.emit(&event).await
}
pub async fn flush(&self) -> Result<(), TelemetryError> {
if self.policy == CollectionPolicy::Disabled {
return Ok(());
}
self.sink.flush().await
}
}
#[must_use]
#[derive(Default)]
pub struct TelemetryContextBuilder {
tool: Option<String>,
tool_version: Option<String>,
salt: Option<String>,
sink: Option<Arc<dyn TelemetrySink>>,
policy: CollectionPolicy,
}
impl TelemetryContextBuilder {
pub fn tool(mut self, tool: impl Into<String>) -> Self {
self.tool = Some(tool.into());
self
}
pub fn tool_version(mut self, version: impl Into<String>) -> Self {
self.tool_version = Some(version.into());
self
}
pub fn salt(mut self, salt: impl Into<String>) -> Self {
self.salt = Some(salt.into());
self
}
pub fn sink(mut self, sink: Arc<dyn TelemetrySink>) -> Self {
self.sink = Some(sink);
self
}
pub const fn policy(mut self, policy: CollectionPolicy) -> Self {
self.policy = policy;
self
}
#[must_use]
pub fn build(self) -> TelemetryContext {
let tool = self.tool.expect("TelemetryContextBuilder: tool is required");
let tool_version =
self.tool_version.expect("TelemetryContextBuilder: tool_version is required");
let sink = self.sink.unwrap_or_else(|| Arc::new(NoopSink));
let policy = self.policy;
let machine_id = match policy {
CollectionPolicy::Disabled => String::new(),
CollectionPolicy::Enabled => {
let salt = self
.salt
.expect("TelemetryContextBuilder: salt is required when policy is Enabled");
MachineId::derive(&salt)
}
};
TelemetryContext {
tool: Arc::new(tool),
tool_version: Arc::new(tool_version),
machine_id: Arc::new(machine_id),
sink,
policy,
}
}
}