Change PubSub usage and add make use of global connection object

Author: pedro-psb

pulpcore/tasking/pubsub.py

@@ -2,30 +2,35 @@
 from pulpcore.constants import TASK_PUBSUB
 import os
 import logging
+from django.db import connection
 from contextlib import suppress
 
 logger = logging.getLogger(__name__)
 
 
 class BasePubSubBackend:
     # Utils
-    def wakeup_worker(self, reason="unknown"):
-        self.publish(TASK_PUBSUB.WAKEUP_WORKER, reason)
+    @classmethod
+    def wakeup_worker(cls, reason="unknown"):
+        cls.publish(TASK_PUBSUB.WAKEUP_WORKER, reason)
 
-    def cancel_task(self, task_pk):
-        self.publish(TASK_PUBSUB.CANCEL_TASK, str(task_pk))
+    @classmethod
+    def cancel_task(cls, task_pk):
+        cls.publish(TASK_PUBSUB.CANCEL_TASK, str(task_pk))
 
-    def record_worker_metrics(self, now):
-        self.publish(TASK_PUBSUB.WORKER_METRICS, str(now))
+    @classmethod
+    def record_worker_metrics(cls, now):
+        cls.publish(TASK_PUBSUB.WORKER_METRICS, str(now))
 
     # Interface
-    def subscribe(self, channel, callback):
+    def subscribe(self, channel):
         raise NotImplementedError()
 
     def unsubscribe(self, channel):
         raise NotImplementedError()
 
-    def publish(self, channel, message=None):
+    @staticmethod
+    def publish(channel, payload=None):
         raise NotImplementedError()
 
     def fileno(self):
@@ -53,19 +58,19 @@ def drain_non_blocking_fd(fd):
 
 class PostgresPubSub(BasePubSubBackend):
 
-    def __init__(self, connection):
-        self.cursor = connection.cursor()
-        self.connection = connection.connection
-        assert self.cursor.connection is self.connection
+    def __init__(self):
         self.subscriptions = []
         self.message_buffer = []
-        self.connection.add_notify_handler(self._store_messages)
+        # Ensures a connection is established
+        if not connection.connection:
+            with connection.cursor():
+                pass
+        connection.connection.add_notify_handler(self._store_messages)
         # Handle message readiness
         # We can use os.evenfd in python >= 3.10
         self.sentinel_r, self.sentinel_w = os.pipe()
         os.set_blocking(self.sentinel_r, False)
         os.set_blocking(self.sentinel_w, False)
-        logger.debug(f"Initialized pubsub. Conn={self.connection}")
 
     def _store_messages(self, notification):
         self.message_buffer.append(
@@ -74,20 +79,25 @@ def _store_messages(self, notification):
 
     def subscribe(self, channel):
         self.subscriptions.append(channel)
-        self.connection.execute(f"LISTEN {channel}")
+        with connection.cursor() as cursor:
+            cursor.execute(f"LISTEN {channel}")
 
     def unsubscribe(self, channel):
         self.subscriptions.remove(channel)
         for i in range(0, len(self.message_buffer), -1):
             if self.message_buffer[i].channel == channel:
                 self.message_buffer.pop(i)
-        self.connection.execute(f"UNLISTEN {channel}")
-
-    def publish(self, channel, message=None):
-        if not message:
-            self.cursor.execute(f"NOTIFY {channel}")
+        with connection.cursor() as cursor:
+            cursor.execute(f"UNLISTEN {channel}")
+
+    @staticmethod
+    def publish(channel, payload=None):
+        if not payload:
+            with connection.cursor() as cursor:
+                cursor.execute(f"NOTIFY {channel}")
         else:
-            self.cursor.execute("SELECT pg_notify(%s, %s)", (channel, message))
+            with connection.cursor() as cursor:
+                cursor.execute("SELECT pg_notify(%s, %s)", (channel, str(payload)))
 
     def fileno(self) -> int:
         if self.message_buffer:
@@ -97,18 +107,25 @@ def fileno(self) -> int:
         return self.sentinel_r
 
     def fetch(self) -> list[PubsubMessage]:
-        self.connection.execute("SELECT 1").fetchone()
+        with connection.cursor() as cursor:
+            cursor.execute("SELECT 1").fetchone()
         result = self.message_buffer.copy()
         self.message_buffer.clear()
         return result
 
     def close(self):
         os.close(self.sentinel_r)
         os.close(self.sentinel_w)
-        self.cursor.close()
+        self.message_buffer.clear()
+        connection.connection.remove_notify_handler(self._store_messages)
+        for channel in self.subscriptions:
+            self.unsubscribe(channel)
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
         self.close()
+
+
+backend = PostgresPubSub

pulpcore/tasking/tasks.py

@@ -304,8 +304,7 @@ def dispatch(
                 task.set_canceling()
                 task.set_canceled(TASK_STATES.CANCELED, "Resources temporarily unavailable.")
     if notify_workers:
-        with pubsub.PostgresPubSub(connection) as pubsub_client:
-            pubsub_client.wakeup_worker(reason=TASK_WAKEUP_UNBLOCK)
+        pubsub.backend.wakeup_worker(reason=TASK_WAKEUP_UNBLOCK)
     return task
 
 
@@ -339,9 +338,8 @@ 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 pubsub.PostgresPubSub(connection) as pubsub_client:
-        pubsub_client.cancel_task(task_pk=task.pk)
-        pubsub_client.wakeup_worker()
+    pubsub.backend.cancel_task(task_pk=task.pk)
+    pubsub.backend.wakeup_worker()
     return task
 
 

pulpcore/tasking/worker.py

@@ -13,7 +13,6 @@
 from packaging.version import parse as parse_version
 
 from django.conf import settings
-from django.db import connection
 from django.db.models import Case, Count, F, Max, Value, When
 from django.db.models.functions import Random
 from django.utils import timezone
@@ -83,7 +82,7 @@ def __init__(self, auxiliary=False):
             int(WORKER_CLEANUP_INTERVAL / 10), WORKER_CLEANUP_INTERVAL
         )
         # Pubsub handling
-        self.pubsub_client = pubsub.PostgresPubSub(connection)
+        self.pubsub_client = pubsub.backend()
         self.pubsub_channel_callback = {}
 
         # Add a file descriptor to trigger select on signals
@@ -576,9 +575,7 @@ def metric_callback(message):
         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)
