use std::io;
use std::sync::OnceLock;
use opentelemetry::KeyValue;
use opentelemetry::global;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_otlp::SpanExporter;
use opentelemetry_sdk::Resource;
use opentelemetry_sdk::propagation::TraceContextPropagator;
use opentelemetry_sdk::trace::SdkTracerProvider;
use tokio::sync::broadcast;
use tracing_subscriber::fmt::MakeWriter;
use tracing_subscriber::{EnvFilter, fmt, prelude::*};
const LOG_BROADCAST_CAPACITY: usize = 512;
static LOG_BROADCAST: OnceLock<broadcast::Sender<String>> = OnceLock::new();
#[must_use]
pub fn subscribe_logs() -> Option<broadcast::Receiver<String>> {
LOG_BROADCAST.get().map(broadcast::Sender::subscribe)
}
#[derive(Clone)]
struct BroadcastTee {
tx: broadcast::Sender<String>,
}
impl<'a> MakeWriter<'a> for BroadcastTee {
type Writer = TeeLine;
fn make_writer(&'a self) -> Self::Writer {
TeeLine {
buf: Vec::new(),
tx: self.tx.clone(),
}
}
}
struct TeeLine {
buf: Vec<u8>,
tx: broadcast::Sender<String>,
}
impl io::Write for TeeLine {
fn write(&mut self, data: &[u8]) -> io::Result<usize> {
io::stderr().write_all(data)?;
self.buf.extend_from_slice(data);
Ok(data.len())
}
fn flush(&mut self) -> io::Result<()> {
io::stderr().flush()
}
}
impl Drop for TeeLine {
fn drop(&mut self) {
if self.buf.is_empty() {
return;
}
let line = String::from_utf8_lossy(&self.buf).trim_end().to_owned();
if !line.is_empty() {
let _ = self.tx.send(line);
}
}
}
pub struct ShutdownGuard {
provider: Option<SdkTracerProvider>,
}
impl Drop for ShutdownGuard {
fn drop(&mut self) {
if let Some(provider) = self.provider.take() {
if let Err(err) = provider.shutdown() {
eprintln!("opentelemetry shutdown failed: {err}");
}
}
}
}
#[must_use]
pub fn init(service_name: &str) -> ShutdownGuard {
let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
let json = std::env::var("RUST_LOG_FORMAT").is_ok_and(|v| v.eq_ignore_ascii_case("json"));
let (tracer, provider) = build_otel_tracer(service_name);
let otel_layer = tracer.map(|t| tracing_opentelemetry::layer().with_tracer(t));
global::set_text_map_propagator(TraceContextPropagator::new());
let (log_tx, _seed) = broadcast::channel(LOG_BROADCAST_CAPACITY);
let _ = LOG_BROADCAST.set(log_tx.clone());
let tee = BroadcastTee { tx: log_tx };
let registry = tracing_subscriber::registry().with(filter).with(otel_layer);
if json {
let fmt_layer = fmt::layer().json().flatten_event(true).with_writer(tee);
registry.with(fmt_layer).init();
} else {
let fmt_layer = fmt::layer().with_writer(tee);
registry.with(fmt_layer).init();
}
install_panic_hook();
ShutdownGuard { provider }
}
fn build_otel_tracer(
service_name: &str,
) -> (
Option<opentelemetry_sdk::trace::Tracer>,
Option<SdkTracerProvider>,
) {
if std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").is_err() {
return (None, None);
}
let resolved_name = std::env::var("OTEL_SERVICE_NAME")
.ok()
.unwrap_or_else(|| service_name.to_owned());
let exporter = match SpanExporter::builder().with_tonic().build() {
Ok(exp) => exp,
Err(err) => {
eprintln!("opentelemetry OTLP exporter build failed: {err}; continuing without traces");
return (None, None);
}
};
let resource = Resource::builder()
.with_attribute(KeyValue::new("service.name", resolved_name.clone()))
.build();
let provider = SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_resource(resource)
.build();
let tracer = provider.tracer(resolved_name);
global::set_tracer_provider(provider.clone());
(Some(tracer), Some(provider))
}
fn install_panic_hook() {
let previous = std::panic::take_hook();
std::panic::set_hook(Box::new(move |info| {
let location = info
.location()
.map_or_else(|| "unknown".to_owned(), ToString::to_string);
tracing::error!(panic = %info, location = %location, "process panicked");
previous(info);
}));
}