naia_shared/
key_generator.rs

1use std::{collections::VecDeque, marker::PhantomData, time::Duration};
2
3use naia_socket_shared::Instant;
4
5/// Simple implementation of a store that manages a recycling pool of u16 keys
6pub struct KeyGenerator<K: From<u16> + Into<u16> + Copy> {
7    recycling_keys: VecDeque<(u16, Instant)>,
8    recycled_keys: VecDeque<u16>,
9    recycle_timeout: Duration,
10    next_new_key: u16,
11    phantom: PhantomData<K>,
12}
13
14impl<K: From<u16> + Into<u16> + Copy> KeyGenerator<K> {
15    /// Creates a `KeyGenerator` that holds recycled keys for at least `recycle_timeout` before reissuing them.
16    pub fn new(recycle_timeout: Duration) -> Self {
17        Self {
18            recycle_timeout,
19            recycling_keys: VecDeque::new(),
20            recycled_keys: VecDeque::new(),
21            next_new_key: 0,
22            phantom: PhantomData,
23        }
24    }
25    /// Get a new, unused key
26    pub fn generate(&mut self) -> K {
27        let now = Instant::now();
28
29        // Check whether we can recycle any keys
30        loop {
31            let Some((_, instant)) = self.recycling_keys.front() else {
32                break;
33            };
34            if instant.elapsed(&now) < self.recycle_timeout {
35                break;
36            }
37            let (key, _) = self.recycling_keys.pop_front().unwrap();
38            self.recycled_keys.push_back(key);
39        }
40
41        // Check whether we can return a recycled key
42        if !self.recycled_keys.is_empty() {
43            let key = self.recycled_keys.pop_front().unwrap();
44            return K::from(key);
45        }
46
47        // Create a new key
48        let output = self.next_new_key;
49        self.next_new_key = self.next_new_key.checked_add(1).expect("KeyGenerator exhausted: all u16 keys are in use");
50        K::from(output)
51    }
52
53    /// Recycle a used key, freeing it up
54    pub fn recycle_key(&mut self, key: &K) {
55        let key_u16: u16 = Into::<u16>::into(*key);
56        self.recycling_keys.push_back((key_u16, Instant::now()));
57    }
58}
naia_shared/key_generator.rs

naia_shared/
key_generator.rs
