use core::fmt;
use std::io::{stdout, Write};
use futures_util::future::BoxFuture;
use ts_opentelemetry_api::trace::{TraceError, TraceResult};
use ts_opentelemetry_sdk::export::{self, trace::ExportResult};
use crate::trace::transform::SpanData;
type Encoder = Box<dyn Fn(&mut dyn Write, SpanData) -> TraceResult<()> + Send + Sync>;
pub struct SpanExporter {
writer: Option<Box<dyn Write + Send + Sync>>,
encoder: Encoder,
}
impl fmt::Debug for SpanExporter {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("SpanExporter")
}
}
impl SpanExporter {
pub fn builder() -> SpanExporterBuilder {
SpanExporterBuilder::default()
}
}
impl Default for SpanExporter {
fn default() -> Self {
SpanExporterBuilder::default().build()
}
}
impl ts_opentelemetry_sdk::export::trace::SpanExporter for SpanExporter {
fn export(&mut self, batch: Vec<export::trace::SpanData>) -> BoxFuture<'static, ExportResult> {
let res = if let Some(writer) = &mut self.writer {
(self.encoder)(writer, crate::trace::SpanData::from(batch)).and_then(|_| {
writer
.write_all(b"\n")
.map_err(|err| TraceError::Other(Box::new(err)))
})
} else {
Err("exporter is shut down".into())
};
Box::pin(std::future::ready(res))
}
fn shutdown(&mut self) {
self.writer.take();
}
}
#[derive(Default)]
pub struct SpanExporterBuilder {
writer: Option<Box<dyn Write + Send + Sync>>,
encoder: Option<Encoder>,
}
impl fmt::Debug for SpanExporterBuilder {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("SpanExporterBuilder")
}
}
impl SpanExporterBuilder {
pub fn with_writer(mut self, writer: impl Write + Send + Sync + 'static) -> Self {
self.writer = Some(Box::new(writer));
self
}
pub fn with_encoder(
mut self,
writer: impl Fn(&mut dyn Write, SpanData) -> TraceResult<()> + Send + Sync + 'static,
) -> Self {
self.encoder = Some(Box::new(writer));
self
}
pub fn build(self) -> SpanExporter {
SpanExporter {
writer: Some(self.writer.unwrap_or_else(|| Box::new(stdout()))),
encoder: self.encoder.unwrap_or_else(|| {
Box::new(|writer, spans| {
serde_json::to_writer(writer, &spans)
.map_err(|err| TraceError::Other(Box::new(err)))
})
}),
}
}
}