refactor(profiling): add type annotations to native-heap arming observability

Annotate the arming-observability additions: the
_NATIVE_HEAP_SIZE_THRESHOLD_BYTES / _NATIVE_HEAP_PARTITION_ARMED_METRIC
module globals and the _emit_native_heap_partition_armed_gauge local
(profiler.py), plus the locals in the new dogstatsd-gauge/WARNING tests.

Author: vlad-scherbich

ddtrace/profiling/profiler.py

@@ -1,6 +1,7 @@
 # -*- encoding: utf-8 -*-
 import json
 import logging
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
 from typing import Mapping
@@ -32,17 +33,21 @@
 from ddtrace.profiling.collector import threading
 
 
+if TYPE_CHECKING:
+    from ddtrace.vendor.dogstatsd import DogStatsd
+
+
 LOG = logging.getLogger(__name__)
 
 # pymalloc's small-request threshold (see CPython Objects/obmalloc.c). Requests
 # larger than this are delegated to glibc malloc, where the native-heap gotter
 # owns them, so the in-process sampler skips them when the partition is armed.
-_NATIVE_HEAP_SIZE_THRESHOLD_BYTES = 512
+_NATIVE_HEAP_SIZE_THRESHOLD_BYTES: int = 512
 # Dogstatsd gauge surfacing the native-heap ownership-partition arming decision
 # (1=armed, 0=not). Structured JSON loggers in real deploys drop ddtrace stdlib
 # records and pods/exec is RBAC-blocked, so the log line alone is not always
 # observable; this gauge gives an out-of-band, queryable signal at service start.
-_NATIVE_HEAP_PARTITION_ARMED_METRIC = "profiling.native_heap.partition_armed"
+_NATIVE_HEAP_PARTITION_ARMED_METRIC: str = "profiling.native_heap.partition_armed"
 
 
 class Profiler(object):
@@ -428,7 +433,7 @@ def _emit_native_heap_partition_armed_gauge(self, armed: bool) -> None:
             from ddtrace.internal.dogstatsd import get_dogstatsd_client
             from ddtrace.internal.settings._agent import config as agent_config
 
-            client = get_dogstatsd_client(agent_config.dogstatsd_url)
+            client: DogStatsd = get_dogstatsd_client(agent_config.dogstatsd_url)
             client.gauge(
                 _NATIVE_HEAP_PARTITION_ARMED_METRIC,
                 1 if armed else 0,

tests/profiling/test_native_heap_gotter.py

@@ -49,7 +49,7 @@ def _only_arming_warning(err: str) -> bool:
     harness would otherwise reject as unexpected stderr. Runs in the parent test
     process against the subprocess's decoded stderr.
     """
-    lines = [line for line in err.splitlines() if line.strip()]
+    lines: list[str] = [line for line in err.splitlines() if line.strip()]
     return all("native heap ownership partition:" in line for line in lines)
 
 
@@ -248,24 +248,24 @@ def test_profiler_start_emits_partition_armed_gauge_when_armed() -> None:
 
     profiling_config.native_heap.enabled = True  # pyright: ignore[reportAttributeAccessIssue]
 
-    client = mock.Mock()
+    client: mock.Mock = mock.Mock()
     with mock.patch.object(heap_gotter, "install", return_value=True):
         with mock.patch.object(heap_gotter, "live_heap_enabled", return_value=False):
             with mock.patch.object(ddtrace.internal.dogstatsd, "get_dogstatsd_client", return_value=client):
                 with mock.patch.object(profiler_mod.LOG, "warning") as warning:
-                    prof = profiler_mod.Profiler()
+                    prof: profiler_mod.Profiler = profiler_mod.Profiler()
                     prof.start()
                     try:
                         assert client.gauge.call_count == 1
                         args, kwargs = client.gauge.call_args
                         assert args[0] == "profiling.native_heap.partition_armed"
                         assert args[1] == 1, "gauge value must be 1 when armed"
-                        tags = kwargs["tags"]
+                        tags: list[str] = kwargs["tags"]
                         assert "domains:OBJ_MEM" in tags
                         assert "size_threshold_bytes:512" in tags
 
                         assert warning.called, "arming decision must be logged at WARNING"
-                        msg = warning.call_args[0][0]
+                        msg: str = warning.call_args[0][0]
                         assert "native heap ownership partition" in msg
                         assert warning.call_args[0][1] is True, "WARNING must report armed=True"
                     finally:
@@ -286,11 +286,11 @@ def test_profiler_start_emits_partition_armed_gauge_zero_when_not_armed() -> Non
 
     profiling_config.native_heap.enabled = True  # pyright: ignore[reportAttributeAccessIssue]
 
-    client = mock.Mock()
+    client: mock.Mock = mock.Mock()
     with mock.patch.object(heap_gotter, "install", return_value=False):
         with mock.patch.object(ddtrace.internal.dogstatsd, "get_dogstatsd_client", return_value=client):
             with mock.patch.object(profiler_mod.LOG, "warning") as warning:
-                prof = profiler_mod.Profiler()
+                prof: profiler_mod.Profiler = profiler_mod.Profiler()
                 prof.start()
                 try:
                     assert client.gauge.call_count == 1
@@ -322,7 +322,7 @@ def test_profiler_start_survives_partition_armed_gauge_error() -> None:
             with mock.patch.object(
                 ddtrace.internal.dogstatsd, "get_dogstatsd_client", side_effect=RuntimeError("no agent")
             ):
-                prof = profiler_mod.Profiler()
+                prof: profiler_mod.Profiler = profiler_mod.Profiler()
                 prof.start()  # must not raise
                 try:
                     assert prof.status.value == "running"