Refactor listen/notify as a pubsub backend

Author: pedro-psb

pulpcore/app/models/task.py

@@ -89,7 +89,7 @@ class Task(BaseModel, AutoAddObjPermsMixin):
 
     The transitions to CANCELING (marked with *) are the only ones allowed to happen without
     holding the tasks advisory lock. Canceling is meant to be initiated asyncronously by a sparate
-    process before signalling the worker via Postgres LISTEN.
+    process before signalling the worker via a pubsub notification (e.g, Postgres LISTEN).
 
     Fields:
 

pulpcore/constants.py

@@ -16,6 +16,13 @@
 TASK_WAKEUP_UNBLOCK = "unblock"
 TASK_WAKEUP_HANDLE = "handle"
 
+#: All valid tasking pubsub channels
+TASK_PUBSUB = SimpleNamespace(
+    WAKEUP_WORKER="pulp_worker_wakeup",
+    CANCEL_TASK="pulp_worker_cancel",
+    WORKER_METRICS="pulp_worker_metrics_heartbeat",
+)
+
 #: All valid task states.
 TASK_STATES = SimpleNamespace(
     WAITING="waiting",

pulpcore/tasking/pubsub.py

@@ -1,28 +1,24 @@
 from typing import NamedTuple
+from pulpcore.constants import TASK_PUBSUB
 import os
 import logging
 from contextlib import suppress
 
 logger = logging.getLogger(__name__)
 
 
-def wakeup_worker(pubsub_backend, reason="unknown"):
-    pubsub_backend.publish(BasePubSubBackend.WORKER_WAKEUP, reason)
-
-
-def cancel_task(task_pk, pubsub_backend):
-    pubsub_backend.publish(BasePubSubBackend.TASK_CANCELLATION, str(task_pk))
-
-
-def record_worker_metrics(pubsub_backend, now):
-    pubsub_backend.publish(BasePubSubBackend.WORKER_METRIC, str(now))
+class BasePubSubBackend:
+    # Utils
+    def wakeup_worker(self, reason="unknown"):
+        self.publish(TASK_PUBSUB.WAKEUP_WORKER, reason)
 
+    def cancel_task(self, task_pk):
+        self.publish(TASK_PUBSUB.CANCEL_TASK, str(task_pk))
 
-class BasePubSubBackend:
-    WORKER_WAKEUP = "pulp_worker_wakeup"
-    TASK_CANCELLATION = "pulp_worker_cancel"
-    WORKER_METRIC = "pulp_worker_metrics_heartbeat"
+    def record_worker_metrics(self, now):
+        self.publish(TASK_PUBSUB.WORKER_METRICS, str(now))
 
+    # Interface
     def subscribe(self, channel, callback):
         raise NotImplementedError()
 

pulpcore/tasking/tasks.py

@@ -28,6 +28,7 @@
     TASK_WAKEUP_UNBLOCK,
 )
 from pulpcore.middleware import x_task_diagnostics_var
+from pulpcore.tasking import pubsub
 from pulpcore.tasking.kafka import send_task_notification
 
 _logger = logging.getLogger(__name__)
@@ -48,12 +49,6 @@ def _validate_and_get_resources(resources):
     return list(resource_set)
 
 
-def wakeup_worker(reason="unknown"):
-    # Notify workers
-    with connection.connection.cursor() as cursor:
-        cursor.execute("SELECT pg_notify('pulp_worker_wakeup', %s)", (reason,))
-
-
 def execute_task(task):
     # This extra stack is needed to isolate the current_task ContextVar
     contextvars.copy_context().run(_execute_task, task)
@@ -309,7 +304,8 @@ def dispatch(
                 task.set_canceling()
                 task.set_canceled(TASK_STATES.CANCELED, "Resources temporarily unavailable.")
     if notify_workers:
-        wakeup_worker(TASK_WAKEUP_UNBLOCK)
+        with pubsub.PostgresPubSub(connection) as pubsub_client:
+            pubsub_client.wakeup_worker(reason=TASK_WAKEUP_UNBLOCK)
     return task
 
 
@@ -343,9 +339,9 @@ def cancel_task(task_id):
     # This is the only valid transition without holding the task lock
     task.set_canceling()
     # Notify the worker that might be running that task and other workers to clean up
-    with connection.cursor() as cursor:
-        cursor.execute("SELECT pg_notify('pulp_worker_cancel', %s)", (str(task.pk),))
-        cursor.execute("NOTIFY pulp_worker_wakeup")
+    with pubsub.PostgresPubSub(connection) as pubsub_client:
+        pubsub_client.cancel_task(task_pk=task.pk)
+        pubsub_client.wakeup_worker()
     return task
 
 

pulpcore/tasking/worker.py

@@ -26,13 +26,15 @@
     TASK_METRICS_HEARTBEAT_LOCK,
     TASK_WAKEUP_UNBLOCK,
     TASK_WAKEUP_HANDLE,
+    TASK_PUBSUB,
 )
 from pulpcore.metrics import init_otel_meter
 from pulpcore.app.apps import pulp_plugin_configs
 from pulpcore.app.models import Worker, Task, ApiAppStatus, ContentAppStatus
 from pulpcore.app.util import PGAdvisoryLock, get_domain
 from pulpcore.exceptions import AdvisoryLockError
 
