Cascade power draw calculation on all downstream power ports

netbox/dcim/models/__init__.py

@@ -708,19 +708,17 @@ class Rack(ChangeLoggedModel, CustomFieldModel):
         """
         Determine the utilization rate of power in the rack and return it as a percentage.
         """
-        power_stats = PowerFeed.objects.filter(
-            rack=self
-        ).annotate(
-            allocated_draw_total=Sum('connected_endpoint__poweroutlets__connected_endpoint__allocated_draw'),
-        ).values(
-            'allocated_draw_total',
-            'available_power'
-        )
+        # Sum up all of the available power from the power feeds assigned to the rack
+        available_power = PowerFeed.objects.filter(rack=self).aggregate(total=Sum('available_power'))
 
-        if power_stats:
-            allocated_draw_total = sum(x['allocated_draw_total'] for x in power_stats)
-            available_power_total = sum(x['available_power'] for x in power_stats)
-            return int(allocated_draw_total / available_power_total * 100) or 0
+        # Get the power draw of the power ports from the power feeds assigned to the rack
+        power_ports = PowerPort.objects.filter(_connected_powerfeed__rack=self)
+        feeds_stats = [x.get_power_draw(leg_stats=False) for x in power_ports]
+
+        if available_power.get('total') and feeds_stats:
+            available_power_total = available_power.get('total')
+            allocated_draw_total = sum([x.get('allocated') or 0 for x in feeds_stats])
+            return round(allocated_draw_total / available_power_total * 100)
         return 0
 
 

netbox/dcim/models/device_components.py

@@ -1,3 +1,4 @@
+from cacheops import cached_as
 from django.contrib.contenttypes.fields import GenericRelation
 from django.core.exceptions import ObjectDoesNotExist, ValidationError
 from django.core.validators import MaxValueValidator, MinValueValidator
@@ -383,37 +384,55 @@ class PowerPort(CableTermination, ComponentModel):
                 "Connected endpoint must be a PowerOutlet or PowerFeed, not {}.".format(type(value))
             )
 
-    def get_power_draw(self):
+    def get_power_draw(self, leg_stats=True):
         """
         Return the allocated and maximum power draw (in VA) and child PowerOutlet count for this PowerPort.
+        If `leg_stats` is True, the `legs` key in the returned dict is populated with a list of per-leg statistics.
         """
+        def get_power_stats(leg=None):
+            """
+            Return tuple of (outlet_count, allocated_draw_total, maximum_draw_total).
+            """
+            if leg:
+                outlets = PowerOutlet.objects.filter(power_port=self, feed_leg=leg)
+            else:
+                outlets = PowerOutlet.objects.filter(power_port=self)
+
+            # The outlets are used as extra to invalidate the cache when an outlet's leg is changed
+            @cached_as(self, extra=outlets)
+            def _stats():
+                # Power ports drawing power from the local outlets
+                return PowerPort.objects.filter(
+                    pk__in=outlets.values_list('downstream_powerports', flat=True),
+                ).aggregate(
+                    Sum('allocated_draw'),
+                    Sum('maximum_draw'),
+                )
+
+            stats = _stats()
+
+            return outlets.count(), stats.get('allocated_draw__sum') or 0, stats.get('maximum_draw__sum') or 0
+
         # Calculate aggregate draw of all child power outlets if no numbers have been defined manually
         if self.allocated_draw is None and self.maximum_draw is None:
-            outlet_ids = PowerOutlet.objects.filter(power_port=self).values_list('pk', flat=True)
-            utilization = PowerPort.objects.filter(_connected_poweroutlet_id__in=outlet_ids).aggregate(
-                maximum_draw_total=Sum('maximum_draw'),
-                allocated_draw_total=Sum('allocated_draw'),
-            )
+            # Global results
+            outlet_count, allocated_draw_total, maximum_draw_total = get_power_stats()
             ret = {
-                'allocated': utilization['allocated_draw_total'] or 0,
-                'maximum': utilization['maximum_draw_total'] or 0,
-                'outlet_count': len(outlet_ids),
+                'allocated': allocated_draw_total,
+                'maximum': maximum_draw_total,
+                'outlet_count': outlet_count,
                 'legs': [],
             }
 
             # Calculate per-leg aggregates for three-phase feeds
