Implement a Redis Delayed Queue to support sorted sets

examples/book/benchmark.py

@@ -134,7 +134,13 @@ def truncate_op(session: sessionmaker, model, nsize: int) -> None:
         case_sensitive=False,
     ),
 )
-def main(config: str, nsize: int, daemon: bool, tg_op: str):
+@click.option(
+    "--weight", "-w", default=0.0, help="Weight for pgsync operations"
+)
+def main(
+    config: str, nsize: int, daemon: bool, tg_op: str, weight: float
+) -> None:
+    """Benchmarking script for Book model operations."""
     show_settings(config)
 
     config: str = get_config(config)
@@ -144,6 +150,9 @@ def main(config: str, nsize: int, daemon: bool, tg_op: str):
         Session = sessionmaker(bind=engine, autoflush=False, autocommit=False)
         session = Session()
 
+        if weight:
+            session.execute(sa.text(f"SET pgsync.weight = {weight}"))
+
         model = Book
         func: dict = {
             INSERT: insert_op,

pgsync/base.py

@@ -73,9 +73,19 @@ class Payload(object):
         new (dict): The new values of the row that was affected by the event (for INSERT and UPDATE operations).
         xmin (int): The transaction ID of the event.
         indices (List[str]): The indices of the affected rows (for UPDATE and DELETE operations).
+        weight (float): The weight of the event.
     """
 
-    __slots__ = ("tg_op", "table", "schema", "old", "new", "xmin", "indices")
+    __slots__ = (
+        "tg_op",
+        "table",
+        "schema",
+        "old",
+        "new",
+        "xmin",
+        "indices",
+        "weight",
+    )
 
     def __init__(
         self,
@@ -86,6 +96,7 @@ def __init__(
         new: t.Optional[t.Dict[str, t.Any]] = None,
         xmin: t.Optional[int] = None,
         indices: t.Optional[t.List[str]] = None,
+        weight: t.Optional[float] = None,
     ):
         self.tg_op: t.Optional[str] = tg_op
         self.table: t.Optional[str] = table
@@ -94,6 +105,7 @@ def __init__(
         self.new: t.Dict[str, t.Any] = new or {}
         self.xmin: t.Optional[int] = xmin
         self.indices: t.List[str] = indices
+        self.weight: float = weight
 
     @property
     def data(self) -> dict:

pgsync/redisqueue.py

@@ -2,6 +2,7 @@
 
 import json
 import logging
+import time
 import typing as t
 
 from redis import Redis
@@ -14,14 +15,18 @@
 )
 from .urls import get_redis_url
 
+# Pick a MULTIPLIER > max timestamp_ms (~1.7e12).
+# 10**13 is safe for now.
+_MULTIPLIER = 10**13
+
+
 logger = logging.getLogger(__name__)
 
 
-class RedisQueue(object):
-    """Simple Queue with Redis Backend."""
+class RedisQueue:
+    """A Redis‐backed queue where items become poppable only once ready is True."""
 
     def __init__(self, name: str, namespace: str = "queue", **kwargs):
-        """Init Simple Queue with Redis Backend."""
         url: str = get_redis_url(**kwargs)
         self.key: str = f"{namespace}:{name}"
         self._meta_key: str = f"{self.key}:meta"
@@ -38,34 +43,50 @@ def __init__(self, name: str, namespace: str = "queue", **kwargs):
 
     @property
     def qsize(self) -> int:
-        """Return the approximate size of the queue."""
-        return self.__db.llen(self.key)
-
-    def pop(self, chunk_size: t.Optional[int] = None) -> t.List[dict]:
-        """Remove and return multiple items from the queue."""
-        chunk_size = chunk_size or REDIS_READ_CHUNK_SIZE
-        if self.qsize > 0:
-            pipeline = self.__db.pipeline()
-            pipeline.lrange(self.key, 0, chunk_size - 1)
-            pipeline.ltrim(self.key, chunk_size, -1)
-            items: t.List = pipeline.execute()
-            logger.debug(f"pop size: {len(items[0])}")
-            return list(map(lambda value: json.loads(value), items[0]))
-
-    def push(self, items: t.List) -> None:
-        """Push multiple items onto the queue."""
-        self.__db.rpush(self.key, *map(json.dumps, items))
+        """Number of items currently in the ZSET (regardless of ready/not)."""
+        return self.__db.zcard(self.key)
+
+    def push(self, items: t.List[dict], weight: float = 0.0) -> None:
+        """
+        Push a batch of items with the given numeric weight.
+
+        - Higher weight -> higher priority.
+        - Among equal weight, FIFO order.
+        """
+        mapping: dict = {}
+        for item in items:
+            now_ms: int = int(time.time() * 1_000)
+            # score = -weight*M + timestamp
+            score = -weight * _MULTIPLIER + now_ms
+            mapping[json.dumps(item)] = score
+        # ZADD will add/update each member's score
+        self.__db.zadd(self.key, mapping)
+
+    def pop(self, chunk_size: int = REDIS_READ_CHUNK_SIZE) -> t.List[dict]:
+        """
+        Pop up to chunk_size highest priority items (by weight, then FIFO).
+        """
+        # ZPOPMIN pulls the entries with the smallest score first
+        popped: t.List[t.Tuple[bytes, float]] = self.__db.zpopmin(
+            self.key, chunk_size
+        )
+        results: t.List[dict] = [
+            json.loads(member) for member, score in popped
+        ]
+        logger.debug(f"popped {len(results)} items (by priority)")
+        return results
 
     def delete(self) -> None:
-        """Delete all items from the named queue."""
-        logger.info(f"Deleting redis key: {self.key}")
+        """Delete all items from the named queue including its metadata."""
+        logger.info(f"Deleting redis key: {self.key} and {self._meta_key}")
         self.__db.delete(self.key)
+        self.__db.delete(self._meta_key)
 
     def set_meta(self, value: t.Any) -> None:
-        """Store an arbitrary JSON-serialisable value in a dedicated key."""
+        """Store an arbitrary JSON‐serializable value in a dedicated key."""
         self.__db.set(self._meta_key, json.dumps(value))
 
     def get_meta(self, default: t.Any = None) -> t.Any:
-        """Retrieve the stored value (or *default* if nothing is set)."""
+        """Retrieve the stored metadata (or *default* if nothing is set)."""
         raw = self.__db.get(self._meta_key)
         return json.loads(raw) if raw is not None else default

pgsync/sync.py

@@ -13,6 +13,7 @@
 import typing as t
 from collections import defaultdict
 from itertools import groupby
+from math import inf
 from pathlib import Path
 
 import click
@@ -1128,7 +1129,9 @@
             self._poll_redis()
 
     async def _async_poll_redis(self) -> None:
-        payloads: list = self.redis.pop()
+        payloads: t.List[t.Dict] = self.redis.pop(
+            settings.REDIS_AUTO_POP_READY_STATE
+        )
         if payloads:
             logger.debug(f"_async_poll_redis: {payloads}")
             self.count["redis"] += len(payloads)
@@ -1144,34 +1147,44 @@
         while True:
             await self._async_poll_redis()
 
+    def _flush_payloads(self, payloads: list[dict]) -> None:
+        if not payloads:
+            return
+
+        # group by weight, default=+inf so missing weights pop first
+        weight_buckets: t.Dict[float, t.List[t.Dict]] = {}
+        for payload in payloads:
+            raw = payload.get("weight")
+            weight: float = float(raw) if raw is not None else inf
+            weight_buckets.setdefault(weight, []).append(payload)
+
+        # push each bucket in descending weight order (highest first)
+        for weight, items in sorted(
+            weight_buckets.items(), key=lambda kv: -kv[0]
+        ):
+            logger.debug(f"Pushing {len(items)} items with weight={weight}")
+            self.redis.push(items, weight=weight)
+
     @threaded
     @exception
     def poll_db(self) -> None:
-        """
-        Producer which polls Postgres continuously.
-
-        Receive a notification message from the channel we are listening on
-        """
         conn = self.engine.connect().connection
         conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-        cursor = conn.cursor()
-        cursor.execute(f'LISTEN "{self.database}"')
+        conn.cursor().execute(f'LISTEN "{self.database}"')
         logger.debug(
             f'Listening to notifications on channel "{self.database}"'
         )
-        payloads: list = []
+
+        payloads: t.List[t.Dict] = []
 
         while True:
-            # NB: consider reducing POLL_TIMEOUT to increase throughput
             if select.select([conn], [], [], settings.POLL_TIMEOUT) == (
                 [],
                 [],
                 [],
             ):
-                # Catch any hanging items from the last poll
-                if payloads:
-                    self.redis.push(payloads)
-                    payloads = []
+                self._flush_payloads(payloads)
+                payloads = []
                 continue
 
             try:
@@ -1182,28 +1195,32 @@
 
             while conn.notifies:
                 if len(payloads) >= settings.REDIS_WRITE_CHUNK_SIZE:
-                    self.redis.push(payloads)
+                    self._flush_payloads(payloads)
                     payloads = []
+
                 notification: t.AnyStr = conn.notifies.pop(0)
-                if notification.channel == self.database:
+                if notification.channel != self.database:
+                    continue
 
-                    try:
-                        payload = json.loads(notification.payload)
-                    except json.JSONDecodeError as e:
-                        logger.exception(
-                            f"Error decoding JSON payload: {e}\n"
-                            f"Payload: {notification.payload}"
-                        )
-                        continue
-                    if (
-                        payload["indices"]
-                        and self.index in payload["indices"]
-                        and payload["schema"] in self.tree.schemas
-                    ):
-                        payloads.append(payload)
-                        logger.debug(f"poll_db: {payload}")
-                        with self.lock:
-                            self.count["db"] += 1
+                try:
+                    payload = json.loads(notification.payload)
+                except json.JSONDecodeError:
+                    logger.exception("Invalid JSON in notification, skipping")
+                    continue
+
+                if (
+                    payload.get("indices")
+                    and self.index in payload["indices"]
+                    and payload.get("schema") in self.tree.schemas
+                ):
+                    payloads.append(payload)
+                    logger.debug(f"Queued payload: {payload}")
+                    with self.lock:
+                        self.count["db"] += 1
+
+            # flush anything left after draining notifications
+            self._flush_payloads(payloads)
+            payloads = []
 
     @exception
     def async_poll_db(self) -> None:
@@ -1220,16 +1237,29 @@
 
         while self.conn.notifies:
             notification: t.AnyStr = self.conn.notifies.pop(0)
-            if notification.channel == self.database:
+            if notification.channel != self.database:
+                continue
+
+            try:
                 payload = json.loads(notification.payload)
-                if (
-                    payload["indices"]
-                    and self.index in payload["indices"]
-                    and payload["schema"] in self.tree.schemas
-                ):
-                    self.redis.push([payload])
-                    logger.debug(f"async_poll: {payload}")
-                    self.count["db"] += 1
+            except json.JSONDecodeError as e:
+                logger.exception(f"Error decoding JSON payload: {e}")
+                continue
+
+            if (
+                payload.get("indices")
+                and self.index in payload["indices"]
+                and payload.get("schema") in self.tree.schemas
+            ):
+                # extract numeric weight (missing +inf for highest priority)
+                raw_w = payload.get("weight")
+                weight = float(raw_w) if raw_w is not None else inf
+
+                # push via priority queue
+                self.redis.push([payload], weight=weight)
+
+                logger.debug(f"async_poll: {payload} (weight={weight})")
+                self.count["db"] += 1
 
     def refresh_views(self) -> None:
         self._refresh_views()
@@ -1331,7 +1361,9 @@
             if polling:
                 return
             else:
-                raise
+                raise Exception(
+                    f"Error while pulling logical slot changes: {e}"
+                ) from e
         self.checkpoint: int = txmax or self.txid_current
         self._truncate = True
 
@@ -1350,9 +1382,32 @@
             await asyncio.sleep(settings.REPLICATION_SLOT_CLEANUP_INTERVAL)
 
     def _truncate_slots(self) -> None:
-        if self._truncate:
-            logger.debug(f"Truncating replication slot: {self.__name}")
-            self.logical_slot_get_changes(self.__name, upto_nchanges=None)
+        if not self._truncate:
+            return
+
+        """
+        Handle eventual consistency of the logical replication slot.
+        We retry logical_slot_changes a few times in case of replication slot in use error.
+        """
+        retries: int = 3
+        backoff: int = 1
+        txmax: int = self.txid_current
+        upto_lsn: str = self.current_wal_lsn
+
+        for attempt in range(1, retries + 1):
+            try:
+                logger.debug(f"Truncating replication slot: {self.__name}")
+                self.logical_slot_changes(txmax=txmax, upto_lsn=upto_lsn)
+                logger.debug("Truncation successful.")
+                break
+            except Exception as e:
+                logger.warning(f"Attempt {attempt} failed with {e}")
+                if attempt == retries:
+                    logger.error("Max retries reached, raising exception.")
+                    raise
+                sleep_time: int = backoff * (2 ** (attempt - 1))
+                logger.debug(f"Retrying in {sleep_time} seconds...")
+                time.sleep(sleep_time)
 
     @threaded
     @exception