broad-announce/scripts/m10_lib.py

"""
m10_lib.py — shared assertion library for M10 smoke scripts.
Used by both m10_smoke.py (5k/s soak) and m10_bench_smoke.py (50k/s broker ceiling).

All Prometheus queries use URL-encoding-safe Python urllib to avoid the
bashquoting bugs that plagued m9_smoke.sh.
"""
import urllib.request
import urllib.parse
import time
import sys

PROM = "http://localhost:9090"
PROM_TIMEOUT = 10  # seconds


def scrape(query: str) -> list[dict]:
    """
    Run a Prometheus query and return the result vector.
    Returns [] on error or no data.
    """
    url = f"{PROM}/api/v1/query?query={urllib.parse.quote(query, safe='')}"
    try:
        with urllib.request.urlopen(url, timeout=PROM_TIMEOUT) as r:
            data = r.read()
        d = __import__("json").loads(data)
        if d.get("status") != "success":
            return []
        return d.get("data", {}).get("result", [])
    except Exception:
        return []


def assert_ingest_p99_under(threshold: float, window_seconds: int = 60,
                             label_filter: str = "") -> float:
    """
    Assert p99 ingest latency (ba_ingestd_publish_latency_seconds) is under threshold.
    Uses histogram_quantile over the given sliding window.
    Returns the actual p99 value, or raises AssertionError.
    """
    filter_expr = f',{{{label_filter}}}' if label_filter else ''
    query = (
        f'histogram_quantile(0.99, '
        f'rate(ba_ingestd_publish_latency_seconds_bucket{{{filter_expr}}}[{window_seconds}s]))'
    )
    results = scrape(query)
    if not results:
        raise AssertionError(
            f"p99 query returned no data. Is ingestd up and generating load? "
            f"(query: {query})"
        )
    p99 = float(results[0]["value"][1])
    if p99 > threshold:
        raise AssertionError(
            f"p99 ingest latency {p99:.3f}s exceeds threshold {threshold}s "
            f"(window={window_seconds}s, filter='{label_filter}')"
        )
    return p99


def per_source_p99(source_id: str, window_seconds: int = 60) -> float:
    """
    Return the p99 publish latency for a specific source_id.
    Returns 0.0 if no data.
    """
    query = (
        f'histogram_quantile(0.99, '
        f'rate(ba_ingestd_publish_latency_seconds_bucket{{source_id="{source_id}"}}[{window_seconds}s]))'
    )
    results = scrape(query)
    if not results:
        return 0.0
    return float(results[0]["value"][1])


def assert_dlq_count_equals(expected: int, window_seconds: int) -> int:
    """
    Assert the DLQ row count (delta over window) equals expected.
    Uses Prometheus counter delta (ba_deliverd_dlq_total).
    Returns the actual count.
    """
    query = (
        f'increase(ba_deliverd_dlq_total[{window_seconds}s])'
    )
    results = scrape(query)
    if not results:
        # No DLQ entries at all
        actual = 0
    else:
        actual = int(float(results[0]["value"][1]))
    if actual != expected:
        raise AssertionError(
            f"DLQ count {actual} != expected {expected} "
            f"(window={window_seconds}s). Run 'docker exec ... psql ... "
            f"SELECT count(*) FROM deliveries_dlq' for precise count."
        )
    return actual


def assert_rate_near(target: float, tolerance: float = 0.05,
                      window_seconds: int = 30) -> float:
    """
    Assert the cluster-wide accept rate is within tolerance of target.
    tolerance=0.05 means ±5%.
    Returns the actual rate.
    """
    query = (
        f'rate(ba_ingestd_alerts_received_total{{result="accepted"}}[{window_seconds}s])'
    )
    results = scrape(query)
    if not results:
        raise AssertionError(
            f"rate query returned no data. Is ingestd receiving alerts? (query: {query})"
        )
    # Sum all series (multiple source_ids)
    actual = sum(float(r["value"][1]) for r in results)
    lower = target * (1 - tolerance)
    upper = target * (1 + tolerance)
    if not (lower <= actual <= upper):
        raise AssertionError(
            f"rate {actual:.1f}/s is outside ±{tolerance*100:.0f}% "
            f"tolerance of target {target}/s (range: {lower:.1f}–{upper:.1f})"
        )
    return actual


def assert_no_rate_limit_storm(window_seconds: int = 30) -> float:
    """
    Assert rate-limited hits are < 5% of total accepted.
    Returns the rate-limited fraction.
    """
    accepted_q = f'rate(ba_ingestd_alerts_received_total{{result="accepted"}}[{window_seconds}s])'
    rl_q = f'rate(ba_ingestd_alerts_received_total{{result="rate_limited"}}[{window_seconds}s])'
    accepted_results = scrape(accepted_q)
    rl_results = scrape(rl_q)
    accepted = sum(float(r["value"][1]) for r in accepted_results) if accepted_results else 0
    rl = sum(float(r["value"][1]) for r in rl_results) if rl_results else 0
    if accepted == 0:
        return 0.0
    frac = rl / accepted
    if frac > 0.05:
        raise AssertionError(
            f"rate-limited fraction {frac:.1%} exceeds 5% threshold "
            f"(accepted={accepted:.1f}/s, rate_limited={rl:.1f}/s)"
        )
    return frac


def wait_until(fn, timeout_seconds: float, interval: float = 1.0, name: str = ""):
    """
    Poll fn() every interval until it returns truthy or timeout.
    Returns fn() result on success, raises on timeout.
    """
    deadline = time.time() + timeout_seconds
    last_err = None
    while time.time() < deadline:
        try:
            result = fn()
            if result:
                return result
        except Exception as e:
            last_err = e
        time.sleep(interval)
    raise AssertionError(
        f"{name or 'condition'} did not become true within {timeout_seconds}s"
        + (f": {last_err}" if last_err else "")
    )


def get_prometheus_targets() -> dict[str, str]:
    """
    Return {job_name: up_value} for all Prometheus scrape targets.
    up_value is "1" (up) or "0" (down).
    """
    results = scrape('up')
    return {r["metric"]["job"]: r["value"][1] for r in results}


def get_metric_series_count(metric_name: str) -> int:
    """Return number of time series for a metric."""
    results = scrape(metric_name)
    return len(results)


def get_circuit_breaker_state(component: str = "nats") -> int:
    """Return current circuit breaker state: 0=CLOSED, 1=HALF-OPEN, 2=OPEN."""
    results = scrape(f'ba_ingestd_circuit_breaker_state{{component="{component}"}}')
    if not results:
        return -1  # not initialized
    return int(float(results[0]["value"][1]))