22851673222

Committed 09 Mar 2026 11:38AM UTC coverage: 97.564%. First build

Build # 22851673222

Build Type

Pull #2

github

Committed by

web-flow

Commit Message

Merge 8d49fc598 into d04f85da2

Pull Request Pull Request #2: Improve task management with Redis Streams, leader election and management API

Coverage Stats

3164 of 3246 new or added lines in 32 files covered. (97.47%)

3324 of 3407 relevant lines covered (97.56%)

4.86 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.86

/src/fastapi_task_manager/coordinator.py

"""Coordinator module - evaluates cron and publishes to stream (leader only).

This module provides the Coordinator class which is responsible for:
- Evaluating cron expressions for all registered tasks
- Publishing ready tasks to the Redis Stream
- Updating next_run timestamps for scheduled tasks

The Coordinator only performs its duties when the associated worker
is the leader. When not the leader, it periodically attempts to acquire
leadership via the LeaderElector.
"""

import asyncio
import logging
import time
from datetime import datetime, timezone
from typing import TYPE_CHECKING

from cronsim import CronSim
from redis.asyncio import Redis

from fastapi_task_manager.async_utils import interruptible_sleep
from fastapi_task_manager.leader_election import LeaderElector
from fastapi_task_manager.redis_keys import RedisKeyBuilder
from fastapi_task_manager.schema.task import Task
from fastapi_task_manager.task_group import TaskGroup

if TYPE_CHECKING:
    from fastapi_task_manager import TaskManager

logger = logging.getLogger("fastapi.task-manager.coordinator")

INITIAL_LOCK_TTL: int = 15

# Safety-net cap for stream length. Messages are deleted after ACK,
# so this limit is only reached if deletes fail or consumers are down.
STREAM_MAX_LEN: int = 10_000


