use axum::http::{Request, Response};
use axum::response::IntoResponse;
use std::cell::RefCell;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::task::{Context, Poll};
use tower::Service;
use crate::container::App;
use crate::container::ScopeId;
thread_local! {
static CURRENT_SCOPE_ID: RefCell<ScopeId> = const { RefCell::new(0) };
}
pub fn current_scope_id() -> Option<ScopeId> {
CURRENT_SCOPE_ID.with(|c| {
let id = *c.borrow();
if id != 0 {
Some(id)
} else {
None
}
})
}
static SCOPE_ID_COUNTER: AtomicU64 = AtomicU64::new(1);
fn generate_scope_id() -> ScopeId {
SCOPE_ID_COUNTER.fetch_add(1, Ordering::Relaxed)
}
#[derive(Clone, Default)]
pub struct RequestScopeLayer;
impl RequestScopeLayer {
pub fn new() -> Self {
Self
}
}
impl<S> tower::Layer<S> for RequestScopeLayer {
type Service = RequestScopeService<S>;
fn layer(&self, inner: S) -> Self::Service {
RequestScopeService { inner }
}
}
#[derive(Clone)]
pub struct RequestScopeService<S> {
inner: S,
}
impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RequestScopeService<S>
where
S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
S::Future: Send + 'static,
ReqBody: Send + 'static,
ResBody: axum::body::HttpBody + Send + 'static,
ResBody::Data: Send,
ResBody::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
{
type Response = Response<ResBody>;
type Error = S::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
let scope_id = generate_scope_id();
CURRENT_SCOPE_ID.with(|c| *c.borrow_mut() = scope_id);
let future = self.inner.clone().call(req);
Box::pin(async move {
let response = future.await?;
CURRENT_SCOPE_ID.with(|c| *c.borrow_mut() = 0);
Ok(response)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use axum::body::Body;
use axum::http::{Method, Request, StatusCode};
use http_body_util::BodyExt;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use tower::{Layer, ServiceExt};
#[derive(Clone)]
struct EchoService;
impl Service<Request<Body>> for EchoService {
type Response = Response<Body>;
type Error = std::convert::Infallible;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, _req: Request<Body>) -> Self::Future {
Box::pin(async {
Ok(Response::builder()
.status(StatusCode::OK)
.body(Body::from("ok"))
.unwrap())
})
}
}
#[tokio::test]
async fn test_p1_arch_di_02_scope_layer_sets_and_clears_scope() {
assert!(current_scope_id().is_none(), "请求外 current_scope_id 应为 None");
let layer = RequestScopeLayer::new();
let mut service = layer.layer(EchoService);
let req = Request::builder()
.method(Method::GET)
.uri("/")
.body(Body::empty())
.unwrap();
let resp = service.call(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
assert!(current_scope_id().is_none(), "请求结束后 current_scope_id 应恢复为 None");
}
#[tokio::test]
async fn test_p1_arch_di_02_scope_id_unique_per_request() {
let layer = RequestScopeLayer::new();
let mut service = layer.layer(EchoService);
let mut seen_ids = Vec::new();
for _ in 0..5 {
let req = Request::builder().method(Method::GET).uri("/").body(Body::empty()).unwrap();
let id1 = generate_scope_id();
let id2 = generate_scope_id();
assert_ne!(id1, id2, "连续生成的 scope_id 应不同");
seen_ids.push(id1);
}
for (i, &id1) in seen_ids.iter().enumerate() {
for &id2 in seen_ids.iter().skip(i + 1) {
assert_ne!(id1, id2, "scope_id 应全局唯一");
}
}
}
fn assert_send<T: Send>() {}
#[test]
fn test_p1_arch_di_02_service_is_send() {
assert_send::<RequestScopeService<EchoService>>();
}
}