Skip to content

Commit db9478a

Browse files
tylfinP403n1x87
tylfin
authored and
P403n1x87
committed
feat(debugger): introduce lazy trampoline for first call
1 parent 3926499 commit db9478a

File tree

2 files changed

+229
-66
lines changed

2 files changed

+229
-66
lines changed

‎ddtrace/debugging/_origin/span.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -223,8 +223,8 @@ def instrument_view(cls, f):
223223

224224
_f = t.cast(FunctionType, f)
225225
if not EntrySpanWrappingContext.is_wrapped(_f):
226-
log.debug("Patching entrypoint %r for code origin", f)
227-
EntrySpanWrappingContext(cls.__uploader__, _f).wrap()
226+
log.debug("Lazy wrapping entrypoint %r for code origin", f)
227+
EntrySpanWrappingContext(cls.__uploader__, _f).wrap_lazy()
228228

229229
@classmethod
230230
def enable(cls):

‎ddtrace/internal/wrapping/context.py‎

Lines changed: 227 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from inspect import iscoroutinefunction
44
from inspect import isgeneratorfunction
55
import sys
6+
from types import CodeType
67
from types import FrameType
78
from types import FunctionType
89
from types import TracebackType
@@ -24,6 +25,117 @@
2425

2526
T = t.TypeVar("T")
2627

28+
# ============================================================================
29+
# Lazy wrapping machinery: allows deferring expensive bytecode instrumentation
30+
# until the first time a function is actually called.
31+
# ============================================================================
32+
33+
_lazy_registry: t.Dict[CodeType, "LazyMeta"] = {}
34+
35+
36+
class LazyMeta:
37+
"""Metadata for a lazily-wrapped function."""
38+
39+
__slots__ = ("func", "builder", "lock", "initialized")
40+
41+
def __init__(self, func: FunctionType, builder: t.Callable[[FunctionType], CodeType]):
42+
self.func = func # the original function object
43+
self.builder = builder # callable: builder(func) -> CodeType
44+
self.initialized = False
45+
46+
47+
def __lazy_trampoline_entry(*args, **kwargs):
48+
"""
49+
Called from the trampoline code on first invocation.
50+
Uses the current frame's code object to find the right LazyMeta,
51+
builds the heavy bytecode, and swaps it in place.
52+
"""
53+
# Get the caller's frame (the trampoline frame)
54+
code = sys._getframe(1).f_code
55+
meta = _lazy_registry[code]
56+
57+
if not meta.initialized:
58+
if not meta.initialized:
59+
# Build heavy code from the original pre-wrapped function
60+
new_code = meta.builder(meta.func)
61+
# Swap the code on the SAME function object frameworks already hold
62+
meta.func.__code__ = new_code
63+
meta.initialized = True
64+
# Drop the registry entry to free memory
65+
_lazy_registry.pop(code, None)
66+
67+
# Now call the function again; this time it runs the heavy code
68+
return meta.func(*args, **kwargs)
69+
70+
71+
def _make_trampoline_code(template_code: CodeType) -> CodeType:
72+
"""
73+
Build a tiny code object that:
74+
- takes *args, **kwargs
75+
- calls the global __lazy_trampoline_entry(*args, **kwargs)
76+
- returns its value
77+
78+
Note: This only works for functions with 0 freevars.
79+
"""
80+
bc = Bytecode()
81+
bc.name = template_code.co_name
82+
bc.filename = template_code.co_filename
83+
bc.first_lineno = template_code.co_firstlineno
84+
85+
# Function signature: accepts *args, **kwargs
86+
bc.argcount = 0
87+
bc.posonlyargcount = 0
88+
bc.kwonlyargcount = 0
89+
bc.flags = bytecode.CompilerFlags.VARARGS | bytecode.CompilerFlags.VARKEYWORDS
90+
91+
# Declare locals for *args, **kwargs
92+
bc.argnames = ["args", "kwargs"]
93+
94+
# No freevars or cellvars
95+
bc.freevars = []
96+
bc.cellvars = []
97+
98+
# Call the global entry: __lazy_trampoline_entry(*args, **kwargs)
99+
if sys.version_info >= (3, 13):
100+
# Python 3.13+
101+
bc.extend(
102+
[
103+
bytecode.Instr("LOAD_GLOBAL", (True, "__lazy_trampoline_entry")), # (True = NULL + func)
104+
bytecode.Instr("LOAD_FAST", "args"),
105+
bytecode.Instr("BUILD_MAP", 0),
106+
bytecode.Instr("LOAD_FAST", "kwargs"),
107+
bytecode.Instr("DICT_MERGE", 1),
108+
bytecode.Instr("CALL_FUNCTION_EX", 1),
109+
bytecode.Instr("RETURN_VALUE"),
110+
]
111+
)
112+
elif sys.version_info >= (3, 11):
113+
# Python 3.11-3.12
114+
bc.extend(
115+
[
116+
bytecode.Instr("PUSH_NULL"),
117+
bytecode.Instr("LOAD_GLOBAL", (False, "__lazy_trampoline_entry")),
118+
bytecode.Instr("LOAD_FAST", "args"),
119+
bytecode.Instr("LOAD_FAST", "kwargs"),
120+
bytecode.Instr("CALL_FUNCTION_EX", 1),
121+
bytecode.Instr("RETURN_VALUE"),
122+
]
123+
)
124+
else:
125+
# Python 3.10 and earlier
126+
bc.extend(
127+
[
128+
bytecode.Instr("LOAD_GLOBAL", "__lazy_trampoline_entry"),
129+
bytecode.Instr("LOAD_FAST", "args"),
130+
bytecode.Instr("LOAD_FAST", "kwargs"),
131+
bytecode.Instr("CALL_FUNCTION_EX", 1),
132+
bytecode.Instr("RETURN_VALUE"),
133+
]
134+
)
135+
136+
return bc.to_code()
137+
138+
27139
# This module implements utilities for wrapping a function with a context
28140
# manager. The rough idea is to re-write the function's bytecode to look like
29141
# this:
@@ -405,6 +517,10 @@ def extract(cls, f: FunctionType) -> "WrappingContext":
405517
def wrap(self) -> None:
406518
t.cast(_UniversalWrappingContext, _UniversalWrappingContext.wrapped(self.__wrapped__)).register(self)
407519

