AsyncClient locks up for good after a node restart under load: every call fails with PoolTimeout while the cluster sits idle #146

"""
Load generator for the typesense-python retry-storm reproducer.

Runs inside the docker network next to nginx (lb:8000, per-node 8001..8003).

  python3 loadgen.py seed <docs>
  python3 loadgen.py calibrate <port> <concurrency> <seconds>
  python3 loadgen.py run <arm> <rate_per_s> <seconds> <out.jsonl>

Arms:
  stock    typesense 2.0.0 as shipped
  patched  2.0.0 with the pause between retries honoured and the node that
           actually answered marked healthy (what 0.21.0 did: PR #141 plus the
           fix for the "wrong node marked healthy" bug)
  noretry  2.0.0 with num_retries=0 (the floor: no client retries at all)
  capped   2.0.0 as shipped, but the app lets at most CAP (50) searches into
           the client at once; the rest wait in the app, so the client never
           has more requests than its 100 connection pool can hold

"run" is open loop: calls start on a fixed schedule whatever the latency, the
way web traffic arrives, so in-flight calls pile up when the cluster slows.
Every app call and every client attempt is written to out.jsonl with epoch
timestamps so analyze.py can line them up with the nginx access log.
"""

import asyncio
import contextlib
import contextvars
import json
import os
import random
import sys
import threading
import time

import typesense
from typesense.async_.api_call import AsyncApiCall, _SERVER_ERRORS
from typesense.exceptions import TypesenseClientError

LB = {"host": "nginx", "port": "8000", "protocol": "http"}
NODES = [{"host": "nginx", "port": p, "protocol": "http"} for p in ("8001", "8002", "8003")]
COLL = "products"
CAP = 50  # "capped" arm: searches allowed inside the client at once

rnd = random.Random(42)
VOCAB = [f"w{i}" for i in range(3000)]
WEIGHTS = [1.0 / (i + 1) for i in range(len(VOCAB))]  # Zipf: a few words are everywhere


def words(n):
    return " ".join(rnd.choices(VOCAB, WEIGHTS, k=n))


def make_client(arm, timeout=3.0):
    cfg = {"api_key": "xyz", "nearest_node": LB, "nodes": NODES, "connection_timeout_seconds": timeout}
    if arm == "noretry":
        cfg["num_retries"] = 0
    client = typesense.AsyncClient(cfg)
    if arm == "patched":
        client.api_call._execute_request = _patched_execute.__get__(client.api_call, AsyncApiCall)
    return client


async def _patched_execute(self, method, endpoint, entity_type, as_json=True, last_exception=None, num_retries=0, **kwargs):
    if num_retries > self.config.num_retries:
        if last_exception:
            raise last_exception
        raise TypesenseClientError("All nodes are unhealthy")
    if num_retries > 0:
        await asyncio.sleep(self.config.retry_interval_seconds)
    node, url, request_kwargs = self._prepare_request_params(endpoint, **kwargs)
    try:
        resp = await self.request_handler.make_request(
            method=method, url=url, as_json=as_json, entity_type=entity_type, client=self._client, **request_kwargs)
        self.node_manager.set_node_health(node, is_healthy=True)
        return resp
    except _SERVER_ERRORS as err:
        self.node_manager.set_node_health(node, is_healthy=False)
        return await _patched_execute(self, method, endpoint, entity_type, as_json,
                                      last_exception=err, num_retries=num_retries + 1, **kwargs)


async def seed(n):
    client = make_client("stock", timeout=120)
    try:
        await client.collections[COLL].delete()
    except Exception:
        pass
    await client.collections.create({
        "name": COLL,
        "fields": [
            {"name": "title", "type": "string"},
            {"name": "description", "type": "string"},
            {"name": "category", "type": "string", "facet": True},
            {"name": "brand", "type": "string", "facet": True},
            {"name": "price", "type": "float"},
            {"name": "popularity", "type": "int32"},
        ],
        "default_sorting_field": "popularity",
    })
    batch, done = [], 0
    for i in range(n):
        batch.append({
            "id": str(i), "title": words(rnd.randint(6, 12)), "description": words(rnd.randint(30, 60)),
            "category": f"c{rnd.randint(0, 59)}", "brand": f"b{rnd.randint(0, 799)}",
            "price": round(rnd.uniform(1, 500), 2), "popularity": rnd.randint(0, 100000),
        })
        if len(batch) == 5000 or i == n - 1:
            res = await client.collections[COLL].documents.import_(batch, {"action": "create"})
            bad = [r for r in res if not r.get("success")]
            if bad:
                print("import errors:", bad[:2], flush=True)
            done += len(batch)
            batch = []
            if done % 50000 == 0 or done == n:
                print(f"seeded {done}/{n}", flush=True)
    await client.api_call.aclose()


def query():
    return {
        "q": words(rnd.choice((1, 1, 2))),
        "query_by": "title,description",
        "facet_by": "category,brand",
        "max_facet_values": 20,
        "per_page": 20,
    }


