austin-python 2.2.0__py3-none-any.whl → 2.4.0__py3-none-any.whl

austin/events.py

@@ -43,6 +43,63 @@ class AustinMetadata(AustinEvent):
     value: str
 
 
+@dataclass(frozen=True)
+class AustinTask:
+    """A node in an asyncio task tree (3.14+ only).
+
+    ``frames`` is this task's last-known position: its coroutine chain if it
+    was suspended, or the frames it was running if it was running.
+    ``awaiting`` are the tasks this task is itself awaiting, i.e. the tasks
+    for which this task is the direct waiter.
+
+    This is the structure of the tree only. The time tasks spend at their
+    positions is carried by AustinTaskReport (see AustinSample.task_reports)
+    and, while they are running, by the samples they are running at (see
+    AustinSample.running_task).
+    """
+
+    task_id: int
+    name: t.Optional[str]
+    frames: t.Tuple[AustinFrame, ...] = ()
+    awaiting: t.Tuple["AustinTask", ...] = ()
+
+
+@dataclass(frozen=True)
+class AustinTaskReport:
+    """A report of a task's position (3.14+ only).
+
+    Austin only reports a task's position when it has changed. ``frames`` is
+    the task's new position, and ``elapsed`` is how long the task was
+    *suspended* at its ``previous`` one, which this report replaces -- the
+    earliest point at which that duration could be known at all. Austin
+    doesn't report it: readers work it out as the time of the task's thread's
+    samples while the task held that position without running. It is await
+    time only: the time a task spends running is the time of the samples it
+    is running at (see AustinSample.running_task). Always 0 in CPU mode, where
+    a suspended task uses no CPU.
+
+    ``previous`` is None on a task's first report, which has no previous
+    position (and so ``elapsed`` is 0). A report with no ``frames`` closes the
+    task: it is gone, and ``elapsed`` is the time at its last position. At the
+    end of a stream, each task still alive gets a report with ``frames`` equal
+    to ``previous``, carrying the time it has spent at its position since.
+
+    Each report is attached to exactly one sample, so summing ``elapsed``
+    over every sample's reports counts each duration once.
+    """
+
+    task_id: int
+    name: t.Optional[str]
+    frames: t.Tuple[AustinFrame, ...]
+    previous: t.Optional[t.Tuple[AustinFrame, ...]]
+    elapsed: MicroSeconds
+
+    @property
+    def closing(self) -> bool:
+        """Whether this report closes the task."""
+        return not self.frames
+
+
 @dataclass(frozen=True)
 class AustinSample(AustinEvent):
     """Austin sample."""
@@ -63,6 +120,13 @@ class AustinSample(AustinEvent):
     frames: t.Optional[t.Tuple[AustinFrame, ...]] = None
     gc: t.Optional[bool] = None
     idle: t.Optional[bool] = None
