[fix] Record monitoring metrics for sessions closed via bulk updates

docs/developer/utils.rst

@@ -27,6 +27,33 @@ completes successfully, just before the response is returned.
 The ``view`` argument can also be used to access the ``request`` object
 i.e. ``view.request``.
 
+``radius_accounting_closed``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+**Path**: ``openwisp_radius.signals.radius_accounting_closed``
+
+**Arguments**:
+
+- ``sender``: ``RadiusAccounting`` model class
+- ``instance``: closed ``RadiusAccounting`` instance
+
+This signal is emitted when a ``RadiusAccounting`` session is closed. For
+regular ``save()`` paths, it is emitted only when ``stop_time`` changes
+from ``None`` to any value, or when a closed session is created directly.
+Editing an already closed session does not emit this signal again.
+
+The signal is emitted after the database transaction is committed.
+
+Integrations which need to react to closed accounting sessions should
+listen to this signal. For example, the monitoring integration uses this
+signal to write RADIUS traffic snapshots for sessions closed by regular
+``post_save`` paths, ``Accounting-On`` packets and automatic stale-session
+cleanup.
+
+When closing multiple sessions with ``bulk_update()``, use
+``RadiusAccounting.emit_radius_accounting_closed()`` after the database
+update to emit this signal once for each closed session.
+
 .. _radius_captive_portal_mock_views:
 
 Captive portal mock views

openwisp_radius/base/models.py

@@ -5,6 +5,7 @@
 import string
 from datetime import timedelta
 from io import StringIO
+from typing import Iterable
 
 import django
 import phonenumbers
@@ -52,6 +53,7 @@
     BATCH_MAIL_SUBJECT,
     DEFAULT_PASSWORD_RESET_URL,
 )
+from ..signals import radius_accounting_closed
 from ..utils import (
     SmsMessage,
     decode_byte_data,
@@ -529,20 +531,66 @@ class AbstractRadiusAccounting(OrgMixin, models.Model):
         blank=True,
     )
 
-    def save(self, *args, **kwargs):
-        if not self.start_time:
-            self.start_time = now()
-        super(AbstractRadiusAccounting, self).save(*args, **kwargs)
-
     class Meta:
         db_table = "radacct"
         verbose_name = _("accounting")
         verbose_name_plural = _("accountings")
         abstract = True
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # used for radius_accounting_closed signal
+        self._set_initial_stop_time()
+
+    def refresh_from_db(self, *args, **kwargs):
+        super().refresh_from_db(*args, **kwargs)
+        fields = kwargs.get("fields")
+        self._set_initial_stop_time(fields=fields)
+
+    def save(self, *args, **kwargs):
+        created = self._state.adding
+        update_fields = kwargs.get("update_fields")
+        if not self.start_time:
+            self.start_time = now()
+            if update_fields is not None:
+                update_fields = set(update_fields) | {"start_time"}
+                kwargs["update_fields"] = update_fields
+        super(AbstractRadiusAccounting, self).save(*args, **kwargs)
+        self._emit_radius_accounting_closed(
+            created=created, update_fields=update_fields
+        )
+        # reset after save
+        self._set_initial_stop_time(update_fields)
+
+    def _set_initial_stop_time(self, fields=None):
+        if fields is None or "stop_time" in fields:
+            self._initial_stop_time = self.stop_time
+
+    def _emit_radius_accounting_closed(self, created, update_fields=None):
+        """Detect whether this save closed the session and emit the signal."""
+        if update_fields is not None and "stop_time" not in update_fields:
+            return
+        being_closed = self.stop_time is not None and (
+            created or self._initial_stop_time is None
+        )
+        if being_closed:
+            self.emit_radius_accounting_closed([self])
+
     def __str__(self):
         return self.unique_id
 
+    @classmethod
+    def emit_radius_accounting_closed(
+        cls, sessions: Iterable["AbstractRadiusAccounting"]
+    ) -> None:
+        """Emit radius_accounting_closed after commit for closed sessions."""
+        for session in sessions:
+            transaction.on_commit(
+                lambda session=session: radius_accounting_closed.send(
+                    sender=session.__class__, instance=session
+                )
+            )
+
     @classmethod
     def close_stale_sessions(cls, days=None, hours=None):
         if hours:
