|
5 | 5 | import string |
6 | 6 | from datetime import timedelta |
7 | 7 | from io import StringIO |
| 8 | +from typing import Iterable |
8 | 9 |
|
9 | 10 | import django |
10 | 11 | import phonenumbers |
|
52 | 53 | BATCH_MAIL_SUBJECT, |
53 | 54 | DEFAULT_PASSWORD_RESET_URL, |
54 | 55 | ) |
| 56 | +from ..signals import radius_accounting_closed |
55 | 57 | from ..utils import ( |
56 | 58 | SmsMessage, |
57 | 59 | decode_byte_data, |
@@ -529,20 +531,66 @@ class AbstractRadiusAccounting(OrgMixin, models.Model): |
529 | 531 | blank=True, |
530 | 532 | ) |
531 | 533 |
|
532 | | - def save(self, *args, **kwargs): |
533 | | - if not self.start_time: |
534 | | - self.start_time = now() |
535 | | - super(AbstractRadiusAccounting, self).save(*args, **kwargs) |
536 | | - |
537 | 534 | class Meta: |
538 | 535 | db_table = "radacct" |
539 | 536 | verbose_name = _("accounting") |
540 | 537 | verbose_name_plural = _("accountings") |
541 | 538 | abstract = True |
542 | 539 |
|
| 540 | + def __init__(self, *args, **kwargs): |
| 541 | + super().__init__(*args, **kwargs) |
| 542 | + # used for radius_accounting_closed signal |
| 543 | + self._set_initial_stop_time() |
| 544 | + |
| 545 | + def refresh_from_db(self, *args, **kwargs): |
| 546 | + super().refresh_from_db(*args, **kwargs) |
| 547 | + fields = kwargs.get("fields") |
| 548 | + self._set_initial_stop_time(fields=fields) |
| 549 | + |
| 550 | + def save(self, *args, **kwargs): |
| 551 | + created = self._state.adding |
| 552 | + update_fields = kwargs.get("update_fields") |
| 553 | + if not self.start_time: |
| 554 | + self.start_time = now() |
| 555 | + if update_fields is not None: |
| 556 | + update_fields = set(update_fields) | {"start_time"} |
| 557 | + kwargs["update_fields"] = update_fields |
| 558 | + super(AbstractRadiusAccounting, self).save(*args, **kwargs) |
| 559 | + self._emit_radius_accounting_closed( |
| 560 | + created=created, update_fields=update_fields |
| 561 | + ) |
| 562 | + # reset after save |
| 563 | + self._set_initial_stop_time(update_fields) |
| 564 | + |
| 565 | + def _set_initial_stop_time(self, fields=None): |
| 566 | + if fields is None or "stop_time" in fields: |
| 567 | + self._initial_stop_time = self.stop_time |
| 568 | + |
| 569 | + def _emit_radius_accounting_closed(self, created, update_fields=None): |
| 570 | + """Detect whether this save closed the session and emit the signal.""" |
| 571 | + if update_fields is not None and "stop_time" not in update_fields: |
| 572 | + return |
| 573 | + being_closed = self.stop_time is not None and ( |
| 574 | + created or self._initial_stop_time is None |
| 575 | + ) |
| 576 | + if being_closed: |
| 577 | + self.emit_radius_accounting_closed([self]) |
| 578 | + |
543 | 579 | def __str__(self): |
544 | 580 | return self.unique_id |
545 | 581 |
|
| 582 | + @classmethod |
| 583 | + def emit_radius_accounting_closed( |
| 584 | + cls, sessions: Iterable["AbstractRadiusAccounting"] |
| 585 | + ) -> None: |
| 586 | + """Emit radius_accounting_closed after commit for closed sessions.""" |
| 587 | + for session in sessions: |
| 588 | + transaction.on_commit( |
| 589 | + lambda session=session: radius_accounting_closed.send( |
| 590 | + sender=session.__class__, instance=session |
| 591 | + ) |
| 592 | + ) |
| 593 | + |
546 | 594 | @classmethod |
547 | 595 | def close_stale_sessions(cls, days=None, hours=None): |
548 | 596 | if hours: |
@@ -579,13 +627,39 @@ def _close_stale_sessions_on_nas_boot(cls, called_station_id): |
579 | 627 | """ |
580 | 628 | if not called_station_id: |
581 | 629 | return 0 |
582 | | - stale_sessions = cls.objects.filter( |
583 | | - called_station_id=called_station_id, |
584 | | - stop_time__isnull=True, |
585 | | - ) |
586 | | - closed_count = stale_sessions.update( |
587 | | - stop_time=now(), terminate_cause="NAS-Reboot" |
588 | | - ) |
| 630 | + stop_time = now() |
| 631 | + closed_count = 0 |
| 632 | + batch_size = 1000 |
| 633 | + has_more_sessions = True |
| 634 | + while has_more_sessions: |
| 635 | + with transaction.atomic(): |
| 636 | + closed_sessions = list( |
| 637 | + cls.objects.select_for_update() |
| 638 | + .filter( |
| 639 | + called_station_id=called_station_id, |
| 640 | + stop_time__isnull=True, |
| 641 | + ) |
| 642 | + .only( |
| 643 | + "unique_id", |
| 644 | + "username", |
| 645 | + "organization_id", |
| 646 | + "input_octets", |
| 647 | + "output_octets", |
| 648 | + "calling_station_id", |
| 649 | + "called_station_id", |
| 650 | + "stop_time", |
| 651 | + )[:batch_size] |
| 652 | + ) |
| 653 | + has_more_sessions = len(closed_sessions) == batch_size |
| 654 | + if not closed_sessions: |
| 655 | + continue |
| 656 | + for session in closed_sessions: |
| 657 | + session.stop_time = stop_time |
| 658 | + session.terminate_cause = "NAS-Reboot" |
| 659 | + closed_count += cls.objects.bulk_update( |
| 660 | + closed_sessions, fields=["stop_time", "terminate_cause"] |
| 661 | + ) |
| 662 | + cls.emit_radius_accounting_closed(closed_sessions) |
589 | 663 | return closed_count |
590 | 664 |
|
591 | 665 |
|
|
0 commit comments