+    # The live tasks owned by this sample's thread, as of this sample.
+    tasks: t.Tuple[AustinTask, ...] = ()
+    # The task reports that came with this sample, each one exactly once.
+    task_reports: t.Tuple[AustinTaskReport, ...] = ()
+    # The task this sample's thread was running, if any. The sample's metrics
+    # then belong to that task too, at its latest position (see AustinTask).
+    running_task: t.Optional[int] = None
 
     def key(
         self,

austin/format/mojo.py

@@ -12,10 +12,17 @@ from austin.events import AustinFrame
 from austin.events import AustinMetadata
 from austin.events import AustinMetrics
 from austin.events import AustinSample
+from austin.events import AustinTask
+from austin.events import AustinTaskReport
 from austin.events import InterpreterId
 from austin.events import ProcessId
 from austin.events import ThreadName
 
+# Guards against a corrupted/cyclic waiter graph in the source stream -- a
+# real await chain never comes anywhere close to this. Mirrors austin's own
+# WHERE_MAX_TREE_DEPTH (src/events.c).
+MAX_TASK_TREE_DEPTH = 64
+
 
 def to_varint(n: int) -> bytes:
     """Convert an integer to a variable-length integer."""
@@ -67,6 +74,9 @@ class MojoEvents:
     STRING = 11
     STRING_REF = 12
     STACK_REPEAT = 13
+    TASK_STACK = 14
+    TASK_WAITER = 15
+    TASK_RUNNING = 16
 
 
 class MojoEventHandler:
@@ -197,6 +207,52 @@ class MojoStackRepeat(MojoEvent):
     EVENT_ID = MojoEvents.STACK_REPEAT
 
 
+@dataclass(frozen=True, eq=True)
+class MojoTaskWaiter(MojoEvent):
+    """MOJO task waiter event.
+
+    One edge of a task's waiter DAG: ``task_id`` is awaited by
+    ``waiter_id``.
+    """
+
+    EVENT_ID = MojoEvents.TASK_WAITER
+
+    task_id: int
+    waiter_id: int
+
+
+@dataclass(frozen=True, eq=True)
+class MojoTaskRunning(MojoEvent):
+    """MOJO task running event.
+
+    The task the enclosing sample's thread was running at that sample. The
+    sample's metrics then belong to that task too, at its latest position.
+    """
+
+    EVENT_ID = MojoEvents.TASK_RUNNING
+
+    task_id: int
+
+
+@dataclass(frozen=True, eq=True)
+class MojoTaskStack(MojoEvent):
+    """MOJO task stack event.
+
+    Introduces a task's new position, keyed by an opaque task id and its
+    (possibly cached) name key (0 if unresolved, and on a closing flush). The
+    next ``frame_count`` frame references are the task's frames; none means
+    the task is gone. The report carries no time: the time a task spends at a
+    position is the time of its thread's samples while it holds it -- see
+    AustinTaskReport.
+    """
+
+    EVENT_ID = MojoEvents.TASK_STACK
+
+    task_id: int
+    name_key: int
+    frame_count: int
+
+
 @dataclass(frozen=True, eq=True)
 class MojoFrame(MojoEvent):
     """MOJO frame."""
@@ -264,6 +320,33 @@ class _RunningSample:
     metrics: t.Dict[MojoMetricType, MojoMetric] = field(default_factory=dict)
     gc: t.Optional[bool] = None
     idle: t.Optional[bool] = None
+    running_task: t.Optional[int] = None
+    task_reports: t.List[AustinTaskReport] = field(default_factory=list)
+
+
+@dataclass
+class _RunningTaskStack:
+    task_id: int
+    name_key: int
+    frame_count: int
+    frames: t.List[MojoFrame] = field(default_factory=list)
+
+
+@dataclass
+class _TaskInfo:
+    """Best-known state for one task_id.
+
+    Built up from MOJO_TASK_STACK events as they arrive. ``owner`` is the
+    (pid, thread) of whichever sample's wire bracket this task's most recent
+    MOJO_TASK_STACK fell within -- see BaseMojoStreamReader's task-stack
+    handling for why that bracket is the only ownership signal that exists.
+    """
+
+    name: t.Optional[str]
+    owner: t.Tuple[ProcessId, ThreadName]
+    frames: t.Tuple[AustinFrame, ...] = ()
+    # Time spent at ``frames`` so far, i.e. since the task's last report.
+    pending: int = 0
 
 
 def int_reader() -> t.Generator[t.Optional[int], bytes, int]:
@@ -340,6 +423,26 @@ class BaseMojoStreamReader(AustinEventIterator):
         # last fully-finalised sample for that thread.
         self._prev_frames: t.Dict[t.Tuple[int, str], t.List[MojoFrame]] = {}
 
+        # Suspended-task coroutine-chain state (3.14+ asyncio). A task stack's
+        # frames/terminating time metric are ordinary MOJO_FRAME/MOJO_FRAME_REF/
+        # MOJO_METRIC_TIME events on the wire, indistinguishable by event ID
+        # alone from those of the enclosing regular stack sample -- so while
+        # this is not None, get_frame_ref/get_time_metric route to it instead
+        # of the running sample.
+        self._running_task_stack: t.Optional[_RunningTaskStack] = None
+
+        # Best-known state per task_id, keyed by remote TaskObj address.
+        self._task_info: t.Dict[int, _TaskInfo] = {}
+
+        # Best-known waiter DAG edges: (task_id, waiter_id) pairs, meaning
+        # task_id is awaited by waiter_id. The wire re-emits a task's
+        # *complete* current waiter set whenever it changes (never a diff --
+        # see py_asyncio.c's fingerprint comment), so a fresh batch for a
+        # task_id replaces its previous edges; batch boundaries are inferred
+        # from consecutive same-task_id events (see get_task_waiter).
+        self._task_edges: t.Set[t.Tuple[int, int]] = set()
+        self._last_waiter_task_id: t.Optional[int] = None
+
         # Austin events
         self.metadata: t.Dict[str, str] = {}
         self.samples: t.List[AustinSample] = []
@@ -384,10 +487,61 @@ class BaseMojoStreamReader(AustinEventIterator):
     def gc(self) -> t.Optional[str]:
         return self.metadata.get("gc")
 
-    def _finalize_sample(self) -> AustinSample:
-        """Finalize the current sample."""
+    def _accrue_task_time(self, sample: _RunningSample) -> None:
+        """Add the time of a sample to the tasks its thread holds.
+
+        A task's time at a position is the time of its thread's samples while
+        it holds it: the running task's go to it as running (see
+        AustinSample.running_task), every other task's as time at its current
+        position, which its next report carries. In CPU mode a task that isn't
+        running isn't using the CPU, so it gets nothing.
+        """
+        if self.mode == "cpu":
+            return
+
+        time_metric = sample.metrics.get(MojoMetricType.TIME)
+        if time_metric is None or not time_metric.value:
+            return
+
+        owner = (sample.pid, sample.thread)
+        for task_id, info in self._task_info.items():
+            if info.owner == owner and task_id != sample.running_task:
+                info.pending += time_metric.value
+
+    def _end_task_reports(self) -> t.List[AustinTaskReport]:
+        """Report the time the tasks still alive have spent at their position.
+
+        Nothing else would: a task's time at a position is only reported when
+        the task moves on, or goes. These reports leave the position unchanged
+        (``frames`` and ``previous`` are the same).
+        """
+        reports = [
+            AustinTaskReport(
+                task_id=task_id,
+                name=info.name,
+                frames=info.frames,
+                previous=info.frames,
+                elapsed=info.pending,
+            )
+            for task_id, info in self._task_info.items()
+            if info.pending
+        ]
+        for info in self._task_info.values():
+            info.pending = 0
+        return reports
+
+    def _finalize_sample(self, last: bool = False) -> AustinSample:
+        """Finalize the current sample.
+
+        The last sample of the stream also carries the end-of-stream task
+        reports (see _end_task_reports).
+        """
         assert self._running_sample is not None, self._running_sample
 
+        self._accrue_task_time(self._running_sample)
+        if last:
+            self._running_sample.task_reports.extend(self._end_task_reports())
+
         self.samples.append(
             sample := AustinSample(
                 pid=self._running_sample.pid,
@@ -416,6 +570,11 @@ class BaseMojoStreamReader(AustinEventIterator):
                 ),
                 gc=self._running_sample.gc,
                 idle=self._running_sample.idle,
+                running_task=self._running_sample.running_task,
+                task_reports=tuple(self._running_sample.task_reports),
+                tasks=self.get_tasks(
+                    self._running_sample.pid, self._running_sample.thread
+                ),
             )
         )
 
