ntex_io/
ops.rs

1#![allow(clippy::cast_possible_truncation)]
2use std::collections::{BTreeMap, VecDeque};
3use std::{cell::RefCell, mem, num::NonZeroUsize, ops, rc::Rc, time::Duration, time::Instant};
4
5use ntex_rt::Arbiter;
6use ntex_util::time::{Millis, Seconds, now, sleep};
7use ntex_util::{HashSet, spawn};
8use slab::Slab;
9
10use crate::IoRef;
11
12const CAP: usize = 32;
13/// Timer sets with a larger capacity are dropped instead of cached.
14const MAX_CACHED_SET_CAP: usize = 512;
15
16thread_local! {
17    static MANAGER: RefCell<Option<IoManager>> = const { RefCell::new(None) };
18}
19
20#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Default)]
21/// Opaque identifier assigned to a registered I/O stream.
22///
23/// The identifier is meaningful only within the current thread's I/O manager.
24pub struct Id(Option<NonZeroUsize>);
25
26#[derive(Copy, Clone, Default, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
27/// Handle to an I/O dispatcher timer.
28///
29/// Handles represent second-granularity deadlines managed by the current
30/// thread's I/O manager, a timer started with a `t` seconds timeout expires
31/// after at least `t` and less than `t + 1` seconds. The timer clock starts
32/// a new second when a timer starts while no other timers are pending, such
33/// a timer expires after `t` seconds. A handle becomes stale
34/// after its timer is stopped or replaced and must not be used as an
35/// independent cancellation token.
36pub struct TimerHandle(u32);
37
38impl TimerHandle {
39    /// A handle that does not refer to an active timer.
40    pub const ZERO: TimerHandle = TimerHandle(0);
41
42    /// Returns `true` if this handle refers to a timer deadline.
43    ///
44    /// This does not guarantee that the timer is still registered or has not
45    /// already elapsed.
46    pub fn is_set(&self) -> bool {
47        self.0 != 0
48    }
49
50    /// Returns the whole seconds remaining until this handle's deadline.
51    ///
52    /// Returns zero if the deadline has elapsed. The remaining time is
53    /// measured from the cached [`now()`](ntex_util::time::now), which lags
54    /// the clock, and is rounded up.
55    pub fn remains(&self) -> Seconds {
56        let rem = self.instant().saturating_duration_since(now());
57        let secs = rem.as_secs() + u64::from(rem.subsec_nanos() != 0);
58        Seconds(secs.min(u64::from(u16::MAX)) as u16)
59    }
60
61    /// Returns the instant represented by this handle.
62    ///
63    /// The instant is based on the current thread's I/O timer clock. For
64    /// [`ZERO`](Self::ZERO), this is the clock's base instant.
65    pub fn instant(&self) -> Instant {
66        IoManager::with(|mgr| mgr.timers.base + Duration::from_secs(u64::from(self.0)))
67    }
68
69    pub(crate) fn update(self, timeout: Seconds, io: &IoRef) -> TimerHandle {
70        IoManager::with(|mgr| {
71            let new_hnd = mgr.timers.deadline(timeout);
72            if self.0 == new_hnd || self.0 == new_hnd + 1 {
73                self
74            } else {
75                mgr.timers.unregister(self, io);
76                mgr.timers.register(timeout, io)
77            }
78        })
79    }
80
81    pub(crate) fn unregister(self, io: &IoRef) {
82        IoManager::with(|manager| manager.timers.unregister(self, io));
83    }
84
85    pub(crate) fn register(timeout: Seconds, io: &IoRef) -> TimerHandle {
86        IoManager::with(move |mgr| mgr.timers.register(timeout, io))
87    }
88}
89
90impl ops::Add<Seconds> for TimerHandle {
91    type Output = TimerHandle;
92
93    #[inline]
94    fn add(self, other: Seconds) -> TimerHandle {
95        TimerHandle(self.0 + u32::from(other.0))
96    }
97}
98
99struct TimerStorage {
100    running: bool,
101    base: Instant,
102    /// Whole seconds elapsed since `base`, the keys up to it have expired.
103    current: u32,
104    cache: VecDeque<HashSet<Id>>,
105    notifications: BTreeMap<u32, HashSet<Id>>,
106}
107
108impl TimerStorage {
109    fn unregister(&mut self, hnd: TimerHandle, io: &IoRef) {
110        if let Some(items) = self.notifications.get_mut(&hnd.0) {
111            items.remove(&io.id());
112            if items.is_empty() {
113                // the timer stops once no timers are left
114                let items = self.notifications.remove(&hnd.0).unwrap();
115                self.recycle(items);
116            }
117        }
118    }
119
120    /// Returns an empty timer set to the cache, a set grown by a burst of
121    /// timers is dropped so its capacity is not retained.
122    fn recycle(&mut self, items: HashSet<Id>) {
123        debug_assert!(items.is_empty());
124        if self.cache.len() < CAP && items.capacity() <= MAX_CACHED_SET_CAP {
125            self.cache.push_back(items);
126        }
127    }
128
129    /// Updates `current` from the clock.
130    fn update_current(&mut self) -> u32 {
131        self.current = self.base.elapsed().as_secs() as u32;
132        self.current
133    }
134
135    /// Returns the key of a timer started now, it expires after at least
136    /// `timeout` and less than `timeout + 1` seconds.
137    ///
138    /// The start is the cached [`now()`](ntex_util::time::now), a timer
139    /// expires early by the age of the cached time.
140    fn deadline(&mut self, timeout: Seconds) -> u32 {
141        let elapsed = now().saturating_duration_since(self.base);
142        let secs = elapsed.as_secs() as u32;
143        if secs < self.current {
144            // the cached time lags the ticker, keys up to `current` expired
145            return self.current + 1 + u32::from(timeout.0);
146        }
147        self.current = secs;
148
149        if self.notifications.is_empty() {
150            // no timers are pending, start the clock second now so a timer
151            // expires after exactly `timeout`. `current` does not change,
152            // stale handles cannot match a new key
153            self.base += Duration::new(0, elapsed.subsec_nanos());
154            self.current + u32::from(timeout.0)
155        } else {
156            let partial = u32::from(elapsed.subsec_nanos() != 0);
157            self.current + partial + u32::from(timeout.0)
158        }
159    }
160
161    fn register(&mut self, timeout: Seconds, io: &IoRef) -> TimerHandle {
162        let hnd = self.deadline(timeout);
163        if let Some(items) = self.notifications.get_mut(&hnd) {
164            items.insert(io.id());
165        } else {
166            let mut items = self.cache.pop_front().unwrap_or_default();
167            items.insert(io.id());
168            self.notifications.insert(hnd, items);
169        }
170
171        self.run_timer();
172
173        TimerHandle(hnd)
174    }
175
176    fn run_timer(&mut self) {
177        if self.running {
178            return;
179        }
180        self.running = true;
181
182        spawn(async move {
183            let guard = TimerGuard;
184            // one wheel timer is reset for every tick
185            let timer = sleep(Self::next_tick());
186            loop {
187                timer.wait().await;
188
189                let stop = IoManager::with(|mgr| {
190                    let current = mgr.timers.update_current();
191
192                    // notify io dispatchers of all expired timers
193                    while let Some(entry) = mgr.timers.notifications.first_entry() {
194                        if *entry.key() > current {
195                            break;
196                        }
197                        let mut items = entry.remove();
198                        for id in items.drain() {
199                            if let Some(io) = mgr.get(id) {
200                                io.notify_timeout();
201                            }
202                        }
203                        mgr.timers.recycle(items);
204                    }
205
206                    if mgr.timers.notifications.is_empty() {
207                        mgr.timers.running = false;
208                        true
209                    } else {
210                        false
211                    }
212                });
213
214                if stop {
215                    break;
216                }
217                timer.reset(Self::next_tick());
218            }
219            drop(guard);
220        });
221    }
222
223    /// Returns the delay until the next whole second of the timer clock, a
224    /// late tick does not delay the later ones.
225    fn next_tick() -> Millis {
226        let next = IoManager::with(|mgr| {
227            let t = &mgr.timers;
228            let next = t.base + Duration::from_secs(u64::from(t.current) + 1);
229            next.saturating_duration_since(Instant::now())
230        });
231        Millis(next.as_millis() as u32 + 1)
232    }
233}
234
235struct TimerGuard;
236
237impl Drop for TimerGuard {
238    fn drop(&mut self) {
239        IoManager::with(|mgr| {
240            mgr.timers.running = false;
241            mgr.timers.notifications.clear();
242        });
243    }
244}
245
246pub(crate) struct IoManager {
247    storage: Slab<Option<IoRef>>,
248    timers: TimerStorage,
249    pub(crate) iops: Iops,
250}
251
252impl Default for IoManager {
253    fn default() -> IoManager {
254        let mut storage = Slab::new();
255        assert_eq!(storage.insert(None), 0);
256
257        IoManager {
258            storage,
259            timers: TimerStorage {
260                running: false,
261                base: Instant::now(),
262                current: 0,
263                cache: VecDeque::with_capacity(CAP),
264                notifications: BTreeMap::default(),
265            },
266            iops: Iops {
267                running: false,
268                ops: Vec::with_capacity(32),
269            },
270        }
271    }
272}
273
274impl IoManager {
275    /// Calls `f` with the current thread's manager.
276    ///
277    /// The manager is created on first use and dropped when the arbiter shuts
278    /// down, so that its state does not carry over to the next runtime on the
279    /// same thread. If the thread-local storage has already been destroyed
280    /// because the thread is exiting, `f` receives a temporary manager.
281    fn with<F, R>(f: F) -> R
282    where
283        F: FnOnce(&mut IoManager) -> R,
284    {
285        let mut f = Some(f);
286        MANAGER
287            .try_with(|cell| {
288                let mut mgr = cell.borrow_mut();
289                let mgr = mgr.get_or_insert_with(|| {
290                    Arbiter::on_shutdown(IoManager::reset);
291                    IoManager::default()
292                });
293                (f.take().unwrap())(mgr)
294            })
295            .unwrap_or_else(|_| (f.take().unwrap())(&mut IoManager::default()))
296    }
297
298    fn reset() {
299        // dropped outside of the borrow, the registered streams it holds may
300        // unregister themselves
301        let mgr = MANAGER
302            .try_with(|cell| cell.borrow_mut().take())
303            .ok()
304            .flatten();
305        drop(mgr);
306    }
307
308    fn get(&self, id: Id) -> Option<&IoRef> {
309        if let Some(id) = id.0 {
310            self.storage.get(id.get()).and_then(|item| item.as_ref())
311        } else {
312            None
313        }
314    }
315
316    pub(crate) fn register(io: &IoRef) -> Id {
317        IoManager::with(|manager| {
318            let entry = manager.storage.vacant_entry();
319            let id = Id(NonZeroUsize::new(entry.key()));
320            entry.insert(Some(io.clone()));
321            id
322        })
323    }
324
325    pub(crate) fn unregister(io: &IoRef) {
326        if let Some(id) = io.id().0 {
327            io.0.id.set(Id(None));
328            IoManager::with(|manager| {
329                // the manager may have been reset since the stream registered,
330                // the id can belong to another stream then
331                if let Some(Some(item)) = manager.storage.get(id.get())
332                    && Rc::ptr_eq(&item.0, &io.0)
333                {
334                    manager.storage.remove(id.get());
335                }
336            });
337        }
338    }
339}
340
341pub(crate) struct Iops {
342    running: bool,
343    pub(crate) ops: Vec<Id>,
344}
345
346impl Iops {
347    pub(crate) fn schedule_write(id: Id) {
348        IoManager::with(|mgr| {
349            mgr.iops.ops.push(id);
350
351            if !mgr.iops.running {
352                mgr.iops.running = true;
353                spawn(async move { Iops::run() });
354            }
355        });
356    }
357
358    pub(crate) fn run() {
359        IoManager::with(|mgr| {
360            mgr.iops.running = false;
361
362            let mut ops = mem::take(&mut mgr.iops.ops);
363            for id in ops.drain(..) {
364                if let Some(io) = mgr.get(id) {
365                    io.ops_send_buf();
366                }
367            }
368            let _ = mem::replace(&mut mgr.iops.ops, ops);
369        });
370    }
371
372    #[cfg(test)]
373    pub(crate) fn is_registered(io: &IoRef) -> bool {
374        IoManager::with(|mgr| mgr.iops.ops.contains(&io.id()))
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    use ntex::rt::{DefaultRuntime, System};
381
382    use super::*;
383
384    async fn wait_timeout(io: &crate::Io) -> Duration {
385        let start = Instant::now();
386        let st = std::future::poll_fn(|cx| io.poll_status_update(cx)).await;
387        assert!(matches!(st, crate::IoStatusUpdate::Timeout));
388        start.elapsed()
389    }
390
391    #[test]
392    fn large_timer_sets_are_not_cached() {
393        let mut timers = TimerStorage {
394            running: false,
395            base: Instant::now(),
396            current: 0,
397            cache: VecDeque::new(),
398            notifications: BTreeMap::new(),
399        };
400
401        let mut small = HashSet::default();
402        small.insert(Id(NonZeroUsize::new(1)));
403        small.clear();
404        timers.recycle(small);
405        assert_eq!(timers.cache.len(), 1);
406
407        // a set grown by a burst of timers is dropped
408        let mut large = HashSet::default();
409        for i in 1..=10_000 {
410            large.insert(Id(NonZeroUsize::new(i)));
411        }
412        large.clear();
413        timers.recycle(large);
414        assert_eq!(timers.cache.len(), 1);
415        assert!(timers.cache[0].capacity() <= MAX_CACHED_SET_CAP);
416
417        // the cache holds at most `CAP` sets
418        for _ in 0..2 * CAP {
419            timers.recycle(HashSet::default());
420        }
421        assert_eq!(timers.cache.len(), CAP);
422    }
423
424    /// A timer started while no timers are pending expires after its
425    /// timeout, the ticker does not add a second.
426    #[ntex::test]
427    async fn timer_expires_after_timeout() {
428        use ntex_service::cfg::SharedCfg;
429
430        use crate::{Io, testing::IoTest};
431
432        let (_client, server) = IoTest::create();
433        let io = Io::new(server, SharedCfg::new("T"));
434
435        // the clock is not aligned to the start of the runtime
436        sleep(Millis(500)).await;
437
438        io.start_timer(Seconds(1));
439        let elapsed = wait_timeout(&io).await;
440        assert!(
441            elapsed >= Duration::from_secs(1) && elapsed < Duration::from_millis(1300),
442            "elapsed: {elapsed:?}"
443        );
444    }
445
446    /// A timer started while other timers are pending expires within a
447    /// second after its timeout.
448    #[ntex::test]
449    async fn timer_expires_within_second_after_timeout() {
450        use ntex_service::cfg::SharedCfg;
451
452        use crate::{Io, testing::IoTest};
453
454        let (_client1, server1) = IoTest::create();
455        let (_client2, server2) = IoTest::create();
456        let io1 = Io::new(server1, SharedCfg::new("T"));
457        let io2 = Io::new(server2, SharedCfg::new("T"));
458
459        io1.start_timer(Seconds(10));
460        sleep(Millis(500)).await;
461
462        // expires at the second clock second
463        io2.start_timer(Seconds(1));
464        let elapsed = wait_timeout(&io2).await;
465        assert!(
466            elapsed >= Duration::from_secs(1) && elapsed < Duration::from_millis(1800),
467            "elapsed: {elapsed:?}"
468        );
469    }
470
471    fn has_manager() -> bool {
472        MANAGER.with(|mgr| mgr.borrow().is_some())
473    }
474
475    /// The manager must not carry over to the next runtime on the thread, a
476    /// write scheduled by a runtime that has stopped would block the writes of
477    /// the next one.
478    #[test]
479    fn manager_reset_on_shutdown() {
480        std::thread::spawn(|| {
481            System::new("test", DefaultRuntime).block_on(async {
482                Iops::schedule_write(Id(None));
483                assert!(has_manager());
484            });
485            assert!(!has_manager());
486
487            System::build()
488                .build(DefaultRuntime)
489                .run(|| {
490                    Iops::schedule_write(Id(None));
491                    System::current().stop();
492                    Ok(())
493                })
494                .unwrap();
495            assert!(!has_manager());
496        })
497        .join()
498        .unwrap();
499    }
500}
ntex_io/ops.rs

ntex_io/
ops.rs
