Skip to main content

naia_shared/world/host/
host_entity_generator.rs

1use std::{
2    collections::{HashMap, VecDeque},
3    time::Duration,
4};
5
6use naia_socket_shared::Instant;
7
8use crate::world::local::local_entity::{HostEntity, RemoteEntity};
9use crate::world::local::local_entity_map::LocalEntityMap;
10use crate::{GlobalEntity, KeyGenerator};
11
12/// Issues and recycles wire-level [`HostEntity`] identifiers for a single connected user.
13pub struct HostEntityGenerator {
14    user_key: u64,
15    generator: KeyGenerator<u16>,
16    static_generator: KeyGenerator<u16>,
17    reserved_host_entities: HashMap<GlobalEntity, HostEntity>,
18    reserved_host_entity_ttl: Duration,
19    reserved_host_entities_ttls: VecDeque<(Instant, GlobalEntity)>,
20}
21
22impl HostEntityGenerator {
23    /// Creates a generator bound to `user_key` with fresh entity and static-entity ID pools.
24    pub fn new(user_key: u64) -> Self {
25        Self {
26            user_key,
27            generator: KeyGenerator::new(Duration::from_secs(60)),
28            static_generator: KeyGenerator::new(Duration::from_secs(60)),
29            reserved_host_entities: HashMap::new(),
30            reserved_host_entity_ttl: Duration::from_secs(60),
31            reserved_host_entities_ttls: VecDeque::new(),
32        }
33    }
34
35    // Host entities
36
37    /// Allocates a [`HostEntity`] for `global_entity` before it has been sent, expiring reservations that have timed out.
38    pub fn host_reserve_entity(
39        &mut self,
40        entity_map: &mut LocalEntityMap,
41        global_entity: &GlobalEntity,
42    ) -> HostEntity {
43        self.process_reserved_entity_timeouts();
44
45        if self.reserved_host_entities.contains_key(global_entity) {
46            panic!("Global Entity has already reserved Local Entity!");
47        }
48        let host_entity = self.generate_host_entity();
49        entity_map.insert_with_host_entity(*global_entity, host_entity);
50        self.reserved_host_entities
51            .insert(*global_entity, host_entity);
52        host_entity
53    }
54
55    fn process_reserved_entity_timeouts(&mut self) {
56        let now = Instant::now();
57
58        loop {
59            let Some((timeout, _)) = self.reserved_host_entities_ttls.front() else {
60                break;
61            };
62            if timeout.elapsed(&now) < self.reserved_host_entity_ttl {
63                break;
64            }
65            let (_, global_entity) = self.reserved_host_entities_ttls.pop_front().unwrap();
66            let Some(_) = self.reserved_host_entities.remove(&global_entity) else {
67                panic!("Reserved Entity does not exist!");
68            };
69        }
70    }
71
72    /// Removes and returns the reserved [`HostEntity`] for `global_entity`, if one exists.
73    pub fn host_remove_reserved_entity(
74        &mut self,
75        global_entity: &GlobalEntity,
76    ) -> Option<HostEntity> {
77        self.reserved_host_entities.remove(global_entity)
78    }
79
80    pub(crate) fn generate_host_entity(&mut self) -> HostEntity {
81        HostEntity::new(self.generator.generate())
82    }
83
84    // Static entities use a separate counter so their wire IDs (0, 1, 2 …)
85    // never collide with dynamic entity wire IDs (also 0, 1, 2 …).
86    // The is_static bit in the wire format keeps them distinguishable.
87    pub(crate) fn generate_static_host_entity(&mut self) -> HostEntity {
88        HostEntity::new_static(self.static_generator.generate())
89    }
90
91    pub(crate) fn remove_by_global_entity(
92        &mut self,
93        entity_map: &mut LocalEntityMap,
94        global_entity: &GlobalEntity,
95    ) {
96        let record = entity_map
97            .remove_by_global_entity(global_entity)
98            .expect("Attempting to despawn entity which does not exist!");
99        if record.is_host_owned() {
100            let host_entity = record.host_entity();
101            if host_entity.is_static() {
102                self.static_generator.recycle_key(&host_entity.value());
103            } else {
104                self.generator.recycle_key(&host_entity.value());
105            }
106        }
107    }
108
109    pub(crate) fn remove_by_host_entity(
110        &mut self,
111        converter: &mut LocalEntityMap,
112        host_entity: &HostEntity,
113    ) {
114        // The mapping may already have been cleared at send time by
115        // `LocalWorldManager::despawn_entity` (see [entity-delegation-15]).
116        // In that case, just recycle the id slot — the entity_map removal
117        // already happened. Otherwise (legacy path, or non-despawn cleanup),
118        // do the full remove.
119        if let Some(global_entity) = converter.global_entity_from_host(host_entity).copied() {
120            self.remove_by_global_entity(converter, &global_entity);
121        } else {
122            // Send-time cleanup already removed the mapping; just free the id.
123            if host_entity.is_static() {
124                self.static_generator.recycle_key(&host_entity.value());
125            } else {
126                self.generator.recycle_key(&host_entity.value());
127            }
128        }
129    }
130
131    /// Removes the entity identified by `remote_entity` from `entity_map`, recycles its host ID, and returns its [`GlobalEntity`].
132    pub fn remove_by_remote_entity(
133        &mut self,
134        entity_map: &mut LocalEntityMap,
135        remote_entity: &RemoteEntity,
136    ) -> GlobalEntity {
137        let global_entity = *(entity_map
138            .global_entity_from_remote(remote_entity)
139            .expect("Attempting to despawn entity which does not exist!"));
140        let record = entity_map
141            .remove_by_global_entity(&global_entity)
142            .expect("Attempting to despawn entity which does not exist!");
143        if record.is_host_owned() {
144            let host_entity = record.host_entity();
145            if host_entity.is_static() {
146                self.static_generator.recycle_key(&host_entity.value());
147            } else {
148                self.generator.recycle_key(&host_entity.value());
149            }
150        }
151        global_entity
152    }
153
154    // Misc
155
156    /// Returns the user key this generator was created for.
157    pub fn get_user_key(&self) -> &u64 {
158        &self.user_key
159    }
160}