@@ -459,6 +618,10 @@ class BaseMojoStreamReader(AustinEventIterator):
         if self._running_sample is not None:
             self._finalize_sample()
 
+        # A task's waiter set is a run of consecutive MOJO_TASK_WAITER events,
+        # which never spans a sample boundary.
+        self._last_waiter_task_id = None
+
         self._running_sample = _RunningSample(
             pid=pid,
             iid=iid,
@@ -499,11 +662,184 @@ class BaseMojoStreamReader(AustinEventIterator):
         """Parse a frame reference."""
         frame = self._frame_map[self.ref(ref)]
 
-        assert self._running_sample is not None, self._running_sample
-        self._running_sample.frames.append(frame)
+        if self._running_task_stack is not None:
+            self._running_task_stack.frames.append(frame)
+            if (
+                len(self._running_task_stack.frames)
+                >= self._running_task_stack.frame_count
+            ):
+                self._finalize_task_stack()
+        else:
+            assert self._running_sample is not None, self._running_sample
+            self._running_sample.frames.append(frame)
 
         return MojoFrameReference(frame)
 
+    def get_task_waiter(self, task_id: int, waiter_id: int) -> MojoTaskWaiter:
+        """Parse a task waiter edge."""
+        if self._last_waiter_task_id != task_id:
+            # Starting a fresh batch for this task_id: the wire re-emits the
+            # complete current waiter set whenever it changes, so drop
+            # whatever we knew about this task_id's waiters before.
+            self._task_edges = {e for e in self._task_edges if e[0] != task_id}
+            self._last_waiter_task_id = task_id
+
+        # A waiter of 0 is no task: the task's waiter set is now empty.
+        if waiter_id:
+            self._task_edges.add((task_id, waiter_id))
+
+        return MojoTaskWaiter(task_id, waiter_id)
+
+    def get_task_running(self, task_id: int) -> MojoTaskRunning:
+        """Parse the task the current sample's thread was running."""
+        assert self._running_sample is not None, self._running_sample
+        self._running_sample.running_task = task_id
+
+        return MojoTaskRunning(task_id)
+
+    def get_task_stack(
+        self, task_id: int, name_key: int, frame_count: int
+    ) -> MojoTaskStack:
+        """Parse the header of a task report.
+
+        Its frames arrive as the next ``frame_count`` frame references -- see
+        get_frame_ref. A report with none closes the task right away.
+        """
+        self._running_task_stack = _RunningTaskStack(task_id, name_key, frame_count)
+        self._last_waiter_task_id = None
+
+        if not frame_count:
+            self._finalize_task_stack()
+
+        return MojoTaskStack(task_id, name_key, frame_count)
+
+    def _finalize_task_stack(self) -> None:
+        """Finalize the running task stack.
+
+        Triggered by its last frame reference (see get_frame_ref), or right
+        away by a header with no frames (see get_task_stack). The report's
+        frames are the task's new position, and its elapsed time is the time
+        the task spent at its previous one -- see AustinTaskReport and
+        _accrue_task_time. The report is attached to the sample
+        currently being assembled, and the task's owner is recorded as that
+        sample's thread: this event only ever arrives bracketed inside that
+        sample's own MOJO_STACK...next-MOJO_STACK window, so there's nothing
+        else it could belong to.
+
+        An empty ``ts.frames`` is a closing flush -- the task is gone
+        (completed, or otherwise), and this only carries its final time. It
+        has no name of its own, nor a meaningful owner (by the time eviction
+        runs, every thread for this cycle has already had its turn, so
+        "whichever sample is currently being assembled" is arbitrary here).
+        The task is forgotten, along with every waiter edge it is part of.
+        """
+        ts = self._running_task_stack
+        assert ts is not None, ts
+        assert self._running_sample is not None, self._running_sample
+
+        self._running_task_stack = None
+
+        existing = self._task_info.get(ts.task_id)
+        previous = existing.frames if existing is not None else None
+        elapsed = existing.pending if existing is not None else 0
+
+        if ts.name_key:
+            name: t.Optional[str] = self._lookup_string(ts.name_key).value
+        else:
+            name = existing.name if existing is not None else None
+
+        frames = tuple(
+            AustinFrame(
+                filename=mf.filename.value,
+                function=mf.scope.value,
+                line=mf.line,
+                line_end=mf.line_end,
+                column=mf.column,
+                column_end=mf.column_end,
+            )
+            for mf in ts.frames
+        )
+
+        self._running_sample.task_reports.append(
+            AustinTaskReport(
+                task_id=ts.task_id,
+                name=name,
+                frames=frames,
+                previous=previous,
+                elapsed=elapsed,
+            )
+        )
+
+        if not frames:
+            self._task_info.pop(ts.task_id, None)
+            self._task_edges = {e for e in self._task_edges if ts.task_id not in e}
+            return
+
+        self._task_info[ts.task_id] = _TaskInfo(
+            name=name,
+            owner=(self._running_sample.pid, self._running_sample.thread),
+            frames=frames,
+        )
+
+    def get_tasks(self, pid: ProcessId, thread: ThreadName) -> t.Tuple[AustinTask, ...]:
+        """Return the current best-known root tasks owned by (pid, thread).
+
+        Every root task it owns (nobody awaits it), and everything it
+        transitively awaits, as one AustinTask per root -- a thread can own
+        more than one at once (e.g. several siblings under one gather()).
+
+        Mirrors austin's own where_event_handler__render_tree: only live
+        tasks that have been reported at least once appear. A waiter we know
+        nothing about (never reported, or already gone) doesn't count as a
+        parent, so the tasks it awaited still show up, as roots. Like that
+        C-side counterpart, this is a no-op (returns empty) when (pid,
+        thread) owns no tasks, so callers can call it unconditionally.
+        """
+
+        def has_parent(task_id: int) -> bool:
+            return any(
+                t_id == task_id and w_id in self._task_info
+                for t_id, w_id in self._task_edges
+            )
+
+        def build(task_id: int, depth: int) -> AustinTask:
+            info = self._task_info[task_id]
+            awaiting = (
+                tuple(
+                    build(awaited_id, depth + 1)
+                    for awaited_id, waiter_id in self._task_edges
+                    if waiter_id == task_id and awaited_id in self._task_info
+                )
+                if depth < MAX_TASK_TREE_DEPTH
+                else ()
+            )
+            return AustinTask(
+                task_id=task_id,
+                name=info.name,
+                frames=info.frames,
+                awaiting=awaiting,
+            )
+
+        owner = (pid, thread)
+        roots = [
+            task_id
+            for task_id, info in self._task_info.items()
+            if info.owner == owner and not has_parent(task_id)
+        ]
+
+        if not roots:
+            # Best-effort fallback, matching austin's own render_tree: if
+            # every one of this owner's tasks appears to have a parent (a
+            # cyclic/bogus edge), there's no true root -- show each as its
+            # own top-level tree rather than silently showing nothing.
+            roots = [
+                task_id
+                for task_id, info in self._task_info.items()
+                if info.owner == owner
+            ]
+
+        return tuple(build(task_id, 0) for task_id in roots)
+
     def get_kernel_frame(self, name: str) -> MojoKernelFrame:
         """Parse kernel frame."""
         return MojoKernelFrame(name)
@@ -630,7 +966,7 @@ class MojoStreamReader(BaseMojoStreamReader):
     @handles(MojoEvents.METRIC_TIME)
     def parse_time_metric(self) -> MojoMetric:
         """Parse time metric."""
-        return self._parse_metric(MojoMetricType.TIME)
+        return self.get_time_metric(self.read_int())
 
     @handles(MojoEvents.METRIC_MEMORY)
     def parse_memory_metric(self) -> MojoMetric:
@@ -667,6 +1003,21 @@ class MojoStreamReader(BaseMojoStreamReader):
         """Parse a stack repeat event."""
         return self.get_stack_repeat()
 
+    @handles(MojoEvents.TASK_WAITER)
+    def parse_task_waiter(self) -> MojoTaskWaiter:
+        """Parse a task waiter edge."""
+        return self.get_task_waiter(self.read_int(), self.read_int())
+
+    @handles(MojoEvents.TASK_RUNNING)
+    def parse_task_running(self) -> MojoTaskRunning:
+        """Parse the running task of a sample."""
+        return self.get_task_running(self.read_int())
+
+    @handles(MojoEvents.TASK_STACK)
+    def parse_task_stack(self) -> MojoTaskStack:
+        """Parse the header of a task report."""
+        return self.get_task_stack(self.read_int(), self.read_int(), self.read_int())
+
     def parse_event(self) -> t.Optional[MojoEvent]:
         """Parse a single event."""
         try:
@@ -724,10 +1075,13 @@ class MojoStreamReader(BaseMojoStreamReader):
                 if self.samples:
                     yield self.samples[-1]
         if self._running_sample is not None:
-            yield self._finalize_sample()
+            yield self._finalize_sample(last=True)
 
     def hexdump(
-        self, start: int, end: int, highlight: t.Set[int] = set()  # noqa: B006
+        self,
+        start: int,
+        end: int,
+        highlight: t.Set[int] = set(),  # noqa: B006
     ) -> None:
         """Print a hexdump of the MOJO file."""
         self.mojo.seek(start)
@@ -818,7 +1172,7 @@ class AsyncMojoStreamReader(BaseMojoStreamReader):
     @handles(MojoEvents.METRIC_TIME)
     async def parse_time_metric(self) -> MojoMetric:
         """Parse time metric."""
-        return await self._parse_metric(MojoMetricType.TIME)
+        return self.get_time_metric(await self.read_int())
 
     @handles(MojoEvents.METRIC_MEMORY)
     async def parse_memory_metric(self) -> MojoMetric:
@@ -857,6 +1211,23 @@ class AsyncMojoStreamReader(BaseMojoStreamReader):
         """Parse a stack repeat event."""
         return self.get_stack_repeat()
 
+    @handles(MojoEvents.TASK_WAITER)
+    async def parse_task_waiter(self) -> MojoTaskWaiter:
+        """Parse a task waiter edge."""
+        return self.get_task_waiter(await self.read_int(), await self.read_int())
+
+    @handles(MojoEvents.TASK_RUNNING)
+    async def parse_task_running(self) -> MojoTaskRunning:
+        """Parse the running task of a sample."""
+        return self.get_task_running(await self.read_int())
+
+    @handles(MojoEvents.TASK_STACK)
+    async def parse_task_stack(self) -> MojoTaskStack:
+        """Parse the header of a task report."""
+        return self.get_task_stack(
+            await self.read_int(), await self.read_int(), await self.read_int()
+        )
+
     async def parse_event(self) -> t.Optional[MojoEvent]:
         """Parse a single event."""
         try:
@@ -914,13 +1285,13 @@ class AsyncMojoStreamReader(BaseMojoStreamReader):
                 if self.samples:
                     yield self.samples[-1]
         if self._running_sample is not None:
-            yield self._finalize_sample()
+            yield self._finalize_sample(last=True)
 
 
 class BaseMojoStreamWriter(abc.ABC):
     """Base class for MOJO stream writers."""
 
-    HEADER = b"MOJ\x03"
+    HEADER = b"MOJ\x04"
 
     def __init__(self, mojo: t.Any) -> None:
         self.mojo = mojo
@@ -936,6 +1307,11 @@ class BaseMojoStreamWriter(abc.ABC):
         self._gc = False
 
         self._new_entries: t.List[MojoEvent] = []
+        self._task_ids: t.Dict[int, int] = {}
+        # The waiter set last written for each task, and the tasks written so
+        # far (see write_tasks).
+        self._written_waiters: t.Dict[int, t.FrozenSet[int]] = {}
+        self._written_tasks: t.Set[int] = set()
 
     def set_metadata(self, metadata: AustinMetadata) -> None:
         self._meta[metadata.name] = metadata.value
@@ -952,6 +1328,15 @@ class BaseMojoStreamWriter(abc.ABC):
             self._new_entries.append(mojo_string)
             return mojo_string
 
+    def resolve_task_id(self, task_id: int) -> int:
+        try:
+            return self._task_ids[task_id]
+        except KeyError:
+            # 0 is reserved: it stands for no task (e.g. an unknown waiter).
+            compact_id = len(self._task_ids) + 1
+            self._task_ids[task_id] = compact_id
+            return compact_id
+
     def resolve_frame(self, frame: AustinFrame) -> MojoFrame:
         try:
             return self._frames[frame]
@@ -968,6 +1353,114 @@ class BaseMojoStreamWriter(abc.ABC):
             self._new_entries.append(mojo_frame)
             return mojo_frame
 
+    def _collect_tasks(
+        self,
+        tasks: t.Sequence[AustinTask],
+        seen: t.Dict[int, AustinTask],
+        edges: t.Dict[int, t.Set[int]],
+    ) -> None:
+        """Flatten a root-task list into its distinct tasks and waiter edges.
+
+        A task can legitimately appear more than once (once per waiter, if
+        it has more than one -- see AustinTask's own docstring on
+        ``awaiting``), so this collects each task_id's own stack exactly
+        once (``seen``, first occurrence wins) and every one of its waiters
+        together (``edges``), rather than re-walking and re-emitting it once
+        per occurrence.
+        """
+        for task in tasks:
+            seen.setdefault(task.task_id, task)
+            for awaited in task.awaiting:
+                edges.setdefault(awaited.task_id, set()).add(task.task_id)
+            self._collect_tasks(task.awaiting, seen, edges)
+
+    def write_tasks(self, sample: AustinSample) -> int:
+        """Write a sample's task reports, then the waiter edges that changed.
+
+        Each report is written once, as it was read, without its elapsed time,
+        which readers work out from the samples. Reports that leave a task
+        where it was, like end-of-stream reports, aren't written: readers
+        produce those themselves. Waiter edges are written as the complete
+        waiter set of every task whose set differs from the last one written
+        for it, contiguously (required for MojoStreamReader.get_task_waiter's
+        own batch-replace logic to reconstruct a multi-waiter task correctly),
+        and an emptied set as a single edge to the reserved 0. See
+        _collect_tasks for why each task is considered exactly once.
+
+        A task in the sample's tree that the stream hasn't introduced yet
+        (e.g. a sample built by hand, with no reports) is introduced with a
+        report of its own, carrying no time: readers only learn of tasks from
+        reports.
+        """
+        size = 0
+
+        seen: t.Dict[int, AustinTask] = {}
+        edges: t.Dict[int, t.Set[int]] = {}
+        self._collect_tasks(sample.tasks, seen, edges)
+
+        reported = {r.task_id for r in sample.task_reports}
+        introductions = tuple(
+            AustinTaskReport(
+                task_id=task.task_id,
+                name=task.name,
+                frames=task.frames,
+                previous=None,
+                elapsed=0,
+            )
+            for task in seen.values()
+            if task.task_id not in self._written_tasks
+            and task.task_id not in reported
+            and task.frames
+        )
+
+        for report in sample.task_reports + introductions:
+            if report.frames and report.frames == report.previous:
+                continue
+
+            name_key = (
+                self.resolve_string(report.name).key
+                if report.name and report.frames
+                else 0
+            )
+            task_frames = [self.resolve_frame(f) for f in report.frames]
+
+            while self._new_entries:
+                size += self.mojo.write(self._new_entries.pop(0).to_bytes())
+
+            size += self.mojo.write(
+                MojoTaskStack(
+                    task_id=self.resolve_task_id(report.task_id),
+                    name_key=name_key,
+                    frame_count=len(task_frames),
+                ).to_bytes()
+            )
+
+            for frame in task_frames:
+                size += self.mojo.write(MojoFrameReference(frame).to_bytes())
+
+            if report.closing:
+                self._written_waiters.pop(report.task_id, None)
+                self._written_tasks.discard(report.task_id)
+            else:
+                self._written_tasks.add(report.task_id)
+
+        for task_id in seen:
+            waiter_ids = frozenset(edges.get(task_id, ()))
+            if waiter_ids == self._written_waiters.get(task_id, frozenset()):
+                continue
+            self._written_waiters[task_id] = waiter_ids
+
+            compact_task_id = self.resolve_task_id(task_id)
+            for waiter_id in waiter_ids or (0,):
+                size += self.mojo.write(
+                    MojoTaskWaiter(
+                        task_id=compact_task_id,
+                        waiter_id=self.resolve_task_id(waiter_id) if waiter_id else 0,
+                    ).to_bytes()
+                )
+
+        return size
+
     @abc.abstractmethod
     def write(self, event: AustinEvent) -> int: ...  # noqa: E704
 
@@ -990,16 +1483,28 @@ class MojoStreamWriter(BaseMojoStreamWriter):
             )
 
         elif isinstance(event, AustinSample):
