From 8891418eee0e96b896662a6bf51e06fae9cb9702 Mon Sep 17 00:00:00 2001
From: "Gabriele N. Tornetta" <phoenix1987@gmail.com>
Date: Sat, 28 Oct 2023 23:37:16 +0100
Subject: [PATCH] feat: add MOJO unwind method

We implement a method to the MojoFile class to unwind the stream
entirely and collect the data in the metadata and samples attributes for
later inspection.
---
 austin/format/mojo.py    | 57 ++++++++++++++++++++++++++++++++++++++--
 test/format/test_mojo.py | 17 ++++++++++++
 2 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/austin/format/mojo.py b/austin/format/mojo.py
index a9925a3..fd67058 100644
--- a/austin/format/mojo.py
+++ b/austin/format/mojo.py
@@ -1,8 +1,12 @@
 import typing as t
 from dataclasses import dataclass
+from dataclasses import field
 from enum import Enum
 from io import BufferedReader
 
+from austin.stats import ProcessId
+from austin.stats import ThreadName
+
 
 __version__ = "0.1.0"
 
@@ -226,6 +230,21 @@ def to_austin(self) -> str:
 
 
 UNKNOWN = MojoString(1, "<unknown>")
+InterpreterId = int
+
+
+@dataclass
+class MojoSample:
+    """Austin sample."""
+
+    pid: ProcessId
+    iid: InterpreterId
+    thread: ThreadName
+    idle: bool = False
+    gc: bool = False
+    metrics: t.List[MojoMetric] = field(default_factory=list)
+    frames: t.List[MojoFrame] = field(default_factory=list)
+    metadata: t.Dict[str, str] = field(default_factory=dict)
 
 
 class MojoFile:
@@ -255,9 +274,12 @@ def __init__(self, mojo: BufferedReader) -> None:
             raise ValueError("Not a MOJO file")
 
         self.mojo_version = self.read_int()
+        self.metadata: t.Dict[str, str] = {}
+        self.samples: t.List[MojoSample] = []
 
         self.header = bytes(self._last_bytes)
         self._last_bytes.clear()
+        self._last_sample: t.Optional[MojoSample] = None
 
     def ref(self, n: int) -> t.Tuple[int, int]:
         """Return a per-process reference key.
@@ -324,6 +346,12 @@ def parse_metadata(self) -> t.Generator[t.Union[MojoEvent, int], None, None]:
         metadata = MojoMetadata(self.read_string(), self.read_string())
         if metadata.key == "mode" and metadata.value == "full":
             self._full_mode = True
+
+        if self._last_sample is not None:
+            self._last_sample.metadata[metadata.key] = metadata.value
+        else:
+            self.metadata[metadata.key] = metadata.value
+
         yield metadata
 
     @handles(MojoEvents.STACK)
@@ -333,8 +361,12 @@ def parse_stack(self) -> t.Generator[t.Union[MojoEvent, int], None, None]:
 
         self._pid = pid = self.read_int()
         iid = self.read_int() if self.mojo_version >= 3 else -1
+        thread = self.read_string()
+
+        self._last_sample = sample = MojoSample(thread=thread, pid=pid, iid=iid)
+        self.samples.append(sample)
 
-        yield MojoStack(pid, iid, self.read_string())
+        yield MojoStack(pid, iid, thread)
 
     def _lookup_string(self) -> MojoString:
         n = self.read_int()
@@ -369,7 +401,12 @@ def parse_frame(self) -> t.Generator[MojoFrame, None, None]:
     @handles(MojoEvents.FRAME_REF)
     def parse_frame_ref(self) -> t.Generator[MojoFrameReference, None, None]:
         """Parse a frame reference."""
-        yield MojoFrameReference(self._frame_map[self.ref(self.read_int())])
+        frame = self._frame_map[self.ref(self.read_int())]
+
+        assert self._last_sample is not None, self._last_sample
+        self._last_sample.frames.append(frame)
+
+        yield MojoFrameReference(frame)
 
     @handles(MojoEvents.FRAME_KERNEL)
     def parse_kernel_frame(self) -> t.Generator[MojoKernelFrame, None, None]:
@@ -386,6 +423,10 @@ def _parse_metric(self, metric_type: MojoMetricType) -> MojoEvent:
             self._metrics.append(metric)
             return MojoEvent()
 
+        assert self._last_sample is not None, self._last_sample
+        self._last_sample.metrics.append(metric)
+        self._last_sample = None
+
         return metric
 
     @handles(MojoEvents.METRIC_TIME)
@@ -407,11 +448,18 @@ def parse_invalid_frame(self) -> t.Generator[MojoSpecialFrame, None, None]:
     def parse_idle(self) -> t.Generator[MojoIdle, None, None]:
         """Parse idle event."""
         self._metrics.append(1)
+
+        assert self._last_sample is not None, self._last_sample
+        self._last_sample.idle = True
+
         yield MojoIdle()
 
     @handles(MojoEvents.GC)
     def parse_gc(self) -> t.Generator[MojoSpecialFrame, None, None]:
         """Parse a GC event."""
+        assert self._last_sample is not None, self._last_sample
+        self._last_sample.gc = True
+
         yield MojoSpecialFrame("GC")
 
     @handles(MojoEvents.STRING)
@@ -461,6 +509,11 @@ def parse(self) -> t.Iterator[MojoEvent]:
                     return
                 yield e
 
+    def unwind(self) -> None:
+        """Read the MOJO file."""
+        for _ in self.parse():
+            pass
+
 
 def main() -> None:
     from argparse import ArgumentParser
diff --git a/test/format/test_mojo.py b/test/format/test_mojo.py
index f1ec883..154c2a2 100644
--- a/test/format/test_mojo.py
+++ b/test/format/test_mojo.py
@@ -167,3 +167,20 @@ def test_mojo_column_info():
 def test_mojo_stack():
     assert MojoStack(1, -1, "noiid").to_austin() == "P1;Tnoiid"
     assert MojoStack(1, 2, "iid").to_austin() == "P1;T2:iid"
+
+
+def test_mojo_data():
+    input = (DATA / "test").with_suffix(".mojo")
+
+    with input.open("rb") as stream:
+        m = MojoFile(stream)
+        m.unwind()
+
+        assert m.metadata == {
+            "austin": "3.4.0",
+            "duration": "1038089",
+            "interval": "100",
+            "mode": "wall",
+        }
+
+        assert len(m.samples) == 13227