aa-killtracker 0.17.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

killtracker/models/webhooks.py

@@ -1,24 +1,20 @@
 """Webhooks models for killtracker."""
 
-import json
-from typing import List, Optional
+from typing import Optional
 
-import dhooks_lite
 from simple_mq import SimpleMQ
 
-from django.core.cache import cache
 from django.db import models
 from django.utils.translation import gettext_lazy as _
 
 from allianceauth.services.hooks import get_extension_logger
 from app_utils.allianceauth import get_redis_client
-from app_utils.json import JSONDateTimeDecoder, JSONDateTimeEncoder
 from app_utils.logging import LoggerAddTag
 from app_utils.urls import static_file_absolute_url
 
-from killtracker import APP_NAME, HOMEPAGE_URL, __title__, __version__
+from killtracker import __title__
 from killtracker.app_settings import KILLTRACKER_WEBHOOK_SET_AVATAR
-from killtracker.exceptions import WebhookTooManyRequests
+from killtracker.core.discord import DiscordMessage, send_message_to_webhook
 from killtracker.managers import WebhookManager
 
 logger = LoggerAddTag(get_extension_logger(__name__), __title__)
@@ -27,8 +23,6 @@ logger = LoggerAddTag(get_extension_logger(__name__), __title__)
 class Webhook(models.Model):
     """A webhook to receive messages"""
 
-    HTTP_TOO_MANY_REQUESTS = 429
-
     class WebhookType(models.IntegerChoices):
         """A webhook type."""
 
@@ -63,8 +57,8 @@ class Webhook(models.Model):
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        self.main_queue = self._create_queue("main")
-        self.error_queue = self._create_queue("error")
+        self._main_queue = self._create_queue("main")
+        self._error_queue = self._create_queue("error")
 
     def __str__(self) -> str:
         return self.name
@@ -78,8 +72,8 @@ class Webhook(models.Model):
         # method to avoid modifying the original state.
         state = self.__dict__.copy()
         # Remove the unpicklable entries.
-        del state["main_queue"]
-        del state["error_queue"]
+        del state["_main_queue"]
+        del state["_error_queue"]
         return state
 
     def __setstate__(self, state):
@@ -87,15 +81,15 @@ class Webhook(models.Model):
         self.__dict__.update(state)
         # Restore the previously opened file's state. To do so, we need to
         # reopen it and read from it until the line count is restored.
-        self.main_queue = self._create_queue("main")
-        self.error_queue = self._create_queue("error")
+        self._main_queue = self._create_queue("main")
+        self._error_queue = self._create_queue("error")
 
     def save(self, *args, **kwargs):
         is_new = self.id is None  # type: ignore
         super().save(*args, **kwargs)
         if is_new:
-            self.main_queue = self._create_queue("main")
-            self.error_queue = self._create_queue("error")
+            self._main_queue = self._create_queue("main")
+            self._error_queue = self._create_queue("error")
 
     def _create_queue(self, suffix: str) -> Optional[SimpleMQ]:
         redis_client = get_redis_client()
