From 3d9b1d4f18b6133c0bc98ebe3801bffeaea7e08c Mon Sep 17 00:00:00 2001
From: Taegyun Kim <taegyun.kim@datadoghq.com>
Date: Wed, 3 Jul 2024 11:58:03 -0400
Subject: [PATCH] fix(profiling): add debug asserts for our stack assumptions
 (#9689)

- [x] Change(s) are motivated and described in the PR description
- [x] Testing strategy is described if automated tests are not included
in the PR
- [x] Risks are described (performance impact, potential for breakage,
maintainability)
- [x] Change is maintainable (easy to change, telemetry, documentation)
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed or label `changelog/no-changelog` is set
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/))
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))
- [x] If this PR changes the public interface, I've notified
`@DataDog/apm-tees`.

- [ ] Title is accurate
- [ ] All changes are related to the pull request's stated goal
- [ ] Description motivates each change
- [ ] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- [ ] Testing strategy adequately addresses listed risks
- [ ] Change is maintainable (easy to change, telemetry, documentation)
- [ ] Release note makes sense to a user of the library
- [ ] Author has acknowledged and discussed the performance implications
of this PR as reported in the benchmarks PR comment
- [ ] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

(cherry picked from commit 076772500ef5f67deafcd4efef0df1f969380693)
---
 ddtrace/profiling/collector/_lock.py | 10 ++++++++--
 ddtrace/settings/profiling.py        |  8 ++++++++
 riotfile.py                          |  3 +++
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/ddtrace/profiling/collector/_lock.py b/ddtrace/profiling/collector/_lock.py
index a2c3f0497dc..1503baf0bd1 100644
--- a/ddtrace/profiling/collector/_lock.py
+++ b/ddtrace/profiling/collector/_lock.py
@@ -260,8 +260,14 @@ def _get_lock_call_loc_with_name(self) -> typing.Optional[str]:
             # 1: _acquire/_release
             # 2: acquire/release (or __enter__/__exit__)
             # 3: caller frame
-            # And we expect additional frame if WRAPT_C_EXT is False
-            frame = sys._getframe(3 if WRAPT_C_EXT else 4)
+            if config.enable_asserts:
+                frame = sys._getframe(1)
+                if frame.f_code.co_name not in {"_acquire", "_release"}:
+                    raise AssertionError("Unexpected frame %s" % frame.f_code.co_name)
+                frame = sys._getframe(2)
+                if frame.f_code.co_name not in {"acquire", "release", "__enter__", "__exit__"}:
+                    raise AssertionError("Unexpected frame %s" % frame.f_code.co_name)
+            frame = sys._getframe(3)
             code = frame.f_code
             call_loc = "%s:%d" % (os.path.basename(code.co_filename), frame.f_lineno)
 
diff --git a/ddtrace/settings/profiling.py b/ddtrace/settings/profiling.py
index 7b42873ae0e..fe8781fd3b1 100644
--- a/ddtrace/settings/profiling.py
+++ b/ddtrace/settings/profiling.py
@@ -188,6 +188,14 @@ class ProfilingConfig(En):
         help="The tags to apply to uploaded profile. Must be a list in the ``key1:value,key2:value2`` format",
     )
 
+    enable_asserts = En.v(
+        bool,
+        "enable_asserts",
+        default=False,
+        help_type="Boolean",
+        help="Whether to enable debug assertions in the profiler code",
+    )
+
     class Stack(En):
         __item__ = __prefix__ = "stack"
 
diff --git a/riotfile.py b/riotfile.py
index 1ccf4d7514d..0d992eb3c78 100644
--- a/riotfile.py
+++ b/riotfile.py
@@ -2677,6 +2677,9 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
             name="profile",
             # NB riot commands that use this Venv must include --pass-env to work properly
             command="python -m tests.profiling.run pytest -v --no-cov --capture=no --benchmark-disable {cmdargs} tests/profiling",  # noqa: E501
+            env={
+                "DD_PROFILING_ENABLE_ASSERTS": "1",
+            },
             pkgs={
                 "gunicorn": latest,
                 #