1717
1818from amazon .opentelemetry .distro .serviceevents .ast_transformation import get_function_info_unlocked
1919
20+ # Hard cap on real call-path frames retained per request. The monitor records a frame for every
21+ # instrumented call (not just sampled ones), so under "auto"/"never" sampling a high-call-volume
22+ # request would otherwise grow this list without bound and could push a serialized IncidentSnapshot
23+ # past the 1 MB CloudWatch OTLP Logs limit. At ~150-250 B/frame, 1024 frames is ~256 KB — well under
24+ # the limit with room for the stack trace, yet far above any legitimate frame count. Java/JS match.
25+ _MAX_CALL_PATH_ENTRIES = 1024
26+ # Marker appended once when the cap is exceeded. duration_ns=0 also trips is_partial (computed from
27+ # timing in the incident collector).
28+ _CALL_PATH_TRUNCATION_SENTINEL = "<call_path_truncated>"
29+
2030# ============================================================================
2131# Sampling Configuration
2232# ============================================================================
2535_sample_tier2_threshold : int = 1000
2636_sample_tier2_rate : int = 10
2737_sample_tier3_rate : int = 100
28- _HOT_ENDPOINT_CYCLES : int = 100
2938
30- _sampling_mode : str = "always" # "auto", "always", "never", or "adaptive "
39+ _sampling_mode : str = "always" # "auto", "always", or "never "
3140_call_counters : Dict [str , int ] = {}
3241_call_counters_lock = threading .Lock ()
3342
43+ # Process-wide count of contexts with an active investigation. The monitor records a call_path
44+ # entry on EVERY instrumented function exit (not just sampled ones), which would otherwise pay a
45+ # ContextVar lookup on the hot path even when no incident investigation is in flight (the common
46+ # case). This counter lets record_call_path_entry skip that lookup when nothing is investigating —
47+ # mirroring the JS distro's _investigationActiveCount and the Java bridge's investigationActiveCount.
48+ # Reads on the hot path are lock-free (a plain int load is atomic under CPython); the infrequent
49+ # begin/clear writes take the lock to avoid lost updates. The counter is global (shared across
50+ # threads/contexts) while the investigation data itself is a per-context ContextVar, exactly like
51+ # JS's global counter + per-context AsyncLocalStorage: when ANY context is investigating, the hot
52+ # path does the ContextVar lookup and finds data only in the contexts that began one.
53+ _investigation_active_count : int = 0
54+ _investigation_active_lock = threading .Lock ()
55+
56+
57+ def _increment_investigation_active () -> None :
58+ """Increment the active-investigation count. Called only when begin_investigation actually
59+ creates new data (so a double-begin in one context doesn't leak the count high)."""
60+ global _investigation_active_count # pylint: disable=global-statement
61+ with _investigation_active_lock :
62+ _investigation_active_count += 1
63+
64+
65+ def _decrement_investigation_active () -> None :
66+ """Decrement the active-investigation count, clamped at zero. Clamping keeps an unbalanced clear
67+ (e.g. a clear with no prior begin) from driving the count negative — a spuriously-high count
68+ only costs an extra hot-path ContextVar lookup (degrades to pre-guard behavior), whereas a
69+ negative count would wrongly disable call_path capture for a genuinely active investigation."""
70+ global _investigation_active_count # pylint: disable=global-statement
71+ with _investigation_active_lock :
72+ if _investigation_active_count > 0 :
73+ _investigation_active_count -= 1
74+
3475
3576def set_sampling_mode (mode : str ) -> None :
36- """Set sampling mode: 'always', 'never', 'auto', or 'adaptive '."""
77+ """Set sampling mode: 'always', 'never', or 'auto '."""
3778 # Singleton module-level sampling state.
3879 global _sampling_mode # pylint: disable=global-statement
39- if mode not in "always" , "never" , "auto" , "adaptive" ):
80+ if mode not in "always" , "never" , "auto" ):
4081 raise ValueError (f"Invalid sampling mode: '{ mode } )
4182 _sampling_mode = mode
4283
@@ -51,17 +92,15 @@ def set_sampling_thresholds(
5192 tier2_threshold : int = 1000 ,
5293 tier2_rate : int = 10 ,
5394 tier3_rate : int = 100 ,
54- hot_endpoint_cycles : int = 100 ,
5595) -> None :
56- """Set sampling thresholds for auto/adaptive mode."""
96+ """Set sampling thresholds for auto mode."""
5797 # Singleton module-level sampling state.
5898 global _sample_tier1_threshold , _sample_tier2_threshold # pylint: disable=global-statement
59- global _sample_tier2_rate , _sample_tier3_rate , _HOT_ENDPOINT_CYCLES # pylint: disable=global-statement
99+ global _sample_tier2_rate , _sample_tier3_rate # pylint: disable=global-statement
60100 _sample_tier1_threshold = tier1_threshold
61101 _sample_tier2_threshold = tier2_threshold
62102 _sample_tier2_rate = tier2_rate
63103 _sample_tier3_rate = tier3_rate
64- _HOT_ENDPOINT_CYCLES = hot_endpoint_cycles
65104
66105
67106def _should_sample (total_calls : int ) -> bool : # pylint: disable=too-many-return-statements
@@ -70,16 +109,7 @@ def _should_sample(total_calls: int) -> bool: # pylint: disable=too-many-return
70109 return True
71110 if _sampling_mode == "never" :
72111 return False
73- if _sampling_mode == "adaptive" :
74- # O(1) cached check - computed once per request, ~50ns per subsequent call
75- cached = _adaptive_sample_cache .get ()
76- if cached is not None :
77- return cached
78- operation = _current_operation .get ()
79- result = operation is not None and operation in _hot_operations
80- _adaptive_sample_cache .set (result )
81- return result
82- # AUTO: adaptive 3-tier sampling. Rates are "1-in-N"; a non-positive N is degenerate
112+ # AUTO: 3-tier sampling. Rates are "1-in-N"; a non-positive N is degenerate
83113 # (it can only arrive via the internal test-config hook, which doesn't validate) and
84114 # would otherwise raise ZeroDivisionError on the modulo. Treat N <= 0 as "sample none
85115 # in this tier" so a misconfigured rate degrades gracefully instead of crashing the
@@ -117,10 +147,6 @@ def _increment_call_counter(function_name: str) -> int:
117147def set_current_operation (operation : str ):
118148 """Set the current operation (e.g., 'POST /users') for the request context."""
119149 _current_operation .set (operation )
120- # Reset adaptive sampling cache so it re-evaluates with the new operation.
121- # This prevents cache poisoning if _should_sample() was called before the
122- # operation was set (e.g., by monitored middleware running before before_request).
123- _adaptive_sample_cache .set (None )
124150
125151
126152def get_current_operation () -> Optional [str ]:
@@ -131,36 +157,6 @@ def get_current_operation() -> Optional[str]:
131157def clear_current_operation ():
132158 """Clear the current operation from the request context."""
133159 _current_operation .set (None )
134- _adaptive_sample_cache .set (None )
135-
136-
137- # ============================================================================
138- # Hot Operation Tracking (for "adaptive" sampling mode)
139- # ============================================================================
140-
141- _hot_operations : Dict [str , int ] = {} # operation -> remaining flush cycles
142- _hot_operations_lock = threading .Lock ()
143-
144- # Per-request cache for adaptive sampling decision (cleared per request)
145- _adaptive_sample_cache : ContextVar [Optional [bool ]] = ContextVar ("serviceevents_adaptive_cache" , default = None )
146-
147-
148- def mark_operation_hot (operation : str ) -> None :
149- """Mark operation as hot after incident. Resets countdown to full cycle count."""
150- with _hot_operations_lock :
151- _hot_operations [operation ] = _HOT_ENDPOINT_CYCLES
152-
153-
154- def tick_hot_operations () -> None :
155- """Decrement hot operation counters. Called once per collector flush cycle."""
156- with _hot_operations_lock :
157- expired = []
158- for op in _hot_operations :
159- _hot_operations [op ] -= 1
160- if _hot_operations [op ] <= 0 :
161- expired .append (op )
162- for op in expired :
163- del _hot_operations [op ]
164160
165161
166162def get_call_stack () -> list :
@@ -187,23 +183,24 @@ def reset_after_fork() -> None:
187183 re-wires.
188184 """
189185 # Singleton module-level sampling state.
190- global _sampling_mode # pylint: disable=global-statement
186+ global _sampling_mode , _investigation_active_count # pylint: disable=global-statement
191187
192188 # Reset sampling mode to default (always)
193189 _sampling_mode = "always"
194190
191+ # The child starts with no in-flight investigations; the singleton's investigation_data is
192+ # cleared below, so the active count must drop to zero to match (avoids a leaked count
193+ # forcing the post-fork hot path to keep doing ContextVar lookups for nothing).
194+ with _investigation_active_lock :
195+ _investigation_active_count = 0
196+
195197 # Clear call counters
196198 with _call_counters_lock :
197199 _call_counters .clear ()
198200
199- # Clear hot operations
200- with _hot_operations_lock :
201- _hot_operations .clear ()
202-
203201 # Clear thread-local state
204202 _call_stack .set ([])
205203 _current_operation .set (None )
206- _adaptive_sample_cache .set (None )
207204
208205 # Clear the existing singleton's mutable state in place. The OTel
209206 # instrument (histogram) and its base attrs are intentionally left intact —
@@ -346,12 +343,19 @@ def record_function_call_metrics(
346343
347344 def begin_investigation (self ):
348345 """Start capturing investigation data for current request."""
346+ # Increment the active count only when we actually create new data, so a re-entrant
347+ # begin in the same context (the analogue of Java's nested servlet dispatch) doesn't
348+ # leak the count high. Mirrors the Java bridge's create-only increment.
349+ if self ._investigation_data .get () is None :
350+ _increment_investigation_active ()
349351 self ._investigation_data .set ({"call_path" : [], "exception" : None , "start_time" : time .time ()})
350352
351353 def get_investigation_data (self ) -> Optional [dict ]:
352354 """Get and clear investigation data."""
353355 data = self ._investigation_data .get ()
354356 self ._investigation_data .set (None )
357+ if data is not None :
358+ _decrement_investigation_active ()
355359 return data
356360
357361 def peek_investigation_data (self ) -> Optional [dict ]:
@@ -369,6 +373,11 @@ def clear_investigation_data(self) -> None:
369373 framework calls this unconditionally in its request-finalize path. Idempotent: a
370374 no-op when already cleared (e.g. an incident was collected this request).
371375 """
376+ # Decrement only when data was actually present, so the unconditional finalize-path call
377+ # (and the incident path, which already consumed via get_investigation_data) each net the
378+ # count correctly — every begin pairs with exactly one decrement.
379+ if self ._investigation_data .get () is not None :
380+ _decrement_investigation_active ()
372381 self ._investigation_data .set (None )
373382
374383 def record_execution_flow (self , caller : Optional [str ], callee : str ):
@@ -391,15 +400,37 @@ def record_call_path_entry(self, function_name: str, caller: Optional[str], dura
391400 caller: Function name that called this one (None if entry point)
392401 duration_ns: Duration in nanoseconds
393402 """
403+ # Hot-path gate: skip the ContextVar lookup entirely when no investigation is in flight
404+ # anywhere (the common case). Mirrors the JS/Java active-count guards. A lock-free int read
405+ # is enough — a stale 0 only happens in the instant before begin_investigation's locked
406+ # increment publishes, which is the same context that hasn't recorded any frames yet.
407+ if _investigation_active_count == 0 :
408+ return
394409 inv_data = self ._investigation_data .get ()
395- if inv_data is not None :
396- inv_data ["call_path" ].append (
410+ if inv_data is None :
411+ return
412+ call_path = inv_data ["call_path" ]
413+ size = len (call_path )
414+ if size > _MAX_CALL_PATH_ENTRIES :
415+ # Already at [MAX real frames + sentinel] — drop everything after.
416+ return
417+ if size == _MAX_CALL_PATH_ENTRIES :
418+ # This frame overflows the cap: keep the first MAX, then a single truncation sentinel.
419+ call_path .append (
397420 {
398- "function_name" : function_name ,
399- "caller_function_name" : caller ,
400- "duration_ns" : duration_ns ,
421+ "function_name" : _CALL_PATH_TRUNCATION_SENTINEL ,
422+ "caller_function_name" : None ,
423+ "duration_ns" : 0 ,
401424 }
402425 )
426+ return
427+ call_path .append (
428+ {
429+ "function_name" : function_name ,
430+ "caller_function_name" : caller ,
431+ "duration_ns" : duration_ns ,
432+ }
433+ )
403434
404435
405436# ============================================================================
@@ -451,7 +482,7 @@ def __enter__(self):
451482 """
452483 try :
453484 # The per-function call counter only drives AUTO-mode tiered sampling. Only AUTO
454- # reads it, so for every other mode (including the default "adaptive ") skip the
485+ # reads it, so for every other mode (including the default "always ") skip the
455486 # lock acquisition and the unbounded counter-dict growth on this hot path.
456487 if _sampling_mode == "auto" :
457488 call_count = _increment_call_counter (self .function_name )
0 commit comments