diff --git a/benchmarks/connection_handshake.py b/benchmarks/connection_handshake.py new file mode 100644 index 0000000000..6d75edd80c --- /dev/null +++ b/benchmarks/connection_handshake.py @@ -0,0 +1,186 @@ +""" +Benchmark for the sync connection handshake (``Connection.on_connect``). + +Opening a connection runs a short sequence of setup commands -- HELLO/AUTH, +CLIENT SETNAME, CLIENT SETINFO (LIB-NAME / LIB-VER) and SELECT. This benchmark +times how long a full connection handshake takes, averaged over many fresh +connections, and reports latency percentiles. + +The handshake cost is dominated by network round-trips, so on localhost (where +RTT is effectively zero) the difference between a sequential and a pipelined +handshake is hard to see. Pass ``--delay-ms`` to inject an artificial per +round-trip latency: the injected delay is charged once per *round-trip* (the +first socket read after a write), so a handshake that issues N sequential +send/read round-trips pays roughly N * delay, while a pipelined handshake that +collapses several commands into one write + batched reads pays it far fewer +times. This makes the round-trip reduction observable without a remote server. + +To measure the improvement of a change, run this script on the baseline commit +and on the change (same arguments) and compare the reported means, e.g.: + + python -m benchmarks.connection_handshake --n 2000 --client-name bench --db 9 --delay-ms 1 + +Assumes a Redis server reachable at --host/--port (defaults to localhost:6379). +Not part of the CI test matrix. +""" + +import argparse +import statistics +import time + +import redis +from redis.maint_notifications import MaintNotificationsConfig + + +class _LatencySocket: + """Socket proxy that charges ``delay`` seconds once per round-trip. + + A round-trip is modeled as: one or more writes, then the first read that + follows them (the wait for the reply to come back). Replies that are already + in flight -- i.e. subsequent reads with no intervening write, as produced by + a pipelined batch -- are served without additional delay. + """ + + def __init__(self, sock, delay): + self._sock = sock + self._delay = delay + self._pending = False + + def sendall(self, data, *args, **kwargs): + self._pending = True + return self._sock.sendall(data, *args, **kwargs) + + def send(self, data, *args, **kwargs): + self._pending = True + return self._sock.send(data, *args, **kwargs) + + def _wait_for_reply(self): + if self._pending: + time.sleep(self._delay) + self._pending = False + + def recv(self, *args, **kwargs): + self._wait_for_reply() + return self._sock.recv(*args, **kwargs) + + def recv_into(self, *args, **kwargs): + self._wait_for_reply() + return self._sock.recv_into(*args, **kwargs) + + def __getattr__(self, name): + # Delegate everything else (settimeout, close, fileno, ...) to the socket. + return getattr(self._sock, name) + + +class _LatencyConnection(redis.Connection): + """Connection whose socket injects an artificial per round-trip latency.""" + + def __init__(self, *args, latency_seconds=0.0, **kwargs): + self._latency_seconds = latency_seconds + super().__init__(*args, **kwargs) + + def _connect(self): + sock = super()._connect() + if self._latency_seconds > 0: + sock = _LatencySocket(sock, self._latency_seconds) + return sock + + +def _percentile(sorted_values, pct): + if not sorted_values: + return float("nan") + k = (len(sorted_values) - 1) * (pct / 100.0) + lo = int(k) + hi = min(lo + 1, len(sorted_values) - 1) + if lo == hi: + return sorted_values[lo] + return sorted_values[lo] + (sorted_values[hi] - sorted_values[lo]) * (k - lo) + + +def run(args): + conn_kwargs = dict( + host=args.host, + port=args.port, + db=args.db, + protocol=args.protocol, + client_name=args.client_name or None, + ) + if args.username: + conn_kwargs["username"] = args.username + if args.password: + conn_kwargs["password"] = args.password + if args.delay_ms > 0: + conn_kwargs["latency_seconds"] = args.delay_ms / 1000.0 + + if args.maint_notifications != "off": + # Maintenance notifications require RESP3; the connection would otherwise + # raise when configuring the (RESP2) parser. + if args.protocol != 3: + raise SystemExit("--maint-notifications requires --protocol 3 (RESP3).") + # enabled="auto" still sends CLIENT MAINT_NOTIFICATIONS (so its cost is + # measured) but tolerates servers that don't support it; enabled=True + # (mode "on") fails the connection if the server rejects the command. + enabled = True if args.maint_notifications == "on" else "auto" + conn_kwargs["maint_notifications_config"] = MaintNotificationsConfig( + enabled=enabled + ) + + connection_cls = _LatencyConnection if args.delay_ms > 0 else redis.Connection + + latencies_ms = [] + for _ in range(args.n): + conn = connection_cls(**conn_kwargs) + start = time.perf_counter() + conn.connect() + latencies_ms.append((time.perf_counter() - start) * 1000.0) + conn.disconnect() + + latencies_ms.sort() + print(f"Connection handshake benchmark ({args.n} connections)") + print( + f" host={args.host}:{args.port} db={args.db} protocol={args.protocol} " + f"client_name={args.client_name or ''!r} delay_ms={args.delay_ms} " + f"maint_notifications={args.maint_notifications}" + ) + print(f" mean : {statistics.fmean(latencies_ms):.3f} ms") + print(f" median : {statistics.median(latencies_ms):.3f} ms") + print(f" p95 : {_percentile(latencies_ms, 95):.3f} ms") + print(f" p99 : {_percentile(latencies_ms, 99):.3f} ms") + print(f" min : {latencies_ms[0]:.3f} ms") + print(f" max : {latencies_ms[-1]:.3f} ms") + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="localhost") + parser.add_argument("--port", type=int, default=6379) + parser.add_argument("--username", default=None) + parser.add_argument("--password", default=None) + parser.add_argument("--protocol", type=int, default=3, choices=(2, 3)) + parser.add_argument("--client-name", default="") + parser.add_argument("--db", type=int, default=0) + parser.add_argument( + "--n", type=int, default=1000, help="number of fresh connections to time" + ) + parser.add_argument( + "--delay-ms", + type=float, + default=0.0, + help="artificial per round-trip latency (ms) to make the handshake " + "round-trip count visible on localhost", + ) + parser.add_argument( + "--maint-notifications", + choices=("off", "auto", "on"), + default="off", + help="enable maintenance notifications during the handshake (requires " + "--protocol 3). 'auto' sends CLIENT MAINT_NOTIFICATIONS but tolerates " + "servers that reject it; 'on' fails the connection if it is rejected. " + "With this change the command is pipelined into the handshake, so " + "enabling it does not add a round-trip.", + ) + return parser.parse_args() + + +if __name__ == "__main__": + run(parse_args()) diff --git a/redis/connection.py b/redis/connection.py index bd2e141240..4236ebf8ad 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -609,21 +609,24 @@ def set_maint_notifications_cluster_handler_for_connection( oss_cluster_maint_notifications_handler.config ) - def activate_maint_notifications_handling_if_enabled(self, check_health=True): - # Send maintenance notifications handshake if RESP3 is active + def _should_enable_maint_notifications(self) -> bool: + # Maintenance notifications are sent only if RESP3 is active # and maintenance notifications are enabled - # and we have a host to determine the endpoint type from - # When the maint_notifications_config enabled mode is "auto", - # we just log a warning if the handshake fails - # When the mode is enabled=True, we raise an exception in case of failure + # and we have a host to determine the endpoint type from. host = getattr(self, "host", None) - if ( + return bool( check_protocol_version(self.get_protocol(), 3) and self.maint_notifications_config and self.maint_notifications_config.enabled and self._maint_notifications_connection_handler and host is not None - ): + ) + + def activate_maint_notifications_handling_if_enabled(self, check_health=True): + # When the maint_notifications_config enabled mode is "auto", + # we just log a warning if the handshake fails + # When the mode is enabled=True, we raise an exception in case of failure + if self._should_enable_maint_notifications(): self._enable_maintenance_notifications( maint_notifications_config=self.maint_notifications_config, check_health=check_health, @@ -631,29 +634,70 @@ def activate_maint_notifications_handling_if_enabled(self, check_health=True): def _enable_maintenance_notifications( self, maint_notifications_config: MaintNotificationsConfig, check_health=True + ): + # Kept for callers that enable maintenance notifications outside of the + # connection handshake. During on_connect the send and the response + # handling are split (see _send_maint_notifications_command / + # _handle_maint_notifications_response) so the reply can be pipelined + # with the rest of the handshake. + self._send_maint_notifications_command( + maint_notifications_config, check_health=check_health + ) + self._handle_maint_notifications_response(maint_notifications_config) + + def _maint_notifications_command_args( + self, maint_notifications_config: MaintNotificationsConfig + ): + host = getattr(self, "host", None) + if host is None: + raise ValueError( + "Cannot enable maintenance notifications for connection" + " object that doesn't have a host attribute." + ) + endpoint_type = maint_notifications_config.get_endpoint_type(host, self) + return ( + "CLIENT", + "MAINT_NOTIFICATIONS", + "ON", + "moving-endpoint-type", + endpoint_type.value, + ) + + def _send_maint_notifications_command( + self, maint_notifications_config: MaintNotificationsConfig, check_health=True + ): + self.send_command( + *self._maint_notifications_command_args(maint_notifications_config), + check_health=check_health, + ) + + def _add_maint_notifications_to_handshake(self, deferred_reads, check_health=True): + # If maintenance notifications are enabled for this connection, send the + # CLIENT MAINT_NOTIFICATIONS command as part of the pipelined handshake tail + # and defer reading its reply (appended to deferred_reads), rather than paying + # its own round-trip. When enabled == "auto" a failure is logged and swallowed; + # when enabled is True it raises. + if not self._should_enable_maint_notifications(): + return + maint_notifications_config = self.maint_notifications_config + self._send_maint_notifications_command( + maint_notifications_config, check_health=check_health + ) + deferred_reads.append( + lambda: self._handle_maint_notifications_response( + maint_notifications_config + ) + ) + + def _handle_maint_notifications_response( + self, maint_notifications_config: MaintNotificationsConfig ): try: - host = getattr(self, "host", None) - if host is None: - raise ValueError( - "Cannot enable maintenance notifications for connection" - " object that doesn't have a host attribute." - ) - else: - endpoint_type = maint_notifications_config.get_endpoint_type(host, self) - self.send_command( - "CLIENT", - "MAINT_NOTIFICATIONS", - "ON", - "moving-endpoint-type", - endpoint_type.value, - check_health=check_health, + response = self.read_response() + if not response or str_if_bytes(response) != "OK": + raise ResponseError( + "The server doesn't support maintenance notifications" ) - response = self.read_response() - if not response or str_if_bytes(response) != "OK": - raise ResponseError( - "The server doesn't support maintenance notifications" - ) except Exception as e: if ( isinstance(e, ResponseError) @@ -1172,54 +1216,80 @@ def on_connect_check_health(self, check_health: bool = True): ): raise ConnectionError("Invalid RESP version") - # Activate maintenance notifications for this connection - # if enabled in the configuration - # This is a no-op if maintenance notifications are not enabled - self.activate_maint_notifications_handling_if_enabled(check_health=check_health) + # The tail of the handshake (optional CLIENT MAINT_NOTIFICATIONS, then + # CLIENT SETNAME / SETINFO / SELECT) does not affect control flow -- the replies + # are only validated or discarded. So we pipeline it: send every command first + # (without blocking on a reply between them), then read the replies back in send + # order. All requests are on the wire before we block on the first read, so the + # whole tail costs a single round-trip instead of one per command. This mirrors + # the async stack (redis/asyncio/connection.py). AUTH/HELLO stays a separate + # round-trip above because its reply drives the RESP2->RESP3 parser upgrade, the + # pre-6.0 AUTH retry, and proto validation. + # + # deferred_reads holds one zero-arg handler per command sent below; each reads + # exactly one reply (in send order) and validates it. Per-command check_health + # reproduces the original behavior: at most one health PING/PONG fires before the + # first tail command (only when no HELLO/AUTH ran, e.g. RESP2 no-auth), and it is + # self-contained so it never desyncs the pipelined replies. + deferred_reads = [] + + # Maintenance notifications (RESP3-only, opt-in) go first when enabled, so their + # reply is pipelined with the rest of the tail. + self._add_maint_notifications_to_handshake(deferred_reads, check_health) # if a client_name is given, set it if self.client_name: self.send_command( - "CLIENT", - "SETNAME", - self.client_name, - check_health=check_health, + "CLIENT", "SETNAME", self.client_name, check_health=check_health ) - if str_if_bytes(self.read_response()) != "OK": - raise ConnectionError("Error setting client name") - # Set the library name and version from driver_info - try: - if self.driver_info and self.driver_info.formatted_name: - self.send_command( - "CLIENT", - "SETINFO", - "LIB-NAME", - self.driver_info.formatted_name, - check_health=check_health, - ) - self.read_response() - except ResponseError: - pass + def _read_setname_response(): + if str_if_bytes(self.read_response()) != "OK": + raise ConnectionError("Error setting client name") - try: - if self.driver_info and self.driver_info.lib_version: - self.send_command( - "CLIENT", - "SETINFO", - "LIB-VER", - self.driver_info.lib_version, - check_health=check_health, - ) + deferred_reads.append(_read_setname_response) + + # Set the library name and version from driver_info. Older servers may not + # support CLIENT SETINFO, so any ResponseError to these replies is swallowed. + def _read_setinfo_response(): + try: self.read_response() - except ResponseError: - pass + except ResponseError: + pass + + if self.driver_info and self.driver_info.formatted_name: + self.send_command( + "CLIENT", + "SETINFO", + "LIB-NAME", + self.driver_info.formatted_name, + check_health=check_health, + ) + deferred_reads.append(_read_setinfo_response) + + if self.driver_info and self.driver_info.lib_version: + self.send_command( + "CLIENT", + "SETINFO", + "LIB-VER", + self.driver_info.lib_version, + check_health=check_health, + ) + deferred_reads.append(_read_setinfo_response) # if a database is specified, switch to it if self.db: self.send_command("SELECT", self.db, check_health=check_health) - if str_if_bytes(self.read_response()) != "OK": - raise ConnectionError("Invalid Database") + + def _read_select_response(): + if str_if_bytes(self.read_response()) != "OK": + raise ConnectionError("Invalid Database") + + deferred_reads.append(_read_select_response) + + # Read the deferred replies in the order the commands were sent. + for read_and_validate_response in deferred_reads: + read_and_validate_response() def disconnect(self, *args, **kwargs): "Disconnects from the Redis server" diff --git a/tests/maint_notifications/test_maint_notifications_handling.py b/tests/maint_notifications/test_maint_notifications_handling.py index 34340ae44e..78ea19064a 100644 --- a/tests/maint_notifications/test_maint_notifications_handling.py +++ b/tests/maint_notifications/test_maint_notifications_handling.py @@ -434,6 +434,8 @@ def _get_client( enable_cache=False, max_connections=10, maint_notifications_config=None, + client_name=None, + db=0, ): """Helper method to create a pool and Redis client with maintenance notifications configuration. @@ -456,12 +458,16 @@ def _get_client( if enable_cache: pool_kwargs = {"cache_config": CacheConfig()} + if client_name is not None: + pool_kwargs["client_name"] = client_name + test_pool = pool_class( connection_class=connection_class, host=DEFAULT_ADDRESS.split(":")[0], port=int(DEFAULT_ADDRESS.split(":")[1]), max_connections=max_connections, protocol=3, # Required for maintenance notifications + db=db, maint_notifications_config=config, **pool_kwargs, ) @@ -534,6 +540,49 @@ def test_handshake_failure_when_enabled(self): finally: test_redis_client.close() + def test_handshake_pipelines_full_tail_with_client_name_and_db(self): + """The whole handshake tail -- CLIENT MAINT_NOTIFICATIONS, SETNAME, both + SETINFO commands and SELECT -- is pipelined (all sent, then the replies read + back in send order). + + Uses enabled="auto" with an internal-ip endpoint so the server rejects + MAINT_NOTIFICATIONS and its error reply is swallowed *mid-tail*. Because the + maint command is first, a following command only succeeds if the swallow + consumed exactly one reply and the remaining SETNAME/SETINFO/SELECT replies + stayed aligned -- this is the reply/command-ordering contract the whole + optimization depends on. + """ + maint_notifications_config = MaintNotificationsConfig( + enabled="auto", endpoint_type=EndpointType.INTERNAL_IP + ) + test_redis_client = self._get_client( + ConnectionPool, + maint_notifications_config=maint_notifications_config, + client_name="myclient", + db=5, + ) + try: + # Post-handshake commands succeed -> replies stayed aligned even though + # the swallowed MAINT_NOTIFICATIONS error sits first in the tail. + assert test_redis_client.set("hello", "world") is True + assert test_redis_client.get("hello") == b"world" + + # HELLO first (its own round-trip), then the tail commands in send order. + handshake_sock = next( + s for s in self.mock_sockets if any(b"HELLO" in w for w in s.sent_data) + ) + wire = b"".join(handshake_sock.sent_data) + assert ( + wire.index(b"HELLO") + < wire.index(b"MAINT_NOTIFICATIONS") + < wire.index(b"SETNAME") + < wire.index(b"LIB-NAME") + < wire.index(b"LIB-VER") + < wire.index(b"SELECT") + ) + finally: + test_redis_client.close() + @pytest.mark.fixed_client class TestMaintenanceNotificationsHandlingSingleProxy(TestMaintenanceNotificationsBase):