DNM: Reproduce RunStorage.read_many crash on redis-py RedisCluster

clients/python/pyproject.toml

@@ -32,6 +32,9 @@ dev = [
     "time-machine>=2.16.0",
     "types-protobuf>=5.27.0.20240626,<6.0.0",
     "datadog>=0.49",
+    # Tests run against redis-py's built-in RedisCluster, which redis-py-cluster (the
+    # `cluster` extra) can't be installed alongside.
+    "redis>=4.1",
 ]
 [project.optional-dependencies]
 # redis<4 does not bundle cluster support; pull in redis-py-cluster (EOL) for it.

clients/python/src/taskbroker_client/scheduler/storage.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Mapping
 from datetime import UTC, datetime
-from typing import TYPE_CHECKING, Protocol
+from typing import TYPE_CHECKING, Protocol, cast
 
 from redis.client import StrictRedis
 
@@ -11,7 +11,7 @@
         # redis>=4.1 ships cluster support natively
         from redis import RedisCluster
     except ImportError:  # pragma: no cover - redis<4 fallback via the `cluster` extra
-        from rediscluster import RedisCluster
+        from rediscluster import RedisCluster  # type: ignore[no-redef]
 
 from taskbroker_client.metrics import MetricsBackend
 
@@ -138,7 +138,7 @@ def read(self, key: str) -> datetime | None:
         Retrieve the last run time of a task
         Returns None if last run time has expired or is unknown.
         """
-        result = self._redis.get(self._make_key(key))
+        result = cast(str | None, self._redis.get(self._make_key(key)))
         if result:
             return datetime.fromisoformat(result)
 
@@ -168,9 +168,15 @@ 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 = cast(
+            list[str | None], self._redis.mget([self._make_key(sk) for sk in storage_keys])
+        )
+        compat_values = cast(
+            list[str | None], self._redis.mget([self._make_key(ck) for ck in compat_keys])
+        )
+        legacy_values = cast(
+            list[str | None], self._redis.mget([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.py

@@ -1,6 +1,15 @@
+from collections.abc import Iterator
 from datetime import UTC, datetime
+from typing import Any
+from unittest import mock
 
-from taskbroker_client.scheduler.storage import VolatileRunStorage
+import pytest
+from redis.cluster import PRIMARY, ClusterNode, NodesManager, RedisCluster
+from redis.crc import REDIS_CLUSTER_HASH_SLOTS
+from redis.exceptions import RedisClusterException
+
+from taskbroker_client.metrics import NoOpMetricsBackend
+from taskbroker_client.scheduler.storage import RunStorage, VolatileRunStorage
 
 from ..conftest import freeze_time
 
@@ -60,3 +69,50 @@ def test_volatile_run_storage_delete() -> None:
     storage.delete("key")
     storage.delete("missing")
     assert storage.read("key") is None
+
+
+class SeerRedisCluster(RedisCluster):
+    """Workaround used by Seer: ``mget`` keys across hash slots, like redis-py-cluster did."""
+
+    def mget(self, keys: Any, *args: Any) -> Any:
+        return self.mget_nonatomic(keys, *args)
+
+
+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  # type: ignore[assignment]
+
+
+@pytest.fixture
+def cluster_storage(request: pytest.FixtureRequest) -> Iterator[RunStorage]:
+    with mock.patch.object(NodesManager, "initialize", _single_server_slots):
+        client = request.param(host="localhost", port=6379, decode_responses=True)
+    storage = RunStorage(metrics=NoOpMetricsBackend(), redis=client)
+    keys = ["first:test:valid:300", "second:test:valid:300"]
+    yield storage
+    for key in keys:
+        storage.delete(key)
+
+
+@pytest.mark.parametrize("cluster_storage", [RedisCluster], indirect=True)
+def test_run_storage_read_many_fails_on_redis_cluster(cluster_storage: RunStorage) -> None:
+    cluster_storage.set("first:test:valid:300", datetime(2099, 1, 1, tzinfo=UTC))
+
+    with pytest.raises(RedisClusterException, match="all keys must map to the same key slot"):
+        cluster_storage.read_many(["first:test:valid:300", "second:test:valid:300"])
+
+
+@pytest.mark.parametrize("cluster_storage", [SeerRedisCluster], indirect=True)
+def test_run_storage_read_many_with_cross_slot_mget(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}

pyproject.toml

@@ -28,6 +28,12 @@ dev = [
 
 [tool.uv]
 environments = ["sys_platform == 'darwin' or sys_platform == 'linux'"]
+conflicts = [
+    [
+        { package = "taskbroker-client", extra = "cluster" },
+        { package = "taskbroker-client", group = "dev" },
+    ],
+]
 
 [[tool.uv.index]]
 name = "sentry-pypi"