refactor(profiling): add type annotations to native-heap ownership partition

Annotate the Phase 2 ownership-partition additions: the
BUILD_NATIVE_HEAP_GOTTER_TEST_SUPPORT build knob (setup.py), the
_arm_native_heap/_start_service locals (profiler.py), and the module-level
constants and remaining local variables in the partition and end-to-end
handoff tests. set_native_heap_partition and the test_hook_hits activator
additions in this PR were already fully annotated.

Author: vlad-scherbich

ddtrace/profiling/profiler.py

@@ -398,7 +398,7 @@ def _arm_native_heap(self) -> bool:
                 # live-heap (ddheap:free + retain flagging) is a build-time
                 # property of the cdylib, not a runtime toggle; report which
                 # mode was actually armed for observability.
-                mode = "live-heap" if heap_gotter.live_heap_enabled() else "allocation-only"
+                mode: str = "live-heap" if heap_gotter.live_heap_enabled() else "allocation-only"
                 LOG.debug("Native heap profiling armed (GOT overrides installed, %s)", mode)
                 return True
             LOG.debug("Native heap profiling requested but GOT overrides were not installed")
@@ -408,7 +408,7 @@ def _arm_native_heap(self) -> bool:
 
     def _start_service(self) -> None:
         """Start the profiler."""
-        native_heap_armed = self._arm_native_heap()
+        native_heap_armed: bool = self._arm_native_heap()
 
         # Native-heap de-duplication (Phase 2 ownership partition): a process
         # must have exactly ONE producer per allocator domain, otherwise the same

setup.py

@@ -145,7 +145,7 @@
 # glibc-malloc tail the in-process sampler drops) without a live eBPF attach.
 # Never set for shipped wheels — it is strictly a test build knob and only has
 # any effect when BUILD_NATIVE_HEAP_GOTTER is also on.