-            frames = (
-                [self.resolve_frame(f) for f in event.frames] if event.frames else []
-            )
-
             size += self.mojo.write(
                 MojoStack(
                     pid=event.pid, iid=event.iid or 0, tid=event.thread
                 ).to_bytes()
             )
 
+            # Task-scan payload (task reports, then waiter edges) is written
+            # before this sample's own frames -- it
+            # rides between the enclosing MOJO_STACK header and the sample's
+            # own frame dump on the real wire too (see austin's own
+            # py_proc.c, _py_proc__sample_threads).
+            size += self.write_tasks(event)
+
+            if event.running_task is not None:
+                size += self.mojo.write(
+                    MojoTaskRunning(self.resolve_task_id(event.running_task)).to_bytes()
+                )
+
+            frames = (
+                [self.resolve_frame(f) for f in event.frames] if event.frames else []
+            )
+
             while self._new_entries:
                 size += self.mojo.write(self._new_entries.pop(0).to_bytes())
 

austin/stats.py

@@ -67,7 +67,8 @@ class HierarchicalStats:
     leaf is given by those instances with an empty ``children`` attribute.
 
     The ``label`` attribute is used for indexing reasons and therefore should
-    be of a hashable type.
+    be of a hashable type. The ``samples`` attribute counts the samples that
+    contributed to the ``total`` of the node.
 
     This class overrides the default ``add`` operator so that one can perform
     operations like ``stats1 + stats2``. Note, however, that instances of this
