mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-09-05 08:02:18 +00:00
123 lines
4.3 KiB
Rust
123 lines
4.3 KiB
Rust
/*
|
|
* Author: Ruben Fiszel
|
|
* Copyright: Windmill Labs, Inc 2022
|
|
* This file and its contents are licensed under the AGPLv3 License.
|
|
* Please see the included NOTICE for copyright information and
|
|
* LICENSE-AGPL for a copy of the license.
|
|
*/
|
|
|
|
use crate::s3_log_batching::{record_s3_log, S3ProxyRequest};
|
|
use ::tracing::{field, Span};
|
|
use axum::extract::Request;
|
|
use axum::middleware::Next;
|
|
use axum::response::Response as AxumResponse;
|
|
use hyper::Response;
|
|
use tower_http::trace::{MakeSpan, OnFailure, OnResponse};
|
|
use uuid::Uuid;
|
|
use windmill_common::log_context::{with_log_context, LogContext};
|
|
|
|
lazy_static::lazy_static! {
|
|
static ref LOG_REQUESTS: bool = std::env::var("LOG_REQUESTS")
|
|
.ok()
|
|
.and_then(|x| x.parse::<bool>().ok())
|
|
.unwrap_or(true);
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct MyOnResponse {}
|
|
|
|
impl<B> OnResponse<B> for MyOnResponse {
|
|
fn on_response(
|
|
self,
|
|
response: &Response<B>,
|
|
latency: std::time::Duration,
|
|
_span: &tracing::Span,
|
|
) {
|
|
if *LOG_REQUESTS {
|
|
if let Some(s3req) = response.extensions().get::<S3ProxyRequest>() {
|
|
let status = response.status().as_u16();
|
|
if response.status().is_success() || response.status().is_redirection() {
|
|
record_s3_log(&s3req.uri, &s3req.method, status, latency.as_millis());
|
|
return;
|
|
}
|
|
// 4xx/5xx errors are rare and shouldn't be batched
|
|
}
|
|
|
|
let latency = latency.as_millis();
|
|
let status = response.status().as_u16();
|
|
if response.status().is_success() || response.status().is_redirection() {
|
|
tracing::info!(latency = latency, status = status, "response")
|
|
} else if response.status().as_u16() == 404 {
|
|
tracing::warn!(latency = latency, status = status, "response")
|
|
} else {
|
|
tracing::error!(latency = latency, status = status, "response")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct MyOnFailure {}
|
|
|
|
impl<B> OnFailure<B> for MyOnFailure {
|
|
fn on_failure(&mut self, _b: B, latency: std::time::Duration, _span: &tracing::Span) {
|
|
tracing::error!(latency = latency.as_millis(), "response failure")
|
|
}
|
|
}
|
|
|
|
lazy_static::lazy_static! {
|
|
static ref TRACING_HEADER: String = std::env::var("TRACING_HEADER")
|
|
.ok().unwrap_or_else(|| "x-tracing-id".to_string());
|
|
}
|
|
#[derive(Clone)]
|
|
pub struct MyMakeSpan {}
|
|
|
|
impl<B> MakeSpan<B> for MyMakeSpan {
|
|
fn make_span(&mut self, request: &hyper::Request<B>) -> Span {
|
|
let tracing_id = request
|
|
.headers()
|
|
.get(TRACING_HEADER.as_str())
|
|
.and_then(|x| x.to_str().map(|x| x.to_string()).ok())
|
|
.unwrap_or(Uuid::new_v4().to_string());
|
|
// ERROR level, not INFO: a span is only recorded by a layer whose filter
|
|
// enables it, so under RUST_LOG=error/warn an INFO span is dropped and the
|
|
// 4xx/5xx `response` events in `MyOnResponse` print with no uri, method or
|
|
// traceId.
|
|
tracing::error_span!(
|
|
"request",
|
|
method = %request.method(),
|
|
uri = %request.uri(),
|
|
username = field::Empty,
|
|
workspace_id = field::Empty,
|
|
traceId = tracing_id,
|
|
email = field::Empty,
|
|
)
|
|
}
|
|
}
|
|
|
|
/// Axum middleware that seeds a per-request `LogContext` with method/uri/
|
|
/// traceId and wraps the downstream chain in a task-local scope. Auth and
|
|
/// workspace-resolution code later mutate this context (via
|
|
/// `update_log_context`) as email/username/workspace_id become known.
|
|
///
|
|
/// Registered at the top of the router layer stack in `windmill-api/src/lib.rs`
|
|
/// so every route — and critically, the `MyOnResponse::on_response` callback
|
|
/// that TraceLayer invokes on the way out — runs inside the scope and thus
|
|
/// flows through to exported OTEL LogRecords via the EE LogContextBridge.
|
|
pub async fn log_context_middleware(request: Request, next: Next) -> AxumResponse {
|
|
let trace_id = request
|
|
.headers()
|
|
.get(TRACING_HEADER.as_str())
|
|
.and_then(|x| x.to_str().ok())
|
|
.map(|s| s.to_string());
|
|
|
|
let ctx = LogContext {
|
|
method: Some(request.method().to_string()),
|
|
uri: Some(request.uri().to_string()),
|
|
trace_id,
|
|
..Default::default()
|
|
};
|
|
|
|
with_log_context(ctx, next.run(request)).await
|
|
}
|