+from pulpcore.tasking import pubsub
 from pulpcore.tasking.storage import WorkerDirectory
 from pulpcore.tasking._util import (
     delete_incomplete_resources,
@@ -72,7 +74,6 @@ def __init__(self, auxiliary=False):
         self.heartbeat_period = timedelta(seconds=settings.WORKER_TTL / 3)
         self.last_metric_heartbeat = timezone.now()
         self.versions = {app.label: app.version for app in pulp_plugin_configs()}
-        self.cursor = connection.cursor()
         self.worker = self.handle_worker_heartbeat()
         # This defaults to immediate task cancellation.
         # It will be set into the future on moderately graceful worker shutdown,
@@ -81,6 +82,9 @@ def __init__(self, auxiliary=False):
         self.worker_cleanup_countdown = random.randint(
             int(WORKER_CLEANUP_INTERVAL / 10), WORKER_CLEANUP_INTERVAL
         )
+        # Pubsub handling
+        self.pubsub_client = pubsub.PostgresPubSub(connection)
+        self.pubsub_channel_callback = {}
 
         # Add a file descriptor to trigger select on signals
         self.sentinel, sentinel_w = os.pipe()
@@ -127,25 +131,6 @@ def _signal_handler(self, thesignal, frame):
             )
         self.shutdown_requested = True
 
-    def _pg_notify_handler(self, notification):
-        if notification.channel == "pulp_worker_wakeup":
-            if notification.payload == TASK_WAKEUP_UNBLOCK:
-                # Auxiliary workers don't do this.
-                self.wakeup_unblock = not self.auxiliary
-            elif notification.payload == TASK_WAKEUP_HANDLE:
-                self.wakeup_handle = True
-            else:
-                _logger.warn("Unknown wakeup call recieved. Reason: '%s'", notification.payload)
-                # We cannot be sure so assume everything happened.
-                self.wakeup_unblock = not self.auxiliary
-                self.wakeup_handle = True
-
-        elif notification.channel == "pulp_worker_metrics_heartbeat":
-            self.last_metric_heartbeat = datetime.fromisoformat(notification.payload)
-        elif self.task and notification.channel == "pulp_worker_cancel":
-            if notification.payload == str(self.task.pk):
-                self.cancel_task = True
-
     def handle_worker_heartbeat(self):
         """
         Create or update worker heartbeat records.
@@ -205,9 +190,6 @@ def beat(self):
                 # to be able to report on a congested tasking system to produce reliable results.
                 self.record_unblocked_waiting_tasks_metric()
 
-    def notify_workers(self, reason="unknown"):
-        self.cursor.execute("SELECT pg_notify('pulp_worker_wakeup', %s)", (reason,))
-
     def cancel_abandoned_task(self, task, final_state, reason=None):
         """Cancel and clean up an abandoned task.
 
@@ -240,7 +222,7 @@ def cancel_abandoned_task(self, task, final_state, reason=None):
         delete_incomplete_resources(task)
         task.set_canceled(final_state=final_state, reason=reason)
         if task.reserved_resources_record:
-            self.notify_workers(TASK_WAKEUP_UNBLOCK)
+            self.pubsub_client.wakeup_worker(reason=TASK_WAKEUP_UNBLOCK)
         return True
 
     def is_compatible(self, task):
@@ -280,7 +262,7 @@ def unblock_tasks(self):
         self.wakeup_unblock_tasks = False
         with contextlib.suppress(AdvisoryLockError), PGAdvisoryLock(TASK_UNBLOCKING_LOCK):
             if count := self._unblock_tasks():
-                self.notify_workers(TASK_WAKEUP_HANDLE)
+                self.pubsub_client.wakeup_worker(TASK_WAKEUP_HANDLE)
         return count
 
     def _unblock_tasks(self):
@@ -405,11 +387,11 @@ def sleep(self):
         _logger.debug(_("Worker %s entering sleep state."), self.name)
         while not self.shutdown_requested and not self.wakeup_handle:
             r, w, x = select.select(
-                [self.sentinel, connection.connection], [], [], self.heartbeat_period.seconds
+                [self.sentinel, self.pubsub_client], [], [], self.heartbeat_period.seconds
             )
             self.beat()
-            if connection.connection in r:
-                connection.connection.execute("SELECT 1")
+            if self.pubsub_client in r:
+                self.pubsub_handle_messages(self.pubsub_client.fetch())
                 if self.wakeup_unblock:
                     self.unblock_tasks()
             if self.sentinel in r:
@@ -447,14 +429,14 @@ def supervise_task(self, task):
                         os.kill(task_process.pid, signal.SIGUSR1)
 
                 r, w, x = select.select(
-                    [self.sentinel, connection.connection, task_process.sentinel],
+                    [self.sentinel, self.pubsub_client, task_process.sentinel],
                     [],
                     [],
                     self.heartbeat_period.seconds,
                 )
                 self.beat()