class Coordinator:
    """Evaluates cron expressions and publishes ready tasks to stream.

    The Coordinator is the "brain" of the task scheduling system in stream mode.
    It runs on all workers but only performs scheduling duties when the worker
    is the leader. This ensures that tasks are only scheduled once, avoiding
    duplicates.

    Responsibilities:
    - Check if tasks are due based on their cron expressions
    - Publish ready tasks to the Redis Stream
    - Update next_run timestamps after scheduling
    - Attempt to acquire leadership when not the leader

    The Coordinator does NOT execute tasks - that's handled by StreamConsumer.
    """

    def __init__(
        self,
        redis_client: Redis,
        key_builder: RedisKeyBuilder,
        leader_elector: LeaderElector,
        task_manager: "TaskManager",
    ) -> None:
        """Initialize the Coordinator.

        Args:
            redis_client: Async Redis client for stream operations.
            key_builder: RedisKeyBuilder instance for key construction.
            leader_elector: LeaderElector instance to check leadership status.
            task_manager: TaskManager instance containing registered task groups.
        """
        self._redis = redis_client
        self._keys = key_builder
        self._leader = leader_elector
        self._task_manager = task_manager
        self._running = False

    async def start(self) -> asyncio.Task:
        """Start the coordinator loop.

        Returns:
            The asyncio Task running the coordinator loop.
        """
        self._running = True
        return asyncio.create_task(self._run(), name="Coordinator")

    async def stop(self) -> None:
        """Stop the coordinator loop."""
        self._running = False

    async def _run(self) -> None:
        """Main coordinator loop.

        This loop continuously:
        1. Checks if this worker is the leader
        2. If not leader, attempts to acquire leadership
        3. If leader, evaluates all tasks and publishes ready ones

        The loop respects the configured poll_interval and leader_retry_interval.
        """
        config = self._task_manager.config

        while self._running:
            try:
                if not self._leader.is_leader:
                    # Not leader, try to acquire leadership
                    await self._leader.try_acquire_leadership()
                    if not self._leader.is_leader:
                        # Still not leader, wait and retry
                        await interruptible_sleep(config.leader_retry_interval, lambda: self._running)
                        continue

                # We are the leader - evaluate and publish tasks
                await self._evaluate_and_publish_all()
                await interruptible_sleep(config.poll_interval, lambda: self._running)

            except asyncio.CancelledError:
                logger.info("Coordinator loop cancelled")
                break
            except Exception:
                logger.exception("Error in coordinator loop")
                await interruptible_sleep(1, lambda: self._running)

    async def _evaluate_and_publish_all(self) -> None:
        """Check all tasks and publish ready ones to stream.

        Iterates through all task groups and their tasks, checking if each
        task is due for execution. Ready tasks are published to the appropriate
        Redis Stream (high or low priority) and their next_run timestamps are updated.
        """
        # Get both stream keys for high and low priority
        stream_key_high = self._keys.task_stream_high_key()
        stream_key_low = self._keys.task_stream_low_key()

        for task_group in self._task_manager.task_groups:
            for task in task_group.tasks:
                try:
                    if await self._is_task_due(task_group, task):
                        # Select stream based on task priority
                        stream_key = stream_key_high if task.high_priority else stream_key_low
                        await self._publish_task(stream_key, task_group, task)
                        await self._update_next_run(task_group, task)
                except Exception:
                    logger.exception(
                        "Error evaluating task %s/%s",
                        task_group.name,
                        task.name,
                    )

    async def _is_task_due(self, task_group: TaskGroup, task: Task) -> bool:
        """Check if a task is due for execution.

        A task is due if:
        1. It's not disabled
        2. Its next_run time is in the past (or not set)

        Args:
            task_group: The TaskGroup containing the task.
            task: The Task to check.

        Returns:
            True if the task should be executed, False otherwise.
        """
        keys = self._keys.get_task_keys(task_group.name, task.name)

        # Check if disabled
        disabled = await self._redis.get(keys.disabled)
        if disabled is not None:
            return False

        # Check if task is in backoff (retry_after timestamp in the future)
        retry_after_raw = await self._redis.get(keys.retry_after)
        if retry_after_raw is not None and time.time() < float(retry_after_raw):
            logger.debug(
                "Task %s/%s in backoff until %.0f, skipping",
                task_group.name,
                task.name,
                float(retry_after_raw),
            )
            return False

        # Check next_run time
        next_run_b = await self._redis.get(keys.next_run)
        if next_run_b is None:
            # Never run before, due immediately
            return True

        next_run = datetime.fromtimestamp(float(next_run_b), tz=timezone.utc)
        return next_run <= datetime.now(timezone.utc)

    async def _publish_task(
        self,
        stream_key: str,
        task_group: TaskGroup,
        task: Task,
    ) -> str:
        """Publish a task to the Redis stream and track it in the scheduled set.

        Creates a message in the appropriate task stream (high or low priority)
        containing the task group name, task name, and scheduling timestamp.
        The stream has a hardcoded safety-net MAXLEN cap; consumed messages
        are deleted individually via XDEL after acknowledgement.

        After publishing, the task identifier is added to the scheduled tracking
        SET so the Reconciler can detect lost tasks.

        Args:
            stream_key: The Redis key for the task stream (high or low).
            task_group: The TaskGroup containing the task.
            task: The Task to publish.

        Returns:
            The message ID assigned by Redis.
        """
        task_id = f"{task_group.name}:{task.name}"

        message_id = await self._redis.xadd(
            stream_key,
            {
                "group": task_group.name,
                "task": task.name,
                "scheduled_at": str(datetime.now(timezone.utc).timestamp()),
            },
            maxlen=STREAM_MAX_LEN,
            approximate=True,  # Use ~ for efficiency
        )

        # Track the published task in the scheduled set for reconciliation
        await self._redis.sadd(self._keys.scheduled_set_key(), task_id)  # ty: ignore[invalid-await]

        priority = "high" if task.high_priority else "low"
        logger.debug(
            "Published task %s/%s to %s priority stream (msg_id: %s)",
            task_group.name,
            task.name,
            priority,
            message_id,
        )
        return message_id

    async def _update_next_run(self, task_group: TaskGroup, task: Task) -> None:
        """Update the next_run timestamp for a task.

        Calculates the next execution time based on the task's cron expression
        and stores it in Redis with an appropriate TTL.

        Args:
            task_group: The TaskGroup containing the task.
            task: The Task to update.
        """
        keys = self._keys.get_task_keys(task_group.name, task.name)
        next_run = next(CronSim(task.expression, datetime.now(timezone.utc)))

        # Calculate TTL: at least INITIAL_LOCK_TTL, or 2x the time until next run
        time_until_next = (next_run - datetime.now(timezone.utc)).total_seconds()
        ttl = max(
            int(time_until_next) * 2,
            INITIAL_LOCK_TTL,
        )
        await self._redis.set(keys.next_run, next_run.timestamp(), ex=ttl)

        logger.debug(
            "Updated next_run for %s/%s to %s",
            task_group.name,
            task.name,
            next_run.isoformat(),
        )

