Only re-run dirty components on PVS re-entry

Robust.Client/GameStates/ClientGameStateManager.cs

@@ -51,6 +51,7 @@ public sealed partial class ClientGameStateManager : IClientGameStateManager
         private StateData[] _toApplySorted = default!;
         private readonly Dictionary<ushort, (IComponent Component, IComponentState? curState, IComponentState? nextState)> _compStateWork = new();
         private readonly Dictionary<EntityUid, HashSet<Type>> _pendingReapplyNetStates = new();
+        private readonly Dictionary<NetEntity, HashSet<ushort>> _detachedDirtyComponents = new();
         private readonly HashSet<NetEntity> _stateEnts = new();
         private readonly List<EntityUid> _toDelete = new();
         private readonly List<IComponent> _toRemove = new();
@@ -76,6 +77,7 @@ private readonly record struct StateData(
 
         private uint _metaCompNetId;
         private uint _xformCompNetId;
+        private uint _containerCompNetId;
 
         [Dependency] private IReplayRecordingManager _replayRecording = default!;
         [Dependency] private IComponentFactory _compFactory = default!;
@@ -191,6 +193,12 @@ public void Initialize()
                 throw new InvalidOperationException("TransformComponent does not have a NetId.");
 
             _xformCompNetId = xformId.Value;
+
+            var containerId = _compFactory.GetRegistration(typeof(ContainerManagerComponent)).NetID;
+            if (!containerId.HasValue)
+                throw new InvalidOperationException("ContainerManagerComponent does not have a NetId.");
+
+            _containerCompNetId = containerId.Value;
         }
 
         private void OnComponentAdded(AddedComponentEventArgs args)
