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    pub fn new(recycle_timeout: Duration) -> Self {
16        Self {
17            recycle_timeout,
18            recycling_keys: VecDeque::new(),
19            recycled_keys: VecDeque::new(),
20            next_new_key: 0,
21            phantom: PhantomData,
22        }
23    }
24    /// Get a new, unused key
25    pub fn generate(&mut self) -> K {
26        let now = Instant::now();
27
28        // Check whether we can recycle any keys
29        loop {
30            let Some((_, instant)) = self.recycling_keys.front() else {
31                break;
32            };
33            if instant.elapsed(&now) < self.recycle_timeout {
34                break;
35            }
36            let (key, _) = self.recycling_keys.pop_front().unwrap();
37            self.recycled_keys.push_back(key);
38        }
39
40        // Check whether we can return a recycled key
41        if self.recycled_keys.len() > 0 {
42            let key = self.recycled_keys.pop_front().unwrap();
43            return K::from(key);
44        }
45
46        // Create a new key
47        let output = self.next_new_key;
48        self.next_new_key = self.next_new_key.wrapping_add(1);
49        K::from(output)
50    }
51
52    /// Recycle a used key, freeing it up
53    pub fn recycle_key(&mut self, key: &K) {
54        let key_u16: u16 = Into::<u16>::into(*key);
55        self.recycling_keys.push_back((key_u16, Instant::now()));
56    }
57}