Refactor listen/notify as a pubsub backend

Author: pedro-psb

pulpcore/tasking/worker.py

@@ -56,6 +56,70 @@
 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
@@ -68,7 +132,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.cursor = connection.cursor()
+        self.pubsub_backend = PostgresPubSub()
         self.worker = self.handle_worker_heartbeat()
         # This defaults to immediate task cancellation.
         # It will be set into the future on moderately graceful worker shutdown,
@@ -123,15 +187,6 @@ def _signal_handler(self, thesignal, frame):
             )
         self.shutdown_requested = True
 
-    def _pg_notify_handler(self, notification):
-        if notification.channel == "pulp_worker_wakeup":
-            self.wakeup = 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.
@@ -190,9 +245,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):
-        self.cursor.execute("NOTIFY pulp_worker_wakeup")
-
     def cancel_abandoned_task(self, task, final_state, reason=None):
         """Cancel and clean up an abandoned task.
 
@@ -225,7 +277,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()
+            self.pubsub_backend.wakeup_workers()
         return True
 
     def is_compatible(self, task):
@@ -369,11 +421,11 @@ def sleep(self):
         _logger.debug(_("Worker %s entering sleep state."), self.name)
         while not self.shutdown_requested and not self.wakeup:
             r, w, x = select.select(
-                [self.sentinel, connection.connection], [], [], self.heartbeat_period.seconds
+                [self.sentinel, self.pubsub_backend], [], [], self.heartbeat_period.seconds
             )
             self.beat()
-            if connection.connection in r:
-                connection.connection.execute("SELECT 1")
+            if self.pubsub_backend in r:
+                self.pubsub_backend.fetch()
             if self.sentinel in r:
                 os.read(self.sentinel, 256)
         self.wakeup = False
@@ -409,14 +461,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_backend, task_process.sentinel],
                     [],
                     [],
                     self.heartbeat_period.seconds,
                 )
                 self.beat()
-                if connection.connection in r:
-                    connection.connection.execute("SELECT 1")
+                if self.pubsub_backend in r:
+                    self.pubsub_backend.fetch()
                     if self.cancel_task:
                         _logger.info(
                             _("Received signal to cancel current task %s in domain: %s."),
@@ -472,7 +524,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()
+            self.pubsub_backend.wakeup_workers()
         self.task = None
 
     def handle_available_tasks(self):
@@ -529,21 +581,38 @@ 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_backend.record_worker_metrics(str(now))
+
+    def pubsub_setup(self):
+        def cancellation_callback(message):
+            if message == str(self.task.pk):
+                self.cancel_task = True
+
+        def wakeup_callback(message):
+            self.wakeup = True
+
+        def metric_callback(message):
+            self.last_metric_heartbeat = datetime.fromisoformat(message)
+
+        self.pubsub_backend.subscribe(self.pubsub_backend.WORKER_WAKEUP, cancellation_callback)
+        self.pubsub_backend.subscribe(self.pubsub_backend.TASK_CANCELLATION, wakeup_callback)
+        self.pubsub_backend.subscribe(self.pubsub_backend.WORKER_METRIC, metric_callback)
+
+    def pubsub_teardown(self):
+        self.pubsub_backend.unsubscribe(self.pubsub_backend.WORKER_WAKEUP)
+        self.pubsub_backend.unsubscribe(self.pubsub_backend.TASK_CANCELLATION)
+        self.pubsub_backend.unsubscribe(self.pubsub_backend.WORKER_METRIC)
 
     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_backend.unsubscribe(self.pubsub_backend.WORKER_WAKEUP)
                 self.handle_available_tasks()
             else:
-                self.cursor.execute("LISTEN pulp_worker_wakeup")
                 while not self.shutdown_requested:
                     # do work
                     if self.shutdown_requested:
@@ -553,7 +622,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()