use crate::{CompletedSpan, ExportMetrics, MetricVisitor, Span, SpanCollector, SpanKind};
use parking_lot::RwLock;
use std::borrow::Cow;
use std::ops::Deref;
use std::sync::Arc;
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct RuntimeConfig {}
#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct MetricScope {
name: Arc<str>,
}
impl MetricScope {
pub fn new(name: impl Into<Arc<str>>) -> Self {
Self { name: name.into() }
}
pub fn name(&self) -> &str {
&self.name
}
}
impl From<&'static str> for MetricScope {
fn from(name: &'static str) -> Self {
Self::new(name)
}
}
impl From<String> for MetricScope {
fn from(name: String) -> Self {
Self::new(Arc::<str>::from(name))
}
}
pub struct Runtime {
config: RuntimeConfig,
registry: RwLock<Vec<RegisteredMetricGroup>>,
span_collector: Arc<SpanCollector>,
}
impl Runtime {
pub fn new(config: RuntimeConfig) -> Arc<Self> {
Arc::new(Self {
config,
registry: RwLock::new(Vec::new()),
span_collector: Arc::new(SpanCollector::new(8, 4096)),
})
}
pub fn config(&self) -> &RuntimeConfig {
&self.config
}
pub fn span_collector(&self) -> &Arc<SpanCollector> {
&self.span_collector
}
pub fn start_span(&self, name: impl Into<Cow<'static, str>>, kind: SpanKind) -> Span {
self.span_collector.start_span(name, kind)
}
pub fn start_span_from_traceparent(
&self,
traceparent: Option<&str>,
name: impl Into<Cow<'static, str>>,
kind: SpanKind,
) -> Span {
self.span_collector
.start_span_from_traceparent(traceparent, name, kind)
}
pub fn flush_local_spans(&self) {
self.span_collector.flush_local();
}
pub fn drain_spans_into(&self, buf: &mut Vec<CompletedSpan>) {
self.span_collector.drain_into(buf);
}
pub fn register_metrics<M>(&self, scope: MetricScope, metrics: M) -> RegisteredMetrics<M>
where
M: ExportMetrics + Send + Sync + 'static,
{
let metrics = Arc::new(metrics);
let erased_metrics: Arc<dyn ErasedExportMetrics> = metrics.clone();
self.registry.write().push(RegisteredMetricGroup {
scope: scope.clone(),
metrics: erased_metrics,
});
RegisteredMetrics { scope, metrics }
}
pub fn visit_metrics<V>(&self, visitor: &mut V)
where
V: MetricVisitor,
{
for group in self.registry.read().iter() {
group.metrics.visit_metrics(visitor);
}
}
pub fn visit_metrics_for_scope<V>(&self, scope: &MetricScope, visitor: &mut V)
where
V: MetricVisitor,
{
for group in self.registry.read().iter() {
if &group.scope == scope {
group.metrics.visit_metrics(visitor);
}
}
}
pub fn registered_metrics_len(&self) -> usize {
self.registry.read().len()
}
pub fn scopes(&self) -> Vec<MetricScope> {
self.registry
.read()
.iter()
.map(|group| group.scope.clone())
.collect()
}
}
#[must_use = "keep RegisteredMetrics so hot paths can update direct metric handles"]
pub struct RegisteredMetrics<M> {
scope: MetricScope,
metrics: Arc<M>,
}
impl<M> RegisteredMetrics<M> {
pub fn scope(&self) -> &MetricScope {
&self.scope
}
pub fn metrics(&self) -> &Arc<M> {
&self.metrics
}
pub fn into_metrics(self) -> Arc<M> {
self.metrics
}
}
impl<M> Clone for RegisteredMetrics<M> {
fn clone(&self) -> Self {
Self {
scope: self.scope.clone(),
metrics: Arc::clone(&self.metrics),
}
}
}
impl<M> Deref for RegisteredMetrics<M> {
type Target = M;
fn deref(&self) -> &Self::Target {
&self.metrics
}
}
struct RegisteredMetricGroup {
scope: MetricScope,
metrics: Arc<dyn ErasedExportMetrics>,
}
trait ErasedExportMetrics: Send + Sync {
fn visit_metrics(&self, visitor: &mut dyn MetricVisitor);
}
impl<M> ErasedExportMetrics for M
where
M: ExportMetrics + Send + Sync + 'static,
{
fn visit_metrics(&self, visitor: &mut dyn MetricVisitor) {
ExportMetrics::visit_metrics(self, visitor);
}
}