@@ -110,129 +104,67 @@ class Webhook(models.Model):
         returns number of moved messages.
         """
         counter = 0
-        if self.error_queue and self.main_queue:
+        if self._error_queue and self._main_queue:
             while True:
-                message = self.error_queue.dequeue()
+                message = self._error_queue.dequeue()
                 if message is None:
                     break
 
-                self.main_queue.enqueue(message)
+                self._main_queue.enqueue(message)
                 counter += 1
 
         return counter
 
-    def enqueue_message(
-        self,
-        content: Optional[str] = None,
-        embeds: Optional[List[dhooks_lite.Embed]] = None,
-        tts: Optional[bool] = None,
-        username: Optional[str] = None,
-        avatar_url: Optional[str] = None,
-    ) -> int:
-        """Enqueues a message to be send with this webhook"""
-        if not self.main_queue:
+    def enqueue_message(self, message: DiscordMessage, is_error: bool = False) -> int:
+        """Enqueues a discord message to be send with this webhook.
+
+        Returns the updated number of messages in the main queue.
+        """
+        q = self._error_queue if is_error else self._main_queue
+
+        if not q:
             return 0
 
-        username = __title__ if KILLTRACKER_WEBHOOK_SET_AVATAR else username
-        brand_url = static_file_absolute_url("killtracker/killtracker_logo.png")
-        avatar_url = brand_url if KILLTRACKER_WEBHOOK_SET_AVATAR else avatar_url
-        return self.main_queue.enqueue(
-            self._discord_message_asjson(
-                content=content,
-                embeds=embeds,
-                tts=tts,
-                username=username,
-                avatar_url=avatar_url,
-            )
-        )
+        if KILLTRACKER_WEBHOOK_SET_AVATAR:
+            message.username = __title__
 
-    @staticmethod
-    def _discord_message_asjson(
-        content: Optional[str] = None,
-        embeds: Optional[List[dhooks_lite.Embed]] = None,
-        tts: Optional[bool] = None,
-        username: Optional[str] = None,
-        avatar_url: Optional[str] = None,
-    ) -> str:
-        """Converts a Discord message to JSON and returns it
-
-        Raises ValueError if message is incomplete
-        """
-        if not content and not embeds:
-            raise ValueError("Message must have content or embeds to be valid")
-
-        if embeds:
-            embeds_list = [obj.asdict() for obj in embeds]
-        else:
-            embeds_list = None
-
-        message = {}
-        if content:
-            message["content"] = content
-        if embeds_list:
-            message["embeds"] = embeds_list
-        if tts:
-            message["tts"] = tts
-        if username:
-            message["username"] = username
-        if avatar_url:
-            message["avatar_url"] = avatar_url
-
-        return json.dumps(message, cls=JSONDateTimeEncoder)
-
-    def send_message_to_webhook(self, message_json: str) -> dhooks_lite.WebhookResponse:
-        """Send given message to webhook
-
-        Params
-            message_json: Discord message encoded in JSON
+        if KILLTRACKER_WEBHOOK_SET_AVATAR:
+            brand_url = static_file_absolute_url("killtracker/killtracker_logo.png")
+            message.avatar_url = brand_url
+
+        return q.enqueue(message.to_json())
+
+    def dequeue_message(self, is_error: bool = False) -> Optional[DiscordMessage]:
+        """Dequeues a message from the main queue and return it.
+
+        Returns None if the queue is empty.
         """
-        timeout = cache.ttl(self._blocked_cache_key())  # type: ignore
-        if timeout:
-            raise WebhookTooManyRequests(timeout)
-
-        message = json.loads(message_json, cls=JSONDateTimeDecoder)
-        if message.get("embeds"):
-            embeds = [
-                dhooks_lite.Embed.from_dict(embed_dict)
-                for embed_dict in message.get("embeds")
-            ]
-        else:
-            embeds = None
-        hook = dhooks_lite.Webhook(
-            url=self.url,
-            user_agent=dhooks_lite.UserAgent(
-                name=APP_NAME, url=HOMEPAGE_URL, version=__version__
-            ),
-        )
-        response = hook.execute(
-            content=message.get("content"),
-            embeds=embeds,
-            username=message.get("username"),
-            avatar_url=message.get("avatar_url"),
-            wait_for_response=True,
-            max_retries=0,  # we will handle retries ourselves
-        )
-        logger.debug("headers: %s", response.headers)
-        logger.debug("status_code: %s", response.status_code)
-        logger.debug("content: %s", response.content)
-        if response.status_code == self.HTTP_TOO_MANY_REQUESTS:
-            logger.error(
-                "%s: Received too many requests error from API: %s",
-                self,
-                response.content,
-            )
-            try:
-                retry_after = int(response.headers["Retry-After"]) + 2
-            except (ValueError, KeyError):
-                retry_after = WebhookTooManyRequests.DEFAULT_RESET_AFTER
-            cache.set(
-                key=self._blocked_cache_key(), value="BLOCKED", timeout=retry_after
-            )
-            raise WebhookTooManyRequests(retry_after)
-        return response
-
-    def _blocked_cache_key(self) -> str:
-        return f"{__title__}_webhook_{self.pk}_blocked"
+        q = self._error_queue if is_error else self._main_queue
+        s = q.dequeue()
+        if not s:
+            return None
+
+        return DiscordMessage.from_json(s)
+
+    def messages_queued(self, is_error: bool = False) -> int:
+        """Returns how many message are currently in the queue."""
+        q = self._error_queue if is_error else self._main_queue
+        if not q:
+            return 0
+
+        return q.size()
+
+    def delete_queued_messages(self, is_error: bool = False) -> int:
+        """Deletes all messages in a queue and returns how many messages where deleted."""
+        q = self._error_queue if is_error else self._main_queue
+        if not q:
+            return 0
+
+        return q.clear()
+
+    def send_message(self, message: DiscordMessage) -> int:
+        """Send a message to the webhook."""
+        return send_message_to_webhook(name=self.name, url=self.url, message=message)
 
     @staticmethod
     def create_message_link(name: str, url: str) -> str:

