use crate::error::app_error::TRACE_ID;
use crate::handler::{Req, Resp};
use crate::router::Router;
use crate::{AppError, IntoResponse};
use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use tower::Service;
pub struct RouterSvc<S> {
pub router: Router<S>,
}
impl<S> Clone for RouterSvc<S> {
fn clone(&self) -> Self {
Self {
router: self.router.clone(),
}
}
}
impl<S: Send + Sync + 'static> Service<Req> for RouterSvc<S> {
type Response = Resp;
type Error = AppError;
type Future = Pin<Box<dyn Future<Output = Result<Resp, AppError>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, mut req: Req) -> Self::Future {
let method = req.method().clone();
let path = req.uri().path().to_string();
let result = self.router.find_handler(&method, &path);
let trace_id = extract_or_generate_trace_id(&req);
let trace_id_clone = trace_id.clone();
let start = std::time::Instant::now();
let task_future = async move {
tracing::debug!(
method = %method,
path = %path,
trace_id = %trace_id,
"Request started"
);
let resp_result = match result {
Some((mut handler, params)) => {
req.extensions_mut().insert(params);
handler.call(req).await
}
None => Ok(AppError::NotFound("404 Not Found".to_string()).into_response()),
};
let elapsed = start.elapsed();
match &resp_result {
Ok(response) => {
tracing::debug!(
method = %method,
path = %path,
trace_id = %trace_id,
status = %response.status(),
elapsed_ms = elapsed.as_millis(),
"Request completed"
);
}
Err(err) => {
tracing::debug!(
method = %method,
path = %path,
trace_id = %trace_id,
status = "500",
elapsed_ms = elapsed.as_millis(),
err = ?err,
"Request completed with error"
);
}
}
Ok(resp_result.unwrap_or_else(|e| e.into_response()))
};
Box::pin(TRACE_ID.scope(trace_id_clone, task_future))
}
}
fn extract_or_generate_trace_id(req: &Req) -> String {
req.headers()
.get("x-trace-id")
.or_else(|| req.headers().get("x-request-id"))
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string())
.unwrap_or_else(generate_trace_id)
}
fn generate_trace_id() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_micros();
let thread_id = std::thread::current().id();
format!("trace-{:x}-{:?}", timestamp, thread_id)
.chars()
.filter(|c| c.is_alphanumeric() || *c == '-')
.collect()
}