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    /// STATEMENT: Keep-alive calls pass through until inactivity causes a readiness error.
227    ///
228    /// HOW:
229    /// - Create and inspect a keep-alive factory and service, then make spaced calls.
230    /// - Wait for the readiness dispatcher to report expiry and verify the error.
231    #[ntex::test]
232    async fn test_ka() {
233        // the keep-alive is much longer than the sleep between calls, so a
234        // delayed wakeup does not expire it
235        let factory = factory::<_, (), usize>(KeepAlive::new(Millis(500), || TestErr));
236        assert!(format!("{factory:?}").contains("KeepAlive"));
237        let _ = factory.clone();
238
239        let svc = factory.create(&()).await.unwrap();
240        assert!(format!("{svc:?}").contains("KeepAliveService"));
241
242        let p = Pipeline::new((), boxed::service(svc));
243        assert_eq!(p.call(1usize).await, Ok(1usize));
244        let svc = p.bind();
245
246        let (tx, rx) = oneshot::channel();
247        drop(spawn(Dispatcher { p, tx: Some(tx) }));
248
249        sleep(Millis(25)).await;
250        assert_eq!(svc.call(1usize).await, Ok(1usize));
251        sleep(Millis(100)).await;
252
253        let res = rx.await;
254        assert_eq!(res, Ok(()));
255        assert_eq!(svc.ready().await, Err(TestErr));
256    }
257
258    /// STATEMENT: A timeout callback does not expire a keep-alive with submillisecond time left.
259    ///
260    /// HOW:
261    /// - Set the last activity so half a millisecond remains and invoke the callback.
262    /// - Verify readiness succeeds, the expired flag stays clear, and the timer is armed.
263    #[ntex::test]
264    async fn test_ka_sub_millis() {
265        let svc = Rc::new(KeepAliveService::new(Millis(100), || TestErr));
266        let p = Pipeline::<usize, usize, TestErr>::new((), svc.clone()).bind();
267        svc.shared.timer.cancel();
268
269        // less than millisecond is left before expiration
270        svc.shared
271            .expire
272            .set(now().checked_sub(svc.shared.dur).unwrap() + time::Duration::from_micros(500));
273        svc.shared.on_timeout();
274        assert_eq!(p.ready().await, Ok(()));
275        assert!(!svc.shared.expired.get());
276        assert!(!svc.shared.timer.is_elapsed());
277    }
278
279    fn notification(p: &Pipeline<usize, usize, TestErr>) -> oneshot::Receiver<()> {
280        let (tx, rx) = oneshot::channel();
281        let tx = Cell::new(Some(tx));
282        p.on_notify(&ReadyCallback::new(move |_| {
283            if let Some(tx) = tx.take() {
284                let _ = tx.send(());
285            }
286        }));
287        rx
288    }
289
290    /// STATEMENT: Idle expiry invalidates cached readiness for every pipeline sharing a service.
291    ///
292    /// HOW:
293    /// - Prime two pipelines' readiness and wait for an expiry notification.
294    /// - Check both caches are invalidated and both readiness and a subsequent call fail.
295    #[ntex::test]
296    async fn idle_timer_invalidates_completed_readiness_for_all_roots() {
297        let svc = Rc::new(KeepAliveService::new(Millis(25), || TestErr));
298        let first = Pipeline::new((), svc.clone());
299        let second = Pipeline::<usize, usize, TestErr>::new((), svc);
300        assert_eq!(first.ready().await, Ok(()));
301        assert_eq!(second.ready().await, Ok(()));
302        let expired = notification(&first);
303
304        crate::time::timeout(Millis(1_000), expired)
305            .await
306            .unwrap()
307            .unwrap();
308        assert!(!first.is_ready());
309        assert!(!second.is_ready());
310        assert_eq!(first.ready().await, Err(TestErr));
311        assert_eq!(second.ready().await, Err(TestErr));
312        assert_eq!(first.call(1).await, Err(TestErr));
313    }
314
315    /// STATEMENT: A call extends the deadline before an overdue timer callback is processed.
316    ///
317    /// HOW:
318    /// - Force the old timer due, make a call, and verify no premature expiry or notification.
319    /// - Await the extended deadline and check one notification and subsequent failures.
320    #[ntex::test]
321    async fn call_extends_deadline_before_old_timer_is_processed() {
322        let svc = Rc::new(KeepAliveService::new(Millis(250), || TestErr));
323        let p = Pipeline::new((), svc.clone());
324        let notifications = Rc::new(Cell::new(0));
325        let count = notifications.clone();
326        p.on_notify(&ReadyCallback::new(move |_| count.set(count.get() + 1)));
327        assert_eq!(p.ready().await, Ok(()));
328        svc.shared.timer.reset(Millis::ZERO);
329        assert_eq!(p.call(1).await, Ok(1));
330        sleep(Millis(50)).await;
331        assert!(p.is_ready());
332        assert_eq!(notifications.get(), 0);
333        assert!(!svc.shared.expired.get());
334        assert!(!svc.shared.timer.is_elapsed());
335
336        crate::time::timeout(Millis(1_000), notification(&p))
337            .await
338            .unwrap()
339            .unwrap();
340        assert_eq!(notifications.get(), 1);
341        assert_eq!(p.ready().await, Err(TestErr));
342        assert_eq!(p.call(2).await, Err(TestErr));
343    }
344
345    struct Bypass<S>(S);
346
347    impl<S: Service<(), usize, Res = usize, Error = TestErr>> Service<(), usize> for Bypass<S> {
348        type Res = usize;
349        type Error = TestErr;
350
351        #[inline]
352        fn bind(&self, ctx: BindCtx<'_>) {
353            ctx.bind(&self.0);
354        }
355
356        // `ready` is not forwarded, so calls reuse this scope's cached readiness
357        async fn call(&self, req: usize, ctx: Ctx<'_, Self>) -> Result<usize, TestErr> {
358            ctx.call(&self.0, req).await
359        }
360    }
361
362    /// STATEMENT: A call from a parent that skips readiness rearms an expired keep-alive timer.
363    ///
364    /// HOW:
365    /// - Wait for expiry, then call through a wrapper that does not check readiness.
366    /// - Check readiness recovers and the rearmed timer later expires again.
367    #[ntex::test]
368    async fn bypassed_call_rearms_an_expired_timer() {
369        let svc = Rc::new(KeepAliveService::new(Millis(25), || TestErr));
370        let p = Pipeline::new((), svc.clone());
371        let bypass = Pipeline::new((), Bypass(svc.clone()));
372        p.ready().await.unwrap();
373        crate::time::timeout(Millis(1_000), notification(&p))
374            .await
375            .unwrap()
376            .unwrap();
377        assert_eq!(p.ready().await, Err(TestErr));
378
379        assert_eq!(bypass.call(1).await, Ok(1));
380        assert!(!svc.shared.timer.is_elapsed());
381        assert_eq!(p.ready().await, Ok(()));
382        crate::time::timeout(Millis(1_000), notification(&p))
383            .await
384            .unwrap()
385            .unwrap();
386        assert_eq!(p.ready().await, Err(TestErr));
387    }
388
389    /// STATEMENT: Shutdown and drop cancel timer callbacks without retaining keep-alive state.
390    ///
391    /// HOW:
392    /// - Exercise drop alone and shutdown followed by drop while counting notifications.
393    /// - Verify no delayed notifications, no rearming after shutdown, and release of shared state.
394    #[ntex::test]
395    async fn timer_callback_is_canceled_on_shutdown_and_drop() {
396        for shutdown in [false, true] {
397            let svc = Rc::new(KeepAliveService::new(Millis(25), || TestErr));
398            let state = Rc::downgrade(&svc.shared);
399            let p = Pipeline::<usize, usize, TestErr>::new((), svc.clone());
400            let notifications = Rc::new(Cell::new(0));
401            let count = notifications.clone();
402            p.on_notify(&ReadyCallback::new(move |_| count.set(count.get() + 1)));
403            p.ready().await.unwrap();
404            if shutdown {
405                p.shutdown().await;
406                assert!(svc.shared.timer.is_elapsed());
407                sleep(Millis(75)).await;
408                assert_eq!(notifications.get(), 0);
409
410                let bypass = Pipeline::new((), Bypass(svc.clone()));
411                assert_eq!(bypass.call(1).await, Ok(1));
412                assert!(svc.shared.timer.is_elapsed());
413                let before = notifications.get();
414                sleep(Millis(75)).await;
415                assert_eq!(notifications.get(), before);
416            }
417            let before = notifications.get();
418            drop(p);
419            drop(svc);
420            assert!(state.upgrade().is_none());
421            sleep(Millis(75)).await;
422            assert_eq!(notifications.get(), before);
423        }
424    }
425
426    /// STATEMENT: A zero keep-alive duration rejects readiness and calls immediately.
427    ///
428    /// HOW:
429    /// - Construct a zero-duration service and assert its configured error from both operations.
430    #[ntex::test]
431    async fn test_ka_zero() {
432        let p = Pipeline::<usize, usize, TestErr>::new(
433            (),
434            KeepAliveService::new(Millis::ZERO, || TestErr),
435        );
436        assert_eq!(p.ready().await, Err(TestErr));
437        assert_eq!(p.call(1).await, Err(TestErr));
438    }
439}
ntex_util/services/keepalive.rs

ntex_util/services/
keepalive.rs