@@ -216,6 +224,7 @@ private void OnComponentAdded(AddedComponentEventArgs args)
         public void Reset()
         {
             _processor.Reset();
+            _detachedDirtyComponents.Clear();
             _timing.CurTick = GameTick.Zero;
             _timing.LastRealTick = GameTick.Zero;
             _lastProcessedInput = 0;
@@ -641,7 +650,13 @@ public void ResetPredictedEntities()
                         if (_sawmill.Level <= LogLevel.Debug)
                             _sawmill.Debug($"  A component was dirtied: {comp.GetType()}");
 
-                        if ((meta.Flags & MetaDataFlags.Detached) == 0 && compState != null)
+                        // If we dirty a component outside of PVS range (god forbid) then we need to ensure
+                        // it gets states re-run for these components if it comes back into PVS range.
+                        if ((meta.Flags & MetaDataFlags.Detached) != 0)
+                        {
+                            _detachedDirtyComponents.GetOrNew(meta.NetEntity).Add(netId);
+                        }
+                        else if (compState != null)
                         {
                             var handleState = new ComponentHandleState(compState, null);
                             _entities.EventBus.RaiseComponentEvent(entity, comp, ref handleState);
@@ -1146,6 +1161,7 @@ public void PartialStateReset(
             }
 
             _sawmill.Info($"Resetting all entity states to tick {state.ToSequence}.");
+            _detachedDirtyComponents.Clear();
 
             // Construct hashset for set.Contains() checks.
             _stateEnts.Clear();
@@ -1230,6 +1246,7 @@ private void ProcessDeletions(ReadOnlySpan<NetEntity> delSpan,
             {
                 // Don't worry about this for later.
                 _entities.PendingNetEntityStates.Remove(netEntity);
+                _detachedDirtyComponents.Remove(netEntity);
 
                 if (!_entities.TryGetEntity(netEntity, out var id))
                     continue;
@@ -1386,20 +1403,46 @@ private void HandleEntityState(in StateData data, IEventBus bus, GameTick toTick
             if (data.EnteringPvs)
             {
                 // last-server state has already been updated with new information from curState
-                // --> simply reset to the most recent server state.
+                // --> simply reset required structural components and any components dirtied during prediction while
+                // detached to the most recent server state.
                 //
                 // as to why we need to reset: because in the process of detaching to null-space, we will have dirtied
                 // the entity. most notably, all entities will have been ejected from their containers.
-                foreach (var (id, state) in _processor.GetLastServerStates(data.NetEntity))
+                var lastState = _processor.GetLastServerStates(data.NetEntity);
+                var netComponents = data.Meta.NetComponents;
+                var uid = data.Uid;
+                var meta = data.Meta;
+
+                void AddCompState(ushort id)
                 {
-                    if (!data.Meta.NetComponents.TryGetValue(id, out var comp))
+                    if (!lastState.TryGetValue(id, out var state))
+                        return;
+
+                    if (!netComponents.TryGetValue(id, out var comp))
                     {
                         comp = _compFactory.GetComponent(id);
-                        _entities.AddComponent(data.Uid, comp, true, metadata: data.Meta);
+                        _entities.AddComponent(uid, comp, true, metadata: meta);
                     }
 
                     _compStateWork[id] = (comp, state, null);
                 }
+
+                AddCompState((ushort) _metaCompNetId);
+                AddCompState((ushort) _xformCompNetId);
+                AddCompState((ushort) _containerCompNetId);
+
+                if (_detachedDirtyComponents.Remove(data.NetEntity, out var dirtyComponents))
+                {
+                    foreach (var netId in dirtyComponents)
+                    {
+                        AddCompState(netId);
+                    }
+                }
+
+                foreach (var compChange in data.CurState!.ComponentChanges.Span)
+                {
+                    AddCompState(compChange.NetID);
+                }
             }
             else if (data.CurState != null)
             {

Robust.Server.IntegrationTests/GameStates/PvsReEntryTest.cs

@@ -7,6 +7,7 @@
 using Robust.Shared;
 using Robust.Shared.Configuration;
 using Robust.Shared.GameObjects;
+using Robust.Shared.GameStates;
 using Robust.Shared.Map;
 using Robust.Shared.Player;
 using Robust.Shared.Timing;
@@ -167,5 +168,148 @@ await client.WaitPost(() =>
 
     }
 #endif
+
+    /// <summary>
+    /// Checks that PVS re-entry only reapplies component states for components dirtied while the entity was out of PVS.
+    /// </summary>
+    [Test]
+    public async Task TestReEntryOnlyReplaysDirtyComponents()
+    {
+        await using var pair = await StartConnectedPair(
+            new ServerIntegrationOptions { Pool = false },
+            new ClientIntegrationOptions { Pool = false });
+
+        var (client, server) = pair;
+
+        var sEntMan = server.ResolveDependency<IEntityManager>();
+        var confMan = server.ResolveDependency<IConfigurationManager>();
+        var sPlayerMan = server.ResolveDependency<ISharedPlayerManager>();
+        var xforms = sEntMan.System<SharedTransformSystem>();
+
+        var cEntMan = client.ResolveDependency<IEntityManager>();
+        PvsReEntryReplayTestSystem? cSystem = null;
+
+        server.Post(() => confMan.SetCVar(CVars.NetPVS, true));
+
+        await RunTicksSync(server, client, 10);
+
+        EntityUid map = default;
+        NetEntity entity = default;
+        NetEntity dirtyEntity = default;
+        NetEntity player = default;
+        EntityCoordinates coords = default;
+
+        await server.WaitPost(() =>
+        {
+            map = server.System<SharedMapSystem>().CreateMap();
+            coords = new EntityCoordinates(map, default);
+
+            var playerUid = sEntMan.SpawnEntity(null, coords);
+            var entUid = sEntMan.SpawnEntity(null, coords);
+            var comp = sEntMan.EnsureComponent<PvsReEntryReplayTestComponent>(entUid);
+            comp.Value = 1;
+            sEntMan.Dirty(entUid, comp);
+
+            var dirtyEntUid = sEntMan.SpawnEntity(null, coords);
+            comp = sEntMan.EnsureComponent<PvsReEntryReplayTestComponent>(dirtyEntUid);
+            comp.Value = 1;
+            sEntMan.Dirty(dirtyEntUid, comp);
+
+            entity = sEntMan.GetNetEntity(entUid);
+            dirtyEntity = sEntMan.GetNetEntity(dirtyEntUid);
+            player = sEntMan.GetNetEntity(playerUid);
+
+            var session = sPlayerMan.Sessions.First();
+            server.PlayerMan.SetAttachedEntity(session, playerUid);
+            sPlayerMan.JoinGame(session);
+        });
+
+        await RunTicksSync(server, client, 10);
+
+        MetaDataComponent? meta = null;
+        MetaDataComponent? dirtyMeta = null;
+        var initialHandleCount = 0;
+        await client.WaitPost(() =>
+        {
+            cSystem = client.System<PvsReEntryReplayTestSystem>();
+            Assert.That(cEntMan.TryGetEntityData(entity, out _, out meta), Is.True);
+            Assert.That(cEntMan.TryGetEntityData(dirtyEntity, out _, out dirtyMeta), Is.True);
+            Assert.That(meta!.Flags & MetaDataFlags.Detached, Is.EqualTo(MetaDataFlags.None));
+            Assert.That(dirtyMeta!.Flags & MetaDataFlags.Detached, Is.EqualTo(MetaDataFlags.None));
+            initialHandleCount = cSystem.HandleStateCount;
+            Assert.That(initialHandleCount, Is.GreaterThanOrEqualTo(2));
+        });
+
+        var farAway = new EntityCoordinates(map, new Vector2(100, 100));
+        await server.WaitPost(() => xforms.SetCoordinates(sEntMan.GetEntity(player), farAway));
+        await RunTicksSync(server, client, 10);
+
+        await client.WaitPost(() =>
+        {
+            Assert.That(meta!.Flags & MetaDataFlags.Detached, Is.EqualTo(MetaDataFlags.Detached));
+            Assert.That(dirtyMeta!.Flags & MetaDataFlags.Detached, Is.EqualTo(MetaDataFlags.Detached));
+            Assert.That(cSystem!.HandleStateCount, Is.EqualTo(initialHandleCount));
+        });
+
+        await server.WaitPost(() =>
+        {
+            var comp = sEntMan.GetComponent<PvsReEntryReplayTestComponent>(sEntMan.GetEntity(dirtyEntity));
+            comp.Value = 2;
+            sEntMan.Dirty(sEntMan.GetEntity(dirtyEntity), comp);
+        });
+
+        await RunTicksSync(server, client, 5);
+
+        await client.WaitPost(() =>
+        {
+            Assert.That(cSystem!.HandleStateCount, Is.EqualTo(initialHandleCount));
+        });
+
+        // Re-enter after dirtying one of the test components. Only the dirtied component should get replayed.
+        await server.WaitPost(() => xforms.SetCoordinates(sEntMan.GetEntity(player), coords));
+        await RunTicksSync(server, client, 10);
+
+        await client.WaitPost(() =>
+        {
+            var uid = cEntMan.GetEntity(entity);
+            var dirtyUid = cEntMan.GetEntity(dirtyEntity);
+            Assert.That(meta!.Flags & MetaDataFlags.Detached, Is.EqualTo(MetaDataFlags.None));
+            Assert.That(dirtyMeta!.Flags & MetaDataFlags.Detached, Is.EqualTo(MetaDataFlags.None));
+            Assert.That(cSystem!.HandleStateCount, Is.EqualTo(initialHandleCount + 1));
+            Assert.That(cEntMan.GetComponent<PvsReEntryReplayTestComponent>(uid).Value, Is.EqualTo(1));
+            Assert.That(cEntMan.GetComponent<PvsReEntryReplayTestComponent>(dirtyUid).Value, Is.EqualTo(2));
+        });
+    }
+}
+
+[RegisterComponent, NetworkedComponent]
+public sealed partial class PvsReEntryReplayTestComponent : Component
+{
+    public int Value;
+}
+
+public sealed partial class PvsReEntryReplayTestSystem : EntitySystem
+{
+    public int HandleStateCount;
+
+    public override void Initialize()
+    {
+        SubscribeLocalEvent<PvsReEntryReplayTestComponent, ComponentGetState>(OnGetState);
+        SubscribeLocalEvent<PvsReEntryReplayTestComponent, ComponentHandleState>(OnHandleState);
+    }
+
+    private void OnGetState(Entity<PvsReEntryReplayTestComponent> ent, ref ComponentGetState args)
+    {
+        args.State = new MetaDataComponentState(null, ent.Comp.Value.ToString(), null, null);
+    }
+
+    private void OnHandleState(Entity<PvsReEntryReplayTestComponent> ent, ref ComponentHandleState args)
+    {
+        if (args.Current is not MetaDataComponentState state)
+            return;
+
+        ent.Comp.Value = int.Parse(state.Description!);
+        HandleStateCount++;
+    }
 }