Update all systemmonitor sensors in one executor call (#48689)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/systemmonitor/sensor.py

@@ -1,8 +1,14 @@
 """Support for monitoring the local system."""
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+import datetime
 import logging
 import os
 import socket
 import sys
+from typing import Any, Callable, cast
 
 import psutil
 import voluptuous as vol
@@ -10,22 +16,30 @@ import voluptuous as vol
 from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
 from homeassistant.const import (
     CONF_RESOURCES,
+    CONF_SCAN_INTERVAL,
     CONF_TYPE,
     DATA_GIBIBYTES,
     DATA_MEBIBYTES,
     DATA_RATE_MEGABYTES_PER_SECOND,
     DEVICE_CLASS_TIMESTAMP,
+    EVENT_HOMEASSISTANT_STOP,
     PERCENTAGE,
     STATE_OFF,
     STATE_ON,
     TEMP_CELSIUS,
 )
+from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect,
+    async_dispatcher_send,
+)
+from homeassistant.helpers.entity_component import DEFAULT_SCAN_INTERVAL
+from homeassistant.helpers.event import async_track_time_interval
+from homeassistant.helpers.typing import ConfigType
 from homeassistant.util import slugify
 import homeassistant.util.dt as dt_util
 
-# mypy: allow-untyped-defs, no-check-untyped-defs
-
 _LOGGER = logging.getLogger(__name__)
 
 CONF_ARG = "arg"
@@ -35,71 +49,80 @@ if sys.maxsize > 2 ** 32:
 else:
     CPU_ICON = "mdi:cpu-32-bit"
 
+SENSOR_TYPE_NAME = 0
+SENSOR_TYPE_UOM = 1
+SENSOR_TYPE_ICON = 2
+SENSOR_TYPE_DEVICE_CLASS = 3
+SENSOR_TYPE_MANDATORY_ARG = 4
+
+SIGNAL_SYSTEMMONITOR_UPDATE = "systemmonitor_update"
+
 # Schema: [name, unit of measurement, icon, device class, flag if mandatory arg]
