From 2b1907a56efbc0b62d35dc1204204d50d306558e Mon Sep 17 00:00:00 2001 From: Hector Date: Wed, 4 Nov 2026 16:33:46 +0000 Subject: [PATCH] test(client): Reproduce RunStorage.read_many crash on redis-py RedisCluster On redis>=4.1, RedisCluster.mget raises unless every key shares a hash slot; redis-py-cluster split the call per node. RunStorage.read_many mgets one key per schedule, so the scheduler crashes on startup with more than one schedule. Run the client tests on redis-py 5 (the cluster extra now conflicts with the dev group), add a test reproducing the crash, and one showing that routing mget through mget_nonatomic fixes it. Type-only changes in storage.py are needed for mypy against redis-py's bundled types. Co-Authored-By: Claude Opus 5.5 --- clients/python/pyproject.toml | 3 + .../taskbroker_client/scheduler/storage.py | 18 ++++-- .../python/tests/scheduler/test_storage.py | 58 ++++++++++++++++++- pyproject.toml | 6 ++ uv.lock | 38 +++++++++++- 5 files changed, 113 insertions(+), 10 deletions(-) diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index f92163fd..23a17e20 100644 --- a/clients/python/pyproject.toml +++ b/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. diff --git a/clients/python/src/taskbroker_client/scheduler/storage.py b/clients/python/src/taskbroker_client/scheduler/storage.py index 92fca58a..054a4c73 100644 --- a/clients/python/src/taskbroker_client/scheduler/storage.py +++ b/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( diff --git a/clients/python/tests/scheduler/test_storage.py b/clients/python/tests/scheduler/test_storage.py index db24de8b..8e765684 100644 --- a/clients/python/tests/scheduler/test_storage.py +++ b/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} diff --git a/pyproject.toml b/pyproject.toml index 8d2ba01f..073a8136 100644 --- a/pyproject.toml +++ b/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" diff --git a/uv.lock b/uv.lock index e6cd60e2..5cbbca18 100644 --- a/uv.lock +++ b/uv.lock @@ -7,6 +7,10 @@ resolution-markers = [ supported-markers = [ "sys_platform == 'darwin' or sys_platform == 'linux'", ] +conflicts = [[ + { package = "taskbroker-client", extra = "cluster" }, + { package = "taskbroker-client", group = "dev" }, +]] [manifest] members = [ @@ -15,6 +19,14 @@ members = [ "taskbroker-integration-tests", ] +[[package]] +name = "async-timeout" +version = "4.0.3" +source = { registry = "https://pypi.devinfra.sentry.io/simple" } +wheels = [ + { url = "https://pypi.devinfra.sentry.io/wheels/async_timeout-4.0.3-py3-none-any.whl", hash = "sha256:7405140ff1230c310e51dc27b3145b9092d659ce68ff733fb0cefe3ee42be028" }, +] + [[package]] name = "black" version = "26.3.1" @@ -573,7 +585,7 @@ name = "pytest-cov" version = "4.1.0" source = { registry = "https://pypi.devinfra.sentry.io/simple" } dependencies = [ - { name = "coverage", extra = ["toml"] }, + { name = "coverage", extra = ["toml"], marker = "extra == 'group-17-taskbroker-client-dev'" }, { name = "pytest" }, ] wheels = [ @@ -636,16 +648,33 @@ wheels = [ name = "redis" version = "3.5.3" source = { registry = "https://pypi.devinfra.sentry.io/simple" } +resolution-markers = [ + "sys_platform == 'darwin' or sys_platform == 'linux'", +] wheels = [ { url = "https://pypi.devinfra.sentry.io/wheels/redis-3.5.3-py2.py3-none-any.whl", hash = "sha256:432b788c4530cfe16d8d943a09d40ca6c16149727e4afe8c2c9d5580c59d9f24" }, ] +[[package]] +name = "redis" +version = "5.0.1" +source = { registry = "https://pypi.devinfra.sentry.io/simple" } +resolution-markers = [ + "sys_platform == 'darwin' or sys_platform == 'linux'", +] +dependencies = [ + { name = "async-timeout", marker = "(python_full_version <= '3.11.2' and extra != 'extra-17-taskbroker-client-cluster') or (extra == 'extra-17-taskbroker-client-cluster' and extra == 'group-17-taskbroker-client-dev')" }, +] +wheels = [ + { url = "https://pypi.devinfra.sentry.io/wheels/redis-5.0.1-py3-none-any.whl", hash = "sha256:ed4802971884ae19d640775ba3b03aa2e7bd5e8fb8dfaed2decce4d0fc48391f" }, +] + [[package]] name = "redis-py-cluster" version = "2.1.3" source = { registry = "https://pypi.devinfra.sentry.io/simple" } dependencies = [ - { name = "redis" }, + { name = "redis", version = "3.5.3", source = { registry = "https://pypi.devinfra.sentry.io/simple" } }, ] wheels = [ { url = "https://pypi.devinfra.sentry.io/wheels/redis_py_cluster-2.1.3-py2.py3-none-any.whl", hash = "sha256:38f08850fde469ffd76bced7309721114acc487e52b76f374a0502c34c69b4ec" }, @@ -806,7 +835,8 @@ dependencies = [ { name = "msgpack" }, { name = "prometheus-client" }, { name = "protobuf" }, - { name = "redis" }, + { name = "redis", version = "3.5.3", source = { registry = "https://pypi.devinfra.sentry.io/simple" }, marker = "extra == 'extra-17-taskbroker-client-cluster'" }, + { name = "redis", version = "5.0.1", source = { registry = "https://pypi.devinfra.sentry.io/simple" }, marker = "extra != 'extra-17-taskbroker-client-cluster' or extra == 'group-17-taskbroker-client-dev'" }, { name = "sentry-arroyo" }, { name = "sentry-protos" }, { name = "sentry-sdk", extra = ["http2"] }, @@ -836,6 +866,7 @@ dev = [ { name = "pre-commit" }, { name = "pytest" }, { name = "pytest-cov" }, + { name = "redis", version = "5.0.1", source = { registry = "https://pypi.devinfra.sentry.io/simple" } }, { name = "sentry-devenv" }, { name = "time-machine" }, { name = "types-protobuf" }, @@ -873,6 +904,7 @@ dev = [ { name = "pre-commit", specifier = ">=4.2.0" }, { name = "pytest", specifier = ">=8.3.3" }, { name = "pytest-cov", specifier = ">=4.0.0" }, + { name = "redis", specifier = ">=4.1" }, { name = "sentry-devenv", specifier = ">=1.22.2" }, { name = "time-machine", specifier = ">=2.16.0" }, { name = "types-protobuf", specifier = ">=5.27.0.20240626,<6.0.0" },