sz_rust_middleware_facade/
request_scope.rs1use axum::http::{Request, Response};
41use std::cell::RefCell;
42use std::future::Future;
43use std::pin::Pin;
44use std::sync::atomic::{AtomicU64, Ordering};
45use std::task::{Context, Poll};
46use tower::Service;
47
48use crate::ScopeId;
49
50thread_local! {
54 static CURRENT_SCOPE_ID: RefCell<ScopeId> = const { RefCell::new(0) };
55}
56
57pub fn current_scope_id() -> Option<ScopeId> {
64 CURRENT_SCOPE_ID.with(|c| {
65 let id = *c.borrow();
66 if id != 0 {
67 Some(id)
68 } else {
69 None
70 }
71 })
72}
73
74static SCOPE_ID_COUNTER: AtomicU64 = AtomicU64::new(1);
76
77fn generate_scope_id() -> ScopeId {
79 SCOPE_ID_COUNTER.fetch_add(1, Ordering::Relaxed)
80}
81
82#[derive(Clone, Default)]
90pub struct RequestScopeLayer;
91
92impl RequestScopeLayer {
93 pub fn new() -> Self {
95 Self
96 }
97}
98
99impl<S> tower::Layer<S> for RequestScopeLayer {
100 type Service = RequestScopeService<S>;
101
102 fn layer(&self, inner: S) -> Self::Service {
103 RequestScopeService { inner }
104 }
105}
106
107#[derive(Clone)]
113pub struct RequestScopeService<S> {
114 inner: S,
115}
116
117impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RequestScopeService<S>
118where
119 S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
120 S::Future: Send + 'static,
121 ReqBody: Send + 'static,
122 ResBody: axum::body::HttpBody + Send + 'static,
123 ResBody::Data: Send,
124 ResBody::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
125{
126 type Response = Response<ResBody>;
127 type Error = S::Error;
128 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
129
130 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
131 self.inner.poll_ready(cx)
132 }
133
134 fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
135 let scope_id = generate_scope_id();
137
138 CURRENT_SCOPE_ID.with(|c| *c.borrow_mut() = scope_id);
140
141 let future = self.inner.clone().call(req);
143
144 Box::pin(async move {
145 let response = future.await?;
146 CURRENT_SCOPE_ID.with(|c| *c.borrow_mut() = 0);
148 Ok(response)
149 })
150 }
151}
152
153#[cfg(test)]
158mod tests {
159 use super::*;
160 use axum::body::Body;
161 use axum::http::{Method, Request, StatusCode};
162 use tower::Layer;
163
164 #[derive(Clone)]
166 struct EchoService;
167
168 impl Service<Request<Body>> for EchoService {
169 type Response = Response<Body>;
170 type Error = std::convert::Infallible;
171 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
172
173 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
174 Poll::Ready(Ok(()))
175 }
176
177 fn call(&mut self, _req: Request<Body>) -> Self::Future {
178 Box::pin(async {
179 Ok(Response::builder()
180 .status(StatusCode::OK)
181 .body(Body::from("ok"))
182 .unwrap())
183 })
184 }
185 }
186
187 #[tokio::test]
188 async fn test_p1_arch_di_02_scope_layer_sets_and_clears_scope() {
189 assert!(
191 current_scope_id().is_none(),
192 "请求外 current_scope_id 应为 None"
193 );
194
195 let layer = RequestScopeLayer::new();
196 let mut service = layer.layer(EchoService);
197
198 let req = Request::builder()
199 .method(Method::GET)
200 .uri("/")
201 .body(Body::empty())
202 .unwrap();
203 let resp = service.call(req).await.unwrap();
204 assert_eq!(resp.status(), StatusCode::OK);
205
206 assert!(
208 current_scope_id().is_none(),
209 "请求结束后 current_scope_id 应恢复为 None"
210 );
211 }
212
213 #[tokio::test]
214 async fn test_p1_arch_di_02_scope_id_unique_per_request() {
215 let mut seen_ids = Vec::new();
216 for _ in 0..5 {
217 let id1 = generate_scope_id();
219 let id2 = generate_scope_id();
220 assert_ne!(id1, id2, "连续生成的 scope_id 应不同");
221 seen_ids.push(id1);
222 }
223
224 for (i, &id1) in seen_ids.iter().enumerate() {
226 for &id2 in seen_ids.iter().skip(i + 1) {
227 assert_ne!(id1, id2, "scope_id 应全局唯一");
228 }
229 }
230 }
231
232 fn assert_send<T: Send>() {}
234
235 #[test]
236 fn test_p1_arch_di_02_service_is_send() {
237 assert_send::<RequestScopeService<EchoService>>();
238 }
239}