From 3795c836d1f990f8bbf08f87d2ab0384904c057d Mon Sep 17 00:00:00 2001 From: Matias Bordese Date: Wed, 31 Jan 2024 11:26:14 -0300 Subject: [PATCH] Add transaction block and callbacks when triggering tasks (#3779) Related to https://github.com/grafana/oncall/issues/3729 --- .../alerts/models/alert_group_log_record.py | 1 + engine/apps/alerts/paging.py | 36 +++++++++++-------- .../apps/alerts/tasks/acknowledge_reminder.py | 9 ++--- .../alerts/tests/test_acknowledge_reminder.py | 11 ++++-- engine/apps/alerts/tests/test_paging.py | 9 +++-- 5 files changed, 42 insertions(+), 24 deletions(-) diff --git a/engine/apps/alerts/models/alert_group_log_record.py b/engine/apps/alerts/models/alert_group_log_record.py index 77d15ab9..dfcabfa7 100644 --- a/engine/apps/alerts/models/alert_group_log_record.py +++ b/engine/apps/alerts/models/alert_group_log_record.py @@ -65,6 +65,7 @@ class AlertGroupLogRecord(models.Model): TYPE_DELETED, TYPE_REGISTERED, # set on creation before having an alert assigned, skip to avoid retries TYPE_ROUTE_ASSIGNED, # set on creation before having an alert assigned, skip to avoid retries + TYPE_ACK_REMINDER_TRIGGERED, # set on acknowledged reminder, no updates to the alert group log ) TYPES_FOR_LICENCE_CALCULATION = ( diff --git a/engine/apps/alerts/paging.py b/engine/apps/alerts/paging.py index 6ff266d4..6a3cadd3 100644 --- a/engine/apps/alerts/paging.py +++ b/engine/apps/alerts/paging.py @@ -1,4 +1,5 @@ import typing +from functools import partial from uuid import uuid4 from django.db import transaction @@ -148,22 +149,27 @@ def direct_paging( title = _construct_title(from_user, team, users) # create alert group if needed - if alert_group is None: - alert_group = _trigger_alert(organization, team, message, title, source_url, grafana_incident_id, from_user) + with transaction.atomic(): + if alert_group is None: + alert_group = _trigger_alert(organization, team, message, title, source_url, grafana_incident_id, from_user) - for u, important in users: - alert_group.log_records.create( - type=AlertGroupLogRecord.TYPE_DIRECT_PAGING, - author=from_user, - reason=f"{from_user.username} paged user {u.username}", - step_specific_info={ - "user": u.public_primary_key, - "important": important, - }, - ) - notify_user_task.apply_async( - (u.pk, alert_group.pk), {"important": important, "notify_even_acknowledged": True, "notify_anyway": True} - ) + for u, important in users: + alert_group.log_records.create( + type=AlertGroupLogRecord.TYPE_DIRECT_PAGING, + author=from_user, + reason=f"{from_user.username} paged user {u.username}", + step_specific_info={ + "user": u.public_primary_key, + "important": important, + }, + ) + transaction.on_commit( + partial( + notify_user_task.apply_async, + (u.pk, alert_group.pk), + {"important": important, "notify_even_acknowledged": True, "notify_anyway": True}, + ) + ) return alert_group diff --git a/engine/apps/alerts/tasks/acknowledge_reminder.py b/engine/apps/alerts/tasks/acknowledge_reminder.py index be6efb05..e0e41500 100644 --- a/engine/apps/alerts/tasks/acknowledge_reminder.py +++ b/engine/apps/alerts/tasks/acknowledge_reminder.py @@ -72,10 +72,11 @@ def acknowledge_reminder_task(alert_group_pk: int, unacknowledge_process_id: str (alert_group.pk, unacknowledge_process_id), countdown=acknowledge_reminder_timeout ) - log_record = alert_group.log_records.create( - type=AlertGroupLogRecord.TYPE_ACK_REMINDER_TRIGGERED, author=alert_group.acknowledged_by_user - ) - transaction.on_commit(partial(send_alert_group_signal.delay, log_record.pk)) + with transaction.atomic(): + log_record = alert_group.log_records.create( + type=AlertGroupLogRecord.TYPE_ACK_REMINDER_TRIGGERED, author=alert_group.acknowledged_by_user + ) + transaction.on_commit(partial(send_alert_group_signal.delay, log_record.pk)) @shared_dedicated_queue_retry_task(autoretry_for=(Exception,), retry_backoff=True, max_retries=MAX_RETRIES) diff --git a/engine/apps/alerts/tests/test_acknowledge_reminder.py b/engine/apps/alerts/tests/test_acknowledge_reminder.py index 759c8c1d..46ba9412 100644 --- a/engine/apps/alerts/tests/test_acknowledge_reminder.py +++ b/engine/apps/alerts/tests/test_acknowledge_reminder.py @@ -159,12 +159,19 @@ def test_acknowledge_reminder_task_skip( @patch.object(acknowledge_reminder_task, "apply_async") @pytest.mark.django_db def test_acknowledge_reminder_task_reschedules_itself( - mock_acknowledge_reminder_task, mock_unacknowledge_timeout_task, ack_reminder_test_setup + mock_acknowledge_reminder_task, + mock_unacknowledge_timeout_task, + ack_reminder_test_setup, + django_capture_on_commit_callbacks, ): organization, alert_group, user = ack_reminder_test_setup( unacknowledge_timeout=Organization.UNACKNOWLEDGE_TIMEOUT_NEVER ) - acknowledge_reminder_task(alert_group.pk, TASK_ID) + with django_capture_on_commit_callbacks(execute=True) as callbacks: + acknowledge_reminder_task(alert_group.pk, TASK_ID) + + # send_alert_group_signal task is queued after commit + assert len(callbacks) == 1 mock_unacknowledge_timeout_task.assert_not_called() mock_acknowledge_reminder_task.assert_called_once_with( diff --git a/engine/apps/alerts/tests/test_paging.py b/engine/apps/alerts/tests/test_paging.py index 2c0ae9f8..05a45720 100644 --- a/engine/apps/alerts/tests/test_paging.py +++ b/engine/apps/alerts/tests/test_paging.py @@ -54,15 +54,16 @@ def test_user_is_oncall(make_organization, make_user_for_organization, make_sche @pytest.mark.django_db -def test_direct_paging_user(make_organization, make_user_for_organization): +def test_direct_paging_user(make_organization, make_user_for_organization, django_capture_on_commit_callbacks): organization = make_organization() user = make_user_for_organization(organization) other_user = make_user_for_organization(organization) from_user = make_user_for_organization(organization) msg = "Fire" - with patch("apps.alerts.paging.notify_user_task") as notify_task: - direct_paging(organization, from_user, msg, users=[(user, False), (other_user, True)]) + with django_capture_on_commit_callbacks(execute=True) as callbacks: + with patch("apps.alerts.paging.notify_user_task") as notify_task: + direct_paging(organization, from_user, msg, users=[(user, False), (other_user, True)]) # alert group created alert_groups = AlertGroup.objects.all() @@ -73,6 +74,8 @@ def test_direct_paging_user(make_organization, make_user_for_organization): assert alert.title == f"{from_user.username} is paging {user.username} and {other_user.username} to join escalation" assert alert.message == msg + # callbacks: distribute_alert + 2 notify_user tasks + assert len(callbacks) == 3 # notifications sent for u, important in ((user, False), (other_user, True)): assert notify_task.apply_async.called_with(