refactor(profiling): type pprof exporter and fix type comparison (#2987) (#3062)

There was an ongoing issue wih the pprof encode mixing type between strings,
ints and None when dealing with labels. This has been fixed in multiple places
and in multiple occasion, but there were still holes.

When encoding Events to pprof, they are grouped by event fields. Those fields
might be str, int or something else. In order to group those events, all the
fields used for grouping must be:

1. hashable (because we index in a dict)
2. sortable (because we sort event to have reproducible)
3. string (because we store those as labels in pprof)

The problem is that most value type used in events are hashable (None, int or
str) but sortability often fails when one of the field is None: you can't
compare `int` with `None`, making sorted() call fails like in #2962.

As we need string in the end, the fix is to convert everything to strings as
soon as possible when grouping events: that makes sure conditions 1, 2 and 3
are validated ASAP and everything works.

To do this, the code has been updated and typed to catch any future error.

Fixes #2962

Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
(cherry picked from commit 6520212)

Co-authored-by: Julien Danjou <[email protected]>

Author: mergify[bot]

ddtrace/profiling/_build.pyx

@@ -1,4 +1,4 @@
 import typing
 
 
-compiled_with: typing.Tuple[int, int, int] = (PY_MAJOR_VERSION, PY_MINOR_VERSION, PY_MICRO_VERSION)
+compiled_with: typing.Tuple[int, int, int] = (PY_MAJOR_VERSION, PY_MINOR_VERSION, PY_MICRO_VERSION)  # type: ignore[name-defined]

ddtrace/profiling/collector/_traceback.pyi

@@ -1,10 +1,9 @@
 import types
 import typing
 
-# (filename, line number, function name)
-FrameType = typing.Tuple[str, int, str]
+from .. import event
 
 def traceback_to_frames(
     traceback: types.TracebackType, max_nframes: int
-) -> typing.Tuple[typing.List[FrameType], int]: ...
-def pyframe_to_frames(frame: types.FrameType, max_nframes: int) -> typing.Tuple[typing.List[FrameType], int]: ...
+) -> typing.Tuple[typing.List[event.FrameType], int]: ...
+def pyframe_to_frames(frame: types.FrameType, max_nframes: int) -> typing.Tuple[typing.List[event.FrameType], int]: ...

ddtrace/profiling/collector/memalloc.py

@@ -27,24 +27,24 @@
 class MemoryAllocSampleEvent(event.StackBasedEvent):
     """A sample storing memory allocation tracked."""
 
-    size = attr.ib(default=None)
+    size = attr.ib(default=0, type=int)
     """Allocation size in bytes."""
 
-    capture_pct = attr.ib(default=None)
+    capture_pct = attr.ib(default=None, type=float)
     """The capture percentage."""
 
-    nevents = attr.ib(default=None)
+    nevents = attr.ib(default=0, type=int)
     """The total number of allocation events sampled."""
 
 
 @event.event_class
 class MemoryHeapSampleEvent(event.StackBasedEvent):
     """A sample storing memory allocation tracked."""
 
-    size = attr.ib(default=None)
+    size = attr.ib(default=0, type=int)
     """Allocation size in bytes."""
 
-    sample_size = attr.ib(default=None)
+    sample_size = attr.ib(default=0, type=int)
     """The sampling size."""
 
 

ddtrace/profiling/collector/stack.pyx

@@ -20,6 +20,7 @@ from ddtrace.profiling import event
 from ddtrace.profiling.collector import _task
 from ddtrace.profiling.collector import _threading
 from ddtrace.profiling.collector import _traceback
+from ddtrace.profiling.collector import stack_event
 
 
 # NOTE: Do not use LOG here. This code runs under a real OS thread and is unable to acquire any lock of the `logging`
@@ -138,22 +139,6 @@ ELSE:
             }
 
 
-@event.event_class
-class StackSampleEvent(event.StackBasedEvent):
-    """A sample storing executions frames for a thread."""
-
-    # Wall clock
-    wall_time_ns = attr.ib(default=0)
-    # CPU time in nanoseconds
-    cpu_time_ns = attr.ib(default=0)
-
-
-@event.event_class
-class StackExceptionSampleEvent(event.StackBasedEvent):
-    """A a sample storing raised exceptions and their stack frames."""
-
-    exc_type = attr.ib(default=None)
-
 from cpython.object cimport PyObject
 
 
@@ -319,7 +304,7 @@ cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_tim
 
         frames, nframes = _traceback.pyframe_to_frames(frame, max_nframes)
 
