Fix despawn for pre-spawned entities (#819)

* fix

* remove Rollback trigger

lightyear/src/client/prediction/despawn.rs

@@ -41,15 +41,11 @@ impl Command for PredictionDespawnCommand {
         }
 
         if let Ok(mut entity) = world.get_entity_mut(self.entity) {
-            if entity.get::<PreSpawnedPlayerObject>().is_some() {
-                // if this is a pre-spawned entity, we can despawn it immediately
-                // - if the entity was spawned/despawned within a server replication interval, then the server
-                //   wouldn't even send a spawn, so we did the right thing
-                // - if the server sends the entity afterwards, the matching will fail and we will just spawn a new predicted entity
-                entity.despawn_recursive();
-                return;
-            }
-            if entity.get::<Predicted>().is_some() || entity.get::<ShouldBePredicted>().is_some() {
+            if entity.get::<Predicted>().is_some()
+                || entity.get::<ShouldBePredicted>().is_some()
+                // see https://github.com/cBournhonesque/lightyear/issues/818
+                || entity.get::<PreSpawnedPlayerObject>().is_some()
+            {
                 // if this is a predicted or pre-predicted entity, do not despawn the entity immediately but instead
                 // add a PredictionDespawn component to it to mark that it should be despawned as soon
                 // as the confirmed entity catches up to it
@@ -131,9 +127,11 @@ pub(crate) fn remove_component_for_despawn_predicted<C: SyncComponent>(
 }
 
 // TODO: compare the performance of cloning the component versus popping from the World directly
-/// In case we rollback the despawn, we need to restore the removed components
-/// even if those components are not checking for rollback (SyncComponent != Full)
-/// For those components, we just re-add them from the cache at the start of rollback
+/// In case of a rollback, check if there were any entities that were predicted-despawn
+/// that we need to re-instate. (all the entities that have RemovedCache<C> are in this scenario)
+/// If we didn't need to re-instate them, the Confirmed entity would have been despawned.
+///
+/// Remember to reinstate components if SyncComponent != Full
 pub(crate) fn restore_components_if_despawn_rolled_back<C: SyncComponent>(
     mut commands: Commands,
     mut query: Query<(Entity, &mut RemovedCache<C>), Without<C>>,
@@ -171,12 +169,13 @@ pub(crate) fn remove_despawn_marker(
 mod tests {
     use crate::client::prediction::despawn::PredictionDespawnMarker;
     use crate::client::prediction::resource::PredictionManager;
+    use crate::client::prediction::rollback::RollbackEvent;
     use crate::prelude::client::{Confirmed, PredictionDespawnCommandsExt};
     use crate::prelude::server::SyncTarget;
     use crate::prelude::{client, server, NetworkTarget};
     use crate::tests::protocol::{ComponentSyncModeFull, ComponentSyncModeSimple};
     use crate::tests::stepper::BevyStepper;
-    use bevy::prelude::{default, Component};
+    use bevy::prelude::{default, Component, Trigger};
 
     #[derive(Component, Debug, PartialEq)]
     struct TestComponent(usize);
@@ -242,6 +241,11 @@ mod tests {
         //     .world()
         //     .get::<TestComponent>(predicted_entity)
         //     .is_none());
+        assert!(stepper
+            .client_app
+            .world()
+            .get_entity(predicted_entity)
+            .is_ok());
         assert!(stepper
             .client_app
             .world()

lightyear/src/client/prediction/history.rs

@@ -3,7 +3,7 @@ use bevy::prelude::{Component, Reflect, Resource};
 use bevy::prelude::{ReflectComponent, ReflectResource};
 use std::collections::VecDeque;
 use std::fmt::Debug;
-use tracing::{debug, info};
+use tracing::debug;
 
 /// Stores a past value in the history buffer
 #[derive(Debug, PartialEq, Clone, Default, Reflect)]

lightyear/src/client/prediction/plugin.rs

@@ -295,6 +295,7 @@ pub fn add_prediction_systems<C: SyncComponent>(app: &mut App, prediction_mode:
             // TODO: register type if C is reflect
             // app.register_type::<HistoryState<C>>();
             // app.register_type::<PredictionHistory<C>>();
+
             app.add_observer(apply_component_removal_predicted::<C>);
             app.add_observer(handle_tick_event_prediction_history::<C>);
             app.add_systems(

lightyear/src/client/prediction/predicted_history.rs

@@ -6,7 +6,7 @@ use bevy::prelude::{
     Added, Commands, Component, DetectChanges, Entity, OnRemove, Or, Query, Ref, Res, Trigger,
     With, Without,
 };
-use tracing::{debug, info, trace};
+use tracing::{debug, trace};
 
 use crate::client::components::{ComponentSyncMode, Confirmed, SyncComponent};
 use crate::client::prediction::history::HistoryBuffer;

lightyear/src/client/prediction/prespawn.rs

@@ -486,14 +486,15 @@ mod tests {
     use crate::client::prediction::history::HistoryState;
     use crate::client::prediction::predicted_history::PredictionHistory;
     use crate::client::prediction::resource::PredictionManager;
+    use crate::client::prediction::rollback::RollbackEvent;
     use crate::prelude::client::PredictionDespawnCommandsExt;
     use crate::prelude::client::{Confirmed, Predicted};
     use crate::prelude::server::{Replicate, SyncTarget};
     use crate::prelude::*;
     use crate::tests::protocol::*;
     use crate::tests::stepper::BevyStepper;
     use crate::utils::ready_buffer::ItemWithReadyKey;
-    use bevy::prelude::{default, Entity, With};
+    use bevy::prelude::{default, Entity, Trigger, With};
 
     #[test]
     fn test_compute_hash() {
@@ -734,17 +735,19 @@ mod tests {
 
     /// Client spawns a PreSpawned entity and tries to despawn it locally
     /// before it gets matched to a server entity.
-    /// The entity should just be despawned completely on the client?
-    /// Then when the server entity arrives, it acts as if there's no match and spawns a new
-    /// predicted entity for it
+    /// The entity should be kept around in case of a match, and then cleanup via the cleanup system.
     #[test]
-    fn test_prespawn_local_despawn() {
+    fn test_prespawn_local_despawn_no_match() {
         let mut stepper = BevyStepper::default();
 
         let client_prespawn = stepper
             .client_app
             .world_mut()
-            .spawn(PreSpawnedPlayerObject::new(1))
+            .spawn((
+                PreSpawnedPlayerObject::new(1),
+                ComponentSyncModeFull(1.0),
+                ComponentSyncModeSimple(1.0),
+            ))
             .id();
         stepper.frame_step();
         stepper
@@ -754,10 +757,124 @@ mod tests {
             .entity(client_prespawn)
             .prediction_despawn();
         stepper.frame_step();
+        assert!(stepper
+            .client_app
+            .world()
+            .get_entity(client_prespawn)
+            .is_ok());
+        assert!(stepper
+            .client_app
+            .world()
+            .get::<ComponentSyncModeFull>(client_prespawn)
+            .is_none());
+        assert!(stepper
+            .client_app
+            .world()
+            .get::<ComponentSyncModeSimple>(client_prespawn)
+            .is_none());
+
+        // if enough frames pass without match, the entity gets cleaned
+        stepper.frame_step();
+        stepper.frame_step();
+        stepper.frame_step();
+        stepper.frame_step();
+        stepper.frame_step();
         assert!(stepper
             .client_app
             .world()
             .get_entity(client_prespawn)
             .is_err());
     }
+
+    fn panic_on_rollback(_: Trigger<RollbackEvent>) {
+        panic!("rollback triggered");
+    }
+
+    /// Client spawns a PreSpawned entity and tries to despawn it locally
+    /// before it gets matched to a server entity.
+    /// The match should work normally without causing any rollbacks, since the server components
+    /// on the PreSpawned entity should match the client history when it was spawned.
+    #[test]
+    fn test_prespawn_local_despawn_match() {
+        let mut stepper = BevyStepper::default();
+        stepper.client_app.add_observer(panic_on_rollback);
+
+        let client_tick = stepper.client_tick().0 as usize;
+        let server_tick = stepper.server_tick().0 as usize;
+        let client_prespawn = stepper
+            .client_app
+            .world_mut()
+            .spawn((
+                PreSpawnedPlayerObject::new(1),
+                ComponentSyncModeFull(1.0),
+                ComponentSyncModeSimple(1.0),
+            ))
+            .id();
+
+        stepper.frame_step();
+
+        // do a predicted despawn (we first wait one frame otherwise the components would get removed
+        //  immediately and the prediction-history would be empty)
+        stepper
+            .client_app
+            .world_mut()
+            .commands()
+            .entity(client_prespawn)
+            .prediction_despawn();
+
+        // we want to advance by the tick difference, so that the server prespawned is spawned on the same
+        // tick as the client prespawned
+        // (i.e. entity is spawned on tick client_tick = X on client, and spawned on tick server_tick = X on server, so that
+        // the Histories match)
+        for tick in server_tick + 1..client_tick {
+            stepper.frame_step();
+        }
+        // make sure that the components were removed on the client prespawned
+        assert!(stepper
+            .client_app
+            .world()
+            .get_entity(client_prespawn)
+            .is_ok());
+        assert!(stepper
+            .client_app
+            .world()
+            .get::<ComponentSyncModeFull>(client_prespawn)
+            .is_none());
+        assert!(stepper
+            .client_app
+            .world()
+            .get::<ComponentSyncModeSimple>(client_prespawn)
+            .is_none());
+
+        // spawn the server prespawned entity
+        let server_prespawn = stepper
+            .server_app
+            .world_mut()
+            .spawn((
+                PreSpawnedPlayerObject::new(1),
+                ComponentSyncModeFull(1.0),
+                ComponentSyncModeSimple(1.0),
+                Replicate {
+                    sync: SyncTarget {
+                        prediction: NetworkTarget::All,
+                        ..default()
+                    },
+                    ..default()
+                },
+            ))
+            .id();
+
+        stepper.frame_step();
+        stepper.frame_step();
+
+        // the server entity gets replicated to the client
+        // we should have a match with no rollbacks.
+        // the ComponentSyncMode::Simple components should not be reinstated (they will be only if there is a rollback)
+        stepper.frame_step();
+        let confirmed = stepper
+            .client_app
+            .world()
+            .get::<Predicted>(client_prespawn)
+            .unwrap();
+    }
 }

lightyear/src/client/prediction/rollback.rs

@@ -11,7 +11,7 @@ use bevy::prelude::{
 use bevy::reflect::Reflect;
 use bevy::time::{Fixed, Time};
 use parking_lot::RwLock;
-use tracing::{debug, error, info, trace, trace_span};
+use tracing::{debug, error, trace, trace_span};
 
 use crate::client::components::{Confirmed, SyncComponent};
 use crate::client::config::ClientConfig;
@@ -1014,7 +1014,7 @@ mod integration_tests {
     /// entity-mapped are mapped when rollbacked.
     #[test]
     fn test_rollback_entity_mapping() {
-        #[derive(Component, Serialize, Deserialize, Clone, Copy, PartialEq)]
+        #[derive(Component, Serialize, Deserialize, Clone, Copy, PartialEq, Debug)]
         struct ComponentWithEntity(Entity);
 
         impl MapEntities for ComponentWithEntity {