Skip to content

Commit 34bb995

Browse files
committed
Properly abstract EntityPools
Since it's used in many places, it's better to abstract it to avoid forgetting to call `clear`.
1 parent fac569d commit 34bb995

File tree

4 files changed

+101
-83
lines changed

4 files changed

+101
-83
lines changed

src/server.rs

Lines changed: 9 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
pub mod client_pools;
12
pub mod message;
23
pub mod related_entities;
34
pub(super) mod removal_buffer;
@@ -29,23 +30,20 @@ use crate::{
2930
postcard_utils,
3031
prelude::*,
3132
server::{
32-
replication_messages::mutations::{EntityMutations, MutationsSplit},
33-
visibility::registry::FilterRegistry,
33+
replication_messages::mutations::MutationsSplit, visibility::registry::FilterRegistry,
3434
},
3535
shared::{
3636
backend::channels::ClientChannel,
3737
message::server_message::message_buffer::MessageBuffer,
3838
replication::{
3939
client_ticks::{ClientTicks, EntityTicks},
40-
registry::{
41-
FnsId, ReplicationRegistry, component_mask::ComponentMask, ctx::SerializeCtx,
42-
serde_fns::SerdeFns,
43-
},
40+
registry::{FnsId, ReplicationRegistry, ctx::SerializeCtx, serde_fns::SerdeFns},
4441
rules::{ReplicationRules, component::ComponentRule},
4542
track_mutate_messages::TrackMutateMessages,
4643
},
4744
},
4845
};
46+
use client_pools::ClientPools;
4947
use related_entities::RelatedEntities;
5048
use removal_buffer::{RemovalBuffer, RemovalReader};
5149
use replication_messages::{
@@ -242,7 +240,7 @@ fn cleanup_acks(
242240
let min_timestamp = time.elapsed().saturating_sub(mutations_timeout);
243241
for mut ticks in &mut clients {
244242
ticks.cleanup_older_mutations(min_timestamp, |mutate_info| {
245-
pools.push_entities(mem::take(&mut mutate_info.entities));
243+
pools.recycle_entities(mem::take(&mut mutate_info.entities));
246244
});
247245
}
248246
}
@@ -263,7 +261,7 @@ fn receive_acks(
263261
)
264262
});
265263
if let Some(entities) = ticks.ack_mutate_message(client, mutate_index) {
266-
pools.push_entities(entities);
264+
pools.recycle_entities(entities);
267265
}
268266
}
269267
Err(e) => {
@@ -532,7 +530,7 @@ fn collect_despawns(
532530
// spawn and despawn could happen during the same tick.
533531
trace!("writing despawn for `{entity}` for client `{client_entity}`");
534532
message.add_despawn(entity_range.clone());
535-
pools.push_components(entity_ticks.components);
533+
pools.recycle_components(entity_ticks.components);
536534
}
537535
visibility.remove_despawned(entity);
538536
priority.remove(&entity);
@@ -545,7 +543,7 @@ fn collect_despawns(
545543
trace!("writing visibility lost for `{entity}` for client `{client_entity}`");
546544
let entity_range = serialized.write_entity(entity)?;
547545
message.add_despawn(entity_range);
548-
pools.push_components(entity_ticks.components);
546+
pools.recycle_components(entity_ticks.components);
549547
}
550548
priority.remove(&entity);
551549
}
@@ -753,7 +751,7 @@ fn collect_changes(
753751
entity_ticks.system_tick = change_tick.this_run();
754752
entity_ticks.server_tick = server_tick;
755753
entity_ticks.components |= &components;
756-
pools.push_components(components);
754+
pools.recycle_components(components);
757755
}
758756
Entry::Vacant(entry) => {
759757
entry.insert(EntityTicks {
@@ -960,41 +958,3 @@ pub struct AuthorizedClient;
960958
/// See its documentation for more details.
961959
#[derive(Component, Deref, DerefMut, Debug, Default, Clone)]
962960
pub struct PriorityMap(EntityHashMap<f32>);
963-
964-
/// Pools for various client components to reuse allocated capacity.
965-
///
966-
/// All data is cleared before the insertion.
967-
#[derive(Resource, Default)]
968-
struct ClientPools {
969-
/// Entities with bitvecs for components from
970-
/// [`MutateInfo`](crate::shared::replication::client_ticks::MutateInfo).
971-
entities: Vec<Vec<(Entity, ComponentMask)>>,
972-
/// Bitvecs for components from [`Updates`], [`Mutations`] and
973-
/// [`MutateInfo`](crate::shared::replication::client_ticks::MutateInfo).
974-
///
975-
/// Only heap-allocated instances are stored.
976-
components: Vec<ComponentMask>,
977-
/// Ranges from [`Updates`] and [`Mutations`].
978-
ranges: Vec<Vec<Range<usize>>>,
979-
/// Entities from [`Mutations`].
980-
mutations: Vec<Vec<EntityMutations>>,
981-
}
982-
983-
impl ClientPools {
984-
fn push_entities(&mut self, mut entities: Vec<(Entity, ComponentMask)>) {
985-
for (_, mut components) in entities.drain(..) {
986-
if components.is_heap() {
987-
components.clear();
988-
self.components.push(components);
989-
}
990-
}
991-
self.entities.push(entities);
992-
}
993-
994-
fn push_components(&mut self, mut components: ComponentMask) {
995-
if components.is_heap() {
996-
components.clear();
997-
self.components.push(components);
998-
}
999-
}
1000-
}

src/server/client_pools.rs

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
use core::ops::Range;
2+
3+
use bevy::prelude::*;
4+
5+
use super::replication_messages::mutations::EntityMutations;
6+
use crate::shared::replication::registry::component_mask::ComponentMask;
7+
8+
/// Pools for various client components to reuse allocated capacity.
9+
///
10+
/// All data is cleared before the insertion.
11+
#[derive(Resource, Default)]
12+
pub(super) struct ClientPools {
13+
/// Entities with bitvecs for components from
14+
/// [`MutateInfo`](crate::shared::replication::client_ticks::MutateInfo).
15+
entities: Vec<Vec<(Entity, ComponentMask)>>,
16+
/// Bitvecs for components from [`Updates`], [`Mutations`] and
17+
/// [`MutateInfo`](crate::shared::replication::client_ticks::MutateInfo).
18+
///
19+
/// Only heap-allocated instances are stored.
20+
components: Vec<ComponentMask>,
21+
/// Ranges from [`Updates`] and [`Mutations`].
22+
ranges: Vec<Vec<Range<usize>>>,
23+
/// Entities from [`Mutations`].
24+
mutations: Vec<Vec<EntityMutations>>,
25+
}
26+
27+
impl ClientPools {
28+
pub(super) fn recycle_entities(&mut self, mut entities: Vec<(Entity, ComponentMask)>) {
29+
for (_, mut components) in entities.drain(..) {
30+
if components.is_heap() {
31+
components.clear();
32+
self.components.push(components);
33+
}
34+
}
35+
self.entities.push(entities);
36+
}
37+
38+
pub(super) fn recycle_components(&mut self, mut components: ComponentMask) {
39+
if components.is_heap() {
40+
components.clear();
41+
self.components.push(components);
42+
}
43+
}
44+
45+
pub(super) fn recycle_ranges(&mut self, ranges: impl Iterator<Item = Vec<Range<usize>>>) {
46+
self.ranges.extend(ranges.map(|mut ranges| {
47+
ranges.clear();
48+
ranges
49+
}));
50+
}
51+
52+
pub(super) fn recycle_mutations(
53+
&mut self,
54+
mutations: impl Iterator<Item = Vec<EntityMutations>>,
55+
) {
56+
self.mutations.extend(mutations.map(|mut mutations| {
57+
mutations.clear();
58+
mutations
59+
}));
60+
}
61+
62+
pub(super) fn take_entities(&mut self) -> Vec<(Entity, ComponentMask)> {
63+
self.entities.pop().unwrap_or_default()
64+
}
65+
66+
pub(super) fn take_components(&mut self) -> ComponentMask {
67+
self.components.pop().unwrap_or_default()
68+
}
69+
70+
pub(super) fn take_ranges(&mut self) -> Vec<Range<usize>> {
71+
self.ranges.pop().unwrap_or_default()
72+
}
73+
74+
pub(super) fn take_mutations(&mut self) -> Vec<EntityMutations> {
75+
self.mutations.pop().unwrap_or_default()
76+
}
77+
}

src/server/replication_messages/mutations.rs

Lines changed: 7 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -65,15 +65,14 @@ impl Mutations {
6565
graph_index: Option<usize>,
6666
entity_range: Range<usize>,
6767
) {
68-
let components = pools.ranges.pop().unwrap_or_default();
6968
let mutations = EntityMutations {
7069
entity,
7170
ranges: ChangeRanges {
7271
entity: entity_range,
7372
components_len: 0,
74-
components,
73+
components: pools.take_ranges(),
7574
},
76-
components: pools.components.pop().unwrap_or_default(),
75+
components: pools.take_components(),
7776
};
7877

7978
match graph_index {
@@ -155,7 +154,7 @@ impl Mutations {
155154
server_tick,
156155
system_tick,
157156
timestamp,
158-
entities: pools.entities.pop().unwrap_or_default(),
157+
entities: pools.take_entities(),
159158
};
160159
let mut mutate_index = ticks.next_mutate_index();
161160
let mut chunks = EntityChunks::new(&mut self.related, &mut self.standalone);
@@ -185,7 +184,7 @@ impl Mutations {
185184
server_tick,
186185
system_tick,
187186
timestamp,
188-
entities: pools.entities.pop().unwrap_or_default(),
187+
entities: pools.take_entities(),
189188
};
190189
chunks_range.start = chunks_range.end;
191190
header_size = metadata_size + serialized_size(&mutate_index)?; // Recalculate since the mutate index changed.
@@ -263,11 +262,7 @@ impl Mutations {
263262
.iter_mut()
264263
.chain(iter::once(&mut self.standalone))
265264
{
266-
let ranges = entities.drain(..).map(|mut mutations| {
267-
mutations.ranges.components.clear();
268-
mutations.ranges.components
269-
});
270-
pools.ranges.extend(ranges);
265+
pools.recycle_ranges(entities.drain(..).map(|m| m.ranges.components));
271266
// We don't take component masks because they are moved to `MutateInfo` during sending.
272267
}
273268
}
@@ -279,14 +274,8 @@ impl Mutations {
279274
match self.related.len().cmp(&graphs_count) {
280275
Ordering::Less => self
281276
.related
282-
.resize_with(graphs_count, || pools.mutations.pop().unwrap_or_default()),
283-
Ordering::Greater => {
284-
let mutations = self.related.drain(graphs_count..).map(|mut mutations| {
285-
mutations.clear();
286-
mutations
287-
});
288-
pools.mutations.extend(mutations);
289-
}
277+
.resize_with(graphs_count, || pools.take_mutations()),
278+
Ordering::Greater => pools.recycle_mutations(self.related.drain(graphs_count..)),
290279
Ordering::Equal => (),
291280
}
292281
}

src/server/replication_messages/updates.rs

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use core::{mem, ops::Range};
1+
use core::{iter, mem, ops::Range};
22

33
use bevy::prelude::*;
44
use postcard::experimental::serialized_size;
@@ -120,7 +120,7 @@ impl Updates {
120120
self.removals.push(RemovalRanges {
121121
entity,
122122
ids_len: 0,
123-
ids: pools.ranges.pop().unwrap_or_default(),
123+
ids: pools.take_ranges(),
124124
});
125125
self.removals_entity_added = true;
126126
}
@@ -159,7 +159,7 @@ impl Updates {
159159
self.changes.push(ChangeRanges {
160160
entity,
161161
components_len: 0,
162-
components: pools.ranges.pop().unwrap_or_default(),
162+
components: pools.take_ranges(),
163163
});
164164
self.changed_entity_added = true;
165165
}
@@ -183,19 +183,18 @@ impl Updates {
183183
/// Takes last mutated entity with its component chunks from the mutate message.
184184
pub(crate) fn take_added_entity(&mut self, pools: &mut ClientPools, mutations: &mut Mutations) {
185185
debug_assert!(mutations.entity_added());
186-
let mut entity_mutations = mutations.pop().expect("entity should be written");
186+
let entity_mutations = mutations.pop().expect("entity should be written");
187187

188188
if !self.changed_entity_added {
189189
self.changes.push(entity_mutations.ranges);
190190
} else {
191191
let changes = self.changes.last_mut().expect("entity should be written");
192192
debug_assert_eq!(entity_mutations.ranges.entity, changes.entity);
193193
changes.extend(&entity_mutations.ranges);
194-
entity_mutations.ranges.components.clear();
195-
pools.ranges.push(entity_mutations.ranges.components);
194+
pools.recycle_ranges(iter::once(entity_mutations.ranges.components));
196195

197196
self.changed_components |= &entity_mutations.components;
198-
pools.push_components(entity_mutations.components);
197+
pools.recycle_components(entity_mutations.components);
199198
}
200199
}
201200

@@ -358,15 +357,8 @@ impl Updates {
358357
self.despawns.clear();
359358
self.despawns_len = 0;
360359

361-
let change_ranges = self.changes.drain(..).map(|mut changes| {
362-
changes.components.clear();
363-
changes.components
364-
});
365-
let removal_ranges = self.removals.drain(..).map(|mut removals| {
366-
removals.ids.clear();
367-
removals.ids
368-
});
369-
pools.ranges.extend(change_ranges.chain(removal_ranges));
360+
pools.recycle_ranges(self.changes.drain(..).map(|c| c.components));
361+
pools.recycle_ranges(self.removals.drain(..).map(|c| c.ids));
370362
}
371363
}
372364

0 commit comments

Comments
 (0)