fix(tasks): optimisation

Author: GefMar

taskiq/_broker_resources.py

@@ -0,0 +1,38 @@
+"""Internal validation for resources owned by broker lifecycle coordinators."""
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:  # pragma: no cover
+    from taskiq.abc.broker import AsyncBroker
+
+
+def validate_unshared_broker_resources(
+    brokers: Iterable["AsyncBroker"],
+    *,
+    error_type: type[ValueError] = ValueError,
+) -> None:
+    """Reject middleware and result backends owned by multiple brokers."""
+    middleware_owners: dict[int, str] = {}
+    backend_owners: dict[int, str] = {}
+    for broker in brokers:
+        for middleware in broker.middlewares:
+            previous_owner = middleware_owners.setdefault(
+                id(middleware),
+                broker.broker_name,
+            )
+            if previous_owner != broker.broker_name:
+                raise error_type(
+                    "One middleware instance cannot be lifecycle-owned by "
+                    f"brokers {previous_owner!r} and {broker.broker_name!r}.",
+                )
+
+        previous_owner = backend_owners.setdefault(
+            id(broker.result_backend),
+            broker.broker_name,
+        )
+        if previous_owner != broker.broker_name:
+            raise error_type(
+                "One result backend instance cannot be lifecycle-owned by "
+                f"brokers {previous_owner!r} and {broker.broker_name!r}.",
+            )

tests/scheduler/test_scheduling_identity.py

@@ -0,0 +1,95 @@
+from datetime import datetime, timezone
+from typing import Any, Literal
+
+import pytest
+
+from taskiq.abc.schedule_source import ScheduleSource
+from taskiq.kicker import AsyncKicker
+from taskiq.scheduler.scheduled_task import ScheduledTask
+from taskiq.scheduler.scheduler import TaskiqScheduler
+from tests.utils import RecordingBroker
+
+ScheduleKind = Literal["cron", "interval", "time"]
+
+
+class RecordingScheduleSource(ScheduleSource):
+    """Store schedules created by a kicker."""
+
+    def __init__(self) -> None:
+        self.schedules: list[ScheduledTask] = []
+
+    async def get_schedules(self) -> list[ScheduledTask]:
+        return self.schedules
+
+    async def add_schedule(self, schedule: ScheduledTask) -> None:
+        self.schedules.append(schedule)
+
+
+async def create_schedule(
+    kind: ScheduleKind,
+    kicker: AsyncKicker[Any, Any],
+    source: RecordingScheduleSource,
+) -> ScheduledTask:
+    """Create one schedule through the selected public kicker method."""
+    if kind == "cron":
+        created = await kicker.schedule_by_cron(source, "* * * * *")
+    elif kind == "interval":
+        created = await kicker.schedule_by_interval(source, 10)
+    else:
+        created = await kicker.schedule_by_time(
+            source,
+            datetime(2030, 1, 1, tzinfo=timezone.utc),
+        )
+    return created.task
+
+
+@pytest.mark.parametrize("kind", ["cron", "interval", "time"])
+async def test_all_schedule_methods_preserve_custom_task_id(
+    kind: ScheduleKind,
+) -> None:
+    broker = RecordingBroker()
+    source = RecordingScheduleSource()
+    kicker: AsyncKicker[Any, Any] = AsyncKicker("demo.task", broker, {})
+
+    scheduled = await create_schedule(
+        kind,
+        kicker.with_task_id("custom-task-id"),
+        source,
+    )
+
+    assert scheduled.task_id == "custom-task-id"
+    assert source.schedules == [scheduled]
+
+
+async def test_interval_without_custom_task_id_keeps_deferred_generation() -> None:
+    broker = RecordingBroker()
+    source = RecordingScheduleSource()
+    kicker: AsyncKicker[Any, Any] = AsyncKicker("demo.task", broker, {})
+
+    scheduled = await create_schedule("interval", kicker, source)
+
+    assert scheduled.task_id is None
+
+
+@pytest.mark.parametrize(
+    ("custom_task_id", "expected_task_id"),
+    [
+        pytest.param("custom-task-id", "custom-task-id", id="custom"),
+        pytest.param(None, "generated-task-id", id="generated-at-dispatch"),
+    ],
+)
+async def test_scheduler_dispatch_uses_stored_or_generated_interval_task_id(
+    custom_task_id: str | None,
+    expected_task_id: str,
+) -> None:
+    broker = RecordingBroker()
+    broker.with_id_generator(lambda: "generated-task-id")
+    source = RecordingScheduleSource()
+    kicker: AsyncKicker[Any, Any] = AsyncKicker("demo.task", broker, {})
+    kicker.with_task_id(custom_task_id)
+    scheduled = await create_schedule("interval", kicker, source)
+    scheduler = TaskiqScheduler(broker, [source])
+
+    await scheduler.on_ready(source, scheduled)
+
+    assert broker.sent[0][0].task_id == expected_task_id