@@ -79,6 +80,7 @@ class HierarchicalStats:
     own: int
     total: int
     children: Dict[Any, "HierarchicalStats"] = field(default_factory=dict)
+    samples: int = 0
 
     def __lshift__(self, other: "HierarchicalStats") -> "HierarchicalStats":
         """Merge the RHS into the LHS."""
@@ -87,6 +89,7 @@ class HierarchicalStats:
 
         self.own += other.own
         self.total += other.total
+        self.samples += other.samples
 
         for frame, child in other.children.items():
             try:
@@ -270,8 +273,11 @@ class AustinStats:
         by using :func:`Sample.parse` on a sample string passed by Austin to
         the sample callback.
         """
-        if self.stats_type in {AustinStatsType.WALL, AustinStatsType.CPU}:
+        if self.stats_type is AustinStatsType.WALL:
             metric = sample.metrics.time
+        elif self.stats_type is AustinStatsType.CPU:
+            # In full mode, idle samples contribute wall time but no CPU time
+            metric = None if sample.idle else sample.metrics.time
         elif (
             self.stats_type is AustinStatsType.MEMORY_ALLOC
             and sample.metrics.memory is not None
@@ -292,13 +298,15 @@ class AustinStats:
 
         pid = sample.pid
         thread_info = ThreadInfo(sample.thread, sample.iid or 0)
-        thread_stats = ThreadStats(thread_info, own=0, total=metric)
+        thread_stats = ThreadStats(thread_info, own=0, total=metric, samples=1)
 
         # Convert the list of frames into a nested FrameStats instance
         stats: HierarchicalStats = thread_stats
         container = thread_stats.children
         for height, frame in enumerate(sample.frames or []):
-            stats = FrameStats(label=frame, height=height, own=0, total=metric)
+            stats = FrameStats(
+                label=frame, height=height, own=0, total=metric, samples=1
+            )
             container[frame] = stats
             container = stats.children
         stats.own = stats.total

{austin_python-2.2.0.dist-info → austin_python-2.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: austin-python
-Version: 2.2.0
+Version: 2.4.0
 Summary: Python wrapper for Austin, the CPython frame stack sampler.
 Project-URL: homepage, https://github.com/P403n1x87/austin-python
 Project-URL: documentation, https://austin-python.readthedocs.io

{austin_python-2.2.0.dist-info → austin_python-2.4.0.dist-info}/RECORD

@@ -4,14 +4,14 @@ austin/base.py,sha256=Xq734VKzE2ggShs6UThsGznBGazg-gCw_nZEJDJKMlM,7553
 austin/cli.py,sha256=--PKm2I4NBBSSSBV5l-YC0mcmLjrvb1FHzxgY4dkby4,7826
 austin/config.py,sha256=PEwdG2Cy8ay1SrFr2N6YfQSZwseYBA1BA5dMfocKR3A,2091
 austin/errors.py,sha256=lCDSTANW_95Az_eDXbn__JRBllB9yw7tGnFo1-tw09Q,1046
-austin/events.py,sha256=u04uArj15b-CkQF7_WYnd1qZ41SFquAnxgpZN9KBskQ,2159
+austin/events.py,sha256=afe6SYCXD_47qn6fVajKAgq4J-ivnEBtSOGqRAWFoA4,4910
 austin/simple.py,sha256=iTmJGAjSBeUv82SgdMTwdNlAQwuK33RyU8FkBCE8rRg,5579
-austin/stats.py,sha256=xY4BZtiE1Z4Bfucah1da6jwkEREvd1jLfmW4WuCgbP0,11084
+austin/stats.py,sha256=YAE641Uwjt0Ty-CbhsHPeCy2B8OgShY_oo-BW-sJ_7Q,11465
 austin/threads.py,sha256=AvcoOkr3FeFAGBcqgXukNIt9xOZUipv7pXMqPIL_Qts,3649
 austin/format/__init__.py,sha256=ZtBJCLpTiZaLFwCN1UbMXLQvYjOSpwmtyGwqgK8yruQ,405
 austin/format/collapsed_stack.py,sha256=vD0CFkUgnRqHaA7dYsh4QYRZubE5IxyhVFnuPa37BSs,7245
 austin/format/compress.py,sha256=by6i79o0pEmGf2Py8edC6RWHVdr5RGsG2FWT6xQjUvk,3797
-austin/format/mojo.py,sha256=iqfPmJQ3tqoP-fclKDd0OiOSLiIQaWRjddA8L-7Huyw,31414
+austin/format/mojo.py,sha256=YWFoturPDxihzF0ACTF_DFaLni5FyE8BJZa9GKzCa2U,51325
 austin/format/mojo_compress.py,sha256=qhzXsenUThxwAz1lX1Z8xtGgBpPEIbD3cqZz2Iar3ws,3711
 austin/format/speedscope.py,sha256=hIuB1A8wbGFFLQUpZqdM7Tlf4dqFwyy01zFHDybs2Iw,8324
 austin/format/pprof/__init__.py,sha256=E2JJDaw_k_m18OZeITKtepUUIo3n09UHdkpvabSSABc,6804
@@ -21,10 +21,10 @@ austin/format/pprof/profile_pb2.py,sha256=i_uM8lDExR95yz7qnEVikhpkpNvetKaYzoiH7A
 austin/tools/diff.py,sha256=5yrHlniofIeMSqzOUD55l2h6Ps0GqEpfisFRj65gpdQ,5008
 austin/tools/mojodbg.py,sha256=l8Lsa4rTQLXfkau8RyGziyz8ZyPEgacFVXafbHT9FN4,2816
 austin/tools/resolve.py,sha256=bTXkll68nOtXqa0XHW79mp_seAJW8_bhd6ej6wa2_Oc,7816
-austin_python-2.2.0.dist-info/METADATA,sha256=EAzdTQnAVmxg6NqXea2G5DzYcqIFwPksjByks3Qvn2s,7537
-austin_python-2.2.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-austin_python-2.2.0.dist-info/entry_points.txt,sha256=xmPVC_6UkSOYu7-YMYm3BwIcMFYEhRw0qlF05Cq5gf0,378
-austin_python-2.2.0.dist-info/licenses/AUTHORS,sha256=pARRmpJQYH8iWy_0SonXsxOiV02AFitowle4AtURxUI,45
-austin_python-2.2.0.dist-info/licenses/COPYING,sha256=jOtLnuWt7d5Hsx6XXB2QxzrSe2sWWh3NgMfFRetluQM,35147
-austin_python-2.2.0.dist-info/licenses/LICENSE.md,sha256=edVDhmEiPf2W5lWzqLkrj6EYLSrCeYg9MJy28YsiMFU,34937
-austin_python-2.2.0.dist-info/RECORD,,
+austin_python-2.4.0.dist-info/METADATA,sha256=ELAptpJDAMuBwv7PBHYsyG4GAYKvT3vD4k436ECt9wo,7537
+austin_python-2.4.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+austin_python-2.4.0.dist-info/entry_points.txt,sha256=xmPVC_6UkSOYu7-YMYm3BwIcMFYEhRw0qlF05Cq5gf0,378
+austin_python-2.4.0.dist-info/licenses/AUTHORS,sha256=pARRmpJQYH8iWy_0SonXsxOiV02AFitowle4AtURxUI,45
+austin_python-2.4.0.dist-info/licenses/COPYING,sha256=jOtLnuWt7d5Hsx6XXB2QxzrSe2sWWh3NgMfFRetluQM,35147
+austin_python-2.4.0.dist-info/licenses/LICENSE.md,sha256=edVDhmEiPf2W5lWzqLkrj6EYLSrCeYg9MJy28YsiMFU,34937
+austin_python-2.4.0.dist-info/RECORD,,

{austin_python-2.2.0.dist-info → austin_python-2.4.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.29.0
+Generator: hatchling 1.32.4
 Root-Is-Purelib: true
 Tag: py3-none-any