rama_http/service/web/
k8s.rs

1//! k8s web service
2
3use crate::{
4    Request, Response, StatusCode, matcher::HttpMatcher,
5    service::web::endpoint::response::IntoResponse,
6};
7use rama_core::{
8    Context, Service,
9    service::{BoxService, service_fn},
10};
11use std::{convert::Infallible, fmt, marker::PhantomData, sync::Arc};
12
13use super::match_service;
14
15/// create a k8s web health service builder
16pub fn k8s_health_builder<S>() -> K8sHealthServiceBuilder<(), (), S> {
17    K8sHealthServiceBuilder::new()
18}
19
20/// create a default k8s web health service
21pub fn k8s_health<S>() -> impl Service<S, Request, Response = Response, Error = Infallible> + Clone
22where
23    S: Clone + Send + Sync + 'static,
24{
25    k8s_health_builder().build()
26}
27
28/// builder to easily create a k8s web service
29///
30/// by default its endpoints will always return 200 (OK),
31/// but this can be made conditional by providing
32/// a ready condition ([`K8sHealthServiceBuilder::ready`], liveness)
33/// and/or alive condition ([`K8sHealthServiceBuilder::alive`], readiness).
34///
35/// In case a conditional is provided and it returns `false`,
36/// a 503 (Service Unavailable) will be returned instead.
37pub struct K8sHealthServiceBuilder<A, R, S> {
38    alive: A,
39    ready: R,
40    _phantom: PhantomData<fn(S) -> ()>,
41}
42
43impl<A: fmt::Debug, R: fmt::Debug, S> std::fmt::Debug for K8sHealthServiceBuilder<A, R, S> {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("K8sHealthServiceBuilder")
46            .field("alive", &self.alive)
47            .field("ready", &self.ready)
48            .field(
49                "_phantom",
50                &format_args!("{}", std::any::type_name::<fn(S) -> ()>()),
51            )
52            .finish()
53    }
54}
55
56impl<A: Clone, R: Clone, S> Clone for K8sHealthServiceBuilder<A, R, S> {
57    fn clone(&self) -> Self {
58        Self {
59            alive: self.alive.clone(),
60            ready: self.ready.clone(),
61            _phantom: PhantomData,
62        }
63    }
64}
65
66impl<S> K8sHealthServiceBuilder<(), (), S> {
67    pub(crate) fn new() -> Self {
68        Self {
69            alive: (),
70            ready: (),
71            _phantom: PhantomData,
72        }
73    }
74}
75
76impl<R, S> K8sHealthServiceBuilder<(), R, S> {
77    /// define an alive condition to be used by the k8s health web service for the liveness check
78    pub fn alive<A: Fn() -> bool>(self, alive: A) -> K8sHealthServiceBuilder<A, R, S> {
79        K8sHealthServiceBuilder {
80            alive,
81            ready: self.ready,
82            _phantom: self._phantom,
83        }
84    }
85}
86
87impl<A, S> K8sHealthServiceBuilder<A, (), S> {
88    /// define an ready condition to be used by the k8s health web service for the readiness check
89    pub fn ready<R: Fn() -> bool>(self, ready: R) -> K8sHealthServiceBuilder<A, R, S> {
90        K8sHealthServiceBuilder {
91            alive: self.alive,
92            ready,
93            _phantom: self._phantom,
94        }
95    }
96}
97
98impl<A, R, S> K8sHealthServiceBuilder<A, R, S>
99where
100    A: ToK8sService<S>,
101    R: ToK8sService<S>,
102    S: Clone + Send + Sync + 'static,
103{
104    /// build the k8s health web server
105    pub fn build(
106        self,
107    ) -> impl Service<S, Request, Response = Response, Error = Infallible> + Clone {
108        Arc::new(match_service! {
109            HttpMatcher::get("/k8s/alive") => self.alive.to_k8s_service(),
110            HttpMatcher::get("/k8s/ready") => self.ready.to_k8s_service(),
111            _ => StatusCode::NOT_FOUND,
112        })
113    }
114}
115
116/// Utility internal trait to create service endpoints for the different checks
117pub trait ToK8sService<S>: private::Sealed<S> {}
118
119impl<S: Clone + Send + Sync + 'static> ToK8sService<S> for () {}
120
121impl<S, F> ToK8sService<S> for F
122where
123    F: Fn() -> bool + Clone + Send + Sync + 'static,
124    S: Clone + Send + Sync + 'static,
125{
126}
127
128struct K8sService<F> {
129    f: F,
130}
131
132impl<F: fmt::Debug> std::fmt::Debug for K8sService<F> {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        f.debug_struct("K8sService").field("f", &self.f).finish()
135    }
136}
137
138impl<F: Clone> Clone for K8sService<F> {
139    fn clone(&self) -> Self {
140        Self { f: self.f.clone() }
141    }
142}
143
144impl<F> K8sService<F> {
145    fn new(f: F) -> Self {
146        Self { f }
147    }
148}
149
150impl<F, State> Service<State, Request> for K8sService<F>
151where
152    F: Fn() -> bool + Send + Sync + 'static,
153    State: Clone + Send + Sync + 'static,
154{
155    type Response = Response;
156    type Error = Infallible;
157
158    async fn serve(&self, _: Context<State>, _: Request) -> Result<Self::Response, Self::Error> {
159        Ok(if (self.f)() {
160            StatusCode::OK
161        } else {
162            StatusCode::SERVICE_UNAVAILABLE
163        }
164        .into_response())
165    }
166}
167
168mod private {
169    use super::*;
170
171    pub trait Sealed<S> {
172        /// create a boxed web service by consuming self
173        fn to_k8s_service(self) -> BoxService<S, Request, Response, Infallible>;
174    }
175
176    impl<S> Sealed<S> for () {
177        fn to_k8s_service(self) -> BoxService<S, Request, Response, Infallible> {
178            service_fn(async || Ok(StatusCode::OK.into_response())).boxed()
179        }
180    }
181
182    impl<S: Clone + Send + Sync + 'static, F: Fn() -> bool + Clone + Send + Sync + 'static>
183        Sealed<S> for F
184    {
185        fn to_k8s_service(self) -> BoxService<S, Request, Response, Infallible> {
186            K8sService::new(self).boxed()
187        }
188    }
189}
rama_http/service/web/k8s.rs

rama_http/service/web/
k8s.rs