-            if self._connected_powerfeed and self._connected_powerfeed.phase == PowerFeedPhaseChoices.PHASE_3PHASE:
+            if leg_stats and getattr(self._connected_powerfeed, 'phase', None) == PowerFeedPhaseChoices.PHASE_3PHASE:
                 for leg, leg_name in PowerOutletFeedLegChoices:
-                    outlet_ids = PowerOutlet.objects.filter(power_port=self, feed_leg=leg).values_list('pk', flat=True)
-                    utilization = PowerPort.objects.filter(_connected_poweroutlet_id__in=outlet_ids).aggregate(
-                        maximum_draw_total=Sum('maximum_draw'),
-                        allocated_draw_total=Sum('allocated_draw'),
-                    )
+                    outlet_count, allocated_draw_total, maximum_draw_total = get_power_stats(leg)
                     ret['legs'].append({
                         'name': leg_name,
-                        'allocated': utilization['allocated_draw_total'] or 0,
-                        'maximum': utilization['maximum_draw_total'] or 0,
-                        'outlet_count': len(outlet_ids),
+                        'allocated': allocated_draw_total,
+                        'maximum': maximum_draw_total,
+                        'outlet_count': outlet_count,
                     })
 
             return ret

netbox/dcim/tests/test_models.py

@@ -552,3 +552,173 @@ class CablePathTestCase(TestCase):
         interface2 = Interface.objects.get(pk=self.interface2.pk)
         self.assertIsNone(interface2.connected_endpoint)
         self.assertIsNone(interface2.connection_status)