-                if connection.connection in r:
-                    connection.connection.execute("SELECT 1")
+                if self.pubsub_client in r:
+                    self.pubsub_handle_messages(self.pubsub_client.fetch())
                     if self.cancel_task:
                         _logger.info(
                             _("Received signal to cancel current task %s in domain: %s."),
@@ -506,7 +488,7 @@ def supervise_task(self, task):
             if cancel_state:
                 self.cancel_abandoned_task(task, cancel_state, cancel_reason)
         if task.reserved_resources_record:
-            self.notify_workers(TASK_WAKEUP_UNBLOCK)
+            self.pubsub_client.wakeup_worker(reason=TASK_WAKEUP_UNBLOCK)
         self.task = None
 
     def handle_unblocked_tasks(self):
@@ -559,26 +541,60 @@ def _record_unblocked_waiting_tasks_metric(self):
                     unblocked_tasks_stats["longest_unblocked_waiting_time"].seconds
                 )
 
-                self.cursor.execute(f"NOTIFY pulp_worker_metrics_heartbeat, '{str(now)}'")
+                self.pubsub_client.record_worker_metrics(now)
+
+    def pubsub_handle_messages(self, messages: pubsub.PubsubMessage):
+        for message in messages:
+            callback = self.pubsub_channel_callback[message.channel]
+            callback(message.payload)
+
+    def pubsub_setup(self):
+        def cancellation_callback(message):
+            if self.task and message == str(self.task.pk):
+                self.cancel_task = True
+
+        def wakeup_callback(message):
+            if message == TASK_WAKEUP_UNBLOCK:
+                # Auxiliary workers don't do this.
+                self.wakeup_unblock = not self.auxiliary
+            elif message == TASK_WAKEUP_HANDLE:
+                self.wakeup_handle = True
+            else:
+                _logger.warn("Unknown wakeup call recieved. Reason: '%s'", message)
+                # We cannot be sure so assume everything happened.
+                self.wakeup_unblock = not self.auxiliary
+                self.wakeup_handle = True
+
+        def metric_callback(message):
+            self.last_metric_heartbeat = datetime.fromisoformat(message)
+
+        self.pubsub_client.subscribe(TASK_PUBSUB.WAKEUP_WORKER)
+        self.pubsub_channel_callback[TASK_PUBSUB.WAKEUP_WORKER] = wakeup_callback
+        self.pubsub_client.subscribe(TASK_PUBSUB.CANCEL_TASK)
+        self.pubsub_channel_callback[TASK_PUBSUB.CANCEL_TASK] = cancellation_callback
+        self.pubsub_client.subscribe(TASK_PUBSUB.WORKER_METRICS)
+        self.pubsub_channel_callback[TASK_PUBSUB.WORKER_METRICS] = metric_callback
+
+    def pubsub_teardown(self):
+        self.pubsub_client.unsubscribe(TASK_PUBSUB.WAKEUP_WORKER)
+        self.pubsub_client.unsubscribe(TASK_PUBSUB.CANCEL_TASK)
+        self.pubsub_client.unsubscribe(TASK_PUBSUB.WORKER_METRICS)
 
     def run(self, burst=False):
         with WorkerDirectory(self.name):
             signal.signal(signal.SIGINT, self._signal_handler)
             signal.signal(signal.SIGTERM, self._signal_handler)
             signal.signal(signal.SIGHUP, self._signal_handler)
-            # Subscribe to pgsql channels
-            connection.connection.add_notify_handler(self._pg_notify_handler)
-            self.cursor.execute("LISTEN pulp_worker_cancel")
-            self.cursor.execute("LISTEN pulp_worker_metrics_heartbeat")
+            self.pubsub_setup()
             if burst:
+                self.pubsub_client.unsubscribe(self.pubsub_client.WORKER_WAKEUP)
                 if not self.auxiliary:
                     # Attempt to flush the task queue completely.
                     # Stop iteration if no new tasks were found to unblock.
                     while self.unblock_tasks():
                         self.handle_unblocked_tasks()
                 self.handle_unblocked_tasks()
             else:
-                self.cursor.execute("LISTEN pulp_worker_wakeup")
                 while not self.shutdown_requested:
                     # do work
                     if self.shutdown_requested:
@@ -588,7 +604,5 @@ def run(self, burst=False):
                         break
                     # rest until notified to wakeup
                     self.sleep()
-                self.cursor.execute("UNLISTEN pulp_worker_wakeup")
-            self.cursor.execute("UNLISTEN pulp_worker_metrics_heartbeat")
-            self.cursor.execute("UNLISTEN pulp_worker_cancel")
+            self.pubsub_teardown()
             self.shutdown()

pulpcore/tests/functional/test_pubsub.py

@@ -1,12 +1,7 @@
-import django
-
-django.setup()
-
 from django.db import connection
 from pulpcore.tasking import pubsub
 from types import SimpleNamespace
 import select
-import time
 import pytest