DoAfter Refactor (#13225)

Co-authored-by: DrSmugleaf <drsmugleaf@gmail.com>

Content.Server/Mech/Systems/MechSystem.cs

@@ -1,19 +1,17 @@
 using System.Linq;
-using System.Threading;
 using Content.Server.Atmos.EntitySystems;
 using Content.Server.DoAfter;
 using Content.Server.Mech.Components;
 using Content.Server.Power.Components;
-using Content.Server.Tools;
 using Content.Server.Wires;
 using Content.Shared.ActionBlocker;
 using Content.Shared.Damage;
+using Content.Shared.DoAfter;
 using Content.Shared.FixedPoint;
 using Content.Shared.Interaction;
 using Content.Shared.Mech;
 using Content.Shared.Mech.Components;
 using Content.Shared.Mech.EntitySystems;
-using Content.Shared.Movement.Components;
 using Content.Shared.Movement.Events;
 using Content.Shared.Tools.Components;
 using Content.Shared.Verbs;
@@ -47,12 +45,9 @@ public sealed class MechSystem : SharedMechSystem
         SubscribeLocalEvent<MechComponent, MapInitEvent>(OnMapInit);
         SubscribeLocalEvent<MechComponent, GetVerbsEvent<AlternativeVerb>>(OnAlternativeVerb);
         SubscribeLocalEvent<MechComponent, MechOpenUiEvent>(OnOpenUi);
-        SubscribeLocalEvent<MechComponent, MechEntryFinishedEvent>(OnEntryFinished);
-        SubscribeLocalEvent<MechComponent, MechEntryCanclledEvent>(OnEntryExitCancelled);
-        SubscribeLocalEvent<MechComponent, MechExitFinishedEvent>(OnExitFinished);
-        SubscribeLocalEvent<MechComponent, MechExitCanclledEvent>(OnEntryExitCancelled);
-        SubscribeLocalEvent<MechComponent, MechRemoveBatteryFinishedEvent>(OnRemoveBatteryFinished);
-        SubscribeLocalEvent<MechComponent, MechRemoveBatteryCancelledEvent>(OnRemoveBatteryCancelled);
+        SubscribeLocalEvent<MechComponent, DoAfterEvent<RemoveBatteryEvent>>(OnRemoveBattery);
+        SubscribeLocalEvent<MechComponent, DoAfterEvent<MechEntryEvent>>(OnMechEntry);
+        SubscribeLocalEvent<MechComponent, DoAfterEvent<MechExitEvent>>(OnMechExit);
 
         SubscribeLocalEvent<MechComponent, DamageChangedEvent>(OnDamageChanged);
         SubscribeLocalEvent<MechComponent, MechEquipmentRemoveMessage>(OnRemoveEquipmentMessage);
@@ -87,32 +82,29 @@ public sealed class MechSystem : SharedMechSystem
             return;
         }
 
-        if (component.EntryTokenSource == null &&
-            TryComp<ToolComponent>(args.Used, out var tool) &&
-            tool.Qualities.Contains("Prying") &&
-            component.BatterySlot.ContainedEntity != null)
+        if (TryComp<ToolComponent>(args.Used, out var tool) && tool.Qualities.Contains("Prying") && component.BatterySlot.ContainedEntity != null)
         {
-            component.EntryTokenSource = new();
-            _doAfter.DoAfter(new DoAfterEventArgs(args.User, component.BatteryRemovalDelay, component.EntryTokenSource.Token, uid, args.Target)
+            var removeBattery = new RemoveBatteryEvent();
+
+            var doAfterEventArgs = new DoAfterEventArgs(args.User, component.BatteryRemovalDelay, target: uid, used: args.Target)
             {
                 BreakOnTargetMove = true,
-                BreakOnUserMove = true,
-                TargetFinishedEvent = new MechRemoveBatteryFinishedEvent(),
-                TargetCancelledEvent = new MechRemoveBatteryCancelledEvent()
-            });
+                BreakOnUserMove = true
+            };
+
+            _doAfter.DoAfter(doAfterEventArgs, removeBattery);
         }
     }
 
-    private void OnRemoveBatteryFinished(EntityUid uid, MechComponent component, MechRemoveBatteryFinishedEvent args)
+    private void OnRemoveBattery(EntityUid uid, MechComponent component, DoAfterEvent<RemoveBatteryEvent> args)
     {
-        component.EntryTokenSource = null;
+        if (args.Cancelled || args.Handled)
+            return;
+
         RemoveBattery(uid, component);
         _actionBlocker.UpdateCanMove(uid);
-    }
 
