Power Consumers Rebalance: Simple Dynamic Power Loading (#41961)

* initial commit

* misc additions and fixes

* final tests and additions

* cleanup 1

* fix tests and add a test

* fix tests AGAIN

* abject horror and misery

* cleanup

* cleanup 2

* address some issues

Content.Client/Power/EntitySystems/PowerStateSystem.cs

@@ -0,0 +1,6 @@
+using Content.Shared.Power.EntitySystems;
+
+namespace Content.Client.Power.EntitySystems;
+
+/// <inheritdoc/>
+public sealed class PowerStateSystem : SharedPowerStateSystem;

Content.IntegrationTests/Tests/Power/PowerStatePrototypeTest.cs

@@ -0,0 +1,59 @@
+using System.Linq;
+using Content.Server.Power.Components;
+using Content.Shared.Power.Components;
+using Content.Shared.Power.EntitySystems;
+using Robust.Shared.GameObjects;
+using Robust.Shared.Prototypes;
+
+namespace Content.IntegrationTests.Tests.Power;
+
+[TestFixture, TestOf(typeof(SharedPowerStateSystem))]
+public sealed class PowerStatePrototypeTest
+{
+    /// <summary>
+    /// Asserts that the <see cref="SharedApcPowerReceiverComponent"/>'s load is the same
+    /// as the idle or working power draw from <see cref="PowerStateComponent"/>,
+    /// depending on the current power state.
+    /// </summary>
+    [Test]
+    public async Task AssertApcPowerMatchesPowerState()
+    {
+        await using var pair = await PoolManager.GetServerClient();
+        var server = pair.Server;
+
+        var protoMan = server.ResolveDependency<IPrototypeManager>();
+        var entMan = server.ResolveDependency<IEntityManager>();
+
+        await server.WaitAssertion(() =>
+        {
+            Assert.Multiple(delegate
+            {
+                foreach (var prototype in protoMan.EnumeratePrototypes<EntityPrototype>()
+                             .Where(p => !p.Abstract)
+                             .Where(p => !pair.IsTestPrototype(p)))
+                {
+                    if (!prototype.TryGetComponent<PowerStateComponent>(out var powerStateComp, entMan.ComponentFactory))
+                        continue;
+
+                    // LESSON LEARNED:
+                    // ENSURE THAT THE COMPONENT YOU ARE TRYING TO GET IS THE SERVER-SIDE VARIANT
+                    if (!prototype.TryGetComponent<ApcPowerReceiverComponent>(out var powerReceiverComp, entMan.ComponentFactory))
+                    {
+                        Assert.Fail(
+                            $"Entity prototype '{prototype.ID}' has a PowerStateComponent but is missing the required ApcPowerReceiverComponent.");
+                    }
+
+                    var expectedLoad = powerStateComp.IsWorking
+                        ? powerStateComp.WorkingPowerDraw
+                        : powerStateComp.IdlePowerDraw;
+
+                    Assert.That(powerReceiverComp.Load,
+                        Is.EqualTo(expectedLoad),
+                        $"Entity prototype '{prototype.ID}' has mismatched power draw between PowerStateComponent and SharedApcPowerReceiverComponent.");
+                }
+            });
+        });
+
+        await pair.CleanReturnAsync();
+    }
+}

Content.IntegrationTests/Tests/Power/PowerStateTest.cs

@@ -56,7 +56,7 @@ public sealed class PowerStateTest
                 Assert.That(receiver.Load, Is.EqualTo(powerState.IdlePowerDraw).Within(0.01f));
             });
 
