ntex_util/services/
keepalive.rs

1use std::{cell::Cell, convert::Infallible, fmt, rc::Rc, rc::Weak, time};
2
3use ntex_service::{BindCtx, Ctx, ReadyNotifier, Service, ServiceFactory};
4
5use crate::time::{CallbackTimer, Millis, now};
6
7/// Middleware that fails readiness once the service has been idle too long.
8///
9/// Every call records the current time and passes the request through
10/// unchanged. When no call has been made for the keep-alive duration,
11/// [`ready`](Service::ready) fails with the error produced by `f`.
12/// A timer callback invalidates cached readiness on expiration, even when
13/// readiness is not being polled. Extended deadlines are rearmed without
14/// notifying readiness; active calls only update the activity timestamp.
15pub struct KeepAlive<F, E>
16where
17    F: Fn() -> E + Clone,
18{
19    f: F,
20    ka: Millis,
21}
22
23impl<F, E> KeepAlive<F, E>
24where
25    F: Fn() -> E + Clone,
26{
27    /// Creates keep-alive middleware.
28    ///
29    /// `ka` is the maximum idle time between calls, and `f` creates the error
30    /// returned once it is exceeded.
31    pub fn new(ka: Millis, f: F) -> Self {
32        KeepAlive { f, ka }
33    }
34}
35
36impl<F, E> Clone for KeepAlive<F, E>
37where
38    F: Fn() -> E + Clone,
39{
40    fn clone(&self) -> Self {
41        KeepAlive {
42            f: self.f.clone(),
43            ka: self.ka,
44        }
45    }
46}
47
48impl<F, E> fmt::Debug for KeepAlive<F, E>
49where
50    F: Fn() -> E + Clone,
51{
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        f.debug_struct("KeepAlive")
54            .field("ka", &self.ka)
55            .field("f", &std::any::type_name::<F>())
56            .finish()
57    }
58}
59
60impl<F, E, St, Req> ServiceFactory<St, Req> for KeepAlive<F, E>
61where
62    F: Fn() -> E + Clone,
63{
64    type Res = Req;
65    type Error = E;
66
67    type Service = KeepAliveService<F, E>;
68    type InitError = Infallible;
69
70    #[inline]
71    async fn create(&self, _: &St) -> Result<Self::Service, Self::InitError> {
72        Ok(KeepAliveService::new(self.ka, self.f.clone()))
73    }
74}
75
76/// A service that fails readiness once it has been idle too long.
77///
78/// See [`KeepAlive`] for details.
79pub struct KeepAliveService<F, E>
80where
81    F: Fn() -> E,
82{
83    f: F,
84    shared: Rc<Shared>,
85}
86
87struct Shared {
88    dur: time::Duration,
89    timer: CallbackTimer,
90    expire: Cell<time::Instant>,
91    expired: Cell<bool>,
92    stopped: Cell<bool>,
93    changed: ReadyNotifier,
94}
95
96impl Shared {
97    fn on_timeout(&self) {
98        let now = now();
99        let expire = self.expire.get() + self.dur;
100        if expire <= now {
101            self.expired.set(true);
102            self.changed.notify();
103        } else {
104            let millis: u32 = (expire - now).as_millis().try_into().unwrap_or(u32::MAX);
105            self.timer.reset(Millis(millis.max(1)));
106        }
107    }
108}
109
110impl<F, E> KeepAliveService<F, E>
111where
112    F: Fn() -> E,
113{
114    /// Creates a keep-alive service with the maximum idle time `dur`.
115    ///
116    /// `f` creates the error returned once the service has been idle longer
117    /// than `dur`. The callback timer is canceled on shutdown or drop.
118    pub fn new(dur: Millis, f: F) -> Self {
119        KeepAliveService {
120            f,
121            shared: Rc::new_cyclic(|weak: &Weak<Shared>| {
122                let weak = weak.clone();
123                Shared {
124                    dur: time::Duration::from(dur),
125                    expire: Cell::new(now()),
126                    expired: Cell::new(false),
127                    stopped: Cell::new(false),
128                    changed: ReadyNotifier::new(),
129                    timer: CallbackTimer::new(dur, move || {
130                        if let Some(shared) = weak.upgrade() {
131                            shared.on_timeout();
132                        }
133                    }),
134                }
135            }),
136        }
137    }
138}
139
140impl<F, E> fmt::Debug for KeepAliveService<F, E>
141where
142    F: Fn() -> E,
143{
144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145        f.debug_struct("KeepAliveService")
146            .field("dur", &self.shared.dur)
147            .field("expire", &self.shared.expire)
148            .field("f", &std::any::type_name::<F>())
149            .finish()
150    }
151}
152
153impl<F, E, St, Req> Service<St, Req> for KeepAliveService<F, E>
154where
155    F: Fn() -> E,
156{
157    type Res = Req;
158    type Error = E;
159
160    fn bind(&self, ctx: BindCtx<'_>) {
161        ctx.subscribe(&self.shared.changed);
162    }
163
164    async fn ready(&self, _: Ctx<'_, Self, St>) -> Result<(), Self::Error> {
165        if self.shared.expire.get() + self.shared.dur <= now() {
166            Err((self.f)())
167        } else {
168            Ok(())
169        }
170    }
171
172    #[inline]
173    async fn call(&self, req: Req, _: Ctx<'_, Self, St>) -> Result<Req, E> {
174        let shared = &self.shared;
175        let now = now();
176        let previous = shared.expire.replace(now);
177        let expired = shared.expired.replace(false);
178        if expired || previous + shared.dur <= now {
179            if !shared.stopped.get() {
180                shared.timer.reset(shared.dur.into());
181            }
182            shared.changed.notify();
183        }
184        Ok(req)
185    }
186
187    async fn shutdown(&self, _: Ctx<'_, Self, St>) {
188        self.shared.stopped.set(true);
189        self.shared.timer.cancel();
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use std::{pin::Pin, task::Context, task::Poll, task::ready};
196
197    use ntex_service::{Pipeline, ReadyCallback, boxed, factory};
198
199    use super::*;
200    use crate::{channel::oneshot, spawn, time::sleep};
201
202    #[derive(Debug, PartialEq)]
203    struct TestErr;
204
205    struct Dispatcher {
206        p: Pipeline<usize, usize, TestErr>,
207        tx: Option<oneshot::Sender<()>>,
208    }
209
210    impl Future for Dispatcher {
211        type Output = ();
212
213        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
214            let mut this = self.as_mut();
215            if ready!(this.p.poll_ready(cx)).is_err() {
216                if let Some(tx) = this.tx.take() {
217                    let _ = tx.send(());
218                }
219                Poll::Ready(())
220            } else {
221                Poll::Pending
222            }
223        }
224    }
225
226    #[ntex::test]
227    async fn test_ka() {
228        // the keep-alive is much longer than the sleep between calls, so a
229        // delayed wakeup does not expire it
230        let factory = factory::<_, (), usize>(KeepAlive::new(Millis(500), || TestErr));
231        assert!(format!("{factory:?}").contains("KeepAlive"));
232        let _ = factory.clone();
233
234        let svc = factory.create(&()).await.unwrap();
235        assert!(format!("{svc:?}").contains("KeepAliveService"));
236
237        let p = Pipeline::new((), boxed::service(svc));
238        assert_eq!(p.call(1usize).await, Ok(1usize));
239        let svc = p.bind();
240
241        let (tx, rx) = oneshot::channel();
242        spawn(Dispatcher { p, tx: Some(tx) }).detach();
243
244        sleep(Millis(25)).await;
245        assert_eq!(svc.call(1usize).await, Ok(1usize));
246        sleep(Millis(100)).await;
247
248        let res = rx.await;
249        assert_eq!(res, Ok(()));
250        assert_eq!(svc.ready().await, Err(TestErr));
251    }
252
253    #[ntex::test]
254    async fn test_ka_sub_millis() {
255        let svc = Rc::new(KeepAliveService::new(Millis(100), || TestErr));
256        let p = Pipeline::<usize, usize, TestErr>::new((), svc.clone()).bind();
257        svc.shared.timer.cancel();
258
259        // less than millisecond is left before expiration
260        svc.shared
261            .expire
262            .set(now().checked_sub(svc.shared.dur).unwrap() + time::Duration::from_micros(500));
263        svc.shared.on_timeout();
264        assert_eq!(p.ready().await, Ok(()));
265        assert!(!svc.shared.expired.get());
266        assert!(!svc.shared.timer.is_elapsed());
267    }
268
269    fn notification(p: &Pipeline<usize, usize, TestErr>) -> oneshot::Receiver<()> {
270        let (tx, rx) = oneshot::channel();
271        let tx = Cell::new(Some(tx));
272        p.on_notify(&ReadyCallback::new(move || {
273            if let Some(tx) = tx.take() {
274                let _ = tx.send(());
275            }
276        }));
277        rx
278    }
279
280    #[ntex::test]
281    async fn idle_timer_invalidates_completed_readiness_for_all_roots() {
282        let svc = Rc::new(KeepAliveService::new(Millis(25), || TestErr));
283        let first = Pipeline::new((), svc.clone());
284        let second = Pipeline::<usize, usize, TestErr>::new((), svc);
285        assert_eq!(first.ready().await, Ok(()));
286        assert_eq!(second.ready().await, Ok(()));
287        let expired = notification(&first);
288
289        crate::time::timeout(Millis(1_000), expired)
290            .await
291            .unwrap()
292            .unwrap();
293        assert!(!first.is_ready());
294        assert!(!second.is_ready());
295        assert_eq!(first.ready().await, Err(TestErr));
296        assert_eq!(second.ready().await, Err(TestErr));
297        assert_eq!(first.call(1).await, Err(TestErr));
298    }
299
300    #[ntex::test]
301    async fn call_extends_deadline_before_old_timer_is_processed() {
302        let svc = Rc::new(KeepAliveService::new(Millis(250), || TestErr));
303        let p = Pipeline::new((), svc.clone());
304        let notifications = Rc::new(Cell::new(0));
305        let count = notifications.clone();
306        p.on_notify(&ReadyCallback::new(move || count.set(count.get() + 1)));
307        assert_eq!(p.ready().await, Ok(()));
308        svc.shared.timer.reset(Millis::ZERO);
309        assert_eq!(p.call(1).await, Ok(1));
310        sleep(Millis(50)).await;
311        assert!(p.is_ready());
312        assert_eq!(notifications.get(), 0);
313        assert!(!svc.shared.expired.get());
314        assert!(!svc.shared.timer.is_elapsed());
315
316        crate::time::timeout(Millis(1_000), notification(&p))
317            .await
318            .unwrap()
319            .unwrap();
320        assert_eq!(notifications.get(), 1);
321        assert_eq!(p.ready().await, Err(TestErr));
322        assert_eq!(p.call(2).await, Err(TestErr));
323    }
324
325    struct Bypass<S>(S);
326
327    impl<S: Service<(), usize, Res = usize, Error = TestErr>> Service<(), usize> for Bypass<S> {
328        type Res = usize;
329        type Error = TestErr;
330
331        ntex_service::forward_bind!((), 0);
332
333        #[allow(
334            deprecated,
335            reason = "verify reset after an intentional readiness bypass"
336        )]
337        async fn call(&self, req: usize, ctx: Ctx<'_, Self>) -> Result<usize, TestErr> {
338            ctx.call_nowait(&self.0, req).await
339        }
340    }
341
342    #[ntex::test]
343    async fn bypassed_call_rearms_an_expired_timer() {
344        let svc = Rc::new(KeepAliveService::new(Millis(25), || TestErr));
345        let p = Pipeline::new((), svc.clone());
346        let bypass = Pipeline::new((), Bypass(svc.clone()));
347        p.ready().await.unwrap();
348        crate::time::timeout(Millis(1_000), notification(&p))
349            .await
350            .unwrap()
351            .unwrap();
352        assert_eq!(p.ready().await, Err(TestErr));
353
354        assert_eq!(bypass.call(1).await, Ok(1));
355        assert!(!svc.shared.timer.is_elapsed());
356        assert_eq!(p.ready().await, Ok(()));
357        crate::time::timeout(Millis(1_000), notification(&p))
358            .await
359            .unwrap()
360            .unwrap();
361        assert_eq!(p.ready().await, Err(TestErr));
362    }
363
364    #[ntex::test]
365    async fn timer_callback_is_canceled_on_shutdown_and_drop() {
366        for shutdown in [false, true] {
367            let svc = Rc::new(KeepAliveService::new(Millis(25), || TestErr));
368            let state = Rc::downgrade(&svc.shared);
369            let p = Pipeline::<usize, usize, TestErr>::new((), svc.clone());
370            let notifications = Rc::new(Cell::new(0));
371            let count = notifications.clone();
372            p.on_notify(&ReadyCallback::new(move || count.set(count.get() + 1)));
373            p.ready().await.unwrap();
374            if shutdown {
375                p.shutdown().await;
376                assert!(svc.shared.timer.is_elapsed());
377                sleep(Millis(75)).await;
378                assert_eq!(notifications.get(), 0);
379
380                let bypass = Pipeline::new((), Bypass(svc.clone()));
381                assert_eq!(bypass.call(1).await, Ok(1));
382                assert!(svc.shared.timer.is_elapsed());
383                let before = notifications.get();
384                sleep(Millis(75)).await;
385                assert_eq!(notifications.get(), before);
386            }
387            let before = notifications.get();
388            drop(p);
389            drop(svc);
390            assert!(state.upgrade().is_none());
391            sleep(Millis(75)).await;
392            assert_eq!(notifications.get(), before);
393        }
394    }
395
396    #[ntex::test]
397    async fn test_ka_zero() {
398        let p = Pipeline::<usize, usize, TestErr>::new(
399            (),
400            KeepAliveService::new(Millis::ZERO, || TestErr),
401        );
402        assert_eq!(p.ready().await, Err(TestErr));
403        assert_eq!(p.call(1).await, Err(TestErr));
404    }
405}
ntex_util/services/keepalive.rs

ntex_util/services/
keepalive.rs
