From 21b7c5f93e6938d48988a8938eed27cf3ea69b96 Mon Sep 17 00:00:00 2001
From: metalgearsloth <31366439+metalgearsloth@users.noreply.github.com>
Date: Sun, 22 Oct 2023 16:53:10 +1100
Subject: [PATCH] Cleanup relays on joint deletion (#4497)

---
 .../Physics/Systems/SharedJointSystem.Relay.cs       | 12 ++++++++----
 Robust.Shared/Physics/Systems/SharedJointSystem.cs   |  4 ++++
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/Robust.Shared/Physics/Systems/SharedJointSystem.Relay.cs b/Robust.Shared/Physics/Systems/SharedJointSystem.Relay.cs
index 28b20ba14..03443b8c9 100644
--- a/Robust.Shared/Physics/Systems/SharedJointSystem.Relay.cs
+++ b/Robust.Shared/Physics/Systems/SharedJointSystem.Relay.cs
@@ -51,7 +51,7 @@ public abstract partial class SharedJointSystem
     {
         foreach (var relay in component.Relayed)
         {
-            if (Deleted(relay) || !_jointsQuery.TryGetComponent(relay, out var joint))
+            if (TerminatingOrDeleted(relay) || !_jointsQuery.TryGetComponent(relay, out var joint))
                 continue;
 
             RefreshRelay(relay, component: joint);
@@ -73,13 +73,13 @@ public abstract partial class SharedJointSystem
             relay = container.Owner;
         }
 
-        RefreshRelay(uid, relay, component);
+        SetRelay(uid, relay, component);
     }
 
     /// <summary>
     /// Refreshes the joint relay for this entity.
     /// </summary>
-    public void RefreshRelay(EntityUid uid, EntityUid? relay, JointComponent? component = null)
+    public void SetRelay(EntityUid uid, EntityUid? relay, JointComponent? component = null)
     {
         if (!Resolve(uid, ref component, false))
             return;
@@ -91,7 +91,11 @@ public abstract partial class SharedJointSystem
         {
             if (relayTarget.Relayed.Remove(uid))
             {
-                // TODO: Comp cleanup.
+                if (relayTarget.Relayed.Count == 0)
+                {
+                    RemCompDeferred<JointRelayTargetComponent>(component.Relay.Value);
+                }
+
                 Dirty(component.Relay.Value, relayTarget);
             }
         }
diff --git a/Robust.Shared/Physics/Systems/SharedJointSystem.cs b/Robust.Shared/Physics/Systems/SharedJointSystem.cs
index 9e7be843f..8e8eb1216 100644
--- a/Robust.Shared/Physics/Systems/SharedJointSystem.cs
+++ b/Robust.Shared/Physics/Systems/SharedJointSystem.cs
@@ -90,6 +90,10 @@ public abstract partial class SharedJointSystem : EntitySystem
         {
             RemoveJoint(joint);
         }
+
+        // If we're relaying elsewhere then cleanup our old data.
+        if (component.Relay != null && !TerminatingOrDeleted(component.Relay.Value))
+            SetRelay(uid, null, component);
     }
 
     #endregion