async def calibrate(port, conc, seconds):
    client = typesense.AsyncClient({"api_key": "xyz", "nodes": [{"host": "nginx", "port": port, "protocol": "http"}],
                                    "connection_timeout_seconds": 60, "num_retries": 0})
    lat, stop = [], time.time() + seconds

    async def worker():
        while time.time() < stop:
            t = time.time()
            await client.collections[COLL].documents.search(query())
            lat.append(time.time() - t)

    await asyncio.gather(*(worker() for _ in range(conc)))
    await client.api_call.aclose()
    lat.sort()
    print(json.dumps({"port": port, "conc": conc, "qps": round(len(lat) / seconds, 1),
                      "p50_ms": round(lat[len(lat) // 2] * 1000), "p95_ms": round(lat[int(len(lat) * .95)] * 1000)}), flush=True)


CUR = contextvars.ContextVar("cur")


async def run(arm, rate, seconds, out_path):
    client = make_client(arm)
    rh = client.api_call.request_handler
    orig = rh.make_request

    def counted(**kw):
        coro = orig(**kw)
        port = kw["url"].split(":")[2].split("/")[0]

        async def attempt():
            rec, t = CUR.get(None), time.time()
            try:
                res = await coro
                if rec is not None:
                    rec["att"].append([port, "ok", round(t, 3), round(time.time(), 3)])
                return res
            except BaseException as e:
                if rec is not None:
                    rec["att"].append([port, type(e).__name__, round(t, 3), round(time.time(), 3)])
                raise
        return attempt()

    rh.make_request = counted
    out = open(out_path, "w")
    nm = client.api_call.node_manager
    t0 = time.time()
    out.write(json.dumps({"k": "start", "arm": arm, "rate": rate, "t": t0}) + "\n")
    inflight = set()
    gate = asyncio.Semaphore(CAP) if arm == "capped" else contextlib.nullcontext()

    async def one(i):
        rec = {"k": "call", "i": i, "t": round(time.time(), 3), "att": []}
        CUR.set(rec)
        try:
            async with gate:  # call time includes any wait here
                await client.collections[COLL].documents.search(query())
            rec["out"] = "ok"
        except Exception as e:
            rec["out"] = type(e).__name__
        rec["end"] = round(time.time(), 3)
        out.write(json.dumps(rec) + "\n")

    pool = client.api_call._client._transport._pool

    async def health():
        while True:
            # Pool census: busy connections with no server traffic means leaked
            # slots (httpcore #1093); idle or churning ones with a long queue
            # means the event loop is just too busy to hand them out.
            conns = list(pool.connections)
            out.write(json.dumps({"k": "health", "t": round(time.time(), 3), "inflight": len(inflight),
                                  "lb": client.config.nearest_node.healthy,
                                  "nodes": [n.healthy for n in nm.nodes],
                                  "conns": len(conns),
                                  "busy": sum(1 for c in conns if not c.is_idle() and not c.is_closed()),
                                  "idle": sum(1 for c in conns if c.is_idle()),
                                  "queued": sum(1 for r in list(pool._requests) if r.is_queued())}) + "\n")
            await asyncio.sleep(1)

    htask = asyncio.create_task(health())
    total = int(rate * seconds)
    stop_path = out_path + ".stop"

    def watchdog():
        # A client stuck in its own pool can starve the event loop for many
        # minutes. Once the run is cut, give in-flight calls 60s, then record
        # how many never finished and exit from this thread.
        while not os.path.exists(stop_path):
            time.sleep(1)
        time.sleep(60)
        with open(out_path, "a") as f:
            f.write(json.dumps({"k": "end", "forced": True, "t": time.time(), "stuck_inflight": len(inflight),
                                "lb": client.config.nearest_node.healthy,
                                "nodes": [n.healthy for n in nm.nodes]}) + "\n")
        os._exit(0)

    threading.Thread(target=watchdog, daemon=True).start()
    for i in range(total):
        if i % 20 == 0 and os.path.exists(stop_path):
            break
        delay = t0 + i / rate - time.time()
        if delay > 0:
            await asyncio.sleep(delay)
        task = asyncio.create_task(one(i))
        inflight.add(task)
        task.add_done_callback(inflight.discard)
    if inflight:
        await asyncio.wait(inflight, timeout=60)
    htask.cancel()
    out.write(json.dumps({"k": "end", "t": time.time()}) + "\n")
    out.close()
    await client.api_call.aclose()
    print(f"run {arm} done: {total} calls", flush=True)


if __name__ == "__main__":
    mode = sys.argv[1]
    if mode == "seed":
        asyncio.run(seed(int(sys.argv[2])))
    elif mode == "calibrate":
        asyncio.run(calibrate(sys.argv[2], int(sys.argv[3]), int(sys.argv[4])))
    elif mode == "run":
        asyncio.run(run(sys.argv[2], float(sys.argv[3]), float(sys.argv[4]), sys.argv[5]))

"""
Line up one run's app calls (out/<tag>.jsonl), the nginx access log
(out/<tag>.access.log) and the restart events (out/<tag>.events).

  python3 analyze.py out/<tag>

Server executions are counted from nginx: every backend address nginx sent a
request to, except connection failures (upstream status 502/504, which nginx
makes up). A request the client abandoned at its 3s timeout still counts,
because Typesense keeps running a search after the caller hangs up.
"""

import collections
import json
import shlex
import sys

tag = sys.argv[1]
BUCKET = 10
NODE = {"10.231.0.11:8108": "n1", "10.231.0.12:8108": "n2", "10.231.0.13:8108": "n3"}

calls, health, t0 = [], [], None
for line in open(f"{tag}.jsonl"):
    r = json.loads(line)
    if r["k"] == "start":
        t0, arm, rate = r["t"], r["arm"], r["rate"]
    elif r["k"] == "call":
        calls.append(r)
    elif r["k"] == "health":
        health.append(r)

events = {}
for line in open(f"{tag}.events"):
    name, t = line.split()
    events[name] = float(t) - t0

hits = []  # (start_rel, node, via)
for line in open(f"{tag}.access.log"):
    parts = shlex.split(line)
    msec, port, rt = float(parts[0]), parts[1], float(parts[7])
    if "/search" not in parts[3]:
        continue
    addrs = [a.strip() for a in parts[5].split(",")]
    stats = [s.strip() for s in parts[6].split(",")]
    for a, s in zip(addrs, stats + ["-"] * len(addrs)):
        if a in NODE and s not in ("502", "504"):
            hits.append((msec - rt - t0, NODE[a], "lb" if port == "8000" else "direct"))

end = events.get("cut", max(c["t"] for c in calls) - t0)
nb = int(end // BUCKET) + 1
B = lambda t: int(t // BUCKET)

started = collections.Counter(B(c["t"] - t0) for c in calls)
failed = collections.Counter(B(c["t"] - t0) for c in calls if c["out"] != "ok")
attempts = collections.Counter(B(a[2] - t0) for c in calls for a in c["att"])
lat = collections.defaultdict(list)
for c in calls:
    lat[B(c["t"] - t0)].append(c["end"] - c["t"])
srv = collections.defaultdict(collections.Counter)
for t, n, via in hits:
    if t >= 0:
        srv[B(t)][n] += 1
        srv[B(t)][via] += 1
lb_down = collections.defaultdict(list)
inflight = collections.defaultdict(int)
for h in health:
    b = B(h["t"] - t0)
    lb_down[b].append(not h["lb"])
    inflight[b] = max(inflight[b], h["inflight"])

print(f"== {tag}: arm={arm}, {rate}/s open loop, {len(calls)} app calls ==")
print("events (s from start):", {k: round(v, 1) for k, v in events.items()})
print(f"{'t':>4} {'calls':>5} {'fail':>4} {'att':>4} {'srv':>4} {'x':>5} {'n1':>4} {'n2':>4} {'n3':>4} {'direct':>6} {'p50':>5} {'p95':>5} {'infl':>4} lb")
for b in range(nb):
    s = srv[b]
    tot = s["n1"] + s["n2"] + s["n3"]
    l = sorted(lat[b]) or [0]
    mark = ""
    for k, v in events.items():
        if B(v) == b:
            mark += f" <{k}>"
    lbd = lb_down[b]
    lbs = "DOWN" if lbd and all(lbd) else ("part" if any(lbd) else "")
    x = tot / started[b] if started[b] else 0
    print(f"{b * BUCKET:>4} {started[b]:>5} {failed[b]:>4} {attempts[b]:>4} {tot:>4} {x:>5.2f} {s['n1']:>4} {s['n2']:>4} {s['n3']:>4} "
          f"{s['direct']:>6} {l[len(l) // 2]:>5.1f} {l[int(len(l) * .95)]:>5.1f} {inflight[b]:>4} {lbs}{mark}")

def phase(a, b):
    cs = [c for c in calls if a <= c["t"] - t0 < b]
    hs = [h for h in hits if a <= h[0] < b]
    if not cs:
        return None
    outs = collections.Counter(c["out"] for c in cs)
    l = sorted(c["end"] - c["t"] for c in cs)
    return {"secs": round(b - a), "app_calls": len(cs), "server_execs": len(hs),
            "amplification": round(len(hs) / len(cs), 2),
            "client_attempts": sum(len(c["att"]) for c in cs),
            "failed": {k: v for k, v in outs.items() if k != "ok"},
            "p50_s": round(l[len(l) // 2], 2), "p95_s": round(l[int(len(l) * .95)], 2)}

rb, nh = events.get("restart_begin", end), events.get("node2_healthy", end)
print("\nphases:")
for name, a, b in (("steady", 10, rb), ("node2 away", rb, nh), ("first 60s back", nh, nh + 60), ("after that", nh + 60, end)):
    print(f"  {name:<15}", json.dumps(phase(a, b)))