+
+
+class PowerCalculationTestCase(TestCase):
+
+    def setUp(self):
+        """
+        The power toplogy:
+
+            power_feed --> power_port1
+                power_outlet12 --> power_port21
+                power_outlet13 --> power_port31
+                    power_outlet34 --> power_port43
+
+        The two numbers at the end denote the direction, so power_outlet12 is from device 1 to device 2.
+        """
+
+        site = Site.objects.create(name='Site 1', slug='site-1')
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        device_type = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type 1', slug='device-type-1')
+        device_role = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1', color='ff0000')
+        power_panel = PowerPanel.objects.create(name='Power Panel 1', site=site)
+
+        self.device1 = Device.objects.create(
+            site=site,
+            device_type=device_type,
+            device_role=device_role,
+            name='Device 1'
+        )
+        self.device2 = Device.objects.create(
+            site=site,
+            device_type=device_type,
+            device_role=device_role,
+            name='Device 2'
+        )
+        self.device3 = Device.objects.create(
+            site=site,
+            device_type=device_type,
+            device_role=device_role,
+            name='Device 3'
+        )
+        self.device4 = Device.objects.create(
+            site=site,
+            device_type=device_type,
+            device_role=device_role,
+            name='Device 4'
+        )
+
+        # Power from power feed to device 1
+        self.power_feed = PowerFeed.objects.create(
+            name='Power Feed 1',
+            power_panel=power_panel,
+            phase=PowerFeedPhaseChoices.PHASE_3PHASE
+        )
+        self.power_port1 = PowerPort.objects.create(device=self.device1, name='Power Port 1')
+        self.power_feed_cable = Cable.objects.create(termination_a=self.power_port1, termination_b=self.power_feed)
+
+        # Power from device 1 to device 2
+        self.power_outlet12 = PowerOutlet.objects.create(
+            device=self.device1,
+            name='Power Outlet 12',
+            power_port=self.power_port1,
+            feed_leg=PowerOutletFeedLegChoices.FEED_LEG_A,
+        )
+        self.power_port21 = PowerPort.objects.create(
+            device=self.device2,
+            name='Power Port 21',
+            maximum_draw=25,
+            allocated_draw=10,
+        )
+        Cable.objects.create(termination_a=self.power_outlet12, termination_b=self.power_port21)
+
+        # Power from device 1 to device 3
+        self.power_outlet13 = PowerOutlet.objects.create(
+            device=self.device1,
+            name='Power Outlet 13',
+            power_port=self.power_port1,
+            feed_leg=PowerOutletFeedLegChoices.FEED_LEG_B,
+        )
+        self.power_port31 = PowerPort.objects.create(
+            device=self.device3,
+            name='Power Port 31',
+            maximum_draw=7,
+            allocated_draw=5,
+        )
+        Cable.objects.create(termination_a=self.power_outlet13, termination_b=self.power_port31)
+
+        # Power from device 3 to device 4
+        self.power_outlet34 = PowerOutlet.objects.create(
+            device=self.device3,
+            name='Power Outlet 34',
+            power_port=self.power_port31,
+        )
+        self.power_port43 = PowerPort.objects.create(
+            device=self.device4,
+            name='Power Port 43',
+            maximum_draw=4,
+            allocated_draw=3,
+        )
+        Cable.objects.create(termination_a=self.power_outlet34, termination_b=self.power_port43)
+
+    def test_power_cascade(self):
+        """
+        Ensure the power calculation cascades over all children. The outlet count is local-only.
+        """
+        stats = self.power_port1.get_power_draw()
+
+        # Global stats: all devices/feeds
+        self.assertEqual(stats['maximum'], 25 + 7 + 4)
+        self.assertEqual(stats['allocated'], 10 + 5 + 3)
+        self.assertEqual(stats['outlet_count'], 2)
+
+        # Feed A stats: device 2
+        self.assertEqual(stats['legs'][0]['maximum'], 25)
+        self.assertEqual(stats['legs'][0]['allocated'], 10)
+        self.assertEqual(stats['legs'][0]['outlet_count'], 1)
+
+        # Feed B stats: devices 3 and 4 (through 3)
+        self.assertEqual(stats['legs'][1]['maximum'], 7 + 4)
+        self.assertEqual(stats['legs'][1]['allocated'], 5 + 3)
+        self.assertEqual(stats['legs'][1]['outlet_count'], 1)
+
+        # Feed C stats: no devices
+        self.assertEqual(stats['legs'][2]['maximum'], 0)
+        self.assertEqual(stats['legs'][2]['allocated'], 0)
+        self.assertEqual(stats['legs'][2]['outlet_count'], 0)
+
+    def test_power_loop(self):
+        """
+        Remove the connection between the power feed and device 1. Instead connect device 4 back to 1. The
+        calculation should continue to be correct and the code should not infinitely loop. The children outlets should
+        have their downstream power ports updated.
+
+        The power toplogy becomes:
+
+            power_outlet41 -> power_port1 (loop)
+                power_outlet12 -> power_port21
+                power_outlet13 -> power_port31
+                    power_outlet34 -> power_port43
+                        power_outlet41 -> power_port1 (loop)
+        """
+        # Power from device 4 to device 1
+        self.power_outlet41 = PowerOutlet.objects.create(
+            device=self.device4,
+            name='Power Outlet 43',
+            power_port=self.power_port43,
+        )
+        self.power_feed_cable.delete()
+        loop_cable = Cable.objects.create(termination_a=self.power_outlet41, termination_b=self.power_port1)
+
+        stats = self.power_port1.get_power_draw()
+
+        self.assertEqual(stats['maximum'], 25 + 7 + 4)
+        self.assertEqual(stats['allocated'], 10 + 5 + 3)
+
+        # With a loop in the topology, all of the outlets affected by the loop have the same children. power_outlet12
+        # is not part of the loop and should only have one child, power_port21.
+        self.assertEqual(self.power_outlet12.downstream_powerports.count(), 1)
+        self.assertEqual(self.power_outlet13.downstream_powerports.count(), 4)
+        self.assertEqual(self.power_outlet34.downstream_powerports.count(), 4)
+        self.assertEqual(self.power_outlet41.downstream_powerports.count(), 4)
+
+        # When a loop-causing cable is removed, the downstream_powerports of the other outlets in the loop should be
+        # updated appropriately. This test is necessary because, in a loop, each outlet is upstream and downstream of
+        # every other outlet in that loop.
+        loop_cable.delete()
+
+        self.assertEqual(self.power_outlet12.downstream_powerports.count(), 1)
+        self.assertEqual(self.power_outlet13.downstream_powerports.count(), 2)
+        self.assertEqual(self.power_outlet34.downstream_powerports.count(), 1)
+        self.assertEqual(self.power_outlet41.downstream_powerports.count(), 0)