Feature: Parallel message processing

src/aleph/db/accessors/pending_messages.py

@@ -1,13 +1,13 @@
 import datetime as dt
-from typing import Any, Collection, Dict, Iterable, Optional, Sequence
+from typing import Any, Collection, Dict, Iterable, List, Optional, Sequence, Set
 
 from aleph_message.models import Chain
 from sqlalchemy import delete, func, select, update
 from sqlalchemy.orm import selectinload
 from sqlalchemy.sql import Update
 
 from aleph.db.models import ChainTxDb, PendingMessageDb
-from aleph.types.db_session import DbSession
+from aleph.types.db_session import AsyncDbSession, DbSession
 
 
 def get_next_pending_message(
@@ -76,6 +76,134 @@ def get_pending_messages(
     return session.execute(select_stmt).scalars()
 
 
+def get_next_pending_messages_by_address(
+    session: DbSession,
+    current_time: dt.datetime,
+    fetched: Optional[bool] = None,
+    exclude_item_hashes: Optional[Set[str]] = None,
+    exclude_addresses: Optional[Set[str]] = None,
+    batch_size: int = 100,
+) -> List[PendingMessageDb]:
+    # Step 1: Get the earliest pending message
+    base_stmt = (
+        select(PendingMessageDb)
+        .where(PendingMessageDb.next_attempt <= current_time)
+        .order_by(PendingMessageDb.next_attempt.asc())
+        .options(selectinload(PendingMessageDb.tx))
+        .limit(1)
+    )
+
+    if fetched is not None:
+        base_stmt = base_stmt.where(PendingMessageDb.fetched == fetched)
+
+    if exclude_item_hashes:  # a non-empty set()
+        base_stmt = base_stmt.where(
+            PendingMessageDb.item_hash.not_in(exclude_item_hashes)
+        )
+
+    if exclude_addresses:
+        base_stmt = base_stmt.where(
+            PendingMessageDb.content["address"].astext.not_in(list(exclude_addresses))
+        )
+
+    first_message = session.execute(base_stmt).scalar_one_or_none()
+
+    if (
+        not first_message
+        or not first_message.content
+        or "address" not in first_message.content
+    ):
+        return []
+
+    address = first_message.content["address"]
+
+    # Step 2: Get a batch of messages with that same address in content
+    match_stmt = (
+        select(PendingMessageDb)
+        .where(
+            PendingMessageDb.next_attempt <= current_time,
+            PendingMessageDb.content["address"].astext == address,
+        )
+        .order_by(PendingMessageDb.next_attempt.asc())
+        .limit(batch_size)  # Limit to batch_size to avoid fetching too many at once
+    )
+
+    if fetched is not None:
+        match_stmt = match_stmt.where(PendingMessageDb.fetched == fetched)
+
+    if exclude_item_hashes:
+        match_stmt = match_stmt.where(
+            PendingMessageDb.item_hash.not_in(exclude_item_hashes)
+        )
+
+    return session.execute(match_stmt).scalars().all()
+
+
+async def async_get_next_pending_messages_by_address(
+    session: AsyncDbSession,
+    current_time: dt.datetime,
+    fetched: Optional[bool] = None,
+    exclude_item_hashes: Optional[Set[str]] = None,
+    exclude_addresses: Optional[Set[str]] = None,
+    batch_size: int = 100,
+) -> List[PendingMessageDb]:
+    # Step 1: Get the earliest pending message
+    base_stmt = (
+        select(PendingMessageDb)
+        .where(PendingMessageDb.next_attempt <= current_time)
+        .order_by(PendingMessageDb.next_attempt.asc())
+        .options(selectinload(PendingMessageDb.tx))
+        .limit(1)
+    )
+
+    if fetched is not None:
+        base_stmt = base_stmt.where(PendingMessageDb.fetched == fetched)
+
+    if exclude_item_hashes:
+        base_stmt = base_stmt.where(
+            PendingMessageDb.item_hash.not_in(exclude_item_hashes)
+        )
+
+    if exclude_addresses:
+        base_stmt = base_stmt.where(
+            PendingMessageDb.content["address"].astext.not_in(list(exclude_addresses))
+        )
+
+    result = await session.execute(base_stmt)
+    first_message = result.scalar_one_or_none()
+
+    if (
+        not first_message
+        or not first_message.content
+        or "address" not in first_message.content
+    ):
+        return []
+
+    address = first_message.content["address"]
+
+    # Step 2: Get a batch of messages with that same address in content
+    match_stmt = (
+        select(PendingMessageDb)
+        .where(
+            PendingMessageDb.next_attempt <= current_time,
+            PendingMessageDb.content["address"].astext == address,
+        )
+        .order_by(PendingMessageDb.next_attempt.asc())
+        .limit(batch_size)
+    )
+
+    if fetched is not None:
+        match_stmt = match_stmt.where(PendingMessageDb.fetched == fetched)
+
+    if exclude_item_hashes:
+        match_stmt = match_stmt.where(
+            PendingMessageDb.item_hash.not_in(exclude_item_hashes)
+        )
+
+    result = await session.execute(match_stmt)
+    return result.scalars().all()
+
+
 def get_pending_message(
     session: DbSession, pending_message_id: int
 ) -> Optional[PendingMessageDb]:

src/aleph/db/connection.py

@@ -69,7 +69,7 @@ def make_async_engine(
     application_name: Optional[str] = None,
 ) -> AsyncEngine:
     return create_async_engine(
-        make_db_url(driver="asyncpg", config=config, application_name=application_name),
+        make_db_url(driver="asyncpg", config=config),
         future=True,
         echo=echo,
     )