+        self.pubsub_client.close()
 
     def run(self, burst=False):
         with WorkerDirectory(self.name):

pulpcore/tests/functional/test_pubsub.py

@@ -1,118 +1,116 @@
 from django.db import connection
 from pulpcore.tasking import pubsub
 from types import SimpleNamespace
+from datetime import datetime
 import select
 import pytest
 
 
 def test_postgres_pubsub():
+    """Testing postgres low-level implementation."""
     state = SimpleNamespace()
-    state.got_first_message = False
-    state.got_second_message = False
+    state.got_message = False
     with connection.cursor() as cursor:
         assert connection.connection is cursor.connection
         conn = cursor.connection
+        # Listen and Notify
         conn.execute("LISTEN abc")
         conn.add_notify_handler(lambda notification: setattr(state, "got_message", True))
         cursor.execute("NOTIFY abc, 'foo'")
+        assert state.got_message is True
         conn.execute("SELECT 1")
+        assert state.got_message is True
+
+        # Reset and retry
+        state.got_message = False
         conn.execute("UNLISTEN abc")
-    assert state.got_message is True
+        cursor.execute("NOTIFY abc, 'foo'")
+        assert state.got_message is False
 
 
 M = pubsub.PubsubMessage
 
+PUBSUB_BACKENDS = [
+    pytest.param(pubsub.PostgresPubSub, id="and-using-postgres-backend"),
+]
+
+
+@pytest.mark.parametrize("pubsub_backend", PUBSUB_BACKENDS)
+class TestPublish:
+
+    @pytest.mark.parametrize(
+        "payload",
+        (
+            pytest.param(None, id="none"),
+            pytest.param("", id="empty-string"),
+            pytest.param("payload", id="non-empty-string"),
+            pytest.param(123, id="int"),
+            pytest.param(datetime.now(), id="datetime"),
+            pytest.param(True, id="bool"),
+        ),
+    )
+    def test_with_payload_as(self, pubsub_backend, payload):
+        pubsub_backend.publish("channel", payload=payload)
+
 
+@pytest.mark.parametrize("pubsub_backend", PUBSUB_BACKENDS)
 @pytest.mark.parametrize(
     "messages",
     (
-        [M("channel_a", "A1")],
-        [M("channel_a", "A1"), M("channel_a", "A2")],
-        [M("channel_a", "A1"), M("channel_a", "A2"), M("channel_b", "B1"), M("channel_c", "C1")],
+        pytest.param([M("a", "A1")], id="single-message"),
+        pytest.param([M("a", "A1"), M("a", "A2")], id="two-messages-in-same-channel"),
+        pytest.param(
+            [M("a", "A1"), M("a", "A2"), M("b", "B1"), M("c", "C1")],
+            id="tree-msgs-in-different-channels",
+        ),
     ),
 )