-    private void OnRemoveBatteryCancelled(EntityUid uid, MechComponent component, MechRemoveBatteryCancelledEvent args)
-    {
-        component.EntryTokenSource = null;
+        args.Handled = true;
     }
 
     private void OnMapInit(EntityUid uid, MechComponent component, MapInitEvent args)
@@ -171,16 +163,14 @@ public sealed class MechSystem : SharedMechSystem
                 Text = Loc.GetString("mech-verb-enter"),
                 Act = () =>
                 {
-                    if (component.EntryTokenSource != null)
-                        return;
-                    component.EntryTokenSource = new CancellationTokenSource();
-                    _doAfter.DoAfter(new DoAfterEventArgs(args.User, component.EntryDelay, component.EntryTokenSource.Token, uid)
+                    var mechEntryEvent = new MechEntryEvent();
+                    var doAfterEventArgs = new DoAfterEventArgs(args.User, component.EntryDelay, target: uid)
                     {
                         BreakOnUserMove = true,
                         BreakOnStun = true,
-                        TargetFinishedEvent = new MechEntryFinishedEvent(args.User),
-                        TargetCancelledEvent = new MechEntryCanclledEvent()
-                    });
+                    };
+
+                    _doAfter.DoAfter(doAfterEventArgs, mechEntryEvent);
                 }
             };
             var openUiVerb = new AlternativeVerb //can't hijack someone else's mech
@@ -199,45 +189,46 @@ public sealed class MechSystem : SharedMechSystem
                 Priority = 1, // Promote to top to make ejecting the ALT-click action
                 Act = () =>
                 {
-                    if (component.EntryTokenSource != null)
-                        return;
                     if (args.User == component.PilotSlot.ContainedEntity)
                     {
                         TryEject(uid, component);
                         return;
                     }
 
-                    component.EntryTokenSource = new CancellationTokenSource();
-                    _doAfter.DoAfter(new DoAfterEventArgs(args.User, component.ExitDelay, component.EntryTokenSource.Token, uid)
+                    var mechExitEvent = new MechExitEvent();
+                    var doAfterEventArgs = new DoAfterEventArgs(args.User, component.ExitDelay, target: uid)
                     {
                         BreakOnUserMove = true,
                         BreakOnTargetMove = true,
-                        BreakOnStun = true,
-                        TargetFinishedEvent = new MechExitFinishedEvent(),
-                        TargetCancelledEvent = new MechExitCanclledEvent()
-                    });
+                        BreakOnStun = true
+                    };
+
+                    _doAfter.DoAfter(doAfterEventArgs, mechExitEvent);
                 }
             };
             args.Verbs.Add(ejectVerb);
         }
     }
 
-    private void OnEntryFinished(EntityUid uid, MechComponent component, MechEntryFinishedEvent args)
+    private void OnMechEntry(EntityUid uid, MechComponent component, DoAfterEvent<MechEntryEvent> args)
     {
-        component.EntryTokenSource = null;
-        TryInsert(uid, args.User, component);
+        if (args.Cancelled || args.Handled)
+            return;
+
+        TryInsert(uid, args.Args.User, component);
         _actionBlocker.UpdateCanMove(uid);
+
+        args.Handled = true;
     }
 
-    private void OnExitFinished(EntityUid uid, MechComponent component, MechExitFinishedEvent args)
+    private void OnMechExit(EntityUid uid, MechComponent component, DoAfterEvent<MechExitEvent> args)
     {
-        component.EntryTokenSource = null;
+        if (args.Cancelled || args.Handled)
+            return;
+
         TryEject(uid, component);
-    }
 
-    private void OnEntryExitCancelled(EntityUid uid, MechComponent component, EntityEventArgs args)
-    {
-        component.EntryTokenSource = null;
+        args.Handled = true;
     }
 
     private void OnDamageChanged(EntityUid uid, SharedMechComponent component, DamageChangedEvent args)
@@ -454,4 +445,27 @@ public sealed class MechSystem : SharedMechSystem
         args.Handled = true;
     }
     #endregion
+
+    /// <summary>
+    ///     Event raised when the battery is successfully removed from the mech,
+    ///     on both success and failure
+    /// </summary>
+    private sealed class RemoveBatteryEvent : EntityEventArgs
+    {
+    }
+
+    /// <summary>
+    ///     Event raised when a person enters a mech, on both success and failure
+    /// </summary>
+    private sealed class MechEntryEvent : EntityEventArgs
+    {
+    }
+
+    /// <summary>
+    ///     Event raised when a person removes someone from a mech,
+    ///     on both success and failure
+    /// </summary>
+    private sealed class MechExitEvent : EntityEventArgs
+    {
+    }
 }