diff --git a/clients/python/src/taskbroker_client/scheduler/storage.py b/clients/python/src/taskbroker_client/scheduler/storage.py index 92fca58a..f223b8af 100644 --- a/clients/python/src/taskbroker_client/scheduler/storage.py +++ b/clients/python/src/taskbroker_client/scheduler/storage.py @@ -118,6 +118,14 @@ def __init__(self, metrics: MetricsBackend, redis: RedisCluster | StrictRedis) - def _make_key(self, key: str) -> str: return f"tw:scheduler:{key}" + def _mget_nonatomic(self, keys: list[str]) -> list[str | None]: + # We want mget_nonatomic (spelled mget in redis-py-cluster) when dealing + # with a cluster. For non clusters we want plain mget. + mget_nonatomic = getattr(self._redis, "mget_nonatomic", None) + if mget_nonatomic is not None: + return mget_nonatomic(keys) + return self._redis.mget(keys) + def set(self, key: str, next_runtime: datetime) -> bool: """ Record a spawn time for a task. @@ -168,9 +176,9 @@ def read_many( compat_keys = [sk.split(":", 1)[1] for sk in storage_keys] legacy_keys = [ck.rsplit(":", 1)[0] for ck in compat_keys] - new_values = self._redis.mget([self._make_key(sk) for sk in storage_keys]) - compat_values = self._redis.mget([self._make_key(ck) for ck in compat_keys]) - legacy_values = self._redis.mget([self._make_key(lk) for lk in legacy_keys]) + new_values = self._mget_nonatomic([self._make_key(sk) for sk in storage_keys]) + compat_values = self._mget_nonatomic([self._make_key(ck) for ck in compat_keys]) + legacy_values = self._mget_nonatomic([self._make_key(lk) for lk in legacy_keys]) run_times: dict[str, datetime | None] = {} for storage_key, new_val, compat_val, legacy_val in zip( diff --git a/clients/python/tests/scheduler/test_storage_cluster.py b/clients/python/tests/scheduler/test_storage_cluster.py new file mode 100644 index 00000000..267635ad --- /dev/null +++ b/clients/python/tests/scheduler/test_storage_cluster.py @@ -0,0 +1,42 @@ +from collections.abc import Iterator +from datetime import UTC, datetime +from unittest import mock + +import pytest + +# redis-py only ships RedisCluster from 4.1; the locked dev version may be older. +pytest.importorskip("redis.cluster") + +from redis.cluster import PRIMARY, ClusterNode, NodesManager, RedisCluster # noqa: E402 +from redis.crc import REDIS_CLUSTER_HASH_SLOTS # noqa: E402 + +from taskbroker_client.metrics import NoOpMetricsBackend # noqa: E402 +from taskbroker_client.scheduler.storage import RunStorage # noqa: E402 + + +def _single_server_slots(self: NodesManager) -> None: + # CI's redis isn't a cluster, so route every slot to it. The slot check that + # fails is client-side, so this reproduces what a real cluster does. + node = ClusterNode("localhost", 6379, server_type=PRIMARY) + node.redis_connection = self.create_redis_node("localhost", 6379, **self.connection_kwargs) + self.nodes_cache = {node.name: node} + self.slots_cache = {slot: [node] for slot in range(REDIS_CLUSTER_HASH_SLOTS)} + self.default_node = node + + +@pytest.fixture +def cluster_storage() -> Iterator[RunStorage]: + with mock.patch.object(NodesManager, "initialize", _single_server_slots): + client = RedisCluster(host="localhost", port=6379, decode_responses=True) + client.flushdb() + yield RunStorage(metrics=NoOpMetricsBackend(), redis=client) + + +def test_run_storage_read_many_on_redis_cluster(cluster_storage: RunStorage) -> None: + cluster_storage.set("first:test:valid:300", datetime(2099, 1, 1, tzinfo=UTC)) + first_run = cluster_storage.read("first:test:valid:300") + + result = cluster_storage.read_many(["first:test:valid:300", "second:test:valid:300"]) + + assert first_run is not None + assert result == {"first:test:valid:300": first_run, "second:test:valid:300": None}