-@pytest.mark.parametrize("same_client", (True, False), ids=("same-clients", "different-clients"))
-class TestPubSub:
-
-    def test_subscribe_publish_fetch(self, same_client, messages):
-        """
-        GIVEN a publisher and a subscriber (which may be the same)
-        AND a queue of messages Q with mixed channels and payloads
-        WHEN the subscriber subscribes to all the channels in Q
-        AND the publisher publishes all the messages in Q
-        THEN the subscriber fetch() call returns a queue equivalent to Q
-        AND calling fetch() a second time returns an empty queue
-        """
-        # Given
-        publisher = pubsub.PostgresPubSub(connection)
-        subscriber = publisher if same_client else pubsub.PostgresPubSub(connection)
-
-        # When
-        for message in messages:
-            subscriber.subscribe(message.channel)
-        for message in messages:
-            publisher.publish(message.channel, message=message.payload)
-
-        # Then
-        assert subscriber.fetch() == messages
-        assert subscriber.fetch() == []
-
-    def test_unsubscribe(self, same_client, messages):
-        """
-        GIVEN a publisher and a subscriber (which may be the same)
-        AND a queue of messages Q with mixed channels and payloads
-        WHEN the subscriber subscribes and unsubscribes to all the channels in Q
-        AND the publisher publishes all the messages in Q
-        THEN the subscriber fetch() call returns an empty queue
-        """
-        # Given
-        publisher = pubsub.PostgresPubSub(connection)
-        subscriber = publisher if same_client else pubsub.PostgresPubSub(connection)
-
-        # When
-        for message in messages:
-            subscriber.subscribe(message.channel)
-        for message in messages:
-            subscriber.unsubscribe(message.channel)
-        for message in messages:
-            publisher.publish(message.channel, message=message.payload)
-
-        # Then
-        assert subscriber.fetch() == []
-
-    def test_select_loop(self, same_client, messages):
-        """
-        GIVEN a publisher and a subscriber (which may be the same)
-        AND a queue of messages Q with mixed channels and payloads
-        AND the subscriber is subscribed to all the channels in Q
-        WHEN the publisher has NOT published anything yet
-        THEN the select loop won't detect the subscriber readiness
-        AND the subscriber fetch() call returns an empty queue
-        BUT WHEN the publisher does publish all messages in Q
-        THEN the select loop detects the subscriber readiness
-        AND the subscriber fetch() call returns a queue equivalent to Q
-        """
+class TestSubscribeFetch:
+    def unsubscribe_all(self, channels, subscriber):
+        for channel in channels:
+            subscriber.unsubscribe(channel)
+
+    def subscribe_all(self, channels, subscriber):
+        for channel in channels:
+            subscriber.subscribe(channel)
+
+    def publish_all(self, messages, publisher):
+        for channel, payload in messages:
+            publisher.publish(channel, payload=payload)
+
+    def test_with(
+        self, pubsub_backend: pubsub.BasePubSubBackend, messages: list[pubsub.PubsubMessage]
+    ):
+        channels = {m.channel for m in messages}
+        publisher = pubsub_backend
+        with pubsub_backend() as subscriber:
+            self.subscribe_all(channels, subscriber)
+            self.publish_all(messages, publisher)
+            assert subscriber.fetch() == messages
+
+            self.unsubscribe_all(channels, subscriber)
+            assert subscriber.fetch() == []
+
+    def test_select_readiness_with(
+        self, pubsub_backend: pubsub.BasePubSubBackend, messages: list[pubsub.PubsubMessage]
+    ):
         TIMEOUT = 0.1
-
-        # Given
-        publisher = pubsub.PostgresPubSub(connection)
-        subscriber = publisher if same_client else pubsub.PostgresPubSub(connection)
-
-        # When
-        for message in messages:
-            subscriber.subscribe(message.channel)
-        r, w, x = select.select([subscriber], [], [], TIMEOUT)
-
-        # Then
-        assert subscriber not in r
-        assert subscriber.fetch() == []
-
-        # But When
-        for message in messages:
-            publisher.publish(message.channel, message=message.payload)
-        r, w, x = select.select([subscriber], [], [], TIMEOUT)
-
-        # Then
-        assert subscriber in r
-        assert subscriber.fetch() == messages
-        assert subscriber.fetch() == []
+        channels = {m.channel for m in messages}
+        publisher = pubsub_backend
+        with pubsub_backend() as subscriber:
+            self.subscribe_all(channels, subscriber)
+            r, w, x = select.select([subscriber], [], [], TIMEOUT)
+            assert subscriber not in r
+            assert subscriber.fetch() == []
+
+            self.publish_all(messages, publisher)
+            r, w, x = select.select([subscriber], [], [], TIMEOUT)
+            assert subscriber in r
+            assert subscriber.fetch() == messages
+            assert subscriber.fetch() == []
+
+            self.unsubscribe_all(channels, subscriber)
+            self.publish_all(messages, publisher)
+            r, w, x = select.select([subscriber], [], [], TIMEOUT)
+            assert subscriber not in r
+            assert subscriber.fetch() == []