From 8eef2becb879e65062c6b796d0536b6487c90f4c Mon Sep 17 00:00:00 2001 From: Aljoscha Krettek Date: Mon, 3 Aug 2026 11:26:29 +0000 Subject: [PATCH 1/3] cluster-spec-sheet: recreate the region for envd CPU changes `mz region enable --environmentd-cpu-allocation` on a region that still has durable state is a 0dt rollout, so the new environmentd comes up in read-only mode and only promotes after the caught-up stability period, ten minutes fleet-wide. The sweep soft-disabled the old envd away first, so nothing served SQL while we waited that out and the 300s readiness wait failed. Hard-disable instead: with no state left there is no predecessor generation to catch up with, the new envd promotes as soon as it has booted, and the sweep re-prepares the small amount of state it needs per scale point. Readiness also gets a positive check that we reached the recreated region (no user tables) plus sustained probes, since neither `mz region enable` nor a single query can tell a leftover envd from the new one. --- test/cluster-spec-sheet/mzcompose.py | 192 ++++++++++++++++++--------- 1 file changed, 131 insertions(+), 61 deletions(-) diff --git a/test/cluster-spec-sheet/mzcompose.py b/test/cluster-spec-sheet/mzcompose.py index 476085025a6d3..f2bce62a4e4a8 100644 --- a/test/cluster-spec-sheet/mzcompose.py +++ b/test/cluster-spec-sheet/mzcompose.py @@ -34,7 +34,6 @@ from materialize import MZ_ROOT, buildkite from materialize.mz_env_util import print_environment_id from materialize.mz_version import MzVersion -from materialize.mzcompose import _wait_for_pg from materialize.mzcompose.composition import ( Composition, WorkflowArgumentParser, @@ -3237,9 +3236,11 @@ def measure(self, runner: ScenarioRunner, point: ScalePoint) -> None: class EnvdCpuSweep(Scenario): """Sweep environmentd's CPU allocation at a fixed compute cluster size. - Wraps a `ClusterScalingScenario` workload. On Cloud, we reset envd's - CPU allocation to the default in ``teardown`` regardless of outcome to - avoid accidentally burning credits.""" + Wraps a `ClusterScalingScenario` workload. On Cloud, each scale point + recreates the region, so every point measures a freshly set up + environment. We reset envd's CPU allocation to the default in + ``teardown`` regardless of outcome to avoid accidentally burning + credits.""" # (So far, I haven't seen a difference between 16 and 32 in manual # testing in cloud. When we start seeing a difference, consider @@ -3283,6 +3284,10 @@ def apply(self, runner: ScenarioRunner, point: ScalePoint) -> bool: assert point.envd_cpus is not None assert self._fixed_replica_size is not None reconfigure_envd_cpus(runner.target, point.envd_cpus, runner) + # On Cloud the CPU change recreates the region, which drops the state that + # `prepare` set up. Re-preparing is cheap, and harmless on Docker, where the + # state survives the restart. + runner.connection.retryable(lambda: self.prepare(runner)) fixed_size = self._fixed_replica_size def recreate() -> None: @@ -3301,20 +3306,10 @@ def measure(self, runner: ScenarioRunner, point: ScalePoint) -> None: def teardown(self, runner: ScenarioRunner) -> None: if isinstance(runner.target, CloudTarget): print("--- Resetting Cloud environmentd CPUs to the default") - target = runner.target - version_args = ( - ["--version", target.version] if target.version is not None else [] - ) - target.composition.run( - "mz", - "region", - "enable", - "--environmentd-cpu-allocation", - "2", - *enable_extra_args(target), - *version_args, - rm=True, - ) + # Recreating rather than enabling on top of the existing region keeps the + # region out of the read-only limbo that a 0dt rollover would put it in, so + # whatever runs after this scenario finds a usable region. + cloud_recreate_region_with_envd_cpus(runner.target, 2) class EnvdObjectsSweep(Scenario): @@ -3407,10 +3402,7 @@ def enable_extra_args(target: "CloudTarget") -> list[str]: return _STABILITY_SOAK_OVERRIDE if target.is_staging else [] -def cloud_disable_enable_and_wait( - target: "BenchTarget", - environmentd_cpu_allocation: int | None = None, -) -> None: +def cloud_disable_enable_and_wait(target: "BenchTarget") -> None: """ Soft-disable and then enable the Cloud region, then wait for environmentd readiness. @@ -3424,9 +3416,6 @@ def cloud_disable_enable_and_wait( envd were still serving, both `mz region enable`'s readiness check and our `wait_for_envd` could pass against the old envd. Waiting for the old envd to stop serving first makes a later successful connection prove that the new envd is up. - - When `environmentd_cpu_allocation` is provided, it is passed to `mz region enable` via - `--environmentd-cpu-allocation` to reconfigure environmentd's CPU allocation. """ assert isinstance(target, CloudTarget) @@ -3443,27 +3432,11 @@ def cloud_disable_enable_and_wait( if host is not None: wait_for_envd_down(target, host) - version_args = ( - ["--version", target.version] - if isinstance(target, CloudTarget) and target.version is not None - else [] - ) + version_args = ["--version", target.version] if target.version is not None else [] - if environmentd_cpu_allocation is None: - target.composition.run( - "mz", "region", "enable", *enable_extra_args(target), *version_args, rm=True - ) - else: - target.composition.run( - "mz", - "region", - "enable", - "--environmentd-cpu-allocation", - str(environmentd_cpu_allocation), - *enable_extra_args(target), - *version_args, - rm=True, - ) + target.composition.run( + "mz", "region", "enable", *enable_extra_args(target), *version_args, rm=True + ) time.sleep(10) @@ -3471,6 +3444,64 @@ def cloud_disable_enable_and_wait( wait_for_envd(target) +def cloud_recreate_region_with_envd_cpus(target: "CloudTarget", envd_cpus: int) -> None: + """ + Recreate the Cloud region from scratch, with the given environmentd CPU allocation. + + All state in the region is dropped, so the caller has to set up whatever the scenario + needs again. + + We hard-disable rather than soft-disable on purpose. Enabling a region that still has + durable state is a 0dt rollout: the new environmentd comes up in read-only mode and + only promotes once every cluster has been caught up *and stable* for + `with_0dt_caught_up_check_stability_period`, which is ten minutes fleet-wide. Nothing + serves SQL while we wait that out, because the soft disable already took the old envd + away. A hard disable leaves no predecessor generation to catch up with, so the new envd + promotes as soon as it has booted, in about a minute. + """ + disable_region(target.composition, hard=True) + + version_args = ["--version", target.version] if target.version is not None else [] + + target.composition.run( + "mz", + "region", + "enable", + "--environmentd-cpu-allocation", + str(envd_cpus), + *enable_extra_args(target), + *version_args, + rm=True, + ) + + assert "materialize.cloud" in target.composition.cloud_hostname() + wait_for_envd(target) + assert_region_is_empty(target) + + +def assert_region_is_empty(target: "CloudTarget") -> None: + """ + Assert that the region we can talk to has no user tables. + + Proves that we reached the freshly created region rather than an environmentd that + outlived the recreate: a leftover envd still has the scenario's tables. Neither + `mz region enable` (which only checks that something answers on the SQL port) nor a + successful query on its own can tell the two apart. + """ + conn = target.new_connection() + try: + with conn.cursor() as cur: + cur.execute("SELECT count(*) FROM mz_tables WHERE id LIKE 'u%'") + row = cur.fetchone() + finally: + conn.close() + + assert row is not None, "could not count user tables in recreated region" + assert ( + row[0] == 0 + ), f"recreated region still has {row[0]} user table(s), we are talking to a stale environmentd" + + def reconfigure_envd_cpus( target: "BenchTarget", envd_cpus: int, runner: ScenarioRunner ) -> None: @@ -3479,8 +3510,9 @@ def reconfigure_envd_cpus( - Docker target: recreate the local `materialized` container with a CPU limit equal to envd_cpus, wait for SQL readiness, and force the benchmark connection to reconnect. - - Cloud target: soft-disable/enable the region with the desired envd CPU allocation, wait for - SQL readiness, and force the benchmark connection to reconnect. + - Cloud target: recreate the region with the desired envd CPU allocation, wait for SQL + readiness, and force the benchmark connection to reconnect. The recreate drops all + state in the region, see `cloud_recreate_region_with_envd_cpus`. """ if isinstance(target, DockerTarget): # For Docker target: restart `materialized` with a CPU limit equal to envd_cpus. @@ -3511,9 +3543,10 @@ def reconfigure_envd_cpus( raise UIError(f"failed to apply Docker CPU override for environmentd: {e}") else: # Cloud target: reconfigure environmentd CPUs via `mz region`. + assert isinstance(target, CloudTarget) try: print(f"--- Reconfiguring Cloud environmentd CPUs to {envd_cpus}") - cloud_disable_enable_and_wait(target, environmentd_cpu_allocation=envd_cpus) + cloud_recreate_region_with_envd_cpus(target, envd_cpus) except Exception as e: raise UIError( f"failed to apply Cloud CPU override for environmentd via 'mz region': {e}" @@ -3526,12 +3559,17 @@ def reconfigure_envd_cpus( pass -def wait_for_envd(target: "BenchTarget", timeout_secs: int = 300) -> None: +def wait_for_envd( + target: "BenchTarget", timeout_secs: int = 600, sustain_probes: int = 3 +) -> None: """ Wait until the environmentd SQL endpoint is ready. - Cloud: uses cloud hostname:6875, sslmode=require, and prefers the per-run - app password if available; falls back to MZ_CLI_APP_PASSWORD. + app password if available; falls back to MZ_CLI_APP_PASSWORD. Requires + `sustain_probes` consecutive successful probes, because a single one can land on an + environmentd that is about to go away, for example the previous generation while the + region is still rolling over. - Docker: probes SQL readiness via Composition.sql_query """ if isinstance(target, CloudTarget): @@ -3539,20 +3577,52 @@ def wait_for_envd(target: "BenchTarget", timeout_secs: int = 300) -> None: user = target.username # Prefer the newly created app password when present; fall back to the CLI password. password = target.new_app_password or target.app_password or "" - sslmode = "require" print( f"Waiting for cloud environmentd at {host}:6875 to come up with username {user} ..." ) - _wait_for_pg( - host=host, - user=user, - password=password, - port=6875, - query="SELECT 1", - expected=[(1,)], - timeout_secs=timeout_secs, - dbname="materialize", - sslmode=sslmode, + start = time.time() + deadline = start + timeout_secs + last_report = start + successes = 0 + last_err: Exception | None = None + while time.time() < deadline: + try: + conn = psycopg.connect( + host=host, + port=6875, + user=user, + password=password, + dbname="materialize", + sslmode="require", + connect_timeout=10, + ) + try: + with conn.cursor() as cur: + cur.execute("SELECT 1") + answered = cur.fetchall() == [(1,)] + finally: + conn.close() + successes = successes + 1 if answered else 0 + except Exception as e: + last_err = e + successes = 0 + + if successes >= sustain_probes: + # Elapsed time is worth reporting: a region that only comes back after + # minutes is doing a 0dt rollover, which we don't expect here. + print( + f"Cloud environmentd at {host}:6875 is serving SQL, " + f"after {time.time() - start:.0f}s" + ) + return + + now = time.time() + if now - last_report >= 30: + print(f" still waiting ({now - start:.0f}s elapsed): {last_err}") + last_report = now + time.sleep(2) + raise UIError( + f"cloud environmentd at {host}:6875 did not serve SQL within {timeout_secs}s: {last_err}" ) else: # Docker target: use the composition helper to query the service via the From f1856b2f774c544851e785b6cc7925ae132483d1 Mon Sep 17 00:00:00 2001 From: Aljoscha Krettek Date: Mon, 3 Aug 2026 12:43:39 +0000 Subject: [PATCH 2/3] cluster-spec-sheet: retry `mz region enable` on transient API errors A sweep now does a dozen region API calls and staging returns the occasional 502, which the CLI does not retry for us. Seen in spec-sheet build 29: the sweep got through four scale points and then `mz region enable` died with `status 502 Bad Gateway`. --- test/cluster-spec-sheet/mzcompose.py | 41 +++++++++++++++++++--------- 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/test/cluster-spec-sheet/mzcompose.py b/test/cluster-spec-sheet/mzcompose.py index f2bce62a4e4a8..2a26c86ae139f 100644 --- a/test/cluster-spec-sheet/mzcompose.py +++ b/test/cluster-spec-sheet/mzcompose.py @@ -3444,7 +3444,9 @@ def cloud_disable_enable_and_wait(target: "BenchTarget") -> None: wait_for_envd(target) -def cloud_recreate_region_with_envd_cpus(target: "CloudTarget", envd_cpus: int) -> None: +def cloud_recreate_region_with_envd_cpus( + target: "CloudTarget", envd_cpus: int, attempts: int = 3 +) -> None: """ Recreate the Cloud region from scratch, with the given environmentd CPU allocation. @@ -3459,20 +3461,33 @@ def cloud_recreate_region_with_envd_cpus(target: "CloudTarget", envd_cpus: int) away. A hard disable leaves no predecessor generation to catch up with, so the new envd promotes as soon as it has booted, in about a minute. """ - disable_region(target.composition, hard=True) - version_args = ["--version", target.version] if target.version is not None else [] - target.composition.run( - "mz", - "region", - "enable", - "--environmentd-cpu-allocation", - str(envd_cpus), - *enable_extra_args(target), - *version_args, - rm=True, - ) + for attempt in range(1, attempts + 1): + disable_region(target.composition, hard=True) + try: + target.composition.run( + "mz", + "region", + "enable", + "--environmentd-cpu-allocation", + str(envd_cpus), + *enable_extra_args(target), + *version_args, + rm=True, + ) + break + except UIError as e: + # A sweep does a dozen of these calls and the staging Cloud API returns the + # occasional 502, which `mz region enable` does not retry on its own. We + # disable again before retrying, so that a half-finished enable can't leave + # the region running with the previous allocation. + if attempt == attempts: + raise + print( + f"WARNING: 'mz region enable' failed (attempt {attempt}/{attempts}): {e}" + ) + time.sleep(30) assert "materialize.cloud" in target.composition.cloud_hostname() wait_for_envd(target) From f7aaac390727ff68ea8eafdc171c955aee562748 Mon Sep 17 00:00:00 2001 From: Aljoscha Krettek Date: Mon, 3 Aug 2026 13:17:41 +0000 Subject: [PATCH 3/3] cluster-spec-sheet: factor out the region enable and SQL probe helpers Both region paths were spelling out the `mz region enable` invocation with its version and staging-override args, and the two readiness waits each carried their own psycopg connect block and poll loop. Pull out `enable_region`, `cloud_sql` and `_await_probe` so each of those lives in one place. No behavior change beyond the poll interval for the down-wait and the sustained-probe count for the Docker wait, which now match the cloud wait. --- test/cluster-spec-sheet/mzcompose.py | 307 +++++++++++++-------------- 1 file changed, 147 insertions(+), 160 deletions(-) diff --git a/test/cluster-spec-sheet/mzcompose.py b/test/cluster-spec-sheet/mzcompose.py index 2a26c86ae139f..c02e246a7af7b 100644 --- a/test/cluster-spec-sheet/mzcompose.py +++ b/test/cluster-spec-sheet/mzcompose.py @@ -22,8 +22,9 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Iterable from dataclasses import dataclass +from functools import partial from textwrap import dedent -from typing import Any, TextIO +from typing import Any, LiteralString, TextIO import matplotlib.pyplot as plt import pandas as pd @@ -3364,9 +3365,8 @@ def teardown(self, runner: ScenarioRunner) -> None: self._workload.teardown(runner) -# TODO: We should factor out the below -# `disable_region`, `cloud_disable_enable_and_wait`, `reconfigure_envd_cpus`, `wait_for_envd` -# functions into a separate module. (Similar `disable_region` functions also occur in other tests.) +# TODO: We should factor the region helpers below out into a separate module. +# (Similar `disable_region` functions also occur in other tests.) def disable_region(composition: Composition, hard: bool) -> None: print("Shutting down region ...") @@ -3380,26 +3380,32 @@ def disable_region(composition: Composition, hard: bool) -> None: pass -# Enabling a region whose catalog already exists is a 0dt deployment: the new envd -# boots read-only and only promotes (and starts serving) once the caught-up check -# passes. Since #37255 that check includes a stability soak with a production -# default of 10 minutes, during which a soft-disabled region serves nothing. Keep -# the soak out of the critical path for test regions, like mzcompose does for -# Docker-based tests. -_STABILITY_SOAK_OVERRIDE = [ - "--environmentd-extra-arg=--system-parameter-default=with_0dt_caught_up_check_stability_period=0s", -] - - -def enable_extra_args(target: "CloudTarget") -> list[str]: +def enable_region(target: "CloudTarget", envd_cpus: int | None = None) -> None: """ - Extra `mz region enable` args, applied to staging only. + Run `mz region enable`, pinning environmentd's CPU allocation when one is given. - Production Cloud forbids callers from injecting environmentd args and rejects - `--environmentd-extra-arg` with a 403 Forbidden, so the stability-soak override - cannot be used there. + Enabling a region whose catalog still exists is a 0dt deployment: the new envd boots + read-only and only promotes, and starts serving, once the caught-up check passes. That + check includes a stability soak whose production default is ten minutes, during which a + soft-disabled region serves nothing at all. We turn the soak off for test regions, like + mzcompose does for Docker-based tests. """ - return _STABILITY_SOAK_OVERRIDE if target.is_staging else [] + args = [] + + if envd_cpus is not None: + args += ["--environmentd-cpu-allocation", str(envd_cpus)] + + if target.is_staging: + # Production Cloud forbids callers from injecting environmentd args and rejects + # `--environmentd-extra-arg` with a 403 Forbidden, so this is staging-only. + args += [ + "--environmentd-extra-arg=--system-parameter-default=with_0dt_caught_up_check_stability_period=0s" + ] + + if target.version is not None: + args += ["--version", target.version] + + target.composition.run("mz", "region", "enable", *args, rm=True) def cloud_disable_enable_and_wait(target: "BenchTarget") -> None: @@ -3432,11 +3438,7 @@ def cloud_disable_enable_and_wait(target: "BenchTarget") -> None: if host is not None: wait_for_envd_down(target, host) - version_args = ["--version", target.version] if target.version is not None else [] - - target.composition.run( - "mz", "region", "enable", *enable_extra_args(target), *version_args, rm=True - ) + enable_region(target) time.sleep(10) @@ -3454,28 +3456,15 @@ def cloud_recreate_region_with_envd_cpus( needs again. We hard-disable rather than soft-disable on purpose. Enabling a region that still has - durable state is a 0dt rollout: the new environmentd comes up in read-only mode and - only promotes once every cluster has been caught up *and stable* for - `with_0dt_caught_up_check_stability_period`, which is ten minutes fleet-wide. Nothing - serves SQL while we wait that out, because the soft disable already took the old envd - away. A hard disable leaves no predecessor generation to catch up with, so the new envd - promotes as soon as it has booted, in about a minute. + durable state is a 0dt rollout, and nothing serves SQL while we wait for the new envd to + catch up and promote, because the soft disable already took the old envd away. A hard + disable leaves no predecessor generation to catch up with, so the new envd promotes as + soon as it has booted, in about a minute. """ - version_args = ["--version", target.version] if target.version is not None else [] - for attempt in range(1, attempts + 1): disable_region(target.composition, hard=True) try: - target.composition.run( - "mz", - "region", - "enable", - "--environmentd-cpu-allocation", - str(envd_cpus), - *enable_extra_args(target), - *version_args, - rm=True, - ) + enable_region(target, envd_cpus=envd_cpus) break except UIError as e: # A sweep does a dozen of these calls and the staging Cloud API returns the @@ -3503,18 +3492,13 @@ def assert_region_is_empty(target: "CloudTarget") -> None: `mz region enable` (which only checks that something answers on the SQL port) nor a successful query on its own can tell the two apart. """ - conn = target.new_connection() - try: - with conn.cursor() as cur: - cur.execute("SELECT count(*) FROM mz_tables WHERE id LIKE 'u%'") - row = cur.fetchone() - finally: - conn.close() + host = target.composition.cloud_hostname() + rows = cloud_sql(target, host, "SELECT count(*) FROM mz_tables WHERE id LIKE 'u%'") - assert row is not None, "could not count user tables in recreated region" + assert rows, "could not count user tables in recreated region" assert ( - row[0] == 0 - ), f"recreated region still has {row[0]} user table(s), we are talking to a stale environmentd" + rows[0][0] == 0 + ), f"recreated region still has {rows[0][0]} user table(s), we are talking to a stale environmentd" def reconfigure_envd_cpus( @@ -3574,126 +3558,129 @@ def reconfigure_envd_cpus( pass -def wait_for_envd( - target: "BenchTarget", timeout_secs: int = 600, sustain_probes: int = 3 -) -> None: +def cloud_sql( + target: "CloudTarget", host: str, query: LiteralString +) -> list[tuple[Any, ...]]: """ - Wait until the environmentd SQL endpoint is ready. - - - Cloud: uses cloud hostname:6875, sslmode=require, and prefers the per-run - app password if available; falls back to MZ_CLI_APP_PASSWORD. Requires - `sustain_probes` consecutive successful probes, because a single one can land on an - environmentd that is about to go away, for example the previous generation while the - region is still rolling over. - - Docker: probes SQL readiness via Composition.sql_query + Run `query` against the environmentd at `host` on a fresh connection and return its rows. + + Takes an explicit host rather than looking it up, because a recreated region gets a new + hostname and callers that probe for the *old* envd need to keep talking to the old one. + """ + conn = psycopg.connect( + host=host, + port=6875, + user=target.username, + # Prefer the newly created app password when present, fall back to the CLI one. + password=target.new_app_password or target.app_password or "", + dbname="materialize", + sslmode="require", + connect_timeout=10, + ) + try: + with conn.cursor() as cur: + cur.execute(query) + return cur.fetchall() + finally: + conn.close() + + +def _err_detail(err: Exception | None) -> str: + return f": {err}" if err is not None else "" + + +def _await_probe( + probe: Callable[[], object], + what: str, + *, + until_ok: bool, + timeout_secs: int, + sustain_probes: int = 3, +) -> float: + """ + Poll `probe` until it has either succeeded or failed, per `until_ok`, `sustain_probes` + times in a row, and return how long that took. Raises `UIError` on timeout. + + A single probe decides nothing in either direction. While a region rolls over, a + connection can land on the generation that is about to go away, and a single failure can + be a transient network error against an envd that is still serving fine. + + `what` completes the sentence "waiting for ...", e.g. "environmentd to serve SQL". + """ + start = time.time() + deadline = start + timeout_secs + last_report = start + streak = 0 + last_err: Exception | None = None + + while time.time() < deadline: + try: + probe() + ok = True + last_err = None + except Exception as e: + last_err = e + ok = False + + streak = streak + 1 if ok == until_ok else 0 + if streak >= sustain_probes: + return time.time() - start + + now = time.time() + if now - last_report >= 30: + print( + f" still waiting for {what} ({now - start:.0f}s){_err_detail(last_err)}" + ) + last_report = now + time.sleep(2) + + raise UIError( + f"timed out after {timeout_secs}s waiting for {what}{_err_detail(last_err)}" + ) + + +def wait_for_envd(target: "BenchTarget", timeout_secs: int = 600) -> None: + """ + Wait until environmentd serves SQL, and report how long that took. + + The elapsed time is worth printing: a region that only comes back after minutes is + promoting out of read-only mode, which none of the callers here expect. """ if isinstance(target, CloudTarget): host = target.composition.cloud_hostname() - user = target.username - # Prefer the newly created app password when present; fall back to the CLI password. - password = target.new_app_password or target.app_password or "" - print( - f"Waiting for cloud environmentd at {host}:6875 to come up with username {user} ..." - ) - start = time.time() - deadline = start + timeout_secs - last_report = start - successes = 0 - last_err: Exception | None = None - while time.time() < deadline: - try: - conn = psycopg.connect( - host=host, - port=6875, - user=user, - password=password, - dbname="materialize", - sslmode="require", - connect_timeout=10, - ) - try: - with conn.cursor() as cur: - cur.execute("SELECT 1") - answered = cur.fetchall() == [(1,)] - finally: - conn.close() - successes = successes + 1 if answered else 0 - except Exception as e: - last_err = e - successes = 0 - - if successes >= sustain_probes: - # Elapsed time is worth reporting: a region that only comes back after - # minutes is doing a 0dt rollover, which we don't expect here. - print( - f"Cloud environmentd at {host}:6875 is serving SQL, " - f"after {time.time() - start:.0f}s" - ) - return - - now = time.time() - if now - last_report >= 30: - print(f" still waiting ({now - start:.0f}s elapsed): {last_err}") - last_report = now - time.sleep(2) - raise UIError( - f"cloud environmentd at {host}:6875 did not serve SQL within {timeout_secs}s: {last_err}" - ) + who = f"cloud environmentd at {host}:6875" + probe: Callable[[], object] = partial(cloud_sql, target, host, "SELECT 1") else: - # Docker target: use the composition helper to query the service via the - # host-mapped port on 127.0.0.1; the container hostname "materialized" - # is not resolvable from the host network when using psycopg directly. - print("Waiting for local environmentd (docker) at materialized:6875 ...") - deadline = time.time() + timeout_secs - last_err: Exception | None = None - while time.time() < deadline: - try: - target.composition.sql_query("SELECT 1", service="materialized") - return - except Exception as e: - last_err = e - time.sleep(1) - raise UIError( - f"materialized did not accept SQL connections within {timeout_secs}s after restart: {last_err}" + # The container hostname "materialized" does not resolve from the host network, so + # we go through the composition instead of connecting with psycopg. + who = "local environmentd (docker)" + probe = partial( + target.composition.sql_query, "SELECT 1", service="materialized" ) + print(f"Waiting for {who} to come up ...") + elapsed = _await_probe( + probe, f"{who} to serve SQL", until_ok=True, timeout_secs=timeout_secs + ) + print(f"{who} is serving SQL, after {elapsed:.0f}s") + def wait_for_envd_down( target: "CloudTarget", host: str, timeout_secs: int = 600 ) -> None: """ - Wait until the environmentd SQL endpoint at `host` stops accepting connections. + Wait until the environmentd at `host` stops serving SQL. - Used after a soft `mz region disable` to observe that the teardown has reached the - data plane before the region is enabled again. + Used after a soft `mz region disable` to observe that the teardown has reached the data + plane before the region is enabled again. """ - print(f"Waiting for cloud environmentd at {host}:6875 to go down ...") - password = target.new_app_password or target.app_password or "" - deadline = time.time() + timeout_secs - consecutive_failures = 0 - while time.time() < deadline: - try: - conn = psycopg.connect( - host=host, - port=6875, - user=target.username, - password=password, - dbname="materialize", - sslmode="require", - connect_timeout=10, - ) - conn.close() - consecutive_failures = 0 - except psycopg.Error: - # A single failed attempt could be a transient network error against a - # live envd. Only treat repeated failures as the teardown having landed. - consecutive_failures += 1 - if consecutive_failures >= 3: - return - time.sleep(1) - raise UIError( - f"environmentd at {host}:6875 still accepting SQL connections " - f"{timeout_secs}s after region disable" + who = f"cloud environmentd at {host}:6875" + print(f"Waiting for {who} to go down ...") + _await_probe( + partial(cloud_sql, target, host, "SELECT 1"), + f"{who} to stop serving SQL", + until_ok=False, + timeout_secs=timeout_secs, )