Add cronjob for syncing heartbeats with heartbeat-missing incidents

src/argus/incident/factories.py

@@ -1,5 +1,5 @@
 from zoneinfo import ZoneInfo
-from datetime import datetime
+from datetime import datetime, timedelta
 from random import randint, choice
 from typing import Optional, Any
 
@@ -47,6 +47,54 @@ def update_tags(incident: Incident, *tags):
     return c
 
 
+def create_fake_source(
+    name: Optional[str] = None,
+    type: Optional[str] = None,
+    user: Optional[str] = None,
+    last_seen: Optional[datetime] = None,
+    heartbeat_frequency: Optional[timedelta] = None,
+    commit=True,
+):
+    generated_name = factory.Sequence(lambda s: "source-%s" % s)
+    name = name or type or user or generated_name
+    if not type:
+        type = "type-" + name
+    if not user:
+        user = "user-" + name
+
+    source_user = SourceUserFactory(username=user)
+    source_type = SourceSystemTypeFactory(name=type)
+    source = SourceSystem(
+        name=name,
+        user=source_user,
+        type=source_type,
+        last_seen=last_seen,
+        heartbeat_frequency=heartbeat_frequency,
+    )
+    if commit:
+        source.save()
+    return source
+
+
+def create_dead_source(
+    name,
+    type: str = None,
+    user: str = None,
+    timestamp: Optional[None] = None,
+):
+    timestamp = timestamp if timestamp else timezone.now()
+    heartbeat_frequency = timedelta(minutes=5)
+    last_seen = timestamp - (2 * heartbeat_frequency)
+    source = create_fake_source(
+        name=name,
+        type=type,
+        user=user,
+        last_seen=last_seen,
+        heartbeat_frequency=heartbeat_frequency,
+    )
+    return source, timestamp
+
+
 def create_fake_incident(
     tags: Optional[list[str]] = None,
     description: Optional[str] = None,

src/argus/incident/heartbeat_utils.py

@@ -0,0 +1,112 @@
+import logging
+from datetime import datetime, timedelta
+from typing import Optional
+
+from django.utils import timezone
+
+from argus.incident.factories import create_stateful_incident
+from argus.incident.models import HEARTBEAT_TAG, SOURCE_TAG_KEY, Incident, SourceSystem
+
+
+LOG = logging.getLogger(__name__)
+DEFAULT_FUDGE_FACTOR = 10
+DEFAULT_LEVEL = 4
+_FUDGE = timedelta(seconds=DEFAULT_FUDGE_FACTOR)
+
+__all__ = [
+    "sync_heartbeats_with_heartbeat_incidents",
+]
+
+
+def sync_heartbeats_with_heartbeat_incidents():
+    timestamp = timezone.now()
+    alive_sources = _close_incidents_whose_sources_are_alive_again(timestamp)
+    new_incidents = _create_incidents_for_dead_sources(timestamp)
+    return alive_sources, new_incidents
+
+
+def _close_incidents_whose_sources_are_alive_again(timestamp):
+    argus = SourceSystem.objects.get(name="argus")
+    still_dead_sources = SourceSystem.objects.dead(timestamp - _FUDGE)
+    outdated_incidents = Incident.objects.heartbeat_incidents().open()
+    format_str = 'Source "{}" is again alive and well'
+    sources = []
+    for incident in outdated_incidents:
+        source_ids = incident.get_values_for_tag_key(SOURCE_TAG_KEY)
+        # We need exactly one source-tag.
+        # We don't need to check for multiple *identical* tags,
+        # that's a UniqueError, but we should also only have one source-tag
+        if not source_ids:
+            LOG.error('Heartbeat incident "%s" is missing source-tag!', incident)
+            continue
+        elif len(source_ids) > 1:
+            LOG.error('Heartbeat incident "%s" has multiple source tags!?', incident)
+            continue
+
+        source_id = source_ids[0]
+        if still_dead_sources.filter(pk=source_id).exists():
+            continue
+
+        try:
+            source = SourceSystem.objects.get(pk=source_id)
+        except SourceSystem.DoesNotExist:
+            LOG.error('Heartbeat source "%s" has disappeared!', source_id)
+            continue
+
+        # phew, *now* we can close the incident
+        incident.set_end(argus.user, timestamp, format_str.format(source))
+        sources.append(source)
+    return sources
+
+
+def _create_incidents_for_dead_sources(timestamp: Optional[datetime] = None):
+    # Create existing incidents whose sources have become dead
+    dead_sources = SourceSystem.objects.dead(timestamp - _FUDGE)
+    incident_owner = SourceSystem.objects.get(name="argus")
+    # prevent multiple incidents
+    heartbeat_incidents = Incident.objects.heartbeat_incidents().open()
+    source_ids = (incident.source_id for incident in heartbeat_incidents)
+    new_incidents = []
+    for source in dead_sources.exclude(id__in=source_ids):
+        incident = _get_or_create_incident_for_dead_source(source, incident_owner=incident_owner, timestamp=timestamp)
+        if incident:
+            new_incidents.append(incident)
+    return new_incidents
+
+
+def _get_or_create_incident_for_dead_source(
+    source: SourceSystem,
+    incident_owner: SourceSystem,
+    timestamp: Optional[datetime] = None,
+    level: int = DEFAULT_LEVEL,
+):
+    assert isinstance(source, SourceSystem), "source is not a SourceSystem"
+    if not timestamp:
+        timestamp = timezone.now()
+
+    source_tag = f"{SOURCE_TAG_KEY}={source.pk}"
+    tags = [HEARTBEAT_TAG, source_tag]
+
+    # One last check in case a heartbeat arrived in the meantime
+    source.refresh_from_db()
+    dead = source.is_dead(timestamp - _FUDGE)
+    if dead:
+        # prevent duplicates
+        existing_incidents = Incident.objects.heartbeat_incidents().open().from_tags(source_tag)
+        if existing_incidents.exists():
+            try:
+                return existing_incidents.get()
+            except Incident.MultipleObjectsReturned:
+                # TODO: fix, different function?
+                LOG.error('Source "%s" has multiple open heartbeat incidents', source)
+                return existing_incidents.first()
+        # save new incident
+        incident = create_stateful_incident(
+            f"Missing heartbeat from source {source}, dead?",
+            incident_owner,
+            level,
+            tags=tags,
+            start_time=timestamp,
+        )
+        return incident
+    return None

src/argus/incident/models.py

@@ -20,6 +20,8 @@
 from .validators import validate_key, validate_lowercase
 
 
+SOURCE_TAG_KEY = "source_system_id"
+HEARTBEAT_TAG = "problem_type=missing_heartbeat"
 MINIMUM_DURATION = timedelta(seconds=60)
 MAXIMUM_DURATION = timedelta(days=1)
 LOG = logging.getLogger(__name__)
@@ -48,6 +50,26 @@ def save(self, *args, **kwargs):
         super().save(*args, **kwargs)
 
 
+class SourceSystemQuerySet(models.QuerySet):
+    def has_heartbeat(self):
+        return self.filter(heartbeat_frequency__isnull=False, last_seen__isnull=False)
+
+    def with_next_expected_heartbeat(self, timestamp: Optional[datetime] = None):
+        timestamp = timestamp if timestamp else timezone.now()
+        qs = self.has_heartbeat()
+        qs = qs.annotate(next_heartbeat=F("last_seen") + F("heartbeat_frequency"))
+        return qs
+
+    def dead(self, timestamp: Optional[datetime] = None):
+        """Find sources that have missed heartbeats
+
+        The calculation is done in the database.
+        """
+        timestamp = timestamp if timestamp else timezone.now()
+        qs = self.with_next_expected_heartbeat(timestamp)
+        return qs.filter(next_heartbeat__lt=timestamp)
+
+
 class SourceSystem(models.Model):
     name = models.TextField()
     type = models.ForeignKey(to=SourceSystemType, on_delete=models.PROTECT, related_name="instances")
@@ -67,6 +89,8 @@ class SourceSystem(models.Model):
         ],
     )
 