killtracker/providers.py

@@ -5,7 +5,7 @@ from esi.clients import EsiClientProvider
 from allianceauth.services.hooks import get_extension_logger
 from app_utils.logging import LoggerAddTag
 
-from . import USER_AGENT_TEXT, __title__
+from killtracker import USER_AGENT_TEXT, __title__
 
 logger = LoggerAddTag(get_extension_logger(__name__), __title__)
 

killtracker/signals.py

@@ -0,0 +1,31 @@
+"""Module signals connects to celery signals."""
+
+# pylint: disable=missing-function-docstring
+
+from celery import signals
+
+from allianceauth.services.hooks import get_extension_logger
+from app_utils.logging import LoggerAddTag
+
+from killtracker import __title__
+from killtracker.core import workers
+
+logger = LoggerAddTag(get_extension_logger(__name__), __title__)
+
+
+@signals.worker_ready.connect
+def worker_ready_handler(sender, **kwargs):
+    workers.state_reset(sender.hostname)
+    logger.debug("worker_ready: %s", sender.hostname)
+
+
+@signals.worker_shutting_down.connect
+def worker_shutting_down_handler(sender, **kwargs):
+    workers.state_set(sender)
+    logger.debug("worker_shutting_down: %s", sender)
+
+
+@signals.worker_shutdown.connect
+def worker_shutdown_handler(sender, **kwargs):
+    workers.state_reset(sender.hostname)
+    logger.debug("worker_shutdown: %s", sender.hostname)

killtracker/tasks.py

@@ -1,8 +1,8 @@
 """Tasks for killtracker."""
 
-from datetime import timedelta
+import time
 
-from celery import chain, shared_task
+from celery import Task, chain, shared_task
 
 from django.db import IntegrityError
 from django.utils.timezone import now
@@ -11,58 +11,82 @@ from eveuniverse.tasks import update_unresolved_eve_entities
 
 from allianceauth.services.hooks import get_extension_logger
 from allianceauth.services.tasks import QueueOnce
-from app_utils.caching import cached_queryset
 from app_utils.esi import retry_task_if_esi_is_down
 from app_utils.logging import LoggerAddTag
 
