Fix replication with a removal at the same tick (#264)

Co-authored-by: UkoeHB <37489173+UkoeHB@users.noreply.github.com>

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- Fix replication with a removal at the same tick.
+
 ## [0.25.2] - 2024-05-18
 
 ### Fixed

src/server.rs

@@ -12,7 +12,8 @@ use std::{io::Cursor, mem, time::Duration};
 use bevy::{
     ecs::{
         archetype::ArchetypeEntity,
-        component::{ComponentId, ComponentTicks, StorageType, Tick},
+        component::{ComponentId, ComponentTicks, StorageType},
+        entity::EntityHashSet,
         storage::{SparseSets, Table},
         system::SystemChangeTick,
     },
@@ -210,6 +211,7 @@ impl ServerPlugin {
 
     /// Collects [`ReplicationMessages`] and sends them.
     pub(super) fn send_replication(
+        mut entities_with_removals: Local<EntityHashSet>,
         mut messages: Local<ReplicationMessages>,
         mut replicated_archetypes: Local<ReplicatedArchetypes>,
         change_tick: SystemChangeTick,
@@ -234,15 +236,17 @@ impl ServerPlugin {
 
         collect_mappings(&mut messages, &mut set.p2())?;
         collect_despawns(&mut messages, &mut set.p3())?;
-        collect_removals(&mut messages, &mut set.p4(), change_tick.this_run())?;
+        collect_removals(&mut messages, &mut set.p4(), &mut entities_with_removals)?;
         collect_changes(
             &mut messages,
             &replicated_archetypes,
             &replication_fns,
+            &entities_with_removals,
             set.p0(),
             &change_tick,
             **server_tick,
         )?;
+        entities_with_removals.clear();
 
         let mut client_buffers = mem::take(&mut *set.p5());
         let connected_clients = messages.send(
@@ -299,6 +303,7 @@ fn collect_changes(
     messages: &mut ReplicationMessages,
     replicated_archetypes: &ReplicatedArchetypes,
     replication_fns: &ReplicationFns,
+    entities_with_removals: &EntityHashSet,
     world: &World,
     change_tick: &SystemChangeTick,
     server_tick: RepliconTick,
@@ -404,8 +409,11 @@ fn collect_changes(
                 }
 
                 let new_entity = marker_added || visibility == Visibility::Gained;
-                if new_entity || init_message.entity_data_size() != 0 {
-                    // If there is any insertion or we must initialize, include all updates into init message
+                if new_entity
+                    || init_message.entity_data_size() != 0
+                    || entities_with_removals.contains(&entity.id())
+                {
+                    // If there is any insertion, removal, or we must initialize, include all updates into init message.
                     // and bump the last acknowledged tick to keep entity updates atomic.
                     init_message.take_entity_data(update_message)?;
                     client.set_change_limit(entity.id(), change_tick.this_run());
@@ -487,19 +495,19 @@ fn collect_despawns(
 fn collect_removals(
     messages: &mut ReplicationMessages,
     removal_buffer: &mut RemovalBuffer,
-    tick: Tick,
+    entities_with_removals: &mut EntityHashSet,
 ) -> bincode::Result<()> {
     for (message, _) in messages.iter_mut() {
         message.start_array();
     }
 
     for (entity, remove_ids) in removal_buffer.iter() {
-        for (message, _, client) in messages.iter_mut_with_clients() {
+        for (message, _) in messages.iter_mut() {
             message.start_entity_data(entity);
             for fns_info in remove_ids {
-                client.set_change_limit(entity, tick);
                 message.write_fns_id(fns_info.fns_id())?;
             }
+            entities_with_removals.insert(entity);
             message.end_entity_data(false)?;
         }
     }

src/server/removal_buffer.rs

@@ -184,7 +184,12 @@ impl RemovalBuffer {
                 }
             }
         }
-        self.removals.push((entity, removed_ids));
+
+        if removed_ids.is_empty() {
+            self.ids_buffer.push(removed_ids);
+        } else {
+            self.removals.push((entity, removed_ids));
+        }
     }
 
     /// Clears all removals.

src/server/replication_messages.rs

@@ -346,6 +346,10 @@ impl InitMessage {
         update_message: &mut UpdateMessage,
     ) -> bincode::Result<()> {
         if update_message.entity_data_size != 0 {
+            if self.entity_data_size == 0 {
+                self.write_data_entity()?;
+            }
+
             let slice = update_message.as_slice();
             let offset = update_message.entity_data_size_pos as usize
                 + mem::size_of_val(&update_message.entity_data_size);

tests/changes.rs

@@ -566,6 +566,49 @@ fn with_insertion() {
     assert!(component.0);
 }
 
+#[test]
+fn with_removal() {
+    let mut server_app = App::new();
+    let mut client_app = App::new();
+    for app in [&mut server_app, &mut client_app] {
+        app.add_plugins((
+            MinimalPlugins,
+            RepliconPlugins.set(ServerPlugin {
+                tick_policy: TickPolicy::EveryFrame,
+                ..Default::default()
+            }),
+        ))
+        .replicate::<BoolComponent>()
+        .replicate::<DummyComponent>();
+    }
+
+    server_app.connect_client(&mut client_app);
+
+    let server_entity = server_app
+        .world
+        .spawn((Replicated, BoolComponent(false), DummyComponent))
+        .id();
+
+    server_app.update();
+    server_app.exchange_with_client(&mut client_app);
+    client_app.update();
+    server_app.exchange_with_client(&mut client_app);
+
+    let mut server_entity = server_app.world.entity_mut(server_entity);
+    server_entity.get_mut::<BoolComponent>().unwrap().0 = true;
+    server_entity.remove::<DummyComponent>();
+
+    server_app.update();
+    server_app.exchange_with_client(&mut client_app);
+    client_app.update();
+
+    let component = client_app
+        .world
+        .query_filtered::<&BoolComponent, Without<DummyComponent>>()
+        .single(&client_app.world);
+    assert!(component.0);
+}
+
 #[test]
 fn with_despawn() {
     let mut server_app = App::new();