-SENSOR_TYPES = {
+SENSOR_TYPES: dict[str, tuple[str, str | None, str | None, str | None, bool]] = {
-    "disk_free": ["Disk free", DATA_GIBIBYTES, "mdi:harddisk", None, False],
+    "disk_free": ("Disk free", DATA_GIBIBYTES, "mdi:harddisk", None, False),
-    "disk_use": ["Disk use", DATA_GIBIBYTES, "mdi:harddisk", None, False],
+    "disk_use": ("Disk use", DATA_GIBIBYTES, "mdi:harddisk", None, False),
-    "disk_use_percent": [
+    "disk_use_percent": (
         "Disk use (percent)",
         PERCENTAGE,
         "mdi:harddisk",
         None,
         False,
-    ],
+    ),
-    "ipv4_address": ["IPv4 address", "", "mdi:server-network", None, True],
+    "ipv4_address": ("IPv4 address", "", "mdi:server-network", None, True),
-    "ipv6_address": ["IPv6 address", "", "mdi:server-network", None, True],
+    "ipv6_address": ("IPv6 address", "", "mdi:server-network", None, True),
-    "last_boot": ["Last boot", None, "mdi:clock", DEVICE_CLASS_TIMESTAMP, False],
+    "last_boot": ("Last boot", None, "mdi:clock", DEVICE_CLASS_TIMESTAMP, False),
-    "load_15m": ["Load (15m)", " ", CPU_ICON, None, False],
+    "load_15m": ("Load (15m)", " ", CPU_ICON, None, False),
-    "load_1m": ["Load (1m)", " ", CPU_ICON, None, False],
+    "load_1m": ("Load (1m)", " ", CPU_ICON, None, False),
-    "load_5m": ["Load (5m)", " ", CPU_ICON, None, False],
+    "load_5m": ("Load (5m)", " ", CPU_ICON, None, False),
-    "memory_free": ["Memory free", DATA_MEBIBYTES, "mdi:memory", None, False],
+    "memory_free": ("Memory free", DATA_MEBIBYTES, "mdi:memory", None, False),
-    "memory_use": ["Memory use", DATA_MEBIBYTES, "mdi:memory", None, False],
+    "memory_use": ("Memory use", DATA_MEBIBYTES, "mdi:memory", None, False),
-    "memory_use_percent": [
+    "memory_use_percent": (
         "Memory use (percent)",
         PERCENTAGE,
         "mdi:memory",
         None,
         False,
-    ],
+    ),
-    "network_in": ["Network in", DATA_MEBIBYTES, "mdi:server-network", None, True],
+    "network_in": ("Network in", DATA_MEBIBYTES, "mdi:server-network", None, True),
-    "network_out": ["Network out", DATA_MEBIBYTES, "mdi:server-network", None, True],
+    "network_out": ("Network out", DATA_MEBIBYTES, "mdi:server-network", None, True),
-    "packets_in": ["Packets in", " ", "mdi:server-network", None, True],
+    "packets_in": ("Packets in", " ", "mdi:server-network", None, True),
-    "packets_out": ["Packets out", " ", "mdi:server-network", None, True],
+    "packets_out": ("Packets out", " ", "mdi:server-network", None, True),
-    "throughput_network_in": [
+    "throughput_network_in": (
         "Network throughput in",
         DATA_RATE_MEGABYTES_PER_SECOND,
         "mdi:server-network",
         None,
         True,
-    ],
+    ),
-    "throughput_network_out": [
+    "throughput_network_out": (
         "Network throughput out",
         DATA_RATE_MEGABYTES_PER_SECOND,
         "mdi:server-network",
+        None,
         True,
-    ],
+    ),
-    "process": ["Process", " ", CPU_ICON, None, True],
+    "process": ("Process", " ", CPU_ICON, None, True),
-    "processor_use": ["Processor use (percent)", PERCENTAGE, CPU_ICON, None, False],
+    "processor_use": ("Processor use (percent)", PERCENTAGE, CPU_ICON, None, False),
-    "processor_temperature": [
+    "processor_temperature": (
         "Processor temperature",
         TEMP_CELSIUS,
         CPU_ICON,
         None,
         False,
-    ],
+    ),
-    "swap_free": ["Swap free", DATA_MEBIBYTES, "mdi:harddisk", None, False],
+    "swap_free": ("Swap free", DATA_MEBIBYTES, "mdi:harddisk", None, False),
-    "swap_use": ["Swap use", DATA_MEBIBYTES, "mdi:harddisk", None, False],
+    "swap_use": ("Swap use", DATA_MEBIBYTES, "mdi:harddisk", None, False),
-    "swap_use_percent": ["Swap use (percent)", PERCENTAGE, "mdi:harddisk", None, False],
+    "swap_use_percent": ("Swap use (percent)", PERCENTAGE, "mdi:harddisk", None, False),
 }
 
 
-def check_required_arg(value):
+def check_required_arg(value: Any) -> Any:
     """Validate that the required "arg" for the sensor types that need it are set."""
     for sensor in value:
         sensor_type = sensor[CONF_TYPE]
         sensor_arg = sensor.get(CONF_ARG)
 
-        if sensor_arg is None and SENSOR_TYPES[sensor_type][4]:
+        if sensor_arg is None and SENSOR_TYPES[sensor_type][SENSOR_TYPE_MANDATORY_ARG]:
             raise vol.RequiredFieldInvalid(
                 f"Mandatory 'arg' is missing for sensor type '{sensor_type}'."
             )
@@ -158,176 +181,280 @@ CPU_SENSOR_PREFIXES = [
 ]
 
 
-def setup_platform(hass, config, add_entities, discovery_info=None):
+@dataclass
+class SensorData:
+    """Data for a sensor."""
+
+    argument: Any
+    state: str | None
+    value: Any | None
+    update_time: datetime.datetime | None
+    last_exception: BaseException | None
+
+
+async def async_setup_platform(
+    hass: HomeAssistant,
+    config: ConfigType,
+    async_add_entities: Callable,
+    discovery_info: Any | None = None,
+) -> None:
     """Set up the system monitor sensors."""
-    dev = []
+    entities = []
+    sensor_registry: dict[str, SensorData] = {}
+
     for resource in config[CONF_RESOURCES]:
+        type_ = resource[CONF_TYPE]
         # Initialize the sensor argument if none was provided.
         # For disk monitoring default to "/" (root) to prevent runtime errors, if argument was not specified.
         if CONF_ARG not in resource:
-            resource[CONF_ARG] = ""
+            argument = ""
             if resource[CONF_TYPE].startswith("disk_"):
-                resource[CONF_ARG] = "/"
+                argument = "/"
+        else:
+            argument = resource[CONF_ARG]
 
         # Verify if we can retrieve CPU / processor temperatures.
         # If not, do not create the entity and add a warning to the log
         if (
-            resource[CONF_TYPE] == "processor_temperature"
+            type_ == "processor_temperature"
-            and SystemMonitorSensor.read_cpu_temperature() is None
+            and await hass.async_add_executor_job(_read_cpu_temperature) is None
         ):
             _LOGGER.warning("Cannot read CPU / processor temperature information")
             continue
 
-        dev.append(SystemMonitorSensor(resource[CONF_TYPE], resource[CONF_ARG]))
+        sensor_registry[type_] = SensorData(argument, None, None, None, None)
+        entities.append(SystemMonitorSensor(sensor_registry, type_, argument))
 
-    add_entities(dev, True)
+    scan_interval = config.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
+    await async_setup_sensor_registry_updates(hass, sensor_registry, scan_interval)
+
+    async_add_entities(entities)
+
+
+async def async_setup_sensor_registry_updates(
+    hass: HomeAssistant,
+    sensor_registry: dict[str, SensorData],
+    scan_interval: datetime.timedelta,
+) -> None:
+    """Update the registry and create polling."""
+
+    _update_lock = asyncio.Lock()
+
+    def _update_sensors() -> None:
+        """Update sensors and store the result in the registry."""
+        for type_, data in sensor_registry.items():
+            try:
+                state, value, update_time = _update(type_, data)
+            except Exception as ex:  # pylint: disable=broad-except
+                _LOGGER.exception("Error updating sensor: %s", type_, exc_info=ex)
+                data.last_exception = ex
+            else:
+                data.state = state
+                data.value = value
+                data.update_time = update_time
+                data.last_exception = None
+
+    async def _async_update_data(*_: Any) -> None:
+        """Update all sensors in one executor jump."""
+        if _update_lock.locked():
+            _LOGGER.warning(
+                "Updating systemmonitor took longer than the scheduled update interval %s",
+                scan_interval,
+            )
+            return
+
+        async with _update_lock:
+            await hass.async_add_executor_job(_update_sensors)
+            async_dispatcher_send(hass, SIGNAL_SYSTEMMONITOR_UPDATE)
+
+    polling_remover = async_track_time_interval(hass, _async_update_data, scan_interval)
+
+    @callback
+    def _async_stop_polling(*_: Any) -> None:
+        polling_remover()
+
+    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_stop_polling)
+
+    await _async_update_data()
 
 
 class SystemMonitorSensor(SensorEntity):
     """Implementation of a system monitor sensor."""
 