1	"""Coordinator module - evaluates cron and publishes to stream (leader only).
2
3	This module provides the Coordinator class which is responsible for:
4	- Evaluating cron expressions for all registered tasks
5	- Publishing ready tasks to the Redis Stream
6	- Updating next_run timestamps for scheduled tasks
7
8	The Coordinator only performs its duties when the associated worker
9	is the leader. When not the leader, it periodically attempts to acquire
10	leadership via the LeaderElector.
11	"""
12
13	import asyncio	5✔
14	import logging	5✔
15	import time	5✔
16	from datetime import datetime, timezone	5✔
17	from typing import TYPE_CHECKING	5✔
18
19	from cronsim import CronSim	5✔
20	from redis.asyncio import Redis	5✔
21
22	from fastapi_task_manager.async_utils import interruptible_sleep	5✔
23	from fastapi_task_manager.leader_election import LeaderElector	5✔
24	from fastapi_task_manager.redis_keys import RedisKeyBuilder	5✔
25	from fastapi_task_manager.schema.task import Task	5✔
26	from fastapi_task_manager.task_group import TaskGroup	5✔
27
28	if TYPE_CHECKING:
29	from fastapi_task_manager import TaskManager
30
31	logger = logging.getLogger("fastapi.task-manager.coordinator")	5✔
32
33	INITIAL_LOCK_TTL: int = 15	5✔
34
35	# Safety-net cap for stream length. Messages are deleted after ACK,
36	# so this limit is only reached if deletes fail or consumers are down.
37	STREAM_MAX_LEN: int = 10_000	5✔
38
39
40	class Coordinator:	5✔
41	"""Evaluates cron expressions and publishes ready tasks to stream.
42
43	The Coordinator is the "brain" of the task scheduling system in stream mode.
44	It runs on all workers but only performs scheduling duties when the worker
45	is the leader. This ensures that tasks are only scheduled once, avoiding
46	duplicates.
47
48	Responsibilities:
49	- Check if tasks are due based on their cron expressions
50	- Publish ready tasks to the Redis Stream
51	- Update next_run timestamps after scheduling
52	- Attempt to acquire leadership when not the leader
53
54	The Coordinator does NOT execute tasks - that's handled by StreamConsumer.
55	"""
56
57	def __init__(	5✔
58	self,
59	redis_client: Redis,
60	key_builder: RedisKeyBuilder,
61	leader_elector: LeaderElector,
62	task_manager: "TaskManager",
63	) -> None:
64	"""Initialize the Coordinator.
65
66	Args:
67	redis_client: Async Redis client for stream operations.
68	key_builder: RedisKeyBuilder instance for key construction.
69	leader_elector: LeaderElector instance to check leadership status.
70	task_manager: TaskManager instance containing registered task groups.
71	"""
72	self._redis = redis_client	5✔
73	self._keys = key_builder	5✔
74	self._leader = leader_elector	5✔
75	self._task_manager = task_manager	5✔
76	self._running = False	5✔
77
78	async def start(self) -> asyncio.Task:	5✔
79	"""Start the coordinator loop.
80
81	Returns:
82	The asyncio Task running the coordinator loop.
83	"""
84	self._running = True	5✔
85	return asyncio.create_task(self._run(), name="Coordinator")	5✔
86
87	async def stop(self) -> None:	5✔
88	"""Stop the coordinator loop."""
89	self._running = False	5✔
90
91	async def _run(self) -> None:	5✔
92	"""Main coordinator loop.
93
94	This loop continuously:
95	1. Checks if this worker is the leader
96	2. If not leader, attempts to acquire leadership
97	3. If leader, evaluates all tasks and publishes ready ones
98
99	The loop respects the configured poll_interval and leader_retry_interval.
100	"""
101	config = self._task_manager.config	5✔
102
103	while self._running:	5✔
104	try:	5✔
105	if not self._leader.is_leader:	5✔
106	# Not leader, try to acquire leadership
107	await self._leader.try_acquire_leadership()	5✔
108	if not self._leader.is_leader:	5✔
109	# Still not leader, wait and retry
110	await interruptible_sleep(config.leader_retry_interval, lambda: self._running)	5✔
111	continue	5✔
112
113	# We are the leader - evaluate and publish tasks
114	await self._evaluate_and_publish_all()	5✔
115	await interruptible_sleep(config.poll_interval, lambda: self._running)	5✔
116
NEW 117	except asyncio.CancelledError:	×
NEW 118	logger.info("Coordinator loop cancelled")	×
NEW 119	break	×
NEW 120	except Exception:	×
NEW 121	logger.exception("Error in coordinator loop")	×
NEW 122	await interruptible_sleep(1, lambda: self._running)	×
123
124	async def _evaluate_and_publish_all(self) -> None:	5✔
125	"""Check all tasks and publish ready ones to stream.
126
127	Iterates through all task groups and their tasks, checking if each
128	task is due for execution. Ready tasks are published to the appropriate
129	Redis Stream (high or low priority) and their next_run timestamps are updated.
130	"""
131	# Get both stream keys for high and low priority
132	stream_key_high = self._keys.task_stream_high_key()	5✔
133	stream_key_low = self._keys.task_stream_low_key()	5✔
134
135	for task_group in self._task_manager.task_groups:	5✔
136	for task in task_group.tasks:	5✔
137	try:	5✔
138	if await self._is_task_due(task_group, task):	5✔
139	# Select stream based on task priority
140	stream_key = stream_key_high if task.high_priority else stream_key_low	5✔
141	await self._publish_task(stream_key, task_group, task)	5✔
142	await self._update_next_run(task_group, task)	5✔
143	except Exception:	5✔
144	logger.exception(	5✔
145	"Error evaluating task %s/%s",
146	task_group.name,
147	task.name,
148	)
149
150	async def _is_task_due(self, task_group: TaskGroup, task: Task) -> bool:	5✔
151	"""Check if a task is due for execution.
152
153	A task is due if:
154	1. It's not disabled
155	2. Its next_run time is in the past (or not set)
156
157	Args:
158	task_group: The TaskGroup containing the task.
159	task: The Task to check.
160
161	Returns:
162	True if the task should be executed, False otherwise.
163	"""
164	keys = self._keys.get_task_keys(task_group.name, task.name)	5✔
165
166	# Check if disabled
167	disabled = await self._redis.get(keys.disabled)	5✔
168	if disabled is not None:	5✔
169	return False	5✔
170
171	# Check if task is in backoff (retry_after timestamp in the future)
172	retry_after_raw = await self._redis.get(keys.retry_after)	5✔
173	if retry_after_raw is not None and time.time() < float(retry_after_raw):	5✔
174	logger.debug(	5✔
175	"Task %s/%s in backoff until %.0f, skipping",
176	task_group.name,
177	task.name,
178	float(retry_after_raw),
179	)
180	return False	5✔
181
182	# Check next_run time
183	next_run_b = await self._redis.get(keys.next_run)	5✔
184	if next_run_b is None:	5✔
185	# Never run before, due immediately
186	return True	5✔
187
188	next_run = datetime.fromtimestamp(float(next_run_b), tz=timezone.utc)	5✔
189	return next_run <= datetime.now(timezone.utc)	5✔
190
191	async def _publish_task(	5✔
192	self,
193	stream_key: str,
194	task_group: TaskGroup,
195	task: Task,
196	) -> str:
197	"""Publish a task to the Redis stream and track it in the scheduled set.
198
199	Creates a message in the appropriate task stream (high or low priority)
200	containing the task group name, task name, and scheduling timestamp.
201	The stream has a hardcoded safety-net MAXLEN cap; consumed messages
202	are deleted individually via XDEL after acknowledgement.
203
204	After publishing, the task identifier is added to the scheduled tracking
205	SET so the Reconciler can detect lost tasks.
206
207	Args:
208	stream_key: The Redis key for the task stream (high or low).
209	task_group: The TaskGroup containing the task.
210	task: The Task to publish.
211
212	Returns:
213	The message ID assigned by Redis.
214	"""
215	task_id = f"{task_group.name}:{task.name}"	5✔
216
217	message_id = await self._redis.xadd(	5✔
218	stream_key,
219	{
220	"group": task_group.name,
221	"task": task.name,
222	"scheduled_at": str(datetime.now(timezone.utc).timestamp()),
223	},
224	maxlen=STREAM_MAX_LEN,
225	approximate=True, # Use ~ for efficiency
226	)
227
228	# Track the published task in the scheduled set for reconciliation
229	await self._redis.sadd(self._keys.scheduled_set_key(), task_id) # ty: ignore[invalid-await]	5✔
230
231	priority = "high" if task.high_priority else "low"	5✔
232	logger.debug(	5✔
233	"Published task %s/%s to %s priority stream (msg_id: %s)",
234	task_group.name,
235	task.name,
236	priority,
237	message_id,
238	)
239	return message_id	5✔
240
241	async def _update_next_run(self, task_group: TaskGroup, task: Task) -> None:	5✔
242	"""Update the next_run timestamp for a task.
243
244	Calculates the next execution time based on the task's cron expression
245	and stores it in Redis with an appropriate TTL.
246
247	Args:
248	task_group: The TaskGroup containing the task.
249	task: The Task to update.
250	"""
251	keys = self._keys.get_task_keys(task_group.name, task.name)	5✔
252	next_run = next(CronSim(task.expression, datetime.now(timezone.utc)))	5✔
253
254	# Calculate TTL: at least INITIAL_LOCK_TTL, or 2x the time until next run
255	time_until_next = (next_run - datetime.now(timezone.utc)).total_seconds()	5✔
256	ttl = max(	5✔
257	int(time_until_next) * 2,
258	INITIAL_LOCK_TTL,
259	)
260	await self._redis.set(keys.next_run, next_run.timestamp(), ex=ttl)	5✔
261
262	logger.debug(	5✔
263	"Updated next_run for %s/%s to %s",
264	task_group.name,
265	task.name,
266	next_run.isoformat(),
267	)

Morry98 / fastapi-task-manager / 22851673222

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous