Skip to content

Commit 7730faf

Browse files
ignatzignatz
committed
Minor: streamline logs layer.
1 parent 1b5c957 commit 7730faf

File tree

1 file changed

+66
-86
lines changed

1 file changed

+66
-86
lines changed

‎trailbase-core/src/logging.rs‎

Lines changed: 66 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -60,18 +60,16 @@ pub(super) fn sqlite_logger_make_span(request: &Request<Body>) -> Span {
6060
let host = get_header(headers, "host").unwrap_or("");
6161
let user_agent = get_header(headers, "user-agent").unwrap_or("");
6262
let referer = get_header(headers, "referer").unwrap_or("");
63-
let client_ip = InsecureClientIp::from(headers, request.extensions())
64-
.map(|ip| ip.0.to_string())
65-
.unwrap_or_else(|_| "".to_string());
66-
// let extensions = request.extensions().get::<ConnectInfo<SocketAddr>>();
63+
let client_ip = InsecureClientIp::from(headers, request.extensions()).map(|ip| ip.0.to_string());
6764

65+
// NOTE: "%" means print using fmt::Display, and "?" means fmt::Debug.
6866
let span = tracing::info_span!(
6967
"request",
7068
method = %request.method(),
7169
uri = %request.uri(),
7270
version = ?request.version(),
7371
host,
74-
client_ip,
72+
client_ip = %client_ip.as_ref().map_or("", |s| s.as_str()),
7573
user_agent,
7674
referer,
7775
);
@@ -84,24 +82,26 @@ pub(super) fn sqlite_logger_on_request(_req: &Request<Body>, _span: &Span) {
8482
// request related information into the span
8583
}
8684

85+
#[inline]
8786
fn as_millis_f64(d: &Duration) -> f64 {
8887
const NANOS_PER_MILLI: f64 = 1_000_000.0;
8988
const MILLIS_PER_SEC: u64 = 1_000;
90-
return d.as_secs_f64() * (MILLIS_PER_SEC as f64) + (d.as_nanos() as f64) / (NANOS_PER_MILLI);
89+
return (d.as_secs() as f64) * (MILLIS_PER_SEC as f64)
90+
+ (d.subsec_nanos() as f64) / (NANOS_PER_MILLI);
9191
}
9292

9393
pub(super) fn sqlite_logger_on_response(
9494
response: &Response<Body>,
9595
latency: Duration,
9696
_span: &Span,
9797
) {
98-
let length = get_header(response.headers(), "content-length").unwrap_or("-1");
98+
let length = get_header(response.headers(), "content-length");
9999

100100
tracing::info!(
101101
name: "response",
102102
latency_ms = as_millis_f64(&latency),
103103
status = response.status().as_u16(),
104-
length = length.parse::<i64>().unwrap(),
104+
length = length.and_then(|l| l.parse::<i64>().ok()),
105105
);
106106
}
107107

@@ -123,89 +123,71 @@ impl SqliteLogLayer {
123123
// TODO: We could use recv_many() and batch insert.
124124
fn write_log(&self, log: Log) -> Result<(), trailbase_sqlite::Error> {
125125
return self.conn.call_and_forget(move |conn| {
126-
let result = conn.execute(
126+
let mut stmt = match conn.prepare_cached(
127127
r#"
128128
INSERT INTO
129129
_logs (type, level, status, method, url, latency, client_ip, referer, user_agent)
130130
VALUES
131131
($1, $2, $3, $4, $5, $6, $7, $8, $9)
132132
"#,
133-
rusqlite::params!(
134-
log.r#type,
135-
log.level,
136-
log.status,
137-
log.method,
138-
log.url,
139-
log.latency,
140-
log.client_ip,
141-
log.referer,
142-
log.user_agent
143-
),
144-
);
145-
146-
if let Err(err) = result {
133+
) {
134+
Ok(stmt) => stmt,
135+
Err(err) => {
136+
warn!("Logs stmt failed: {err}");
137+
return;
138+
}
139+
};
140+
141+
if let Err(err) = stmt.execute(rusqlite::params!(
142+
log.r#type,
143+
log.level,
144+
log.status,
145+
log.method,
146+
log.url,
147+
log.latency,
148+
log.client_ip,
149+
log.referer,
150+
log.user_agent
151+
)) {
147152
warn!("logs writing failed: {err}");
148153
}
149154
});
150155
}
151156
}
152157

158+
const BUG_TEXT: &str = "Span not found, this is a bug";
159+
153160
impl<S> Layer<S> for SqliteLogLayer
154161
where
155162
S: tracing::Subscriber,
156163
S: for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
157164
{
158165
fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
159-
let span = ctx.span(id).unwrap();
160-
161-
// let mut fields = BTreeMap::new();
162-
// attrs.record(&mut JsonVisitor(&mut fields));
163-
// span.extensions_mut().insert(CustomFieldStorage(fields));
166+
let span = ctx.span(id).expect(BUG_TEXT);
164167

165168
let mut storage = LogFieldStorage::default();
166169
attrs.record(&mut LogJsonVisitor(&mut storage));
167170
span.extensions_mut().insert(storage);
168171
}
169172

170173
fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
171-
// Get a mutable reference to the data we created in new_span
172-
let span = ctx.span(id).unwrap();
173-
let mut extensions_mut = span.extensions_mut();
174+
let span = ctx.span(id).expect(BUG_TEXT);
174175

175-
// And add to using our old friend the visitor!
176-
// let custom_field_storage = extensions_mut.get_mut::<CustomFieldStorage>().unwrap();
177-
// values.record(&mut JsonVisitor(&mut custom_field_storage.0));
178-
179-
let log_field_storage = extensions_mut.get_mut::<LogFieldStorage>().unwrap();
180-
values.record(&mut LogJsonVisitor(log_field_storage));
176+
let mut extensions = span.extensions_mut();
177+
if let Some(storage) = extensions.get_mut::<LogFieldStorage>() {
178+
values.record(&mut LogJsonVisitor(storage));
179+
}
181180
}
182181

183182
fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
184-
let mut request_storage: Option<LogFieldStorage> = None;
185-
186-
let scope = ctx.event_scope(event).unwrap();
187-
for span in scope.from_root() {
188-
// TODO: we should be merging here to account for multiple spans. Maybe we should have a json
189-
// span representation in the data field.
190-
let extensions = span.extensions();
191-
if let Some(storage) = extensions.get::<LogFieldStorage>() {
192-
request_storage = Some(storage.clone());
193-
}
194-
}
183+
let Some(span) = ctx.event_span(event) else {
184+
return;
185+
};
186+
187+
let mut extensions = span.extensions_mut();
188+
if let Some(s) = extensions.get_mut::<LogFieldStorage>() {
189+
let mut storage = std::mem::take(s);
195190

196-
// The fields of the event
197-
// let mut fields = BTreeMap::new();
198-
// event.record(&mut JsonVisitor(&mut fields));
199-
// let output = json!({
200-
// "target": event.metadata().target(),
201-
// "name": event.metadata().name(),
202-
// "level": format!("{:?}", event.metadata().level()),
203-
// "fields": fields,
204-
// "spans": spans,
205-
// });
206-
// println!("{}", serde_json::to_string_pretty(&output).unwrap());
207-
208-
if let Some(mut storage) = request_storage {
209191
event.record(&mut LogJsonVisitor(&mut storage));
210192

211193
let log = Log {
@@ -222,7 +204,13 @@ where
222204
client_ip: storage.client_ip,
223205
referer: storage.referer,
224206
user_agent: storage.user_agent,
225-
data: Some(json!(storage.fields)),
207+
data: {
208+
if storage.fields.is_empty() {
209+
None
210+
} else {
211+
Some(json!(storage.fields))
212+
}
213+
},
226214
};
227215

228216
if let Err(err) = self.write_log(log) {
@@ -259,77 +247,69 @@ struct LogFieldStorage {
259247
length: i64,
260248

261249
// All other fields.
262-
fields: BTreeMap<String, serde_json::Value>,
250+
fields: BTreeMap<&'static str, serde_json::Value>,
263251
}
264252

265253
struct LogJsonVisitor<'a>(&'a mut LogFieldStorage);
266254

267255
impl tracing::field::Visit for LogJsonVisitor<'_> {
268256
fn record_f64(&mut self, field: &Field, double: f64) {
269-
let name = field.name();
270-
match name {
257+
match field.name() {
271258
"latency_ms" => self.0.latency_ms = double,
272-
_ => {
273-
self.0.fields.insert(name.to_string(), json!(double));
259+
name => {
260+
self.0.fields.insert(name, json!(double));
274261
}
275262
};
276263
}
277264

278265
fn record_i64(&mut self, field: &Field, int: i64) {
279-
let name = field.name();
280-
match name {
266+
match field.name() {
281267
"length" => self.0.length = int,
282-
_ => {
283-
self.0.fields.insert(name.to_string(), json!(int));
268+
name => {
269+
self.0.fields.insert(name, json!(int));
284270
}
285271
};
286272
}
287273

288274
fn record_u64(&mut self, field: &Field, uint: u64) {
289-
let name = field.name();
290-
match name {
275+
match field.name() {
291276
"status" => self.0.status = uint,
292-
_ => {
293-
self.0.fields.insert(name.to_string(), json!(uint));
277+
name => {
278+
self.0.fields.insert(name, json!(uint));
294279
}
295280
};
296281
}
297282

298283
fn record_bool(&mut self, field: &Field, b: bool) {
299-
self.0.fields.insert(field.name().to_string(), json!(b));
284+
self.0.fields.insert(field.name(), json!(b));
300285
}
301286

302287
fn record_str(&mut self, field: &Field, s: &str) {
303-
let name: &str = field.name();
304-
match name {
288+
match field.name() {
305289
"client_ip" => self.0.client_ip = s.to_string(),
306290
"host" => self.0.host = s.to_string(),
307291
"referer" => self.0.referer = s.to_string(),
308292
"user_agent" => self.0.user_agent = s.to_string(),
309293
name => {
310-
self.0.fields.insert(name.to_string(), json!(s));
294+
self.0.fields.insert(name, json!(s));
311295
}
312296
};
313297
}
314298

315299
fn record_debug(&mut self, field: &Field, dbg: &dyn std::fmt::Debug) {
316-
let name = field.name();
317300
let v = format!("{:?}", dbg);
318-
match name {
301+
match field.name() {
319302
"method" => self.0.method = v,
320303
"uri" => self.0.uri = v,
321304
"version" => self.0.version = v,
322305
name => {
323-
self.0.fields.insert(name.to_string(), json!(v));
306+
self.0.fields.insert(name, json!(v));
324307
}
325308
};
326309
}
327310

328311
fn record_error(&mut self, field: &Field, err: &(dyn std::error::Error + 'static)) {
329-
self
330-
.0
331-
.fields
332-
.insert(field.name().to_string(), json!(err.to_string()));
312+
self.0.fields.insert(field.name(), json!(err.to_string()));
333313
}
334314
}
335315

0 commit comments

Comments
 (0)