-    def __init__(self, sensor_type, argument=""):
+    def __init__(
+        self,
+        sensor_registry: dict[str, SensorData],
+        sensor_type: str,
+        argument: str = "",
+    ) -> None:
         """Initialize the sensor."""
-        self._name = "{} {}".format(SENSOR_TYPES[sensor_type][0], argument)
+        self._type: str = sensor_type
-        self._unique_id = slugify(f"{sensor_type}_{argument}")
+        self._name: str = "{} {}".format(self.sensor_type[SENSOR_TYPE_NAME], argument)
-        self.argument = argument
+        self._unique_id: str = slugify(f"{sensor_type}_{argument}")
-        self.type = sensor_type
+        self._sensor_registry = sensor_registry
-        self._state = None
-        self._unit_of_measurement = SENSOR_TYPES[sensor_type][1]
-        self._available = True
-        if sensor_type in ["throughput_network_out", "throughput_network_in"]:
-            self._last_value = None
-            self._last_update_time = None
 
     @property
-    def name(self):
+    def name(self) -> str:
         """Return the name of the sensor."""
         return self._name.rstrip()
 
     @property
-    def unique_id(self):
+    def unique_id(self) -> str:
         """Return the unique ID."""
         return self._unique_id
 
     @property
-    def device_class(self):
+    def device_class(self) -> str | None:
         """Return the class of this sensor."""
-        return SENSOR_TYPES[self.type][3]
+        return self.sensor_type[SENSOR_TYPE_DEVICE_CLASS]  # type: ignore[no-any-return]
 
     @property
-    def icon(self):
+    def icon(self) -> str | None:
         """Icon to use in the frontend, if any."""
