1use std::fmt;
25
26use tracing_subscriber::{
27 fmt::{FmtContext, FormatEvent, FormatFields, format::Writer},
28 layer::Context,
29 registry::LookupSpan,
30};
31
32pub const REQUEST_SPAN_NAME: &str = "request";
35
36#[derive(Clone, Debug)]
40pub struct RequestShortId(pub String);
41
42pub struct RequestSpanLayer;
46
47impl<S> tracing_subscriber::Layer<S> for RequestSpanLayer
48where
49 S: tracing::Subscriber + for<'a> LookupSpan<'a>,
50{
51 fn on_new_span(
52 &self,
53 attrs: &tracing::span::Attributes<'_>,
54 id: &tracing::span::Id,
55 ctx: Context<'_, S>,
56 ) {
57 if attrs.metadata().name() != REQUEST_SPAN_NAME {
58 return;
59 }
60 let mut visitor = IdVisitor(None);
61 attrs.record(&mut visitor);
62 if let (Some(short), Some(span)) = (visitor.0, ctx.span(id)) {
63 span.extensions_mut().insert(RequestShortId(short));
64 }
65 }
66}
67
68struct IdVisitor(Option<String>);
69
70impl tracing::field::Visit for IdVisitor {
71 fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
72 if field.name() == "id" {
73 self.0 = Some(value.to_string());
74 }
75 }
76
77 fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn fmt::Debug) {
78 if field.name() == "id" && self.0.is_none() {
82 self.0 = Some(format!("{:?}", value));
83 }
84 }
85}
86
87pub struct CloudilloFormat;
91
92impl CloudilloFormat {
93 pub fn new() -> Self {
94 Self
95 }
96}
97
98impl Default for CloudilloFormat {
99 fn default() -> Self {
100 Self::new()
101 }
102}
103
104impl<S, N> FormatEvent<S, N> for CloudilloFormat
105where
106 S: tracing::Subscriber + for<'a> LookupSpan<'a>,
107 N: for<'a> FormatFields<'a> + 'static,
108{
109 fn format_event(
110 &self,
111 ctx: &FmtContext<'_, S, N>,
112 mut writer: Writer<'_>,
113 event: &tracing::Event<'_>,
114 ) -> fmt::Result {
115 let now = chrono::Utc::now();
116 write!(writer, "{} ", now.format("%Y-%m-%dT%H:%M:%S%.3fZ"))?;
117
118 let level = event.metadata().level();
119 write!(writer, "{:>5} ", level)?;
120
121 if let Some(id) = find_request_short_id(ctx) {
122 write!(writer, "REQ:{} ", id)?;
123 }
124
125 ctx.field_format().format_fields(writer.by_ref(), event)?;
126 writeln!(writer)
127 }
128}
129
130fn find_request_short_id<S, N>(ctx: &FmtContext<'_, S, N>) -> Option<String>
133where
134 S: tracing::Subscriber + for<'a> LookupSpan<'a>,
135 N: for<'a> FormatFields<'a> + 'static,
136{
137 let scope = ctx.event_scope()?;
138 for span in scope.from_root() {
139 if span.name() == REQUEST_SPAN_NAME
140 && let Some(rsid) = span.extensions().get::<RequestShortId>()
141 {
142 return Some(rsid.0.clone());
143 }
144 }
145 None
146}
147
148#[cfg(test)]
149#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
150mod tests {
151 use super::*;
152 use std::io::Write;
153 use std::sync::{Arc, Mutex};
154 use tracing::{info, info_span};
155 use tracing_subscriber::fmt::MakeWriter;
156 use tracing_subscriber::layer::SubscriberExt;
157
158 #[derive(Clone, Default)]
159 struct BufWriter(Arc<Mutex<Vec<u8>>>);
160
161 struct BufWriterGuard(Arc<Mutex<Vec<u8>>>);
162
163 impl Write for BufWriterGuard {
164 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
165 match self.0.lock() {
166 Ok(mut g) => g.write(buf),
167 Err(_) => Ok(buf.len()),
168 }
169 }
170 fn flush(&mut self) -> std::io::Result<()> {
171 Ok(())
172 }
173 }
174
175 impl<'a> MakeWriter<'a> for BufWriter {
176 type Writer = BufWriterGuard;
177 fn make_writer(&'a self) -> Self::Writer {
178 BufWriterGuard(self.0.clone())
179 }
180 }
181
182 impl BufWriter {
183 fn snapshot(&self) -> String {
184 let g = self.0.lock().unwrap();
185 String::from_utf8(g.clone()).unwrap()
186 }
187 }
188
189 fn build_subscriber(buf: BufWriter) -> impl tracing::Subscriber + Send + Sync {
190 let fmt_layer = tracing_subscriber::fmt::layer()
191 .with_writer(buf)
192 .event_format(CloudilloFormat::new());
193 tracing_subscriber::registry().with(RequestSpanLayer).with(fmt_layer)
194 }
195
196 #[test]
197 fn request_span_prefixes_with_req_id() {
198 let buf = BufWriter::default();
199 let subscriber = build_subscriber(buf.clone());
200
201 tracing::subscriber::with_default(subscriber, || {
202 info_span!("request", id = "abcd").in_scope(|| {
203 info!("hello world");
204 });
205 });
206
207 let out = buf.snapshot();
208 assert_eq!(out.matches(" REQ:abcd ").count(), 1, "output: {out:?}");
209
210 let req_pos = out.find("REQ:").unwrap();
211 let level_pos = out.find("INFO").unwrap();
212 assert!(level_pos < req_pos, "REQ before level: {out:?}");
213 assert!(out.contains("hello world"), "missing message: {out:?}");
214 }
215
216 #[test]
217 fn request_span_with_display_value_captures_correctly() {
218 let buf = BufWriter::default();
219 let subscriber = build_subscriber(buf.clone());
220
221 let short = String::from("abcd");
222 tracing::subscriber::with_default(subscriber, || {
223 info_span!("request", id = %short).in_scope(|| {
224 info!("display value");
225 });
226 });
227
228 let out = buf.snapshot();
229 assert!(out.contains(" REQ:abcd "), "output: {out:?}");
230 }
231
232 #[test]
233 fn no_span_means_no_prefix() {
234 let buf = BufWriter::default();
235 let subscriber = build_subscriber(buf.clone());
236
237 tracing::subscriber::with_default(subscriber, || {
238 info!("startup");
239 });
240
241 let out = buf.snapshot();
242 assert!(!out.contains("REQ:"), "unexpected prefix: {out:?}");
243 assert!(out.contains("startup"), "missing message: {out:?}");
244 }
245
246 #[test]
247 fn other_span_names_do_not_get_captured() {
248 let buf = BufWriter::default();
249 let subscriber = build_subscriber(buf.clone());
250
251 tracing::subscriber::with_default(subscriber, || {
252 info_span!("background", id = "abcd").in_scope(|| {
253 info!("background work");
254 });
255 });
256
257 let out = buf.snapshot();
258 assert!(!out.contains("REQ:"), "unexpected prefix: {out:?}");
259 }
260}
261
262