spacegate_kernel/helper_layers/
reload.rs

1use std::sync::{Arc, OnceLock};
2
3use crossbeam_utils::sync::ShardedLock;
4use tower_layer::Layer;
5
6#[derive(Default, Debug, Clone)]
7pub struct ReloadLayer<S> {
8    pub reloader: Reloader<S>,
9}
10
11impl<S> Layer<S> for ReloadLayer<S>
12where
13    S: Send,
14{
15    type Service = Reload<S>;
16
17    fn layer(&self, inner: S) -> Self::Service {
18        let inner = Arc::new(ShardedLock::new(inner));
19        self.reloader.setup(inner.clone());
20        Reload { service: inner }
21    }
22}
23
24#[derive(Debug, Clone)]
25pub struct Reload<S>
26where
27    S: Send,
28{
29    pub(crate) service: Arc<ShardedLock<S>>,
30}
31
32#[derive(Debug)]
33pub struct Reloader<S> {
34    pub service: Arc<OnceLock<Arc<ShardedLock<S>>>>,
35}
36
37impl<S> Default for Reloader<S> {
38    fn default() -> Self {
39        Self { service: Default::default() }
40    }
41}
42
43impl<S> Clone for Reloader<S> {
44    fn clone(&self) -> Self {
45        Self { service: self.service.clone() }
46    }
47}
48
49impl<S> Reloader<S> {
50    pub fn setup(&self, service: Arc<ShardedLock<S>>) {
51        if self.service.set(service).is_err() {
52            tracing::warn!("reloader already settled");
53        }
54    }
55    pub fn reload(&self, service: S) {
56        if let Some(wg) = self.service.get() {
57            let mut wg = wg.write().expect("should never be poisoned");
58            *wg = service;
59        } else {
60            tracing::warn!("reloader not initialized");
61        }
62    }
63    pub fn into_layer(self) -> ReloadLayer<S> {
64        ReloadLayer { reloader: self }
65    }
66}
67
68impl<Request, S> hyper::service::Service<Request> for Reload<S>
69where
70    Request: Send + Sync + 'static,
71    S: hyper::service::Service<Request> + Send + Sync + 'static,
72    <S as hyper::service::Service<Request>>::Future: std::marker::Send,
73{
74    type Response = S::Response;
75
76    type Error = S::Error;
77
78    type Future = <S as hyper::service::Service<Request>>::Future;
79
80    fn call(&self, req: Request) -> Self::Future {
81        let service = self.service.clone();
82        let rg = service.read().expect("should never be posisoned");
83        let fut = rg.call(req);
84        drop(rg);
85        fut
86    }
87}