-            var system = entManager.System<PowerStateSystem>();
+            var system = entManager.System<SharedPowerStateSystem>();
             system.SetWorkingState((ent, powerState), true);
 
             Assert.Multiple(() =>
@@ -93,7 +93,7 @@ public sealed class PowerStateTest
 
             var receiver = entManager.GetComponent<Server.Power.Components.ApcPowerReceiverComponent>(ent);
             var powerState = entManager.GetComponent<PowerStateComponent>(ent);
-            var system = entManager.System<PowerStateSystem>();
+            var system = entManager.System<SharedPowerStateSystem>();
             Entity<PowerStateComponent> newEnt = (ent, powerState);
 
             Assert.Multiple(() =>
@@ -146,7 +146,7 @@ public sealed class PowerStateTest
 
             var receiver = entManager.GetComponent<Server.Power.Components.ApcPowerReceiverComponent>(ent);
             var powerState = entManager.GetComponent<PowerStateComponent>(ent);
-            var system = entManager.System<PowerStateSystem>();
+            var system = entManager.System<SharedPowerStateSystem>();
             Entity<PowerStateComponent> valueTuple = (ent, powerState);
 
             Assert.Multiple(() =>

Content.Server/Holopad/HolopadSystem.cs

@@ -23,6 +23,7 @@ using Robust.Shared.Containers;
 using Robust.Shared.Timing;
 using Robust.Shared.Utility;
 using System.Linq;
+using Content.Shared.Power.EntitySystems;
 
 namespace Content.Server.Holopad;
 
@@ -40,6 +41,7 @@ public sealed class HolopadSystem : SharedHolopadSystem
     [Dependency] private readonly PopupSystem _popupSystem = default!;
     [Dependency] private readonly IGameTiming _timing = default!;
     [Dependency] private readonly PvsOverrideSystem _pvs = default!;
+    [Dependency] private readonly SharedPowerStateSystem _powerState = default!;
 
     private float _updateTimer = 1.0f;
     private const float UpdateTime = 1.0f;
@@ -548,10 +550,14 @@ public sealed class HolopadSystem : SharedHolopadSystem
         {
             _telephoneSystem.SetSpeakerForTelephone((entity, entityTelephone), (hologramUid, hologramSpeech));
         }
+
+        _powerState.SetWorkingState(entity.Owner, true);
     }
 
     private void DeleteHologram(Entity<HolopadHologramComponent> hologram, Entity<HolopadComponent> attachedHolopad)
     {
+        _powerState.SetWorkingState(attachedHolopad.Owner, false);
+
         attachedHolopad.Comp.Hologram = null;
 
         QueueDel(hologram);

Content.Server/Kitchen/EntitySystems/MicrowaveSystem.cs

@@ -37,6 +37,7 @@ using Content.Shared.Stacks;
 using Content.Server.Construction.Components;
 using Content.Shared.Chat;
 using Content.Shared.Damage.Components;
+using Content.Shared.Power.EntitySystems;
 using Content.Shared.Temperature.Components;
 
 namespace Content.Server.Kitchen.EntitySystems
@@ -64,6 +65,7 @@ namespace Content.Server.Kitchen.EntitySystems
         [Dependency] private readonly IPrototypeManager _prototype = default!;
         [Dependency] private readonly IAdminLogManager _adminLogger = default!;
         [Dependency] private readonly SharedSuicideSystem _suicide = default!;
+        [Dependency] private readonly SharedPowerStateSystem _powerState = default!;
 
         private static readonly EntProtoId MalfunctionSpark = "Spark";
 
@@ -112,6 +114,7 @@ namespace Content.Server.Kitchen.EntitySystems
 
             microwaveComponent.PlayingStream =
                 _audio.PlayPvs(microwaveComponent.LoopingSound, ent, AudioParams.Default.WithLoop(true).WithMaxDistance(5))?.Entity;
+            _powerState.SetWorkingState(ent.Owner, true);
         }
 
         private void OnCookStop(Entity<ActiveMicrowaveComponent> ent, ref ComponentShutdown args)
@@ -121,6 +124,7 @@ namespace Content.Server.Kitchen.EntitySystems
 
             SetAppearance(ent.Owner, MicrowaveVisualState.Idle, microwaveComponent);
             microwaveComponent.PlayingStream = _audio.Stop(microwaveComponent.PlayingStream);
+            _powerState.SetWorkingState(ent.Owner, false);
         }
 
         private void OnActiveMicrowaveInsert(Entity<ActiveMicrowaveComponent> ent, ref EntInsertedIntoContainerMessage args)

Content.Server/Kitchen/EntitySystems/ReagentGrinderSystem.cs

@@ -22,6 +22,7 @@ using Content.Server.Jittering;
 using Content.Shared.Jittering;
 using Content.Shared.Kitchen.EntitySystems;
 using Content.Shared.Power;
+using Content.Shared.Power.EntitySystems;
 
 namespace Content.Server.Kitchen.EntitySystems
 {
@@ -40,6 +41,7 @@ namespace Content.Server.Kitchen.EntitySystems
         [Dependency] private readonly SharedDestructibleSystem _destructible = default!;
         [Dependency] private readonly RandomHelperSystem _randomHelper = default!;
         [Dependency] private readonly JitteringSystem _jitter = default!;
+        [Dependency] private readonly SharedPowerStateSystem _powerState = default!;
 
         public override void Initialize()
         {
@@ -152,11 +154,15 @@ namespace Content.Server.Kitchen.EntitySystems
         private void OnActiveGrinderStart(Entity<ActiveReagentGrinderComponent> ent, ref ComponentStartup args)
         {
             _jitter.AddJitter(ent, -10, 100);
+
+            // Not all grinders need power.
+            _powerState.TrySetWorkingState(ent.Owner, true);
         }
 
         private void OnActiveGrinderRemove(Entity<ActiveReagentGrinderComponent> ent, ref ComponentRemove args)
         {
             RemComp<JitteringComponent>(ent);
+            _powerState.TrySetWorkingState(ent.Owner, false);
         }
 
         private void OnEntRemoveAttempt(Entity<ReagentGrinderComponent> entity, ref ContainerIsRemovingAttemptEvent args)

Content.Server/Lathe/LatheProducingSystem.cs

@@ -0,0 +1,34 @@
+using Content.Server.Lathe.Components;
+using Content.Shared.Power.EntitySystems;
+
+namespace Content.Server.Lathe;
+
+/// <summary>
+/// System for handling lathes that are actively producing items.
+/// The component is used more so as a marker for EntityQueryEnumerator,
+/// however it's also used to set the power state of the lathe when producing.
+/// </summary>
+public sealed class LatheProducingSystem : EntitySystem
+{
+    [Dependency] private readonly SharedPowerStateSystem _powerState = default!;
+
+    public override void Initialize()
+    {
+        base.Initialize();
+
+        SubscribeLocalEvent<LatheProducingComponent, ComponentStartup>(OnComponentStartup);
+        SubscribeLocalEvent<LatheProducingComponent, ComponentShutdown>(OnComponentShutdown);
+    }
+
+    private void OnComponentShutdown(Entity<LatheProducingComponent> ent, ref ComponentShutdown args)
+    {
+        // use the Try variant of this here
+        // or else you get trolled by AllComponentsOneToOneDeleteTest
+        _powerState.TrySetWorkingState(ent.Owner, false);
+    }
+
+    private void OnComponentStartup(Entity<LatheProducingComponent> ent, ref ComponentStartup args)
+    {
+        _powerState.TrySetWorkingState(ent.Owner, true);
+    }
+}

Content.Server/Power/EntitySystems/PowerStateSystem.cs

@@ -0,0 +1,21 @@
+using Content.Server.Power.Components;
+using Content.Shared.Power.Components;
+using Content.Shared.Power.EntitySystems;
+
+namespace Content.Server.Power.EntitySystems;
+
+public sealed class PowerStateSystem : SharedPowerStateSystem
+{
+    public override void Initialize()
+    {
+        base.Initialize();
+
+        SubscribeLocalEvent<PowerStateComponent, ComponentStartup>(OnComponentStartup);
+    }
+
+    private void OnComponentStartup(Entity<PowerStateComponent> ent, ref ComponentStartup args)
+    {
+        EnsureComp<ApcPowerReceiverComponent>(ent);
+        SetWorkingState(ent.Owner, ent.Comp.IsWorking);
+    }
+}

Content.Shared/Chemistry/EntitySystems/SharedSolutionContainerMixerSystem.cs

@@ -2,6 +2,7 @@ using Content.Shared.Chemistry.Components;
 using Content.Shared.Chemistry.Reaction;
 using Content.Shared.Interaction;
 using Content.Shared.Popups;
+using Content.Shared.Power.EntitySystems;
 using Robust.Shared.Audio.Systems;
 using Robust.Shared.Containers;
 using Robust.Shared.Network;
@@ -21,6 +22,7 @@ public abstract class SharedSolutionContainerMixerSystem : EntitySystem
     [Dependency] private readonly SharedContainerSystem _container = default!;
     [Dependency] private readonly SharedPopupSystem _popup = default!;
     [Dependency] private readonly SharedSolutionContainerSystem _solution = default!;
+    [Dependency] private readonly SharedPowerStateSystem _powerState = default!;
 
     /// <inheritdoc/>
     public override void Initialize()
@@ -74,6 +76,7 @@ public abstract class SharedSolutionContainerMixerSystem : EntitySystem
             comp.MixingSoundEntity = _audio.PlayPvs(comp.MixingSound, entity, comp.MixingSound?.Params.WithLoop(true));
         comp.MixTimeEnd = _timing.CurTime + comp.MixDuration;
         _appearance.SetData(entity, SolutionContainerMixerVisuals.Mixing, true);
+        _powerState.SetWorkingState(entity.Owner, true);
         Dirty(uid, comp);
     }
 
@@ -86,6 +89,7 @@ public abstract class SharedSolutionContainerMixerSystem : EntitySystem
         _appearance.SetData(entity, SolutionContainerMixerVisuals.Mixing, false);
         comp.Mixing = false;
         comp.MixingSoundEntity = null;
+        _powerState.SetWorkingState(entity.Owner, false);
         Dirty(uid, comp);
     }
 

Content.Shared/Power/EntitySystems/PowerStateSystem.cs

@@ -7,7 +7,7 @@ namespace Content.Shared.Power.EntitySystems;
 /// Generic system that handles entities with <see cref="PowerStateComponent"/>.
 /// Used for simple machines that only need to switch between "idle" and "working" power states.
 /// </summary>
-public sealed class PowerStateSystem : EntitySystem
+public abstract class SharedPowerStateSystem : EntitySystem
 {
     [Dependency] private readonly SharedPowerReceiverSystem _powerReceiverSystem = default!;
 
@@ -17,16 +17,9 @@ public sealed class PowerStateSystem : EntitySystem
     {
         base.Initialize();
 
-        SubscribeLocalEvent<PowerStateComponent, ComponentStartup>(OnComponentStartup);
-
         _powerStateQuery = GetEntityQuery<PowerStateComponent>();
     }
 
-    private void OnComponentStartup(Entity<PowerStateComponent> ent, ref ComponentStartup args)
-    {
-        SetWorkingState(ent.Owner, ent.Comp.IsWorking);
-    }
-
     /// <summary>
     /// Sets the working state of the entity, adjusting its power draw accordingly.
     /// </summary>
@@ -41,4 +34,22 @@ public sealed class PowerStateSystem : EntitySystem
         _powerReceiverSystem.SetLoad(ent.Owner, working ? ent.Comp.WorkingPowerDraw : ent.Comp.IdlePowerDraw);
         ent.Comp.IsWorking = working;
     }
+
+    /// <summary>
+    /// Tries to set the working state of the entity, adjusting its power draw accordingly.
+    /// Use this for if you're not sure if the entity has a <see cref="PowerStateComponent"/>.
+    /// </summary>
+    /// <param name="ent">The entity to set the working state for.</param>
+    /// <param name="working">Whether the entity should be in the working state.</param>
+    [PublicAPI]
+    public void TrySetWorkingState(Entity<PowerStateComponent?> ent, bool working)
+    {
+        // Sometimes systems calling this API handle generic objects that can or can't consume power,
+        // so to reduce boilerplate we don't log an error. Any entity that *should* have an ApcPowerRecieverComponent
+        // will log an error in tests if someone tries to add an entity that doesn't have one.
+        if (!_powerStateQuery.Resolve(ent, ref ent.Comp, false))
+            return;
+
+        SetWorkingState(ent, working);
+    }
 }

