wip

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/app/pubsub.py

@@ -0,0 +1,87 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def wakeup_worker(pubsub_backend):
+    pubsub_backend.publish(BasePubSubBackend.WORKER_WAKEUP)
+
+
+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:
+    WORKER_WAKEUP = "pulp_worker_wakeup"
+    TASK_CANCELLATION = "pulp_worker_cancel"
+    WORKER_METRIC = "pulp_worker_metrics_heartbeat"
+
+    def subscribe(self, channel, callback):
+        raise NotImplementedError()
+
+    def unsubscribe(self, channel):
+        raise NotImplementedError()
+
+    def publish(self, channel, message=None):
+        raise NotImplementedError()
+
+    def fileno(self):
+        """Add support for being used in select loop."""
+        raise NotImplementedError()
+
+    def fetch(self):
+        """Fetch messages new message, if required."""
+        raise NotImplementedError()
+
+    def close(self):
+        raise NotImplementedError()
+
+
+class PostgresPubSub(BasePubSubBackend):
+
+    def __init__(self, connection):
+        logger.info("Initialized pubsub")
+        self.cursor = connection.cursor()
+        self.connection = connection.connection
+        self.listening_callback = {}
+        self.connection.add_notify_handler(self._notification_handler)
+
+    def _notification_handler(self, notification):
+        logger.info(f"Handling notification: {notification}")
+        callback = self.listening_callback[notification.channel]
+        callback(message=notification.payload)
+
+    def subscribe(self, channel, callback):
+        self.listening_callback[channel] = callback
+        self.cursor.execute(f"LISTEN {channel}")
+
+    def unsubscribe(self, channel):
+        self.cursor.execute(f"UNLISTEN {channel}")
+
+    def publish(self, channel, message=None):
+        logger.info(f"Publish: {channel}: {message}")
+        if not message:
+            self.cursor.execute(f"NOTIFY {channel}")
+        else:
+            self.cursor.execute("SELECT pg_notify(%s, %s)", (channel, message))
+
+    def fileno(self):
+        return self.connection.fileno()
+
+    def fetch(self):
+        logger.info("Fetching")
+        self.connection.execute("SELECT 1").fetchone()
+
+    def close(self):
+        logger.info("Closing")
+        self.cursor.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.cursor.close()

pulpcore/tasking/tasks.py

@@ -19,6 +19,7 @@
 from pulpcore.app.apps import MODULE_PLUGIN_VERSIONS
 from pulpcore.app.models import Task, TaskGroup
 from pulpcore.app.util import current_task, get_domain, get_prn, deprecation_logger
