zon_middleware/
task_local.rs

1use std::future::Future;
2
3use http::{Request, Response};
4use tokio::task::LocalKey;
5use zon_core::{HttpMiddleware, HttpService};
6
7/// Middleware to set tokio task-local data.
8#[derive(Debug, Clone, Copy)]
9pub struct SetTaskLocal<T: 'static> {
10    key: &'static LocalKey<T>,
11    value: T,
12}
13
14impl<T> SetTaskLocal<T>
15where
16    T: Clone + Send + Sync + 'static,
17{
18    pub fn new(key: &'static LocalKey<T>, value: T) -> Self {
19        SetTaskLocal { key, value }
20    }
21}
22
23impl<S, T> HttpMiddleware<S> for SetTaskLocal<T>
24where
25    T: Clone + Send + Sync + 'static,
26{
27    type Service = SetTaskLocalService<S, T>;
28
29    fn apply(self, inner: S) -> Self::Service {
30        SetTaskLocalService::new(inner, self.key, self.value)
31    }
32}
33
34/// Service for [`SetTaskLocal`] middleware.
35#[derive(Debug, Clone, Copy)]
36pub struct SetTaskLocalService<S, T: 'static> {
37    inner: S,
38    key: &'static LocalKey<T>,
39    value: T,
40}
41
42impl<S, T> SetTaskLocalService<S, T>
43where
44    T: Clone + Send + Sync + 'static,
45{
46    pub fn new(inner: S, key: &'static LocalKey<T>, value: T) -> Self {
47        Self { inner, key, value }
48    }
49}
50
51impl<S, T, B> HttpService<B> for SetTaskLocalService<S, T>
52where
53    S: HttpService<B>,
54    T: Clone + Send + Sync + 'static,
55    B: Send,
56{
57    type ResponseBody = S::ResponseBody;
58
59    fn call(
60        &self,
61        request: Request<B>,
62    ) -> impl Future<Output = Response<Self::ResponseBody>> + Send {
63        self.key.scope(self.value.clone(), self.inner.call(request))
64    }
65}
zon_middleware/task_local.rs

zon_middleware/
task_local.rs
