patx/skirr

"""Diagnostic-only FastAPI wrapper: await suspension, app time, GC and loop lag.

Set SKIRR_TAIL_OUTPUT to a JSON path. Never use this wrapper for throughput claims.
"""
import asyncio
import atexit
from collections import Counter
import gc
import json
import os
from pathlib import Path
import time
from apps import fastapi_app

histograms = {}
yields = Counter()
counts = Counter()
started = False
gc_started = {}
clock = time.perf_counter_ns


def record(name, ns):
    histogram = histograms.setdefault(name, Counter())
    histogram[min(ns // 1000, 1000000)] += 1


def dump():
    Path(os.environ['SKIRR_TAIL_OUTPUT']).write_text(json.dumps({
        'pid': os.getpid(), 'diagnostic_only': True, 'counts': counts,
        'await_yields': yields, 'microsecond_histograms': histograms,
    }))


def collection(phase, info):
    generation = info['generation']
    if phase == 'start':
        gc_started[generation] = clock()
    elif generation in gc_started:
        record(f'gc_gen{generation}', clock() - gc_started.pop(generation))


class Observed:
    def __init__(self, operation, kind):
        self.iterator = operation.__await__()
        self.kind = kind

    def __await__(self):
        return self

    def __iter__(self):
        return self

    def __next__(self):
        result = next(self.iterator)
        yields[self.kind] += 1
        return result

    def send(self, value):
        if value is None:
            return next(self)
        result = self.iterator.send(value)
        yields[self.kind] += 1
        return result

    def throw(self, error):
        method = getattr(self.iterator, 'throw', None)
        if method is None:
            raise error
        return method(error)


def initialize():
    gc.callbacks.append(collection)
    atexit.register(dump)
    loop = asyncio.get_running_loop()
    last_dump = clock()

    def heartbeat(due):
        nonlocal last_dump
        now = clock()
        record('loop_lag', max(0, now - due))
        if now - last_dump > 1000000000:
            dump()
            last_dump = now
        loop.call_later(.001, heartbeat, clock() + 1000000)

    loop.call_later(.001, heartbeat, clock() + 1000000)


async def app(scope, receive, send):
    global started
    if not started:
        initialize()
        started = True
    begin = clock()

    async def observed_send(message):
        kind = message['type']
        start = clock()
        try:
            return await Observed(send(message), kind)
        finally:
            record(kind, clock() - start)
            counts[kind] += 1

    try:
        await fastapi_app(scope, receive, observed_send)
    finally:
        record('app', clock() - begin)
        counts['requests'] += 1