-        event = StackSampleEvent(
+        event = stack_event.StackSampleEvent(
             thread_id=thread_id,
             thread_native_id=thread_native_id,
             thread_name=thread_name,
@@ -338,7 +323,7 @@ cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_tim
         if exception is not None:
             exc_type, exc_traceback = exception
             frames, nframes = _traceback.traceback_to_frames(exc_traceback, max_nframes)
-            exc_event = StackExceptionSampleEvent(
+            exc_event = stack_event.StackExceptionSampleEvent(
                 thread_id=thread_id,
                 thread_name=thread_name,
                 thread_native_id=thread_native_id,

ddtrace/profiling/collector/stack_event.py

@@ -0,0 +1,22 @@
+import typing
+
+import attr
+
+from ddtrace.profiling import event
+
+
+@event.event_class
+class StackSampleEvent(event.StackBasedEvent):
+    """A sample storing executions frames for a thread."""
+
+    # Wall clock
+    wall_time_ns = attr.ib(default=0, type=int)
+    # CPU time in nanoseconds
+    cpu_time_ns = attr.ib(default=0, type=int)
+
+
+@event.event_class
+class StackExceptionSampleEvent(event.StackBasedEvent):
+    """A a sample storing raised exceptions and their stack frames."""
+
+    exc_type = attr.ib(default=None, type=typing.Optional[str])

ddtrace/profiling/collector/threading.py

@@ -23,22 +23,22 @@
 class LockEventBase(event.StackBasedEvent):
     """Base Lock event."""
 
-    lock_name = attr.ib(default=None, type=typing.Optional[str])
-    sampling_pct = attr.ib(default=None, type=typing.Optional[int])
+    lock_name = attr.ib(default="<unknown lock name>", type=str)
+    sampling_pct = attr.ib(default=0, type=int)
 
 
 @event.event_class
 class LockAcquireEvent(LockEventBase):
     """A lock has been acquired."""
 
-    wait_time_ns = attr.ib(default=None, type=typing.Optional[int])
+    wait_time_ns = attr.ib(default=0, type=int)
 
 
 @event.event_class
 class LockReleaseEvent(LockEventBase):
     """A lock has been released."""
 
-    locked_for_ns = attr.ib(default=None, type=typing.Optional[int])
+    locked_for_ns = attr.ib(default=0, type=int)
 
 
 def _current_thread():

ddtrace/profiling/event.py

@@ -8,9 +8,13 @@
 
 _T = typing.TypeVar("_T")
 
+# (filename, line number, function name)
+FrameType = typing.Tuple[str, int, typing.Optional[str]]
+StackTraceType = typing.List[FrameType]
+
 
 def event_class(
-        klass  # type: typing.Type[_T]
+    klass,  # type: typing.Type[_T]
 ):
     # type: (...) -> typing.Type[_T]
     return attr.s(slots=True)(klass)
@@ -45,23 +49,24 @@ class SampleEvent(Event):
 
 @event_class
 class StackBasedEvent(SampleEvent):
-    thread_id = attr.ib(default=None)
-    thread_name = attr.ib(default=None)
-    thread_native_id = attr.ib(default=None)
-    task_id = attr.ib(default=None)
-    task_name = attr.ib(default=None)
-    frames = attr.ib(default=None)
-    nframes = attr.ib(default=None)
+    thread_id = attr.ib(default=None, type=typing.Optional[int])
+    thread_name = attr.ib(default=None, type=typing.Optional[str])
+    thread_native_id = attr.ib(default=None, type=typing.Optional[int])
+    task_id = attr.ib(default=None, type=typing.Optional[int])
+    task_name = attr.ib(default=None, type=typing.Optional[str])
+    frames = attr.ib(default=None, type=StackTraceType)
+    nframes = attr.ib(default=0, type=int)
     local_root_span_id = attr.ib(default=None, type=typing.Optional[int])
     span_id = attr.ib(default=None, type=typing.Optional[int])
     trace_type = attr.ib(default=None, type=typing.Optional[str])
-    trace_resource_container = attr.ib(default=None)
+    trace_resource_container = attr.ib(default=None, type=typing.List[str])
 
     def set_trace_info(
         self,
         span,  # type: typing.Optional[ddspan.Span]
         endpoint_collection_enabled,  # type: bool
     ):
+        # type: (...) -> None
         if span:
             self.span_id = span.span_id
             if span._local_root is not None:

ddtrace/profiling/exporter/__init__.py

@@ -1,3 +1,5 @@
+import typing
+
 import attr
 
 
@@ -10,7 +12,7 @@ class Exporter(object):
     """Exporter base class."""
 
     def export(self, events, start_time_ns, end_time_ns):
-        # type: (...) -> None
+        # type: (...) -> typing.Any
         """Export events.
 
         :param events: List of events to export.

ddtrace/profiling/exporter/file.py

@@ -5,16 +5,23 @@
 
 from ddtrace.profiling.exporter import pprof
 
+from .. import recorder
+
 
 @attr.s
 class PprofFileExporter(pprof.PprofExporter):
     """PProf file exporter."""
 
-    prefix = attr.ib()
-    _increment = attr.ib(default=1, init=False, repr=False)
+    prefix = attr.ib(type=str)
+    _increment = attr.ib(default=1, init=False, repr=False, type=int)
 
-    def export(self, events, start_time_ns, end_time_ns):
-        # type: (...) -> None
+    def export(
+        self,
+        events,  # type: recorder.EventsType
+        start_time_ns,  # type: int
+        end_time_ns,  # type: int
+    ):
+        # type: (...) -> pprof.pprof_ProfileType
         """Export events to pprof file.
 
         The file name is based on the prefix passed to init. The process ID number and type of export is then added as a
@@ -26,5 +33,6 @@ def export(self, events, start_time_ns, end_time_ns):
         """
         profile = super(PprofFileExporter, self).export(events, start_time_ns, end_time_ns)
         with gzip.open(self.prefix + (".%d.%d" % (os.getpid(), self._increment)), "wb") as f:
-            f.write(profile.SerializeToString())
+            f.write(profile.SerializeToString())  # type: ignore[attr-defined]
         self._increment += 1
+        return profile