-BUILD_NATIVE_HEAP_GOTTER_TEST_SUPPORT = os.getenv("DD_PROFILING_NATIVE_HEAP_TEST_SUPPORT", "0").lower() in (
+BUILD_NATIVE_HEAP_GOTTER_TEST_SUPPORT: bool = os.getenv("DD_PROFILING_NATIVE_HEAP_TEST_SUPPORT", "0").lower() in (
     "1",
     "yes",
     "on",

tests/profiling/collector/test_memalloc.py

@@ -1585,12 +1585,12 @@ def test_obj_and_mem_domain_coexist(tmp_path: Path) -> None:
 
 # Each buffer is far larger than pymalloc's 512B threshold, so it is served by
 # the raw allocator (glibc malloc).
-_PARTITION_LARGE_ALLOC_BYTES = 1024 * 1024
+_PARTITION_LARGE_ALLOC_BYTES: int = 1024 * 1024
 # Kept <= 256 so the ``range()`` loop variables stay in CPython's small-int
 # cache: ``_allocate_large_buffers`` then makes *only* > 512B allocations, which
 # makes the "large allocations are skipped" assertion deterministic.
-_PARTITION_LARGE_ALLOC_COUNT = 64
-_PARTITION_SMALL_ALLOC_COUNT = 200_000
+_PARTITION_LARGE_ALLOC_COUNT: int = 64
+_PARTITION_SMALL_ALLOC_COUNT: int = 200_000
 
 
 def _allocate_large_buffers(store: list[object]) -> None:
@@ -1605,10 +1605,10 @@ def _allocate_small_objects(store: list[object]) -> None:
 
 def test_native_heap_partition_skips_large_managed_allocations(tmp_path: Path) -> None:
     """With the partition on, > 512B OBJ allocations are skipped but <= 512B kept."""
-    output_filename = _setup_profiling_prelude(tmp_path, "test_native_heap_partition_skips_large")
+    output_filename: str = _setup_profiling_prelude(tmp_path, "test_native_heap_partition_skips_large")
 
     store: list[object] = []
-    mc = memalloc.MemoryCollector(heap_sample_size=64 * 1024)
+    mc: memalloc.MemoryCollector = memalloc.MemoryCollector(heap_sample_size=64 * 1024)
     memalloc.set_native_heap_partition(True)
     try:
         with mc:
@@ -1617,17 +1617,17 @@ def test_native_heap_partition_skips_large_managed_allocations(tmp_path: Path) -
             mc.snapshot()
         ddup.upload()
 
-        profile = pprof_utils.parse_newest_profile(output_filename)
-        heap_samples = pprof_utils.get_samples_with_value_type(profile, "heap-space")
+        profile: "pprof_pb2.Profile" = pprof_utils.parse_newest_profile(output_filename)
+        heap_samples: "list[pprof_pb2.Sample]" = pprof_utils.get_samples_with_value_type(profile, "heap-space")
 
         # Small pool-served allocations (<= 512B) are invisible to the gotter, so
         # the in-process sampler must keep sampling them.
-        small_count = _count_heap_samples_with_function(profile, heap_samples, "_allocate_small_objects")
+        small_count: int = _count_heap_samples_with_function(profile, heap_samples, "_allocate_small_objects")
         assert small_count > 0, "small (<=512B) managed allocations must still be sampled when the partition is on"
 
         # Large allocations (> 512B) hit glibc malloc and are owned by the gotter,
         # so the in-process sampler must skip them entirely.
-        large_count = _count_heap_samples_with_function(profile, heap_samples, "_allocate_large_buffers")
+        large_count: int = _count_heap_samples_with_function(profile, heap_samples, "_allocate_large_buffers")
         assert large_count == 0, (
             f"large (>512B) managed allocations must be skipped when the partition is on (got {large_count})"
         )
@@ -1640,10 +1640,10 @@ def test_native_heap_partition_skips_large_managed_allocations(tmp_path: Path) -
 
 def test_native_heap_partition_disabled_samples_all_sizes(tmp_path: Path) -> None:
     """With the partition off (default, fail-safe), all sizes are sampled."""
-    output_filename = _setup_profiling_prelude(tmp_path, "test_native_heap_partition_disabled")
+    output_filename: str = _setup_profiling_prelude(tmp_path, "test_native_heap_partition_disabled")
 
     store: list[object] = []
-    mc = memalloc.MemoryCollector(heap_sample_size=64 * 1024)
+    mc: memalloc.MemoryCollector = memalloc.MemoryCollector(heap_sample_size=64 * 1024)
     # Default is off; set explicitly in case a prior test left it on.
     memalloc.set_native_heap_partition(False)
     with mc:
@@ -1652,12 +1652,12 @@ def test_native_heap_partition_disabled_samples_all_sizes(tmp_path: Path) -> Non
         mc.snapshot()
     ddup.upload()
 
-    profile = pprof_utils.parse_newest_profile(output_filename)
-    heap_samples = pprof_utils.get_samples_with_value_type(profile, "heap-space")
+    profile: "pprof_pb2.Profile" = pprof_utils.parse_newest_profile(output_filename)
+    heap_samples: "list[pprof_pb2.Sample]" = pprof_utils.get_samples_with_value_type(profile, "heap-space")
 
-    large_count = _count_heap_samples_with_function(profile, heap_samples, "_allocate_large_buffers")
+    large_count: int = _count_heap_samples_with_function(profile, heap_samples, "_allocate_large_buffers")
     assert large_count > 0, "large (>512B) allocations must be sampled when the partition is off"
-    small_count = _count_heap_samples_with_function(profile, heap_samples, "_allocate_small_objects")
+    small_count: int = _count_heap_samples_with_function(profile, heap_samples, "_allocate_small_objects")
     assert small_count > 0, "small (<=512B) allocations must be sampled when the partition is off"
 
     del store

tests/profiling/test_native_heap_gotter.py

@@ -15,10 +15,16 @@
 """
 
 import sys
+from typing import TYPE_CHECKING
 
 import pytest
 
 
+if TYPE_CHECKING:
+    # We need the pyright: ignore because pprof_pb2 does not exist as a real module, only as a pyi.
+    from tests.profiling.collector import pprof_pb2  # pyright: ignore[reportMissingModuleSource]
+
+
 # Evaluated in the PARENT interpreter (subprocess bodies cannot express a skip:
 # an in-body ``pytest.skip`` would surface as a non-zero exit and FAIL the outer
 # test). ``test_hook_hits()`` is a read-only counter query with no side effects —
@@ -170,11 +176,11 @@ def test_profiler_keeps_managed_heap_when_native_heap_armed() -> None:
             with mock.patch.object(memalloc, "set_native_heap_partition") as set_partition:
                 from ddtrace.profiling.profiler import Profiler
 
-                prof = Profiler()
+                prof: Profiler = Profiler()
                 prof.start()
                 try:
                     assert install.called, "the gotter must still be armed when native heap is enabled"
-                    has_mem = any(isinstance(c, memalloc.MemoryCollector) for c in prof._profiler._collectors)
+                    has_mem: bool = any(isinstance(c, memalloc.MemoryCollector) for c in prof._profiler._collectors)
                     assert has_mem, (
                         "in-process managed-heap (OBJ/MEM) collector must stay active when the gotter is armed; "
                         "the gotter owns only the native/raw glibc malloc domain"
@@ -204,10 +210,10 @@ def test_profiler_keeps_managed_heap_when_gotter_not_installed() -> None:
         with mock.patch.object(memalloc, "set_native_heap_partition") as set_partition:
             from ddtrace.profiling.profiler import Profiler
 
-            prof = Profiler()
+            prof: Profiler = Profiler()
             prof.start()
             try:
-                has_mem = any(isinstance(c, memalloc.MemoryCollector) for c in prof._profiler._collectors)
+                has_mem: bool = any(isinstance(c, memalloc.MemoryCollector) for c in prof._profiler._collectors)
                 assert has_mem, "in-process memory collector must stay active when the gotter fails to install"
                 # Not armed -> partition off so ALL sizes keep being sampled.
                 set_partition.assert_called_once_with(False)
@@ -234,11 +240,11 @@ def test_profiler_keeps_managed_heap_when_native_heap_disabled() -> None:
         with mock.patch.object(memalloc, "set_native_heap_partition") as set_partition:
             from ddtrace.profiling.profiler import Profiler
 
-            prof = Profiler()
+            prof: Profiler = Profiler()
             prof.start()
             try:
                 assert not install.called
-                has_mem = any(isinstance(c, memalloc.MemoryCollector) for c in prof._profiler._collectors)
+                has_mem: bool = any(isinstance(c, memalloc.MemoryCollector) for c in prof._profiler._collectors)
                 assert has_mem, "in-process memory collector must run when native heap is disabled"
                 # Feature off -> partition off (all sizes sampled).
                 set_partition.assert_called_once_with(False)
@@ -313,8 +319,8 @@ def test_native_heap_ownership_handoff_end_to_end() -> None:
     assert heap_gotter.install() is True
     assert heap_gotter.is_installed() is True
 
-    prefix = os.path.join(tempfile.mkdtemp(), "handoff")
-    output_filename = prefix + "." + str(os.getpid())
+    prefix: str = os.path.join(tempfile.mkdtemp(), "handoff")
+    output_filename: str = prefix + "." + str(os.getpid())
     ddup.config(
         service="test_native_heap_ownership_handoff",
         version="test",
@@ -324,7 +330,7 @@ def test_native_heap_ownership_handoff_end_to_end() -> None:
     ddup.start()
 
     store: list[object] = []
-    mc = memalloc.MemoryCollector(heap_sample_size=64 * 1024)
+    mc: memalloc.MemoryCollector = memalloc.MemoryCollector(heap_sample_size=64 * 1024)
     memalloc.set_native_heap_partition(True)
     try:
         with mc:
@@ -333,19 +339,19 @@ def test_native_heap_ownership_handoff_end_to_end() -> None:
             # raw malloc (it is NOT sampling-gated), so background allocations
             # can only inflate the delta — never shrink it below the number of
             # large buffers we deliberately allocate.
-            hits_before = heap_gotter.test_hook_hits()
+            hits_before: "int | None" = heap_gotter.test_hook_hits()
             _allocate_large_buffers(store)
-            hits_after = heap_gotter.test_hook_hits()
+            hits_after: "int | None" = heap_gotter.test_hook_hits()
 
             _allocate_small_objects(store)
             mc.snapshot()
         ddup.upload()
 
-        profile = pprof_utils.parse_newest_profile(output_filename)
-        heap_samples = pprof_utils.get_samples_with_value_type(profile, "heap-space")
+        profile: "pprof_pb2.Profile" = pprof_utils.parse_newest_profile(output_filename)
+        heap_samples: "list[pprof_pb2.Sample]" = pprof_utils.get_samples_with_value_type(profile, "heap-space")
 
         # (a) In-process producer dropped the > 512B tail ...
-        large_count = _count_heap_samples_with_function(profile, heap_samples, "_allocate_large_buffers")
+        large_count: int = _count_heap_samples_with_function(profile, heap_samples, "_allocate_large_buffers")
         assert large_count == 0, (
             f"partition ON: > 512B managed allocations must NOT be sampled in-process (got {large_count})"
         )
@@ -354,15 +360,15 @@ def test_native_heap_ownership_handoff_end_to_end() -> None:
         # is a single raw malloc routed through the patched GOT, so the hook-hit
         # counter must advance by at least the number of large buffers.
         assert hits_before is not None and hits_after is not None
-        delta = hits_after - hits_before
+        delta: int = hits_after - hits_before
         assert delta >= _PARTITION_LARGE_ALLOC_COUNT, (
             "native gotter must capture the > 512B raw-malloc tail the in-process sampler dropped "
             f"(hook-hit delta {delta} < {_PARTITION_LARGE_ALLOC_COUNT} large allocations)"
         )
 
         # Control: <= 512B pool-served allocations are invisible to the gotter
         # and must still be sampled in-process — the partition splits by size.
-        small_count = _count_heap_samples_with_function(profile, heap_samples, "_allocate_small_objects")
+        small_count: int = _count_heap_samples_with_function(profile, heap_samples, "_allocate_small_objects")
         assert small_count > 0, "partition ON: <= 512B managed allocations must still be sampled in-process"
     finally:
         # Reset the process-global flag so it cannot bleed into other tests