+from pulpcore.app import pubsub
 from pulpcore.constants import (
     TASK_FINAL_STATES,
     TASK_INCOMPLETE_STATES,
@@ -47,12 +48,6 @@ def _validate_and_get_resources(resources):
     return list(resource_set)
 
 
-def wakeup_worker():
-    # Notify workers
-    with connection.connection.cursor() as cursor:
-        cursor.execute("NOTIFY pulp_worker_wakeup")
-
-
 def execute_task(task):
     # This extra stack is needed to isolate the current_task ContextVar
     contextvars.copy_context().run(_execute_task, task)
@@ -308,10 +303,18 @@ def dispatch(
                 task.set_canceling()
                 task.set_canceled(TASK_STATES.CANCELED, "Resources temporarily unavailable.")
     if notify_workers:
-        wakeup_worker()
+        pubsub_backend = pubsub.PostgresPubSub(connection)
+        pubsub.wakeup_worker(pubsub_backend)
+        pubsub_backend.close()
     return task
 
 
+def wakeup_worker():
+    pubsub_backend = pubsub.PostgresPubSub(connection)
+    pubsub.wakeup_worker(pubsub_backend)
+    pubsub_backend.close()
+
+
 def cancel_task(task_id):
     """
     Cancel the task that is represented by the given task_id.
@@ -342,9 +345,10 @@ 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")
+    pubsub_backend = pubsub.PostgresPubSub(connection)
+    pubsub.cancel_task(pubsub_backend, task_pk=task.pk)
+    pubsub.wakeup_worker(pubsub_backend)
+    pubsub_backend.close()
     return task
 
 

pulpcore/tasking/worker.py

@@ -29,6 +29,7 @@
 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.app import pubsub
 from pulpcore.exceptions import AdvisoryLockError
 
 from pulpcore.tasking.storage import WorkerDirectory
@@ -56,70 +57,6 @@
 THRESHOLD_UNBLOCKED_WAITING_TIME = 5
 
 
-class BasePubSubBackend:
-    WORKER_WAKEUP = "pulp_worker_wakeup"
-    TASK_CANCELLATION = "pulp_worker_cancel"
-    WORKER_METRIC = "pulp_worker_metrics_heartbeat"
-
-    def wakeup_workers(self):
-        self.publish(self.WORKER_WAKEUP)
-
-    def cancel_task(self):
-        self.publish(self.TASK_CANCELLATION)
-
-    def record_worker_metrics(self, now):
-        self.publish(self.WORKER_METRIC, str(now))
-
-    # Specific implementation
-    def subscribe(self, channel, callback):
-        raise NotImplementedError()
-
-    def unsubscribe(self, channel):
-        raise NotImplementedError()
-
-    def publish(self, channel, message=None):
-        raise NotImplementedError()
-
-    def fileno(self):
-        """Add support for being used in select loop."""
-        raise NotImplementedError()
-
-    def fetch(self):
-        """Fetch messages new message, if required."""
-        raise NotImplementedError()
-
-
-class PostgresPubSub(BasePubSubBackend):
-
-    def __init__(self):
-        self.cursor = connection.cursor()
-        self.listening_callback = {}
-
-    def _notification_handler(self, notification):
-        callback = self.listening_callback[notification.channel]
-        callback(message=notification.payload)
-
-    def subscribe(self, channel, callback):
-        self.listening_callback[channel] = callback
-        self.cursor.execute(f"LISTEN {channel}")
-        connection.connection.add_notify_handler(self._notification_handler)
-
-    def unsubscribe(self, channel):
-        self.cursor.execute(f"UNLISTEN {channel}")
-
-    def publish(self, channel, message=None):
-        if not message:
-            self.cursor.execute(f"NOTIFY {channel}")
-        else:
-            self.cursor.execute(f"NOTIFY {channel}, {message}")
-
-    def fileno(self):
-        return connection.connection.fileno()
-
-    def fetch(self):
-        connection.connection.execute("SELECT 1")
-
-
 class PulpcoreWorker:
     def __init__(self):
         # Notification states from several signal handlers
@@ -132,7 +69,7 @@ def __init__(self):
         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.pubsub_backend = PostgresPubSub()
+        self.pubsub_backend = pubsub.PostgresPubSub(connection)
         self.worker = self.handle_worker_heartbeat()
         # This defaults to immediate task cancellation.
         # It will be set into the future on moderately graceful worker shutdown,
@@ -277,7 +214,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.pubsub_backend.wakeup_workers()
+            pubsub.wakeup_worker(self.pubsub_backend)
         return True
 
     def is_compatible(self, task):
@@ -524,7 +461,7 @@ def supervise_task(self, task):
             if cancel_state:
                 self.cancel_abandoned_task(task, cancel_state, cancel_reason)
         if task.reserved_resources_record:
-            self.pubsub_backend.wakeup_workers()
+            pubsub.wakeup_worker(self.pubsub_backend)
         self.task = None
 
     def handle_available_tasks(self):
@@ -581,11 +518,11 @@ def _record_unblocked_waiting_tasks_metric(self):
                     unblocked_tasks_stats["longest_unblocked_waiting_time"].seconds
                 )
 
-                self.pubsub_backend.record_worker_metrics(str(now))
+                pubsub.record_worker_metrics(self.pubsub_backend, now)
 
     def pubsub_setup(self):
         def cancellation_callback(message):
-            if message == str(self.task.pk):
+            if self.task and message == str(self.task.pk):
                 self.cancel_task = True
 
         def wakeup_callback(message):