broad-announce/scripts/m11_lib.py

"""
m11_lib.py — shared assertion library for M11 smoke scripts.
Used by m11_smoke.py.

All Prometheus queries use URL-encoding-safe Python urllib.
"""
import urllib.request
import urllib.parse
import json

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 = json.loads(data)
        if d.get("status") != "success":
            return []
        return d.get("data", {}).get("result", [])
    except Exception:
        return []


def assert_grpc_rate_near(target: float, tolerance: float = 0.10,
                          window_seconds: int = 30) -> float:
    """
    Assert the gRPC ingest rate is within tolerance of target.
    target: expected alerts/sec cluster-wide
    tolerance: fraction (0.10 = ±10%)
    Returns the actual rate, or raises AssertionError.

    Query: sum(rate(ba_ingestd_alerts_received_total{transport="grpc",result="accepted"}[window]))
    """
    query = (
        f'sum(rate(ba_ingestd_alerts_received_total{{transport="grpc",result="accepted"}}[{window_seconds}s]))'
    )
    results = scrape(query)
    if not results:
        raise AssertionError(
            f"gRPC rate query returned no data. Is ingestd gRPC server up and receiving load? "
            f"(query: {query})"
        )
    rate = float(results[0]["value"][1])
    min_rate = target * (1 - tolerance)
    max_rate = target * (1 + tolerance)
    if rate < min_rate:
        raise AssertionError(
            f"gRPC rate {rate:.0f}/s is below target {target:.0f}/s "
            f"(tolerance ±{tolerance*100:.0f}%, min allowed: {min_rate:.0f}/s, "
            f"window={window_seconds}s)"
        )
    return rate


def assert_grpc_p99_under(threshold_ms: float, window_seconds: int = 60) -> float:
    """
    Assert gRPC ack p99 latency is under threshold (in milliseconds).
    Uses histogram_quantile over the given window.
    Returns the actual p99 in milliseconds, or raises AssertionError.

    Query: histogram_quantile(0.99, rate(ba_ingestd_grpc_ack_latency_seconds_bucket[window]))
    """
    query = (
        f'histogram_quantile(0.99, '
        f'rate(ba_ingestd_grpc_ack_latency_seconds_bucket[{window_seconds}s]))'
    )
    results = scrape(query)
    if not results:
        raise AssertionError(
            f"gRPC p99 query returned no data. Is ingestd gRPC server up? "
            f"(query: {query})"
        )
    p99_seconds = float(results[0]["value"][1])
    p99_ms = p99_seconds * 1000
    if p99_ms > threshold_ms:
        raise AssertionError(
            f"gRPC ack p99 {p99_ms:.1f}ms exceeds threshold {threshold_ms:.1f}ms "
            f"(window={window_seconds}s)"
        )
    return p99_ms


def assert_nats_publish_rate_near(target: float, tolerance: float = 0.10,
                                  window_seconds: int = 30) -> float:
    """
    Assert the NATS publish OK rate is within tolerance of target.

    F2 (M11 NATS investigation): receive rate alone is not enough — a broken
    publish path can hide behind a healthy receive metric. The M11 10-min
    soak that "shipped" M11 was a false positive: ba_ingestd_alerts_received_total
    counts gRPC receive, but PublishAsync to NATS was failing because the
    server hit its max_storage cap. The smoke stayed green while the publish
    path was broken end-to-end.

    This assertion catches the next class of this bug by querying
    ba_ingestd_nats_publish_total{result="ok"} and checking it tracks the
    receive rate.

    Query: sum(rate(ba_ingestd_nats_publish_total{result="ok"}[window]))
    """
    query = (
        f'sum(rate(ba_ingestd_nats_publish_total{{result="ok"}}[{window_seconds}s]))'
    )
    results = scrape(query)
    if not results:
        raise AssertionError(
            f"NATS publish OK rate query returned no data. "
            f"Is the F2 metric (ba_ingestd_nats_publish_total) exposed by ingestd? "
            f"(query: {query})"
        )
    rate = float(results[0]["value"][1])
    min_rate = target * (1 - tolerance)
    if rate < min_rate:
        raise AssertionError(
            f"NATS publish OK rate {rate:.0f}/s is below target {target:.0f}/s "
            f"(tolerance ±{tolerance*100:.0f}%, min allowed: {min_rate:.0f}/s, "
            f"window={window_seconds}s). This means the publish path is broken "
            f"even though gRPC receive may be green. See M11_NATS_INVESTIGATION.md."
        )
    return rate


def assert_dlq_count_equals(expected: int, window_seconds: int) -> int:
    """
    Assert the DLQ row count (delta over window) equals expected.
    Returns the actual count.
    """
    query = (
        f'increase(ba_deliverd_dlq_total[{window_seconds}s])'
    )
    results = scrape(query)
    if not results:
        return 0  # No data → assume 0
    count = float(results[0]["value"][1])
    if abs(count - expected) > 0.5:
        raise AssertionError(
            f"DLQ count {count:.0f} over last {window_seconds}s does not equal expected {expected}"
        )
    return int(count)



    """
    Assert the DLQ row count (delta over window) equals expected.
    Returns the actual count.
    """
    query = (
        f'increase(ba_deliverd_dlq_total[{window_seconds}s])'
    )
    results = scrape(query)
    if not results:
        return 0  # No data → assume 0
    count = float(results[0]["value"][1])
    if abs(count - expected) > 0.5:
        raise AssertionError(
            f"DLQ count {count:.0f} over last {window_seconds}s does not equal expected {expected}"
        )
    return int(count)


def get_prometheus_targets() -> dict[str, str]:
    """
    Return a dict mapping service name → health status from Prometheus target health API.
    "1" = healthy, "0" = down, other = unknown state.
    """
    url = f"{PROM}/api/v1/targets?state=active"
    try:
        with urllib.request.urlopen(url, timeout=PROM_TIMEOUT) as r:
            data = json.loads(r.read())
        out = {}
        for t in data.get("data", {}).get("activeTargets", []):
            labels = t.get("labels", {})
            job = labels.get("job", "?")
            # Use the job name (e.g. "ingestd", "loadgen-grpc-1") as the key
            out[job] = "1" if t.get("health") == "up" else str(t.get("health", "0"))
        return out
    except Exception:
        return {}


def get_grpc_streams_active() -> int:
    """
    Return the current number of active gRPC streams on ingestd
    (ba_ingestd_grpc_streams_active{service="ingestd"} gauge).
    Returns 0 if no data.

    NB: this metric is exported by every service that uses the grpcserver
    package (deliverd-fcm, deliverd-telegram, admind, routerd, ingestd),
    all with value 0 except ingestd under load. An unfiltered query returns
    5 series and the first happens to be 0, so we filter by service="ingestd".
    """
    query = 'ba_ingestd_grpc_streams_active{service="ingestd"}'
    results = scrape(query)
    if not results:
        return 0
    return int(float(results[0]["value"][1]))


def get_grpc_rate_limited_total() -> int:
    """
    Return the cumulative ba_ingestd_grpc_rate_limited_total counter on ingestd.
    Returns 0 if no data.

    NB: same multi-service issue as get_grpc_streams_active — filter by service.
    """
    query = 'ba_ingestd_grpc_rate_limited_total{service="ingestd"}'
    results = scrape(query)
    if not results:
        return 0
    return int(float(results[0]["value"][1]))