Buffered insertions (#484)

Originally, we used commands to insert or remove components during
replication. This wasn't optimal because Bevy doesn't batch commands.
However, since Bevy was considering adding support for batching, we
decided to wait.

With the introduction of triggers, this became a real issue: triggers
would fire after each insertion, and inside the observer, the user
might not be able to access all components (depending on the insertion
order).

This change replaces commands with insertion and removal methods on
`DeferredEntity` that buffer all changes and apply them later using
`EntityWorldMut::insert_by_ids` and `EntityWorldMut::remove_by_ids`.
They are flushed after processing each entity. This not only makes
triggers behave as expected, but also significantly improves
performance by avoiding extra archetype moves and lookups.

I also slightly changed the behavior of the `Replicated` component to
simplify the implementation and unlock additional performance.
`Replicated` is no longer automatically inserted into non-replicated
entities spawned from replicated components. This change is consistent
with server behavior, where such entities also do not have the
`Replicated` component.

The commit entry-based API for `ServerEntityMap` that is required for this change.

Author: Shatur

CHANGELOG.md

@@ -21,9 +21,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `ReplicationRule` now stores `Vec<ComponentRule>` instead of `Vec<(ComponentId, FnsId)>`
 - `RuleFns` now available from prelude.
 - Rules created with the same priority now evaluated in their creation order.
+- Component removals and insertions for an entity are now buffered and applied as bundles to avoid triggering observers without all components being inserted or removed. This also significantly improves performance by avoiding extra archetype moves and lookups.
+- The `Replicated` component is no longer automatically inserted into non-replicated entities spawned from replicated components.
+- Replace `ServerEntityMap::get_by_*` and `ServerEntityMap::remove_by_*` with an entry-based API. Use `ServerEntityMap::server_entry` or `ServerEntityMap::client_entry` instead.
+- Print error instead of panic on mapping overwrite in `ServerEntityMap`.
 
 ### Removed
 
+- `WriteCtx::commands`. You can now insert and remove components directly through `DeferredEntity`.
 - Deprecated methods.
 
 ## [0.33.0] - 2025-04-27

src/client.rs

@@ -4,7 +4,7 @@ pub mod diagnostics;
 pub mod event;
 pub mod server_mutate_ticks;
 
-use bevy::{ecs::world::CommandQueue, prelude::*, reflect::TypeRegistry};
+use bevy::{prelude::*, reflect::TypeRegistry};
 use bytes::{Buf, Bytes};
 use log::{debug, trace};
 use postcard::experimental::max_size::MaxSize;
@@ -19,7 +19,7 @@ use crate::shared::{
     replication::{
         Replicated,
         command_markers::{CommandMarkers, EntityMarkers},
-        deferred_entity::DeferredEntity,
+        deferred_entity::{DeferredChanges, DeferredEntity},
         mutate_index::MutateIndex,
         replication_registry::{
             ReplicationRegistry,
@@ -29,7 +29,7 @@ use crate::shared::{
         update_message_flags::UpdateMessageFlags,
     },
     replicon_tick::RepliconTick,
-    server_entity_map::ServerEntityMap,
+    server_entity_map::{EntityEntry, ServerEntityMap},
 };
 use confirm_history::{ConfirmHistory, EntityReplicated};
 use server_mutate_ticks::{MutateTickReceived, ServerMutateTicks};
@@ -103,7 +103,7 @@ fn setup_channels(mut client: ResMut<RepliconClient>, channels: Res<RepliconChan
 /// See also [`ReplicationMessages`](crate::server::replication_messages::ReplicationMessages).
 pub(super) fn receive_replication(
     world: &mut World,
-    mut queue: Local<CommandQueue>,
+    mut changes: Local<DeferredChanges>,
     mut entity_markers: Local<EntityMarkers>,
 ) -> Result<()> {
     world.resource_scope(|world, mut client: Mut<RepliconClient>| {
@@ -119,7 +119,7 @@ pub(super) fn receive_replication(
                                     let mut mutate_ticks =
                                         world.remove_resource::<ServerMutateTicks>();
                                     let mut params = ReceiveParams {
-                                        queue: &mut queue,
+                                        changes: &mut changes,
                                         entity_markers: &mut entity_markers,
                                         entity_map: &mut entity_map,
                                         replicated_events: &mut replicated_events,
@@ -385,7 +385,8 @@ fn apply_despawn(
     let server_entity = entity_serde::deserialize_entity(message)?;
     if let Some(client_entity) = params
         .entity_map
-        .remove_by_server(server_entity)
+        .server_entry(server_entity)
+        .remove()
         .and_then(|entity| world.get_entity_mut(entity).ok())
     {
         let ctx = DespawnCtx { message_tick };
@@ -404,28 +405,31 @@ fn apply_removals(
 ) -> Result<()> {
     let server_entity = entity_serde::deserialize_entity(message)?;
 
-    let client_entity = params
-        .entity_map
-        .get_by_server_or_insert(server_entity, || world.spawn(Replicated).id());
+    let mut client_entity = match params.entity_map.server_entry(server_entity) {
+        EntityEntry::Occupied(entry) => {
+            DeferredEntity::new(world.entity_mut(entry.get()), params.changes)
+        }
+        EntityEntry::Vacant(entry) => {
+            // It's possible to receive a removal when an entity is spawned and has a component removed in the same tick.
+            // We could serialize the size of the removals instead of the total number of removals and just advance the cursor,
+            // but it's a very rare case and not worth optimizing for.
+            let mut client_entity = DeferredEntity::new(world.spawn_empty(), params.changes);
+            client_entity.insert(Replicated);
+            entry.insert(client_entity.id());
+            client_entity
+        }
+    };
 
-    let mut client_entity = DeferredEntity::new(world, client_entity);
-    let mut commands = client_entity.commands(params.queue);
     params
         .entity_markers
         .read(params.command_markers, &*client_entity);
 
-    confirm_tick(
-        &mut commands,
-        &mut client_entity,
-        params.replicated_events,
-        message_tick,
-    );
+    confirm_tick(&mut client_entity, params.replicated_events, message_tick);
 
     let len = apply_array(ArrayKind::Sized, message, |message| {
         let fns_id = postcard_utils::from_buf(message)?;
         let (component_id, component_fns, _) = params.registry.get(fns_id);
         let mut ctx = RemoveCtx {
-            commands: &mut commands,
             message_tick,
             component_id,
         };
@@ -438,7 +442,7 @@ fn apply_removals(
         stats.components_changed += len;
     }
 
-    params.queue.apply(world);
+    client_entity.flush();
 
     Ok(())
 }
@@ -452,32 +456,39 @@ fn apply_changes(
 ) -> Result<()> {
     let server_entity = entity_serde::deserialize_entity(message)?;
 
-    let client_entity = params
-        .entity_map
-        .get_by_server_or_insert(server_entity, || world.spawn(Replicated).id());
+    let world_cell = world.as_unsafe_world_cell();
+    let entities = world_cell.entities();
+    // SAFETY: split into `Entities` and `DeferredEntity`.
+    // The latter won't apply any structural changes until `flush`, and `Entities` won't be used afterward.
+    let world = unsafe { world_cell.world_mut() };
+
+    let mut client_entity = match params.entity_map.server_entry(server_entity) {
+        EntityEntry::Occupied(entry) => {
+            DeferredEntity::new(world.entity_mut(entry.get()), params.changes)
+        }
+        EntityEntry::Vacant(entry) => {
+            let mut client_entity = DeferredEntity::new(world.spawn_empty(), params.changes);
+            client_entity.insert(Replicated);
+            entry.insert(client_entity.id());
+            client_entity
+        }
+    };
 
-    let mut client_entity = DeferredEntity::new(world, client_entity);
-    let mut commands = client_entity.commands(params.queue);
     params
         .entity_markers
         .read(params.command_markers, &*client_entity);
 
-    confirm_tick(
-        &mut commands,
-        &mut client_entity,
-        params.replicated_events,
-        message_tick,
-    );
+    confirm_tick(&mut client_entity, params.replicated_events, message_tick);
 
     let len = apply_array(ArrayKind::Sized, message, |message| {
         let fns_id = postcard_utils::from_buf(message)?;
         let (component_id, component_fns, rule_fns) = params.registry.get(fns_id);
         let mut ctx = WriteCtx {
-            commands: &mut commands,
             entity_map: params.entity_map,
             type_registry: params.type_registry,
             component_id,
             message_tick,
+            entities,
             ignore_mapping: false,
         };
 
@@ -499,7 +510,7 @@ fn apply_changes(
         stats.components_changed += len;
     }
 
-    params.queue.apply(world);
+    client_entity.flush();
 
     Ok(())
 }
@@ -540,17 +551,14 @@ enum ArrayKind {
 }
 
 fn confirm_tick(
-    commands: &mut Commands,
     entity: &mut DeferredEntity,
     replicated_events: &mut Events<EntityReplicated>,
     tick: RepliconTick,
 ) {
     if let Some(mut history) = entity.get_mut::<ConfirmHistory>() {
         history.set_last_tick(tick);
     } else {
-        commands
-            .entity(entity.id())
-            .insert(ConfirmHistory::new(tick));
+        entity.insert(ConfirmHistory::new(tick));
     }
     replicated_events.send(EntityReplicated {
         entity: entity.id(),
@@ -570,15 +578,20 @@ fn apply_mutations(
     let server_entity = entity_serde::deserialize_entity(message)?;
     let data_size: usize = postcard_utils::from_buf(message)?;
 
-    let Some(client_entity) = params.entity_map.get_by_server(server_entity) else {
+    let Some(&client_entity) = params.entity_map.to_client().get(&server_entity) else {
         // Mutation could arrive after a despawn from update message.
         debug!("ignoring mutations received for unknown server's {server_entity:?}");
         message.advance(data_size);
         return Ok(());
     };
 
-    let mut client_entity = DeferredEntity::new(world, client_entity);
-    let mut commands = client_entity.commands(params.queue);
+    let world_cell = world.as_unsafe_world_cell();
+    let entities = world_cell.entities();
+    // SAFETY: split into `Entities` and `DeferredEntity`.
+    // The latter won't apply any structural changes until `flush`, and `Entities` won't be used afterward.
+    let world = unsafe { world_cell.world_mut() };
+
+    let mut client_entity = DeferredEntity::new(world.entity_mut(client_entity), params.changes);
     params
         .entity_markers
         .read(params.command_markers, &*client_entity);
@@ -622,11 +635,11 @@ fn apply_mutations(
         let fns_id = postcard_utils::from_buf(&mut data)?;
         let (component_id, component_fns, rule_fns) = params.registry.get(fns_id);
         let mut ctx = WriteCtx {
-            commands: &mut commands,
             entity_map: params.entity_map,
             type_registry: params.type_registry,
             component_id,
             message_tick,
+            entities,
             ignore_mapping: false,
         };
 
@@ -659,7 +672,7 @@ fn apply_mutations(
         stats.components_changed += components_count;
     }
 
-    params.queue.apply(world);
+    client_entity.flush();
 
     Ok(())
 }
@@ -668,7 +681,7 @@ fn apply_mutations(
 ///
 /// To avoid passing a lot of arguments into all receive functions.
 struct ReceiveParams<'a> {
-    queue: &'a mut CommandQueue,
+    changes: &'a mut DeferredChanges,
     entity_markers: &'a mut EntityMarkers,
     entity_map: &'a mut ServerEntityMap,
     replicated_events: &'a mut Events<EntityReplicated>,

src/shared/replication/command_markers.rs

@@ -84,17 +84,15 @@ pub trait AppMarkerExt {
         if let Some(mut history) = entity.get_mut::<History<C>>() {
             history.insert(ctx.message_tick, component);
         } else {
-            ctx.commands
-                .entity(entity.id())
-                .insert(History([(ctx.message_tick, component)].into()));
+            entity.insert(History([(ctx.message_tick, component)].into()));
         }
 
         Ok(())
     }
 
     /// Removes component `C` and its history.
-    fn remove_history<C: Component>(ctx: &mut RemoveCtx, entity: &mut DeferredEntity) {
-        ctx.commands.entity(entity.id()).remove::<History<C>>().remove::<C>();
+    fn remove_history<C: Component>(_ctx: &mut RemoveCtx, entity: &mut DeferredEntity) {
+        entity.remove::<History<C>>().remove::<C>();
     }
 
     /// If this marker is present on an entity, registered components will be stored in [`History<T>`].