fix(client): Use mget_nonatomic for RunStorage.read_many on RedisCluster

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(

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}