query.py: flame graphs from traces

`query.py flame '<TraceQL>' [since] [out.svg]` folds the matching spans of up to
100 traces into stacks of self time and draws them as a flame graph, with the
folded stacks beside it for speedscope or inferno. A wall-clock picture of where
a route's time goes, waiting included, which a CPU profiler would not show: ipx
waits on Postgres and the disk far more than it computes. The traces are fetched
in one container, a docker run apiece being a second each.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

.claude/skills/ipx-prod-check/SKILL.md

@@ -29,6 +29,8 @@ $Q traces  '<TraceQL>'         [since]   # slowest first
 $Q spans   '<TraceQL>'          [since]   # the matching spans, each with its own time, slowest first
 $Q trace   <trace id>                    # one trace as a tree, with its log lines
 $Q prom    '<PromQL>'          [since]   # Prometheus, now; $range becomes `since`
+$Q flame   '<TraceQL>'  [since] [out.svg] # where the matching spans' time went, summed over up to
+                                         # 100 traces: a flame graph, and out.svg.folded for speedscope
 ```
 
 `since` is `30m`, `1h`, `24h`, `7d`. Default to `24h`; widen it to see whether something is new.

.claude/skills/ipx-prod-check/query.py

@@ -7,11 +7,14 @@
     query.py spans   '<TraceQL query>'      [since]   the matching spans, each with its own time
     query.py trace   <trace id>                       one trace's spans, as a tree
     query.py prom    '<PromQL query>'       [since]   one value per series, now; $range is `since`
+    query.py flame   '<TraceQL query>'      [since] [out.svg]   where the matching spans' time
+                                                      went, summed over up to 100 traces, as a
+                                                      flame graph (and out.svg.folded)
 
 `since` is 1h, 24h, 7d and the like (default 24h). Loki and Tempo publish no query port on the
 host, so each call runs a throwaway alpine container on the monitoring project's network.
 """
-import json, subprocess, sys, time, urllib.parse
+import base64, html, json, subprocess, sys, time, urllib.parse, zlib
 
 NET = "monitoring_default"
 
@@ -24,11 +27,108 @@ def fetch(url):
     return json.loads(out.stdout)
 
 
+def fetch_many(urls):
+    # One container for them all: a docker run apiece is about a second, a hundred traces a minute and a half.
+    out = subprocess.run(["docker", "run", "--rm", "--network", NET, "alpine", "sh", "-c",
+                          'for u; do wget -qO- "$u"; echo; echo @@; done', "--", *urls],
+                         capture_output=True, text=True, timeout=600)
+    return [json.loads(part) for part in out.stdout.split("\n@@\n") if part.strip()]
+
+
+def trace_spans(r):
+    return [sp for b in r.get("batches", r.get("resourceSpans", []))
+            for ss in b.get("scopeSpans", b.get("instrumentationLibrarySpans", [])) for sp in ss["spans"]]
+
+
+def span_hex(sid):
+    # A trace's span ids are base64 and the search's are hex.
+    return sid if len(sid) == 16 and all(c in "0123456789abcdef" for c in sid) else base64.b64decode(sid).hex()
+
+
+def label(sp):
+    attrs = {a["key"]: next(iter(a["value"].values()), None) for a in sp.get("attributes", [])}
+    return f"step:{attrs['of']}" if sp["name"] == "step" and attrs.get("of") else sp["name"]
+
+
+def fold(spans, roots):
+    """Self time in µs for each stack under the spans in `roots` (hex ids), identical stacks summed.
+
+    shortcut: only as fine as the spans; a route's time outside its `step`s is its own frame. Add
+    steps in ipx where that is large. Children run at once (try_join) can add up past their
+    parent, so a frame's self time is floored at 0.
+    """
+    kids, stacks = {}, {}
+    for sp in spans:
+        kids.setdefault(span_hex(sp.get("parentSpanId") or "0" * 16), []).append(sp)
+
+    def walk(sp, path):
+        path = f"{path};{label(sp)}" if path else label(sp)
+        dur = lambda s: (int(s["endTimeUnixNano"]) - int(s["startTimeUnixNano"])) // 1000
+        cs = kids.get(span_hex(sp["spanId"]), [])
+        stacks[path] = stacks.get(path, 0) + max(0, dur(sp) - sum(dur(c) for c in cs))
+        for c in cs:
+            walk(c, path)
+
+    for sp in spans:
+        if span_hex(sp["spanId"]) in roots:
+            walk(sp, "")
+    return stacks
+
+
+def check_fold():
+    def sp(sid, parent, name, start, end, of=None):
+        attrs = [{"key": "of", "value": {"stringValue": of}}] if of else []
+        return {"spanId": sid, "parentSpanId": parent, "name": name, "attributes": attrs,
+                "startTimeUnixNano": str(start * 1000), "endTimeUnixNano": str(end * 1000)}
+    a, b, c, d = (f"{i:016x}" for i in range(1, 5))
+    spans = [sp(a, "", "GET /x", 0, 100), sp(b, a, "step", 0, 60, "q"), sp(c, a, "step", 0, 60, "q"),
+             sp(d, b, "inner", 0, 10)]
+    assert fold(spans, {a}) == {"GET /x": 0, "GET /x;step:q": 110, "GET /x;step:q;inner": 10}, fold(spans, {a})
+    assert fold(spans, {b}) == {"step:q": 50, "step:q;inner": 10}
+    # base64 ids, as a trace sends them, match the search's hex.
+    assert span_hex(base64.b64encode(bytes.fromhex(a)).decode()) == a
+    print("fold: ok")
+
+
+def flame_svg(stacks, n, title):
+    tree = {}  # name -> [total, children]
+    for path, us in stacks.items():
+        node = tree
+        for part in path.split(";"):
+            node = node.setdefault(part, [0, {}])
+            node[0] += us
+            node = node[1]
+    total = sum(v[0] for v in tree.values()) or 1
+    width, row, rects = 1200, 18, []
+
+    def draw(node, x, depth, trail):
+        for name, (us, sub) in sorted(node.items()):
+            w = us / total * width
+            stack = f"{trail};{name}" if trail else name
+            hue = 10 + zlib.crc32(name.encode()) % 45
+            tip = f"{stack}\n{us / 1000:.0f} ms in all, {100 * us / total:.1f}%, {us / 1000 / n:.1f} ms a request"
+            text = html.escape(name[: int(w / 7)]) if w > 30 else ""
+            rects.append(f'<g><title>{html.escape(tip)}</title><rect x="{x:.1f}" y="{30 + depth * row}" '
+                         f'width="{max(w - 0.5, 0.1):.1f}" height="{row - 1}" fill="hsl({hue},85%,62%)"/>'
+                         f'<text x="{x + 3:.1f}" y="{30 + depth * row + 13}">{text}</text></g>')
+            draw(sub, x, depth + 1, stack)
+            x += w
+
+    draw(tree, 0, 0, "")
+    depth = max(p.count(";") for p in stacks) + 1 if stacks else 1
+    head = html.escape(f"{title}: {n} request(s), {total / 1000 / n:.1f} ms each on average")
+    return (f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{40 + depth * row}" '
+            f'font-family="monospace" font-size="11"><rect width="100%" height="100%" fill="#fff"/>'
+            f'<text x="4" y="18" font-size="13">{head}</text>{"".join(rects)}</svg>')
+
+
 def seconds(since):
     return int(since[:-1]) * {"m": 60, "h": 3600, "d": 86400}[since[-1]]
 
 
 def main():
+    if sys.argv[1:] == ["flame", "--check"]:
+        return check_fold()
     if len(sys.argv) < 3:
         sys.exit(__doc__)
     kind, q = sys.argv[1], sys.argv[2]
@@ -82,10 +182,24 @@ def main():
             labels = {k: v for k, v in s["metric"].items() if k not in ("job", "instance")}
             print(f"{float(s['value'][1]):>14g}  {json.dumps(labels) if labels else ''}")
         print(f"-- {len(rows)} series", file=sys.stderr)
+    elif kind == "flame":
+        out = sys.argv[4] if len(sys.argv) > 4 else "flame.svg"
+        r = fetch(f"http://tempo:3200/api/search?q={enc(q)}&start={int(start)}&end={int(now)}&limit=100&spss=20")
+        found = r.get("traces", [])
+        roots = {sp["spanID"] for t in found for ss in t.get("spanSets") or [t.get("spanSet") or {}]
+                 for sp in ss.get("spans", [])}
+        stacks = {}
+        for tr in fetch_many([f"http://tempo:3200/api/traces/{t['traceID']}" for t in found]):
+            for path, us in fold(trace_spans(tr), roots).items():
+                stacks[path] = stacks.get(path, 0) + us
+        n = len(roots) or 1
+        with open(out + ".folded", "w") as f:
+            f.writelines(f"{path} {us}\n" for path, us in sorted(stacks.items()))
+        with open(out, "w") as f:
+            f.write(flame_svg(stacks, n, q))
+        print(f"-- {len(roots)} span(s) in {len(found)} trace(s): {out}, {out}.folded", file=sys.stderr)
     elif kind == "trace":
-        r = fetch(f"http://tempo:3200/api/traces/{q}")
+        spans = trace_spans(fetch(f"http://tempo:3200/api/traces/{q}"))
-        spans = [sp for b in r.get("batches", r.get("resourceSpans", []))
-                 for ss in b.get("scopeSpans", b.get("instrumentationLibrarySpans", [])) for sp in ss["spans"]]
         kids = {}
         for sp in spans:
             kids.setdefault(sp.get("parentSpanId", ""), []).append(sp)