520+
def wrap_lazy(self) -> None:
521+
"""Install lazy wrapping that defers bytecode instrumentation until first call."""
522+
t.cast(_UniversalWrappingContext, _UniversalWrappingContext.wrapped_lazy(self.__wrapped__)).register(self)
523+
408524
def unwrap(self) -> None:
409525
f = self.__wrapped__
410526

@@ -475,6 +591,16 @@ def register(self, context: WrappingContext) -> None:
475591
self._contexts.append(context)
476592
self._contexts.sort(key=lambda c: c.__priority__)
477593

594+
@classmethod
595+
def wrapped_lazy(cls, f: FunctionType) -> "_UniversalWrappingContext":
596+
"""Create a universal wrapping context with lazy bytecode instrumentation."""
597+
if cls.is_wrapped(f):
598+
context = cls.extract(f)
599+
else:
600+
context = cls(f)
601+
context.wrap_lazy()
602+
return context
603+
478604
def unregister(self, context: WrappingContext) -> None:
479605
try:
480606
self._contexts.remove(context)
@@ -541,86 +667,123 @@ def extract(cls, f: FunctionType) -> "_UniversalWrappingContext":
541667
raise ValueError("Function is not wrapped")
542668
return t.cast(_UniversalWrappingContext, t.cast(ContextWrappedFunction, f).__dd_context_wrapped__)
543669

544-
if sys.version_info >= (3, 11):
545-
546-
def wrap(self) -> None:
547-
f = self.__wrapped__
670+
def _build_wrapped_code(self, f: FunctionType) -> CodeType:
671+
"""
672+
Build the heavy instrumented bytecode for a function.
673+
This is extracted from wrap() to enable lazy wrapping.
674+
"""
675+
bc = Bytecode.from_code(f.__code__)
548676

549-
if self.is_wrapped(f):
550-
raise ValueError("Function already wrapped")
677+
# Prefix every return
678+
i = 0
679+
while i < len(bc):
680+
instr = bc[i]
681+
try:
682+
if instr.name == "RETURN_VALUE":
683+
return_code = CONTEXT_RETURN.bind({"context_return": self.__return__}, lineno=instr.lineno)
684+
elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+
685+
return_code = CONTEXT_RETURN_CONST.bind(
686+
{"context_return": self.__return__, "value": instr.arg}, lineno=instr.lineno
687+
)
688+
else:
689+
return_code = []
551690

552-
bc = Bytecode.from_code(f.__code__)
691+
bc[i:i] = return_code
692+
i += len(return_code)
693+
except AttributeError:
694+
# Not an instruction
695+
pass
696+
i += 1
553697

554-
# Prefix every return
698+
# Search for the RESUME instruction
699+
for i, instr in enumerate(bc, 1):
700+
try:
701+
if instr.name == "RESUME":
702+
break
703+
except AttributeError:
704+
# Not an instruction
705+
pass
706+
else:
555707
i = 0
556-
while i < len(bc):
557-
instr = bc[i]
558-
try:
559-
if instr.name == "RETURN_VALUE":
560-
return_code = CONTEXT_RETURN.bind({"context_return": self.__return__}, lineno=instr.lineno)
561-
elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+
562-
return_code = CONTEXT_RETURN_CONST.bind(
563-
{"context_return": self.__return__, "value": instr.arg}, lineno=instr.lineno
564-
)
565-
else:
566-
return_code = []
567708