Content.Shared/Power/EntitySystems/UIPowerStateSystem.cs

@@ -10,7 +10,7 @@ namespace Content.Shared.Power.EntitySystems;
 public sealed class UIPowerStateSystem : EntitySystem
 {
     [Dependency] private readonly SharedUserInterfaceSystem _ui = default!;
-    [Dependency] private readonly PowerStateSystem _powerState = default!;
+    [Dependency] private readonly SharedPowerStateSystem _powerState = default!;
 
     public override void Initialize()
     {

Resources/Maps/Misc/terminal.yml

@@ -1611,6 +1611,7 @@ entities:
     - type: Godmode
     missingComponents:
     - ApcPowerReceiver
+    - PowerState
     - Anchorable
     - Construction
     - Destructible
@@ -1623,6 +1624,7 @@ entities:
     - type: Godmode
     missingComponents:
     - ApcPowerReceiver
+    - PowerState
     - Anchorable
     - Construction
     - Destructible
@@ -7949,6 +7951,7 @@ entities:
     - type: Godmode
     missingComponents:
     - ApcPowerReceiver
+    - PowerState
     - Anchorable
     - Construction
     - Destructible
@@ -7961,6 +7964,7 @@ entities:
     - type: Godmode
     missingComponents:
     - ApcPowerReceiver
+    - PowerState
     - Anchorable
     - Construction
     - Destructible

Resources/Prototypes/Entities/Structures/Machines/Computers/arcades.yml

@@ -5,8 +5,6 @@
   name: arcade
   parent: BaseComputer
   components:
-  - type: ApcPowerReceiver
-    powerLoad: 350
   - type: ExtensionCableReceiver
   - type: PointLight
     radius: 1.8

Resources/Prototypes/Entities/Structures/Machines/Computers/base_structurecomputers.yml

@@ -17,7 +17,11 @@
       - board
   - type: Computer
   - type: ApcPowerReceiver
-    powerLoad: 200
+    powerLoad: 50
+  - type: PowerState
+    idlePowerDraw: 50
+    workingPowerDraw: 500
+  - type: UIPowerState
   - type: ExtensionCableReceiver
   - type: ActivatableUIRequiresPower
   - type: Sprite

Resources/Prototypes/Entities/Structures/Machines/Computers/computers.yml

@@ -528,7 +528,10 @@
       enum.WiresUiKey.Key:
         type: WiresBoundUserInterface
   - type: ApcPowerReceiver
-    powerLoad: 1000
+    powerLoad: 50
+  - type: PowerState
+    idlePowerDraw: 50
+    workingPowerDraw: 1000
   - type: Computer
     board: ResearchComputerCircuitboard
   - type: AccessReader
@@ -579,7 +582,10 @@
       enum.WiresUiKey.Key:
         type: WiresBoundUserInterface
   - type: ApcPowerReceiver
-    powerLoad: 1000
+    powerLoad: 50
+  - type: PowerState
+    idlePowerDraw: 50
+    workingPowerDraw: 1000
   - type: Computer
     board: AnalysisComputerCircuitboard
   - type: PointLight
@@ -662,8 +668,6 @@
   name: body scanner computer
   description: A body scanner.
   components:
-  - type: ApcPowerReceiver
-    powerLoad: 500
   - type: Computer
     board: BodyScannerComputerCircuitboard
   - type: PointLight
@@ -1207,8 +1211,6 @@
       state: generic_keys
     - map: [ "enum.WiresVisualLayers.MaintenancePanel" ]
       state: generic_panel_open
-  - type: ApcPowerReceiver
-    powerLoad: 3100 #We want this to fail first so I transferred most of the scanner and pod's power here. (3500 in total)
   - type: Computer
     board: CloningConsoleComputerCircuitboard
   - type: PointLight
@@ -1578,7 +1580,11 @@
       enum.WiresUiKey.Key:
         type: WiresBoundUserInterface
   - type: ApcPowerReceiver
-    powerLoad: 1000
+    powerLoad: 5
+  - type: PowerState
+    idlePowerDraw: 5
+    workingPowerDraw: 1000
+  - type: UIPowerState
   - type: DeviceNetwork
     deviceNetId: Wireless
     receiveFrequencyId: RoboticsConsole
@@ -1623,7 +1629,10 @@
     channels:
     - Xenoborg
   - type: ApcPowerReceiver
-    powerLoad: 1000
+    powerLoad: 50
+  - type: PowerState
+    idlePowerDraw: 50
+    workingPowerDraw: 1000
   - type: DeviceNetwork
     deviceNetId: Wireless
     receiveFrequencyId: Mothership
@@ -1650,7 +1659,10 @@
     - map: [ "enum.WiresVisualLayers.MaintenancePanel" ]
       state: generic_panel_open
   - type: ApcPowerReceiver
-    powerLoad: 1000
+    powerLoad: 50
+  - type: PowerState
+    idlePowerDraw: 50
+    workingPowerDraw: 1000
   - type: Computer
     board: StationAiUploadCircuitboard
   - type: AccessReader
@@ -1726,7 +1738,10 @@
           True: { visible: false }
           False: { visible: true }
   - type: ApcPowerReceiver
-    powerLoad: 1000
+    powerLoad: 50
+  - type: PowerState
+    idlePowerDraw: 50
+    workingPowerDraw: 1000
   - type: Computer
     board: StationAiFixerCircuitboard
   - type: AccessReader

Resources/Prototypes/Entities/Structures/Machines/Medical/chemistry_machines.yml

@@ -76,6 +76,11 @@
           False: { visible: False }
   - type: Machine
     board: ElectrolysisUnitMachineCircuitboard
+  - type: ApcPowerReceiver
+    powerLoad: 0
+  - type: PowerState
+    idlePowerDraw: 0
+    workingPowerDraw: 1000 # for a lab-grade machine
 
 # TODO centrifuge should spill the vial if the lid is off
 - type: entity
@@ -124,3 +129,8 @@
             - CentrifugeCompatible
   - type: Machine
     board: CentrifugeMachineCircuitboard
+  - type: ApcPowerReceiver
+    powerLoad: 0
+  - type: PowerState
+    idlePowerDraw: 0
+    workingPowerDraw: 500

Resources/Prototypes/Entities/Structures/Machines/holopad.yml

@@ -15,7 +15,10 @@
         mask:
         - Impassable
   - type: ApcPowerReceiver
-    powerLoad: 300
+    powerLoad: 5
+  - type: PowerState
+    idlePowerDraw: 5
+    workingPowerDraw: 300
   - type: StationAiVision
     range: 1
     needsAnchoring: true

Resources/Prototypes/Entities/Structures/Machines/lathe.yml

@@ -56,6 +56,11 @@
       input: AutomationSlotMaterials
       output: null
   - type: TechnologyDatabase
+  - type: ApcPowerReceiver
+    powerLoad: 150
+  - type: PowerState
+    idlePowerDraw: 150
+    workingPowerDraw: 1000
     supportedDisciplines: # DeltaV - don't add it to every map
     - Industrial
     - Biochemical

Resources/Prototypes/Entities/Structures/Machines/microwave.yml

@@ -98,7 +98,10 @@
     canCreateVacuum: false
     deleteAfterExplosion: false
   - type: ApcPowerReceiver
-    powerLoad: 400
+    powerLoad: 5
+  - type: PowerState
+    idlePowerDraw: 5
+    workingPowerDraw: 1000
   - type: Machine
     board: MicrowaveMachineCircuitboard
   - type: ContainerContainer

Resources/Prototypes/Entities/Structures/Machines/reagent_grinder.yml

@@ -39,6 +39,10 @@
     - map: [ "grinder" ]
       state: "grinder_empty"
   - type: ApcPowerReceiver
+    powerLoad: 0
+  - type: PowerState
+    idlePowerDraw: 0
+    workingPowerDraw: 750 # medium power blender
     powerLoad: 300
   - type: AutomationSlots # Goobstation
     slots:

Resources/Prototypes/Entities/Structures/Wallmounts/WallmountMachines/station_map.yml

@@ -52,7 +52,11 @@
       containers:
         board: !type:Container
     - type: ApcPowerReceiver
-      powerLoad: 200
+      powerLoad: 50
+    - type: PowerState
+      idlePowerDraw: 50
+      workingPowerDraw: 200
+    - type: UIPowerState
     - type: Construction
       graph: StationMap
       node: station_map