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;
use tracing_subscriber::fmt::format::{FormatEvent, Writer};
use tracing_subscriber::fmt::{FmtContext, MakeWriter};
use tracing_subscriber::registry::LookupSpan;
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)
.map_event_format(CloudLoggingSeverity::wrap)
.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);
}));
}
struct CloudLoggingSeverity<F> {
inner: F,
}
impl<F> CloudLoggingSeverity<F> {
const fn wrap(inner: F) -> Self {
Self { inner }
}
}
impl<S, N, F> FormatEvent<S, N> for CloudLoggingSeverity<F>
where
S: Subscriber + for<'a> LookupSpan<'a>,
N: for<'a> tracing_subscriber::fmt::FormatFields<'a> + 'static,
F: FormatEvent<S, N>,
{
fn format_event(
&self,
ctx: &FmtContext<'_, S, N>,
mut writer: Writer<'_>,
event: &tracing::Event<'_>,
) -> std::fmt::Result {
let mut formatted = String::new();
self.inner
.format_event(ctx, Writer::new(&mut formatted), event)?;
let line = insert_severity(&formatted, *event.metadata().level());
writer.write_str(&line)
}
}
fn cloud_logging_severity(level: tracing::Level) -> &'static str {
match level.as_str() {
"ERROR" => "ERROR",
"WARN" => "WARNING",
"INFO" => "INFO",
_ => "DEBUG",
}
}
fn insert_severity(formatted: &str, level: tracing::Level) -> String {
let Some(rest) = formatted.strip_prefix('{') else {
return formatted.to_owned();
};
let severity = cloud_logging_severity(level);
let mut out = String::with_capacity(formatted.len() + severity.len() + 14);
out.push('{');
out.push_str("\"severity\":\"");
out.push_str(severity);
out.push_str("\",");
out.push_str(rest);
out
}
#[cfg(test)]
mod tests {
use super::{BroadcastTee, CloudLoggingSeverity, insert_severity};
use tracing_subscriber::fmt;
use tracing_subscriber::prelude::*;
fn json_lines(emit: impl FnOnce()) -> Vec<String> {
let (tx, mut rx) = tokio::sync::broadcast::channel(8);
let layer = fmt::layer()
.json()
.flatten_event(true)
.map_event_format(CloudLoggingSeverity::wrap)
.with_writer(BroadcastTee { tx });
tracing::subscriber::with_default(tracing_subscriber::registry().with(layer), emit);
let mut lines = Vec::new();
while let Ok(line) = rx.try_recv() {
lines.push(line);
}
lines
}
fn top_level_severity(line: &str) -> String {
let value: serde_json::Value = serde_json::from_str(line).expect("json line");
let object = value.as_object().expect("object");
let severities: Vec<_> = object
.iter()
.filter(|(key, _)| *key == "severity")
.map(|(_, value)| value.as_str().expect("severity string").to_owned())
.collect();
assert_eq!(
severities.len(),
1,
"exactly one top-level severity in {line}"
);
severities.into_iter().next().expect("one severity")
}
#[test]
fn json_lines_carry_cloud_logging_severity() {
let lines = json_lines(|| {
tracing::info!(answer = 7, "ready");
tracing::warn!(answer = 7, "careful");
tracing::error!(answer = 7, "failed");
});
assert_eq!(lines.len(), 3);
assert_eq!(top_level_severity(&lines[0]), "INFO");
assert_eq!(top_level_severity(&lines[1]), "WARNING");
assert_eq!(top_level_severity(&lines[2]), "ERROR");
for line in &lines {
let value: serde_json::Value = serde_json::from_str(line).expect("json line");
assert!(value.get("level").is_some(), "keeps level: {line}");
assert_eq!(
value.get("answer").and_then(serde_json::Value::as_i64),
Some(7)
);
assert!(value.get("message").is_some(), "keeps message: {line}");
}
}
#[test]
fn field_text_does_not_add_a_second_severity() {
let note = "has { and \"severity\":\"FAKE\"";
let body = "body with { and \"severity\":\"FAKE\"";
let lines = json_lines(|| {
tracing::warn!(note, "{body}");
});
assert_eq!(lines.len(), 1);
let value: serde_json::Value = serde_json::from_str(&lines[0]).expect("json line");
assert_eq!(top_level_severity(&lines[0]), "WARNING");
assert_eq!(
value.get("level").and_then(serde_json::Value::as_str),
Some("WARN")
);
let note = value
.get("note")
.and_then(serde_json::Value::as_str)
.expect("note");
let message = value
.get("message")
.and_then(serde_json::Value::as_str)
.expect("message");
assert!(note.contains('{'), "{note}");
assert!(note.contains("\"severity\""), "{note}");
assert!(message.contains('{'), "{message}");
assert!(message.contains("\"severity\""), "{message}");
}
#[test]
fn pretty_lines_do_not_gain_a_severity_field() {
let (tx, mut rx) = tokio::sync::broadcast::channel(4);
let layer = fmt::layer().with_writer(BroadcastTee { tx });
tracing::subscriber::with_default(tracing_subscriber::registry().with(layer), || {
tracing::info!(answer = 7, "ready");
});
let line = rx.try_recv().expect("pretty line");
assert!(!line.contains("\"severity\""), "{line}");
assert!(line.contains("ready"), "{line}");
assert!(!line.starts_with('{'), "{line}");
}
#[test]
fn insert_severity_keeps_a_line_that_does_not_start_with_an_object() {
assert_eq!(
insert_severity("not json", tracing::Level::INFO),
"not json"
);
}
}