trace-test.py

import trio
import os
import json
from itertools import count

# https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#

class Trace(trio.abc.Instrument):
    def __init__(self, out):
        self.out = out
        self.out.write("[\n")
        self.ids = count()
        self._task_metadata(-1, "I/O manager")

    def _write(self, **ev):
        ev.setdefault("pid", os.getpid())
        if ev["ph"] != "M":
            ev.setdefault("ts", trio.current_time() * 1e6)
        self.out.write(json.dumps(ev))
        self.out.write(",\n")

    def _task_metadata(self, tid, name):
        self._write(
            name="thread_name",
            ph="M",
            tid=tid,
            args={"name": name},
        )
        self._write(
            name="thread_sort_index",
            ph="M",
            tid=tid,
            args={"sort_index": tid},
        )

    def task_spawned(self, task):
        self._task_metadata(task._counter, task.name)
        self._write(
            name="task lifetime",
            ph="B",
            tid=task._counter,
        )

    def task_exited(self, task):
        self._write(
            name="task lifetime",
            ph="E",
            tid=task._counter,
        )

    def before_task_step(self, task):
        self._write(
            name="running",
            ph="B",
            tid=task._counter,
        )

    def after_task_step(self, task):
        self._write(
            name="running",
            ph="E",
            tid=task._counter,
        )

    def task_scheduled(self, task):
        # This doesn't work... why?
        try:
            waker = trio.hazmat.current_task()
        except RuntimeError:
            pass
        else:
            id = next(self.ids)
            self._write(
                ph="s",
                cat="wakeup",
                id=id,
                tid=waker._counter,
            )
            self._write(
                cat="wakeup",
                ph="f",
                id=id,
                tid=task._counter,
            )

    def before_io_wait(self, timeout):
        self._write(
            name=f"I/O wait",
            ph="B",
            tid=-1,
        )

    def after_io_wait(self, timeout):
        self._write(
            name=f"I/O wait",
            ph="E",
            tid=-1,
        )


async def child1():
    print("  child1: started! sleeping now...")
    await trio.sleep(1)
    print("  child1: exiting!")

async def child2():
    print("  child2: started! sleeping now...")
    await trio.sleep(1)
    print("  child2: exiting!")

async def parent():
    print("parent: started!")
    async with trio.open_nursery() as nursery:
        print("parent: spawning child1...")
        nursery.start_soon(child1)

        print("parent: spawning child2...")
        nursery.start_soon(child2)

        print("parent: waiting for children to finish...")
        # -- we exit the nursery block here --
    print("parent: all done!")

t = Trace(open("/tmp/t.json", "w"))
trio.run(parent, instruments=[t])