Merge pull request taskiq-python#62 from taskiq-python/bugfix/cluster-source

Author: s3rius

taskiq_redis/schedule_source.py

@@ -117,7 +117,6 @@ def __init__(
  self,
  url: str,
  prefix: str = "schedule",
- buffer_size: int = 50,
  serializer: Optional[TaskiqSerializer] = None,
  **connection_kwargs: Any,
  ) -> None:
@@ -126,7 +125,6 @@ def __init__(
  url,
  **connection_kwargs,
  )
- self.buffer_size = buffer_size
  if serializer is None:
  serializer = PickleSerializer()
  self.serializer = serializer
@@ -156,19 +154,14 @@ async def get_schedules(self) -> List[ScheduledTask]:
  :return: list of schedules.
  """
  schedules = []
- buffer = []
  async for key in self.redis.scan_iter(f"{self.prefix}:*"): # type: ignore[attr-defined]
- buffer.append(key)
- if len(buffer) >= self.buffer_size:
- schedules.extend(await self.redis.mget(buffer)) # type: ignore[attr-defined]
- buffer = []
- if buffer:
- schedules.extend(await self.redis.mget(buffer)) # type: ignore[attr-defined]
- return [
- model_validate(ScheduledTask, self.serializer.loadb(schedule))
- for schedule in schedules
- if schedule
- ]
+ raw_schedule = await self.redis.get(key) # type: ignore[attr-defined]
+ parsed_schedule = model_validate(
+ ScheduledTask,
+ self.serializer.loadb(raw_schedule),
+ )
+ schedules.append(parsed_schedule)
+ return schedules
 
  async def post_send(self, task: ScheduledTask) -> None:
  """Delete a task after it's completed."""

tests/test_schedule_source.py

@@ -196,17 +196,31 @@ async def test_cluster_post_run_time(redis_cluster_url: str) -> None:
 
 
 @pytest.mark.anyio
-async def test_cluster_buffer(redis_cluster_url: str) -> None:
+async def test_cluster_get_schedules(redis_cluster_url: str) -> None:
+ """
+ Test of a redis cluster source.
+
+ This test checks that if the schedules are located on different nodes,
+ the source will still be able to get them all.
+
+ To simulate this we set a specific shard key for each schedule.
+ The shard keys are from this gist:
+
+ https://gist.githubusercontent.com/dvirsky/93f43277317f629bb06e858946416f7e/raw/b0438faf6f5a0020c12a0730f6cd6ac4bdc4b171/crc16_slottable.h
+
+ """
  prefix = uuid.uuid4().hex
- source = RedisClusterScheduleSource(redis_cluster_url, prefix=prefix, buffer_size=1)
+ source = RedisClusterScheduleSource(redis_cluster_url, prefix=prefix)
  schedule1 = ScheduledTask(
+ schedule_id=r"id-{06S}",
  task_name="test_task1",
  labels={},
  args=[],
  kwargs={},
  cron="* * * * *",
  )
  schedule2 = ScheduledTask(
+ schedule_id=r"id-{4Rs}",
  task_name="test_task2",
  labels={},
  args=[],