Add latency-aware cluster read balancing

Author: lh0156

benchmarks/README.md

@@ -123,3 +123,20 @@ Stop the example stack when finished:
 ```shell
 docker compose -f docs/examples/opentelemetry/docker-compose.yml down
 ```
+
+## Cluster Latency-Aware Read Balancing
+
+`cluster_latency_load_balancing.py` compares round-robin reads with
+`LoadBalancingStrategy.LATENCY_BASED` on one fixed key slot. It adds a
+configurable delay to one replica's async client command path, then reports
+that replica's selection share and the measured p99 latency:
+
+```shell
+python -m benchmarks.cluster_latency_load_balancing \
+  --host 127.0.0.1 --port 7000 --delay-ms 10 \
+  --requests 2000 --concurrency 32
+```
+
+The cluster must expose at least one replica for the selected key slot and
+must permit read-only connections. Pass `--delayed-node HOST:PORT` to choose a
+specific replica; otherwise the last replica in the slot's topology is used.

benchmarks/cluster_latency_load_balancing.py

@@ -0,0 +1,138 @@
+"""Compare round-robin and latency-aware reads on a Redis cluster.
+
+The benchmark injects delay into one replica's client-side command path. This
+keeps the setup small and repeatable while exercising the same async
+``RedisCluster`` selection and measurement paths used in production.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import math
+import time
+from collections import Counter
+from dataclasses import dataclass, field
+
+from redis.asyncio.cluster import ClusterNode, RedisCluster
+from redis.cluster import LoadBalancingStrategy
+
+
+@dataclass
+class Measurements:
+    counts: Counter[str] = field(default_factory=Counter)
+    latencies: list[float] = field(default_factory=list)
+
+    @property
+    def p99_ms(self) -> float:
+        if not self.latencies:
+            return 0.0
+        ordered = sorted(self.latencies)
+        index = min(len(ordered) - 1, math.ceil(len(ordered) * 0.99) - 1)
+        return ordered[index] * 1000
+
+
+async def run_mode(
+    host: str,
+    port: int,
+    key: str,
+    requests: int,
+    concurrency: int,
+    delay_ms: float,
+    strategy: LoadBalancingStrategy,
+    delayed_node_name: str | None,
+) -> tuple[str, Measurements, str]:
+    client = await RedisCluster(
+        host=host,
+        port=port,
+        load_balancing_strategy=strategy,
+    )
+    slot_nodes = client.nodes_manager.slots_cache[client.keyslot(key)]
+    replicas = slot_nodes[1:]
+    if not replicas:
+        await client.aclose()
+        raise RuntimeError("the selected key slot must have at least one replica")
+
+    delayed_node = next(
+        (node for node in replicas if node.name == delayed_node_name),
+        replicas[-1] if delayed_node_name is None else None,
+    )
+    if delayed_node is None:
+        await client.aclose()
+        raise ValueError(
+            f"node {delayed_node_name!r} is not a replica for key slot "
+            f"{client.keyslot(key)}"
+        )
+
+    measurements = Measurements()
+    original_execute_command = ClusterNode.execute_command
+
+    async def delayed_execute_command(node, *args, **kwargs):
+        started = time.perf_counter()
+        if args and args[0] == "GET":
+            measurements.counts[node.name] += 1
+            if node.name == delayed_node.name:
+                await asyncio.sleep(delay_ms / 1000)
+        try:
+            return await original_execute_command(node, *args, **kwargs)
+        finally:
+            if args and args[0] == "GET":
+                measurements.latencies.append(time.perf_counter() - started)
+
+    ClusterNode.execute_command = delayed_execute_command
+    try:
+        requests_per_worker, remainder = divmod(requests, concurrency)
+
+        async def worker(worker_index: int) -> None:
+            worker_requests = requests_per_worker + (worker_index < remainder)
+            for _ in range(worker_requests):
+                await client.get(key)
+
+        await asyncio.gather(*(worker(i) for i in range(concurrency)))
+    finally:
+        ClusterNode.execute_command = original_execute_command
+        await client.aclose()
+
+    return strategy.value, measurements, delayed_node.name
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=6379)
+    parser.add_argument("--key", default="{latency-benchmark}:key")
+    parser.add_argument("--requests", type=int, default=2000)
+    parser.add_argument("--concurrency", type=int, default=32)
+    parser.add_argument("--delay-ms", type=float, default=10.0)
+    parser.add_argument("--delayed-node")
+    return parser.parse_args()
+
+
+async def main(args: argparse.Namespace) -> None:
+    if args.requests < args.concurrency:
+        raise ValueError("--requests must be at least --concurrency")
+
+    for strategy in (
+        LoadBalancingStrategy.ROUND_ROBIN,
+        LoadBalancingStrategy.LATENCY_BASED,
+    ):
+        name, measurements, delayed_node = await run_mode(
+            host=args.host,
+            port=args.port,
+            key=args.key,
+            requests=args.requests,
+            concurrency=args.concurrency,
+            delay_ms=args.delay_ms,
+            strategy=strategy,
+            delayed_node_name=args.delayed_node,
+        )
+        total = sum(measurements.counts.values())
+        delayed_share = measurements.counts[delayed_node] / total * 100
+        print(
+            f"strategy={name:<14} delayed_node={delayed_node:<24} "
+            f"delayed_share={delayed_share:6.2f}% p99={measurements.p99_ms:8.3f}ms"
+        )
+
+
+if __name__ == "__main__":
+    asyncio.run(main(parse_args()))

docs/clustering.rst

@@ -223,6 +223,20 @@ the primary and its replications in a Round-Robin manner.
 With load_balancing_strategy you can define a custom strategy for
 assigning read commands to the replicas and primary nodes.
 
+``LoadBalancingStrategy.LATENCY_BASED`` uses power-of-two choices and a
+peak-sensitive latency estimate to prefer the less busy of two sampled
+nodes.  The strategy is opt-in; existing round-robin and random strategies
+are unchanged.  Latency measurements are collected separately for each
+cluster node and are shared by regular commands and pipelines.
+
+.. code:: python
+
+   >>> from redis.cluster import LoadBalancingStrategy, RedisCluster
+   >>> rc = RedisCluster(
+   ...     startup_nodes=startup_nodes,
+   ...     load_balancing_strategy=LoadBalancingStrategy.LATENCY_BASED,
+   ... )
+
 READONLY mode can be set at runtime by calling the readonly() method
 with target_nodes=‘replicas’, and read-write access can be restored by
 calling the readwrite() method.

redis/asyncio/cluster.py

@@ -1205,6 +1205,8 @@ async def _execute_command(
 
         while ttl > 0:
             ttl -= 1
+            tracked_node_name = None
+            attempt_start_time = None
             try:
                 if asking:
                     target_node = self.get_node(node_name=redirect_addr)
@@ -1223,6 +1225,15 @@ async def _execute_command(
                     )
                     moved = False
 
+                if (
+                    self.load_balancing_strategy == LoadBalancingStrategy.LATENCY_BASED
+                    and command in READ_COMMANDS
+                ):
+                    tracked_node_name = target_node.name
+                    attempt_start_time = time.monotonic()
+                    self.nodes_manager.read_load_balancer.start_request(
+                        tracked_node_name
+                    )
                 response = await target_node.execute_command(*args, **kwargs)
                 await self._record_command_metric(
                     command_name=command,
@@ -1370,6 +1381,11 @@ async def _execute_command(
                     error=e,
                 )
                 raise
+            finally:
+                if tracked_node_name is not None and attempt_start_time is not None:
+                    self.nodes_manager.read_load_balancer.record_latency(
+                        tracked_node_name, time.monotonic() - attempt_start_time
+                    )
 
         e = ClusterError("TTL exhausted.")
         e.connection = target_node
@@ -2066,7 +2082,10 @@ def get_node_from_slot(
                 # get the server index using the strategy defined in load_balancing_strategy
                 primary_name = self.slots_cache[slot][0].name
                 node_idx = self.read_load_balancer.get_server_index(
-                    primary_name, len(self.slots_cache[slot]), load_balancing_strategy
+                    primary_name,
+                    len(self.slots_cache[slot]),
+                    load_balancing_strategy,
+                    nodes=self.slots_cache[slot],
                 )
                 return self.slots_cache[slot][node_idx]
             return self.slots_cache[slot][0]
@@ -2793,10 +2812,27 @@ async def _execute(
         # Start timing for observability
         start_time = time.monotonic()
 
+        async def execute_node_pipeline(node, commands):
+            track_latency = (
+                client.load_balancing_strategy == LoadBalancingStrategy.LATENCY_BASED
+                and any(command.args[0] in READ_COMMANDS for command in commands)
+            )
+            if not track_latency:
+                return await node.execute_pipeline(commands)
+
+            client.nodes_manager.read_load_balancer.start_request(node.name)
+            start_time = time.monotonic()
+            try:
+                return await node.execute_pipeline(commands)
+            finally:
+                client.nodes_manager.read_load_balancer.record_latency(
+                    node.name, time.monotonic() - start_time
+                )
+
         errors = await asyncio.gather(
             *(
-                asyncio.create_task(node[0].execute_pipeline(node[1]))
-                for node in nodes.values()
+                asyncio.create_task(execute_node_pipeline(node, commands))
+                for node, commands in nodes.values()
             )
         )