568-
bc[i:i] = return_code
569-
i += len(return_code)
570-
except AttributeError:
571-
# Not an instruction
572-
pass
573-
i += 1
709+
bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=f.__code__.co_firstlineno)
574710

575-
# Search for the RESUME instruction
576-
for i, instr in enumerate(bc, 1):
577-
try:
578-
if instr.name == "RESUME":
711+
# Wrap every line outside a try block
712+
except_label = bytecode.Label()
713+
first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True)
714+
715+
i = 0
716+
while i < len(bc):
717+
instr = bc[i]
718+
if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None:
719+
bc.insert(i, bytecode.TryEnd(last_try_begin))
720+
last_try_begin = None
721+
i += 1
722+
elif isinstance(instr, bytecode.TryEnd):
723+
j = i + 1
724+
while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin):
725+
if isinstance(bc[j], bytecode.Instr):
726+
last_try_begin = bytecode.TryBegin(except_label, push_lasti=True)
727+
bc.insert(i + 1, last_try_begin)
579728
break
580-
except AttributeError:
581-
# Not an instruction
582-
pass
583-
else:
584-
i = 0
729+
j += 1
730+
i += 1
731+
i += 1
585732

586-
bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=f.__code__.co_firstlineno)
733+
bc.insert(0, first_try_begin)
587734

588-
# Wrap every line outside a try block
589-
except_label = bytecode.Label()
590-
first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True)
735+
bc.append(bytecode.TryEnd(last_try_begin))
736+
bc.append(except_label)
737+
bc.extend(CONTEXT_FOOT.bind({"context_exit": self._exit}))
591738

592-
i = 0
593-
while i < len(bc):
594-
instr = bc[i]
595-
if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None:
596-
bc.insert(i, bytecode.TryEnd(last_try_begin))
597-
last_try_begin = None
598-
i += 1
599-
elif isinstance(instr, bytecode.TryEnd):
600-
j = i + 1
601-
while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin):
602-
if isinstance(bc[j], bytecode.Instr):
603-
last_try_begin = bytecode.TryBegin(except_label, push_lasti=True)
604-
bc.insert(i + 1, last_try_begin)
605-
break
606-
j += 1
607-
i += 1
608-
i += 1
739+
# Link the function to its original code object so that we can retrieve
740+
# it later if required.
741+
link_function_to_code(f.__code__, f)
609742

610-
bc.insert(0, first_try_begin)
743+
return bc.to_code()
611744

612-
bc.append(bytecode.TryEnd(last_try_begin))
613-
bc.append(except_label)
614-
bc.extend(CONTEXT_FOOT.bind({"context_exit": self._exit}))
745+
def wrap_lazy(self) -> None:
746+
"""
747+
Install lazy wrapping: replace function's __code__ with a trampoline
748+
that defers the expensive bytecode instrumentation until first call.
749+
"""
750+
f = self.__wrapped__
751+
752+
if self.is_wrapped(f):
753+
raise ValueError("Function already wrapped")
754+
755+
# Lazy wrapping doesn't work with closures (functions with freevars)
756+
# because we can't replace __code__ with different freevar counts.
757+
# Fall back to eager wrapping in this case.
758+
if len(f.__code__.co_freevars) > 0:
759+
return self.wrap()
760+
761+
# Mark the function as wrapped immediately (before the trampoline)
762+
t.cast(ContextWrappedFunction, f).__dd_context_wrapped__ = self
763+
764+
# Create a trampoline code object with matching freevars
765+
tramp_code = _make_trampoline_code(f.__code__)
766+
767+
# Register the builder that will be called on first invocation
768+
_lazy_registry[tramp_code] = LazyMeta(f, self._build_wrapped_code)
769+
770+
# Swap in the trampoline (cheap operation)
771+
f.__code__ = tramp_code
772+
773+
if sys.version_info >= (3, 11):
774+
775+
def wrap(self) -> None:
776+
"""Eagerly wrap the function with full bytecode instrumentation."""
777+
f = self.__wrapped__
778+
779+
if self.is_wrapped(f):
780+
raise ValueError("Function already wrapped")
615781

616782
# Mark the function as wrapped by a wrapping context
617783
t.cast(ContextWrappedFunction, f).__dd_context_wrapped__ = self
618784

619-
# Replace the function code with the wrapped code. We also link
620-
# the function to its original code object so that we can retrieve
621-
# it later if required.
622-
link_function_to_code(f.__code__, f)
623-
f.__code__ = bc.to_code()
785+
# Build and install the heavy wrapped code immediately
786+
f.__code__ = self._build_wrapped_code(f)
624787

625788
def unwrap(self) -> None:
626789
f = self.__wrapped__

0 commit comments

Comments
 (0)