@@ -579,13 +627,39 @@ def _close_stale_sessions_on_nas_boot(cls, called_station_id):
         """
         if not called_station_id:
             return 0
-        stale_sessions = cls.objects.filter(
-            called_station_id=called_station_id,
-            stop_time__isnull=True,
-        )
-        closed_count = stale_sessions.update(
-            stop_time=now(), terminate_cause="NAS-Reboot"
-        )
+        stop_time = now()
+        closed_count = 0
+        batch_size = 1000
+        has_more_sessions = True
+        while has_more_sessions:
+            with transaction.atomic():
+                closed_sessions = list(
+                    cls.objects.select_for_update()
+                    .filter(
+                        called_station_id=called_station_id,
+                        stop_time__isnull=True,
+                    )
+                    .only(
+                        "unique_id",
+                        "username",
+                        "organization_id",
+                        "input_octets",
+                        "output_octets",
+                        "calling_station_id",
+                        "called_station_id",
+                        "stop_time",
+                    )[:batch_size]
+                )
+                has_more_sessions = len(closed_sessions) == batch_size
+                if not closed_sessions:
+                    continue
+                for session in closed_sessions:
+                    session.stop_time = stop_time
+                    session.terminate_cause = "NAS-Reboot"
+                closed_count += cls.objects.bulk_update(
+                    closed_sessions, fields=["stop_time", "terminate_cause"]
+                )
+                cls.emit_radius_accounting_closed(closed_sessions)
         return closed_count
 
 

openwisp_radius/integrations/monitoring/apps.py

@@ -4,14 +4,14 @@
 from django.db import models
 from django.db.models import Count, Sum
 from django.db.models.functions import Cast, Round
-from django.db.models.signals import post_save
 from django.utils.translation import gettext_lazy as _
 from openwisp_monitoring.monitoring.configuration import (
     _register_chart_configuration_choice,
     register_metric,
 )
 from swapper import load_model
 
+from openwisp_radius.signals import radius_accounting_closed
 from openwisp_utils.admin_theme import register_dashboard_chart
 
 from .utils import (
@@ -121,12 +121,12 @@ def register_radius_metrics(self):
                 _register_chart_configuration_choice(chart_key, chart_config)
 
     def connect_signal_receivers(self):
-        from .receivers import post_save_radiusaccounting
+        from .receivers import radius_accounting_closed_handler
 
         RadiusAccounting = load_model("openwisp_radius", "RadiusAccounting")
 
-        post_save.connect(
-            post_save_radiusaccounting,
+        radius_accounting_closed.connect(
+            radius_accounting_closed_handler,
             sender=RadiusAccounting,
-            dispatch_uid="post_save_radiusaccounting_radius_acc_metric",
+            dispatch_uid="radius_accounting_closed_radius_acc_metric",
         )

openwisp_radius/integrations/monitoring/management/commands/rebuild_radius_accounting_metrics.py

@@ -0,0 +1,122 @@
+from django.core.management.base import BaseCommand, CommandError
+from django.utils import timezone
+from django.utils.dateparse import parse_datetime
+from openwisp_monitoring.db import timeseries_db
+from swapper import load_model
+
+from openwisp_radius.integrations.monitoring import tasks
+from openwisp_radius.integrations.monitoring.utils import sha1_hash
+
+RadiusAccounting = load_model("openwisp_radius", "RadiusAccounting")
+
+
+class Command(BaseCommand):
+    """
+    Rebuild RADIUS accounting metrics missed by the NAS-Reboot bulk update path.
+
+    See https://github.com/openwisp/openwisp-radius/issues/734.
+    TODO: remove in version 1.4.
+    """
+
+    help = "Private command to rebuild monitoring metrics for closed RADIUS sessions."
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--commit",
+            action="store_true",
+            help="Write metrics. Without this flag the command only reports the count.",
+        )
+        parser.add_argument(
+            "--start",
+            help=(
+                "Only process sessions with stop_time greater than or equal "
+                "to this date."
+            ),
+        )
+        parser.add_argument(
+            "--end",
+            help=(
+                "Only process sessions with stop_time lower than or equal "
+                "to this date."
+            ),
+        )
+        parser.add_argument(
+            "--chunk-size",
+            type=int,
+            default=1000,
+            help="Number of sessions fetched per database batch.",
+        )
+
+    def handle(self, *args, **options):
+        queryset = RadiusAccounting.objects.filter(
+            stop_time__isnull=False,
+            terminate_cause="NAS-Reboot",
+        )
+        if options["start"]:
+            queryset = queryset.filter(
+                stop_time__gte=self._parse_datetime(options["start"], "start")
+            )
+        if options["end"]:
+            queryset = queryset.filter(
+                stop_time__lte=self._parse_datetime(options["end"], "end")
+            )
+        queryset = queryset.order_by("stop_time", "unique_id").only(
+            "unique_id",
+            "username",
+            "organization_id",
+            "input_octets",
+            "output_octets",
+            "calling_station_id",
+            "called_station_id",
+            "stop_time",
+            "terminate_cause",
+        )
+        count = queryset.count()
+        if not options["commit"]:
+            self.stdout.write(f"Dry run: {count} closed sessions would be processed.")
+            return
+        processed = 0
+        chunk_size = options["chunk_size"]
+        self.stdout.write(f"Starting to rebuild {count} accounting metrics.")
+        for session in queryset.iterator(chunk_size=chunk_size):
+            self._delete_radius_accounting_metric(session)
+            tasks.post_save_radiusaccounting(
+                username=session.username,
+                organization_id=str(session.organization_id),
+                input_octets=session.input_octets,
+                output_octets=session.output_octets,
+                calling_station_id=session.calling_station_id,
+                called_station_id=session.called_station_id,
+                time=session.stop_time,
+            )
+            processed += 1
+            if processed % chunk_size == 0 or processed == count:
+                self.stdout.write(
+                    f"Processed {processed} of {count} accounting metrics."
+                )
+
+    def _parse_datetime(self, value, option):
+        parsed = parse_datetime(value)
+        if parsed is None:
+            raise CommandError(f"Invalid --{option} datetime: {value}")
+        if timezone.is_naive(parsed):
+            parsed = timezone.make_aware(parsed)
+        return parsed
+
+    def _delete_radius_accounting_metric(self, session):
+        tags = {
+            "organization_id": str(session.organization_id),
+            "calling_station_id": sha1_hash(session.calling_station_id),
+            "called_station_id": session.called_station_id,
+        }
+        where = " AND ".join(
+            f"\"{key}\" = '{self._escape_tag_value(value)}'"
+            for key, value in tags.items()
+        )
+        timeseries_db.query(
+            "DELETE FROM radius_acc "
+            f"WHERE time = '{session.stop_time.isoformat()}' AND {where}"
+        )
+
+    def _escape_tag_value(self, value):
+        return str(value).replace("'", r"\'")