+    objects = SourceSystemQuerySet.as_manager()
+
     class Meta:
         constraints = [
             models.UniqueConstraint(fields=["name", "type"], name="%(class)s_unique_name_per_type"),
@@ -79,9 +103,16 @@ def update_last_seen(self, timestamp: Optional[datetime] = None):
         self.last_seen = timestamp if timestamp else timezone.now()
         self.save()
 
+    def is_dead(self, timestamp) -> None | bool:
+        """Check if an expected heartbeat of a source is missing"""
+        if not self.heartbeat_frequency:
+            return None
+        dead = self.last_seen + self.heartbeat_frequency < timestamp
+        return dead
+
 
 class TagQuerySet(models.QuerySet):
-    def parse(self, *tags):
+    def parse(self, *tags: str):
         "Return a list of querysets that match `tags`"
         set_dict = defaultdict(set)
         for k, v in (Tag.split(tag) for tag in tags):
@@ -94,6 +125,17 @@ def create_from_tag(self, tag):
         tag, _ = self.get_or_create(key=key, value=value)
         return tag
 
+    def from_tags(self, *tags: str):
+        qss = self.parse(*tags)
+        tagobjs = []
+        for qs in qss:
+            tagobjs.extend(list(qs))
+        return set(tagobjs)
+
+    def from_tag_keys(self, *keys: str):
+        qs = self.filter(key__in=keys)
+        return qs
+
 
 class Tag(models.Model):
     TAG_DELIMITER = "="
@@ -249,6 +291,13 @@ def from_tags(self, *tags):
         qs = reduce(and_, qs)
         return qs.distinct()
 
+    def from_tag_keys(self, *keys):
+        qs = self.filter(incident_tag_relations__tag__key__in=keys)
+        return qs
+
+    def heartbeat_incidents(self):
+        return self.from_tags(HEARTBEAT_TAG)
+
     def is_longer_than_minutes(self, minutes):
         min_duration = timedelta(minutes=minutes)
         open = self.open().annotate(duration=timezone.now() - F("start_time"))
@@ -469,6 +518,27 @@ def acked(self):
 
         return self.events.filter((acks_query & acks_not_expired_query) | ack_is_just_being_created).exists()
 
+    def has_tags(self, *tags: str):
+        tags = Tag.objects.from_tags(*tags)
+        if not tags:
+            return False
+        return bool(set(self.deprecated_tags).issuperset(tags))
+
+    def has_tag_keys(self, *keys: str):
+        return self.incident_tag_relations.only("tag").filter(tag__key__in=keys).exists()
+
+    def is_heartbeat_incident(self):
+        heartbeat_tag = self.has_tags(HEARTBEAT_TAG)
+        source_tag = self.has_tag_keys(SOURCE_TAG_KEY)
+        return heartbeat_tag and source_tag
+
+    def get_values_for_tag_key(self, key):
+        values = []
+        for tag in self.deprecated_tags:
+            if tag.key == key:
+                values.append(tag.value)
+        return values
+
     def is_acked_by(self, group: str) -> bool:
         return group in self.acks.active().group_names()
 
@@ -492,6 +562,7 @@ def create_first_event(self):
 
     # @transaction.atomic
     def set_open(self, actor: User, timestamp: datetime = None, description=""):
+        "Incident reopened by user"
         if not self.stateful:
             raise ValidationError("Cannot set a stateless incident as open")
         if self.open:
@@ -509,6 +580,7 @@ def set_open(self, actor: User, timestamp: datetime = None, description=""):
 
     # @transaction.atomic
     def set_closed(self, actor: User, timestamp: datetime = None, description=""):
+        "Incident closed by user"
         if not self.stateful:
             raise ValidationError("Cannot set a stateless incident as closed")
         if not self.open:
@@ -525,15 +597,22 @@ def set_closed(self, actor: User, timestamp: datetime = None, description=""):
         )
 
     # @transaction.atomic
-    def set_end(self, actor: User):
+    def set_end(self, actor: User, timestamp: datetime = None, description: str = ""):
+        "Incident closed by source"
         if not self.stateful:
             raise ValidationError("Cannot set a stateless incident as ended")
         if not self.open:
             return
 
-        self.end_time = timezone.now()
+        self.end_time = timestamp or timezone.now()
         self.save(update_fields=["end_time"])
-        Event.objects.create(incident=self, actor=actor, timestamp=self.end_time, type=Event.Type.INCIDENT_END)
+        Event.objects.create(
+            incident=self,
+            actor=actor,
+            timestamp=self.end_time,
+            type=Event.Type.INCIDENT_END,
+            description=description,
+        )
 
     # @transaction.atomic
     def create_ack(self, actor: User, timestamp=None, description="", expiration=None):

tests/incident/test_factories.py

@@ -106,7 +106,7 @@ def test_create_fake_incident_does_not_raise_error_on_extra_args(self):
 
 
 def TestCreateStatelessIncident(TestCase):
-    def test_create_stateless_incident(self):
+    def test_it_should_always_return_a_stateless_incident(self):
         source_name = "source_a"
         sst = SourceSystemTypeFactory(name=source_name)
         user = SourceUserFactory(username=source_name)