-        return SENSOR_TYPES[self.type][2]
+        return self.sensor_type[SENSOR_TYPE_ICON]  # type: ignore[no-any-return]
 
     @property
-    def state(self):
+    def state(self) -> str | None:
         """Return the state of the device."""
-        return self._state
+        return self.data.state
 
     @property
-    def unit_of_measurement(self):
+    def unit_of_measurement(self) -> str | None:
         """Return the unit of measurement of this entity, if any."""
-        return self._unit_of_measurement
+        return self.sensor_type[SENSOR_TYPE_UOM]  # type: ignore[no-any-return]
 
     @property
-    def available(self):
+    def available(self) -> bool:
         """Return True if entity is available."""
-        return self._available
+        return self.data.last_exception is None
 
-    def update(self):
+    @property
-        """Get the latest system information."""
+    def should_poll(self) -> bool:
-        if self.type == "disk_use_percent":
+        """Entity does not poll."""
-            self._state = psutil.disk_usage(self.argument).percent
+        return False
-        elif self.type == "disk_use":
+
-            self._state = round(psutil.disk_usage(self.argument).used / 1024 ** 3, 1)
+    @property
-        elif self.type == "disk_free":
+    def sensor_type(self) -> list:
-            self._state = round(psutil.disk_usage(self.argument).free / 1024 ** 3, 1)
+        """Return sensor type data for the sensor."""
-        elif self.type == "memory_use_percent":
+        return SENSOR_TYPES[self._type]  # type: ignore
-            self._state = psutil.virtual_memory().percent
+
-        elif self.type == "memory_use":
+    @property
-            virtual_memory = psutil.virtual_memory()
+    def data(self) -> SensorData:
-            self._state = round(
+        """Return registry entry for the data."""
-                (virtual_memory.total - virtual_memory.available) / 1024 ** 2, 1
+        return self._sensor_registry[self._type]
+
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass."""
+        await super().async_added_to_hass()
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass, SIGNAL_SYSTEMMONITOR_UPDATE, self.async_write_ha_state
             )
-        elif self.type == "memory_free":
+        )
-            self._state = round(psutil.virtual_memory().available / 1024 ** 2, 1)
+
-        elif self.type == "swap_use_percent":
+
-            self._state = psutil.swap_memory().percent
+def _update(
-        elif self.type == "swap_use":
+    type_: str, data: SensorData
-            self._state = round(psutil.swap_memory().used / 1024 ** 2, 1)
+) -> tuple[str | None, str | None, datetime.datetime | None]:
-        elif self.type == "swap_free":
+    """Get the latest system information."""
-            self._state = round(psutil.swap_memory().free / 1024 ** 2, 1)
+    state = None
-        elif self.type == "processor_use":
+    value = None
-            self._state = round(psutil.cpu_percent(interval=None))
+    update_time = None
-        elif self.type == "processor_temperature":
+
-            self._state = self.read_cpu_temperature()
+    if type_ == "disk_use_percent":
-        elif self.type == "process":
+        state = psutil.disk_usage(data.argument).percent
+    elif type_ == "disk_use":
+        state = round(psutil.disk_usage(data.argument).used / 1024 ** 3, 1)
+    elif type_ == "disk_free":
+        state = round(psutil.disk_usage(data.argument).free / 1024 ** 3, 1)
+    elif type_ == "memory_use_percent":
+        state = psutil.virtual_memory().percent
+    elif type_ == "memory_use":
+        virtual_memory = psutil.virtual_memory()
+        state = round((virtual_memory.total - virtual_memory.available) / 1024 ** 2, 1)
+    elif type_ == "memory_free":
+        state = round(psutil.virtual_memory().available / 1024 ** 2, 1)
+    elif type_ == "swap_use_percent":
+        state = psutil.swap_memory().percent
+    elif type_ == "swap_use":
+        state = round(psutil.swap_memory().used / 1024 ** 2, 1)
+    elif type_ == "swap_free":
+        state = round(psutil.swap_memory().free / 1024 ** 2, 1)
+    elif type_ == "processor_use":
+        state = round(psutil.cpu_percent(interval=None))
+    elif type_ == "processor_temperature":
+        state = _read_cpu_temperature()
+    elif type_ == "process":
+        state = STATE_OFF
         for proc in psutil.process_iter():
             try:
-                    if self.argument == proc.name():
+                if data.argument == proc.name():
-                        self._state = STATE_ON
+                    state = STATE_ON
-                        return
+                    break
             except psutil.NoSuchProcess as err:
                 _LOGGER.warning(
                     "Failed to load process with ID: %s, old name: %s",
                     err.pid,
                     err.name,
                 )
-            self._state = STATE_OFF
+    elif type_ in ["network_out", "network_in"]:
-        elif self.type in ["network_out", "network_in"]:
         counters = psutil.net_io_counters(pernic=True)
-            if self.argument in counters:
+        if data.argument in counters:
-                counter = counters[self.argument][IO_COUNTER[self.type]]
+            counter = counters[data.argument][IO_COUNTER[type_]]
-                self._state = round(counter / 1024 ** 2, 1)
+            state = round(counter / 1024 ** 2, 1)
         else:
-                self._state = None
+            state = None
-        elif self.type in ["packets_out", "packets_in"]:
+    elif type_ in ["packets_out", "packets_in"]:
         counters = psutil.net_io_counters(pernic=True)
-            if self.argument in counters:
+        if data.argument in counters:
-                self._state = counters[self.argument][IO_COUNTER[self.type]]
+            state = counters[data.argument][IO_COUNTER[type_]]
         else:
-                self._state = None
+            state = None
-        elif self.type in ["throughput_network_out", "throughput_network_in"]:
+    elif type_ in ["throughput_network_out", "throughput_network_in"]:
         counters = psutil.net_io_counters(pernic=True)
-            if self.argument in counters:
+        if data.argument in counters:
-                counter = counters[self.argument][IO_COUNTER[self.type]]
+            counter = counters[data.argument][IO_COUNTER[type_]]
             now = dt_util.utcnow()
-                if self._last_value and self._last_value < counter:
+            if data.value and data.value < counter:
-                    self._state = round(
+                state = round(
-                        (counter - self._last_value)
+                    (counter - data.value)
                     / 1000 ** 2
-                        / (now - self._last_update_time).seconds,
+                    / (now - (data.update_time or now)).seconds,
                     3,
                 )
             else:
-                    self._state = None
+                state = None
-                self._last_update_time = now
+            update_time = now
-                self._last_value = counter
+            value = counter
         else:
-                self._state = None
+            state = None
-        elif self.type in ["ipv4_address", "ipv6_address"]:
+    elif type_ in ["ipv4_address", "ipv6_address"]:
         addresses = psutil.net_if_addrs()
-            if self.argument in addresses:
+        if data.argument in addresses:
-                for addr in addresses[self.argument]:
+            for addr in addresses[data.argument]:
-                    if addr.family == IF_ADDRS_FAMILY[self.type]:
+                if addr.family == IF_ADDRS_FAMILY[type_]:
-                        self._state = addr.address
+                    state = addr.address
         else:
-                self._state = None
+            state = None
-        elif self.type == "last_boot":
+    elif type_ == "last_boot":
         # Only update on initial setup
-            if self._state is None:
+        if data.state is None:
-                self._state = dt_util.as_local(
+            state = dt_util.as_local(
                 dt_util.utc_from_timestamp(psutil.boot_time())
             ).isoformat()
-        elif self.type == "load_1m":
+        else:
-            self._state = round(os.getloadavg()[0], 2)
+            state = data.state
-        elif self.type == "load_5m":
+    elif type_ == "load_1m":
-            self._state = round(os.getloadavg()[1], 2)
+        state = round(os.getloadavg()[0], 2)
-        elif self.type == "load_15m":
+    elif type_ == "load_5m":
-            self._state = round(os.getloadavg()[2], 2)
+        state = round(os.getloadavg()[1], 2)
+    elif type_ == "load_15m":
+        state = round(os.getloadavg()[2], 2)
 
-    @staticmethod
+    return state, value, update_time
-    def read_cpu_temperature():
+
+
+def _read_cpu_temperature() -> float | None:
     """Attempt to read CPU / processor temperature."""
     temps = psutil.sensors_temperatures()
 
@@ -337,4 +464,6 @@ class SystemMonitorSensor(SensorEntity):
             # construct it ourself here based on the sensor key name.
             _label = f"{name} {i}" if not entry.label else entry.label
             if _label in CPU_SENSOR_PREFIXES:
-                    return round(entry.current, 1)
+                return cast(float, round(entry.current, 1))
+
+    return None