-from . import APP_NAME, __title__
-from .app_settings import (
+from killtracker import __title__
+from killtracker.app_settings import (
     KILLTRACKER_DISCORD_SEND_DELAY,
     KILLTRACKER_GENERATE_MESSAGE_MAX_RETRIES,
     KILLTRACKER_GENERATE_MESSAGE_RETRY_COUNTDOWN,
     KILLTRACKER_MAX_KILLMAILS_PER_RUN,
+    KILLTRACKER_MAX_MESSAGES_SENT_PER_RUN,
     KILLTRACKER_PURGE_KILLMAILS_AFTER_DAYS,
+    KILLTRACKER_RUN_TIMEOUT,
     KILLTRACKER_STORING_KILLMAILS_ENABLED,
     KILLTRACKER_TASK_OBJECTS_CACHE_TIMEOUT,
     KILLTRACKER_TASKS_TIMEOUT,
-    KILLTRACKER_ZKB_REQUEST_DELAY,
 )
-from .core.killmails import Killmail
-from .exceptions import WebhookTooManyRequests
-from .models import EveKillmail, Tracker, Webhook
+from killtracker.core import workers
+from killtracker.core.discord import (
+    DiscordMessage,
+    HTTPError,
+    WebhookRateLimitExhausted,
+)
+from killtracker.core.zkb import Killmail, KillmailDoesNotExist, ZKBTooManyRequestsError
+from killtracker.models import EveKillmail, Tracker, Webhook
 
 logger = LoggerAddTag(get_extension_logger(__name__), __title__)
 
 
-@shared_task(timeout=KILLTRACKER_TASKS_TIMEOUT)
-def run_killtracker(runs: int = 0) -> None:
-    """Main task for running the Killtracker.
+@shared_task(bind=True, base=QueueOnce, timeout=KILLTRACKER_TASKS_TIMEOUT)
+def run_killtracker(self: Task) -> int:
+    """Fetches and processes new killmails from ZKB API
+    and returns how many killmails were processed.
 
-    Will fetch new killmails from ZKB and start running trackers for them
+    This is the main periodic task for running Killtracker.
     """
     if not is_esi_online():
         logger.warning("ESI is currently offline. Aborting")
-        return
+        return 0
 
-    if runs == 0:
-        logger.info("Killtracker run started...")
-        qs = cached_queryset(
-            Webhook.objects.filter(is_enabled=True),
-            key=f"{APP_NAME}_enabled_webhooks",
-            timeout=KILLTRACKER_TASK_OBJECTS_CACHE_TIMEOUT,
-        )
-        for webhook in qs:
-            webhook.reset_failed_messages()
+    for webhook in Webhook.objects.filter(is_enabled=True):
+        webhook.reset_failed_messages()
+
+    started = time.time()
+
+    def is_timed_out() -> bool:
+        elapsed = time.time() - started
+        return KILLTRACKER_RUN_TIMEOUT - elapsed <= 0
 
-    killmail = Killmail.create_from_zkb_redisq()
-    if killmail:
+    killmails_count = 0
+    for _ in range(KILLTRACKER_MAX_KILLMAILS_PER_RUN):
+        if is_timed_out():
+            break
+
+        if workers.is_shutting_down(self):
+            logger.debug("Aborting due to worker shutdown")
+            break
+
+        killmail = None
+        try:
+            killmail = Killmail.create_from_zkb_redisq()
+        except ZKBTooManyRequestsError as ex:
+            seconds = (ex.retry_at - now()).total_seconds()
+            if seconds < 0:
+                break
+
+            logger.warning(
+                "Killtracker has been baned from ZKB API for %f seconds", seconds
+            )
+            raise self.retry(countdown=seconds)
+
+        if not killmail:
+            break
+
+        killmails_count += 1
         killmail.save()
-        qs = cached_queryset(
-            Tracker.objects.filter(is_enabled=True),
-            key=f"{APP_NAME}_enabled_trackers",
-            timeout=KILLTRACKER_TASK_OBJECTS_CACHE_TIMEOUT,
-        )
-        for tracker in qs:
+        for tracker in Tracker.objects.filter(is_enabled=True):
             run_tracker.delay(tracker_pk=tracker.pk, killmail_id=killmail.id)
 
         if KILLTRACKER_STORING_KILLMAILS_ENABLED:
@@ -71,59 +95,64 @@ def run_killtracker(runs: int = 0) -> None:
                 update_unresolved_eve_entities.si(),
             ).delay()
 
-    total_killmails = runs + (1 if killmail else 0)
-    if killmail and total_killmails < KILLTRACKER_MAX_KILLMAILS_PER_RUN:
-        run_killtracker.apply_async(
-            kwargs={"runs": runs + 1},
-            eta=now() + timedelta(milliseconds=KILLTRACKER_ZKB_REQUEST_DELAY),
-        )
-    else:
-        if (
-            KILLTRACKER_STORING_KILLMAILS_ENABLED
-            and KILLTRACKER_PURGE_KILLMAILS_AFTER_DAYS > 0
-        ):
-            delete_stale_killmails.delay()
+    elapsed = time.time() - started
+    logger.info(
+        "Killtracker processed %d new killmails from ZKB in %f seconds",
+        killmails_count,
+        elapsed,
+    )
 
-        logger.info(
-            "Killtracker runs completed. %d killmails received from ZKB",
-            total_killmails,
-        )
+    if (
+        KILLTRACKER_STORING_KILLMAILS_ENABLED
+        and KILLTRACKER_PURGE_KILLMAILS_AFTER_DAYS > 0
+    ):
+        delete_stale_killmails.delay()
+
+    return killmails_count
 
 
 @shared_task(bind=True, max_retries=None)
 def run_tracker(
-    self, tracker_pk: int, killmail_id: int, ignore_max_age: bool = False
+    self: Task, tracker_pk: int, killmail_id: int, ignore_max_age: bool = False
 ) -> None:
     """Run tracker for given killmail and trigger sending if needed."""
     retry_task_if_esi_is_down(self)
-    tracker = Tracker.objects.get_cached(
+    tracker: Tracker = Tracker.objects.get_cached(
         pk=tracker_pk,
         select_related="webhook",
         timeout=KILLTRACKER_TASK_OBJECTS_CACHE_TIMEOUT,
     )
-    logger.info("%s: Started running tracker", tracker)
-    killmail = Killmail.get(killmail_id)
+    try:
+        killmail = Killmail.get(killmail_id)
+    except KillmailDoesNotExist as ex:
+        logger.error("Aborting. %s", ex)
+        return
+
     killmail_new = tracker.process_killmail(
         killmail=killmail, ignore_max_age=ignore_max_age
     )
     if killmail_new:
+        logger.info("%s: Killmail %d matches", tracker, killmail_id)
         killmail_new.save()
         generate_killmail_message.delay(tracker_pk=tracker_pk, killmail_id=killmail_id)
-    elif tracker.webhook.main_queue.size():
+    elif tracker.webhook.messages_queued():
         send_messages_to_webhook.delay(webhook_pk=tracker.webhook.pk)
 
 
 @shared_task(bind=True, max_retries=None)
-def generate_killmail_message(self, tracker_pk: int, killmail_id: int) -> None:
+def generate_killmail_message(self: Task, tracker_pk: int, killmail_id: int) -> None:
     """Generate and enqueue message from given killmail and start sending."""
     retry_task_if_esi_is_down(self)
-    tracker = Tracker.objects.get_cached(
+    tracker: Tracker = Tracker.objects.get_cached(
         pk=tracker_pk,
         select_related="webhook",
         timeout=KILLTRACKER_TASK_OBJECTS_CACHE_TIMEOUT,
     )
-    killmail = Killmail.get(killmail_id)
-    logger.info("%s: Generating message from killmail %s", tracker, killmail.id)
+    try:
+        killmail = Killmail.get(killmail_id)
+    except KillmailDoesNotExist as ex:
+        logger.error("Aborting. %s", ex)
+        return
     try:
         tracker.generate_killmail_message(killmail)
     except Exception as ex:
@@ -142,12 +171,19 @@ def generate_killmail_message(self, tracker_pk: int, killmail_id: int) -> None:
         )
 
     send_messages_to_webhook.delay(webhook_pk=tracker.webhook.pk)
+    logger.info(
+        "%s: Added message from killmail %s to send queue", tracker, killmail.id
+    )
 
 
 @shared_task(timeout=KILLTRACKER_TASKS_TIMEOUT)
 def store_killmail(killmail_id: int) -> None:
-    """stores killmail as EveKillmail object"""
-    killmail = Killmail.get(killmail_id)
+    """Stores killmail as EveKillmail object."""
+    try:
+        killmail = Killmail.get(killmail_id)
+    except KillmailDoesNotExist as ex:
+        logger.error("Aborting. %s", ex)
+        return
     try:
         EveKillmail.objects.create_from_killmail(killmail, resolve_ids=False)
     except IntegrityError:
@@ -160,63 +196,74 @@ def store_killmail(killmail_id: int) -> None:
 
 @shared_task(timeout=KILLTRACKER_TASKS_TIMEOUT)
 def delete_stale_killmails() -> None:
-    """deleted all EveKillmail objects that are considered stale"""
+    """Deletes all EveKillmail objects that are considered stale."""
     _, details = EveKillmail.objects.delete_stale()
     if details:
         logger.info("Deleted %d stale killmails", details["killtracker.EveKillmail"])
 
 
 @shared_task(
-    bind=True,
-    base=QueueOnce,  # celery_once locks stay intact during retries
-    timeout=KILLTRACKER_TASKS_TIMEOUT,
-    retry_backoff=False,
-    max_retries=None,
+    bind=True, base=QueueOnce, timeout=KILLTRACKER_TASKS_TIMEOUT, max_retries=None
 )
-def send_messages_to_webhook(self, webhook_pk: int) -> None:
-    """send all queued messages to given Webhook"""
+def send_messages_to_webhook(self: Task, webhook_pk: int) -> None:
+    """Sends queued messages to a webhook.
 
-    webhook = Webhook.objects.get_cached(
-        pk=webhook_pk,
-        timeout=KILLTRACKER_TASK_OBJECTS_CACHE_TIMEOUT,
-    )
+    Note: This task will retry after processing a set number of messages
+    to avoid running potentially forever.
+    """
+
+    webhook: Webhook = Webhook.objects.get(pk=webhook_pk)
     if not webhook.is_enabled:
         logger.info("%s: Webhook is disabled - aborting", webhook)
         return
 
-    message = webhook.main_queue.dequeue()
-    if message:
-        logger.info("%s: Sending message to webhook", webhook)
+    for _ in range(KILLTRACKER_MAX_MESSAGES_SENT_PER_RUN):
+        if workers.is_shutting_down(self):
+            logger.debug("Aborting due to worker shutdown")
+            return
+
+        message = webhook.dequeue_message()
+        if not message:
+            logger.debug("%s: No more messages to send for webhook", webhook)
+            break
+
         try:
-            response = webhook.send_message_to_webhook(message)
-        except WebhookTooManyRequests as ex:
-            webhook.main_queue.enqueue(message)
+            message_id = webhook.send_message(message)
+
+        except WebhookRateLimitExhausted as ex:
+            webhook.enqueue_message(message)
             logger.warning(
-                "%s: Too many requests for webhook. Blocked for %s seconds. Aborting.",
-                webhook,
-                ex.retry_after,
+                "%s: Webhook temporarily blocked. Retrying at %s.", webhook, ex.retry_at
             )
-            return
+            raise self.retry(eta=ex.retry_at)
 
-        if not response.status_ok:
-            webhook.error_queue.enqueue(message)
+        except HTTPError as ex:
+            webhook.enqueue_message(message, is_error=True)
             logger.warning(
-                "%s: Failed to send message to webhook, will retry. "
-                "HTTP status code: %d, response: %s",
+                "%s: Failed to send message for Killmail %d to webhook, will retry. "
+                "HTTP status code: %d",
                 webhook,
-                response.status_code,
-                response.content,
+                message.killmail_id,
+                ex.status_code,
             )
+            continue
 
-        raise self.retry(countdown=KILLTRACKER_DISCORD_SEND_DELAY)
+        logger.info(
+            "%s: Discord message %s created for killmail %d",
+            webhook,
+            message_id,
+            message.killmail_id,
+        )
 
-    logger.debug("%s: No more messages to send for webhook", webhook)
+    if webhook.messages_queued() > 0:
+        raise self.retry(countdown=KILLTRACKER_DISCORD_SEND_DELAY)
 
 
 @shared_task(timeout=KILLTRACKER_TASKS_TIMEOUT)
 def send_test_message_to_webhook(webhook_pk: int, count: int = 1) -> None:
-    """send a test message to given webhook.
-    Optional inform user about result if user ok is given
+    """Send a test message to given webhook.
+
+    Optional inform user about result if user ok is given.
     """
     try:
         webhook = Webhook.objects.get(pk=webhook_pk)
@@ -224,8 +271,10 @@ def send_test_message_to_webhook(webhook_pk: int, count: int = 1) -> None:
         logger.error("Webhook with pk = %s does not exist", webhook_pk)
         return
 
-    logger.info("Sending %s test messages to webhook %s", count, webhook)
     for num in range(count):
         num_str = f"{num+1}/{count} " if count > 1 else ""
-        webhook.enqueue_message(content=f"Test message {num_str}from {__title__}.")
+        message = DiscordMessage(content=f"Test message {num_str}from {__title__}.")
+        webhook.enqueue_message(message)
+
     send_messages_to_webhook.delay(webhook.pk)
+    logger.info("%s test messages submitted to webhook %s", count, webhook)