diff --git a/netra/__init__.py b/netra/__init__.py index d36b346e..88d25f8a 100644 --- a/netra/__init__.py +++ b/netra/__init__.py @@ -107,9 +107,18 @@ def init( root_instruments: Set of instruments allowed to produce root-level spans. Independent of ``instruments``. Defaults to ``DEFAULT_INSTRUMENTS_FOR_ROOT`` when ``None``. When a root - span is blocked, its entire subtree is discarded. Pass a set - containing ``NetraInstruments.ALL`` to allow all - instrumentations to produce root spans (legacy behaviour). + span comes from an instrumentation outside this set, only that + span is dropped — its children are reparented onto its parent + (a true root's children become the new roots). The peel then + repeats: if a promoted child is itself from a non-root + instrumentation it is dropped too, recursively, until a + surviving span is reached. This keeps LLM/vector spans that + originate under an unwanted server root without emitting the server span itself. + Note: when ``enable_root_span=True``, Netra attaches its own root span and + every auto-instrumentation span becomes its child, so no + reparenting occurs. Pass a set containing + ``NetraInstruments.ALL`` to allow all instrumentations to + produce root spans (legacy behaviour). Returns: None diff --git a/netra/exporters/filtering_span_exporter.py b/netra/exporters/filtering_span_exporter.py index cea41e1c..f730d67e 100644 --- a/netra/exporters/filtering_span_exporter.py +++ b/netra/exporters/filtering_span_exporter.py @@ -1,30 +1,40 @@ import logging -from typing import Any, Dict, List, Sequence +from typing import Any, Dict, List, Sequence, Set from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.sdk.trace.export import ( - SpanExporter, - SpanExportResult, -) - -from netra.exporters.utils import add_blocked_trace_id, get_trace_id, is_trace_id_blocked, is_trial_blocked -from netra.processors.local_filtering_span_processor import ( - BLOCKED_LOCAL_PARENT_MAP, +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + +from netra.exporters.utils import ( + PatternMatcher, + add_blocked_trace_id, + find_root_spans_blocked, + get_span_id, + get_trace_id, + has_local_block_flag, + is_trace_id_blocked, + is_trial_blocked, + normalize_parent, + read_local_block_patterns, + reparent_spans, ) +from netra.processors.local_filtering_span_processor import BLOCKED_LOCAL_PARENT_MAP logger = logging.getLogger(__name__) class FilteringSpanExporter(SpanExporter): # type: ignore[misc] - """ - SpanExporter wrapper that filters out spans by name. - - Matching rules: - - Exact match: pattern "Foo" blocks span.name == "Foo". - - Prefix match: pattern ending with '*' (e.g., "CloudSpanner.*") blocks spans whose - names start with the prefix before '*', e.g., "CloudSpanner.", "CloudSpanner.Query". - - Suffix match: pattern starting with '*' (e.g., "*.Query") blocks spans whose - names end with the suffix after '*', e.g., "DB.Query", "Search.Query". + """SpanExporter wrapper that drops spans by name and by root-instrument policy. + + A span is dropped when any of the following holds: + - its trace ID was blocked while a trial/quota block was active; + - its name matches a globally configured block pattern (see ``PatternMatcher``); + - its name matches a per-span local block pattern set by + ``LocalFilteringSpanProcessor``; + - it is a root-connected span from an instrumentation not allowed to emit + root spans (resolved by ``find_root_spans_blocked``). + + Children of a dropped span are reparented onto the dropped span's parent so + subtrees are never silently discarded. """ def __init__(self, exporter: SpanExporter, patterns: Sequence[str]) -> None: @@ -32,272 +42,176 @@ def __init__(self, exporter: SpanExporter, patterns: Sequence[str]) -> None: Initialize the filtering span exporter. Args: - exporter: The span exporter to use. - patterns: List of patterns to block. + exporter: The underlying span exporter to forward surviving spans to. + patterns: Global name patterns to block. """ self._exporter = exporter - # Normalize once for efficient checks - exact: List[str] = [] - prefixes: List[str] = [] - suffixes: List[str] = [] - for p in patterns: - if not p: - continue - if p.endswith("*") and not p.startswith("*"): - prefixes.append(p[:-1]) - elif p.startswith("*") and not p.endswith("*"): - suffixes.append(p[1:]) - else: - exact.append(p) - self._exact = set(exact) - self._prefixes = prefixes - self._suffixes = suffixes + self._matcher = PatternMatcher(patterns) def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: - """ - Export spans to the exporter. + """Filter *spans*, reparent survivors, and forward them to the wrapped exporter. Args: - spans: List of spans to export. + spans: The batch of spans to export. Returns: - SpanExportResult.SUCCESS if the export was successful. + The wrapped exporter's ``SpanExportResult``, or ``SUCCESS`` when the + batch is fully filtered (nothing left to forward). """ if is_trial_blocked(): logger.debug("Trial/quota exhausted: blocking spans from export") - # Track trace IDs from spans being blocked during blocking period - for span in spans: - trace_id = get_trace_id(span) - if trace_id: - add_blocked_trace_id(trace_id) + self._record_blocked_trace_ids(spans) return SpanExportResult.SUCCESS - filtered: List[ReadableSpan] = [] - blocked_parent_map: Dict[Any, Any] = {} - for span in spans: - trace_id = get_trace_id(span) + # Find unconditional name/local drops first, then resolve which + # root-block candidates are root-connected. Feeding the name/local drops + # into that walk is what stops a disallowed candidate from being promoted + # to a root when its only surviving ancestor is itself name-blocked. + parents_of_spans_blocked_by_name = self._find_spans_blocked_by_name(spans) + root_spans_blocked, parents_of_root_spans_blocked = find_root_spans_blocked( + spans, parents_of_spans_blocked_by_name + ) - # Check if this span belongs to a trace ID that was blocked - if trace_id and is_trace_id_blocked(trace_id): - continue + all_blocked_span_ids: Set[Any] = set(parents_of_spans_blocked_by_name) + all_blocked_span_ids.update(root_spans_blocked) - name = getattr(span, "name", None) - if name is None: - filtered.append(span) - continue + surviving_spans = self._collect_survivors(spans, all_blocked_span_ids) - # Global blocking (configured patterns) - globally_blocked = self._is_blocked(name) - - # Local per-span blocking via attribute set by LocalFilteringSpanProcessor - locally_blocked = False - try: - local_patterns = self._get_local_patterns(span) - if local_patterns: - locally_blocked = self._matches_any_pattern(name, local_patterns) - # Fallback: if processor explicitly marked the span as locally blocked - if not locally_blocked and self._has_local_block_flag(span): - locally_blocked = True - except Exception: - locally_blocked = False - - if not (globally_blocked or locally_blocked): - filtered.append(span) - continue + reparent_map = self._build_reparent_map(parents_of_root_spans_blocked, parents_of_spans_blocked_by_name) + if reparent_map: + reparent_spans(surviving_spans, reparent_map) + if not surviving_spans: + return SpanExportResult.SUCCESS + return self._exporter.export(surviving_spans) - # Collect mapping for reparenting children of the blocked span - span_context = getattr(span, "context", None) - span_id = getattr(span_context, "span_id", None) if span_context else None - if span_id is not None: - blocked_parent_map[span_id] = getattr(span, "parent", None) + def _record_blocked_trace_ids(self, spans: Sequence[ReadableSpan]) -> None: + """Remember the trace IDs seen during a block so their later spans are dropped too. - # Merge with registry of locally blocked spans captured by processor to handle - # cases where children export before their blocked parent (SimpleSpanProcessor) - merged_map: Dict[Any, Any] = {} - try: - if BLOCKED_LOCAL_PARENT_MAP: - merged_map.update(BLOCKED_LOCAL_PARENT_MAP) - except Exception: - pass - merged_map.update(blocked_parent_map) + Args: + spans: The batch of spans being dropped during the block. + """ + for span in spans: + trace_id = get_trace_id(span) + if trace_id: + add_blocked_trace_id(trace_id) - if merged_map: - self._reparent_blocked_children(filtered, merged_map) - if not filtered: - return SpanExportResult.SUCCESS - return self._exporter.export(filtered) + def _find_spans_blocked_by_name(self, spans: Sequence[ReadableSpan]) -> Dict[Any, Any]: + """Find the spans dropped unconditionally by a global or local name rule. - def _is_blocked(self, name: str) -> bool: - """ - Check if a span name is blocked. + Trace-blocked spans are skipped entirely (they are dropped without + reparenting). The result feeds both the root-connected candidate walk — + as transparent dropped ancestors — and the reparent map. Args: - name: The span name to check. + spans: The batch of spans to classify. Returns: - True if the span name is blocked, False otherwise. + A ``{dropped_span_id -> normalized_parent}`` map for name/locally + blocked spans in this batch. """ - if name in self._exact: - return True - for pref in self._prefixes: - if name.startswith(pref): - return True - for suf in self._suffixes: - if name.endswith(suf): - return True - return False + parent_map: Dict[Any, Any] = {} + for span in spans: + trace_id = get_trace_id(span) + if trace_id and is_trace_id_blocked(trace_id): + continue - def _get_local_patterns(self, span: ReadableSpan) -> List[str]: - """ - Fetch local-block patterns from span attributes set by LocalFilteringSpanProcessor. + name = getattr(span, "name", None) + if name is None: + continue - Args: - span: The span to fetch local-block patterns from. + if self._matcher.matches(name) or self._is_locally_blocked(span, name): + span_id = get_span_id(span) + if span_id is not None: + parent_map[span_id] = normalize_parent(getattr(span, "parent", None)) - Returns: - List of local-block patterns. - """ - try: - attrs = getattr(span, "attributes", None) - if not attrs: - return [] - value = None - # Prefer Mapping.get if available - try: - if hasattr(attrs, "get"): - value = attrs.get("netra.local_blocked_spans") - else: - value = attrs["netra.local_blocked_spans"] - except Exception: - value = None - if isinstance(value, (list, tuple)) and all(isinstance(v, str) for v in value): - return [v for v in value if v] - except Exception: - logger.debug("Failed reading local blocked patterns from span", exc_info=True) - return [] + return parent_map - def _matches_any_pattern(self, name: str, patterns: Sequence[str]) -> bool: - """ - Check if a span name matches any of the given patterns. + def _collect_survivors(self, spans: Sequence[ReadableSpan], all_blocked_span_ids: Set[Any]) -> List[ReadableSpan]: + """Return the spans that survive filtering. Args: - name: The span name to check. - patterns: List of patterns to check against. + spans: The batch of spans to filter. + all_blocked_span_ids: Span IDs dropped by name/local rules or by the + root-instrument policy. Returns: - True if the span name matches any of the given patterns, False otherwise. + The spans to forward to the wrapped exporter (before reparenting). """ - for p in patterns: - if not p: + surviving_spans: List[ReadableSpan] = [] + for span in spans: + trace_id = get_trace_id(span) + if trace_id and is_trace_id_blocked(trace_id): continue - if p.endswith("*") and not p.startswith("*"): - if name.startswith(p[:-1]): - return True - elif p.startswith("*") and not p.endswith("*"): - if name.endswith(p[1:]): - return True - else: - if name == p: - return True - return False - - def _has_local_block_flag(self, span: ReadableSpan) -> bool: - """ - Check if a span has a local-block flag. + + span_id = get_span_id(span) + if span_id is not None and span_id in all_blocked_span_ids: + continue + + surviving_spans.append(span) + + return surviving_spans + + def _is_locally_blocked(self, span: ReadableSpan, name: str) -> bool: + """Check whether *span* is blocked by its per-span local rules. Args: - span: The span to check. + span: The span carrying the local-block attributes. + name: The span name to match against local patterns. Returns: - True if the span has a local-block flag, False otherwise. + ``True`` if *name* matches a local pattern or the span carries the + local-block flag; ``False`` on any read error. """ try: - attrs = getattr(span, "attributes", None) - if not attrs: - return False - try: - if hasattr(attrs, "get"): - value = attrs.get("netra.local_blocked") - else: - value = attrs["netra.local_blocked"] - except Exception: - value = None - return bool(value) is True + local_patterns = read_local_block_patterns(span) + if local_patterns and PatternMatcher(local_patterns).matches(name): + return True + return has_local_block_flag(span) except Exception: return False - def _reparent_blocked_children( - self, - spans: Sequence[ReadableSpan], - blocked_parent_map: Dict[Any, Any], - ) -> None: - """ - Reparent blocked children of a span. - - Args: - spans: List of spans to reparent. - blocked_parent_map: Dictionary mapping span IDs to their blocked parent spans. - """ - if not blocked_parent_map: - return + def _build_reparent_map( + self, parents_of_root_spans_blocked: Dict[int, Any], parents_of_spans_blocked_by_name: Dict[Any, Any] + ) -> Dict[Any, Any]: + """Merge the cross-batch registry with this batch's dropped-parent maps. - for span in spans: - parent_context = getattr(span, "parent", None) - if parent_context is None: - continue - - updated_parent = parent_context - visited: set[Any] = set() - changed = False - - while updated_parent is not None: - parent_span_id = getattr(updated_parent, "span_id", None) - if parent_span_id not in blocked_parent_map or parent_span_id in visited: - break - visited.add(parent_span_id) - updated_parent = blocked_parent_map[parent_span_id] - changed = True - - if changed: - self._set_span_parent(span, updated_parent) - - def _set_span_parent(self, span: ReadableSpan, parent: Any) -> None: - """ - Set the parent of a span. + Ordering matters: the process-global registry captured by processors is + the base, overlaid by root-blocked parents, then this batch's name/local + blocks — so in-batch parents win on conflict. Args: - span: The span to set the parent of. - parent: The parent to set. + parents_of_root_spans_blocked: Root-blocked span IDs to their parents. + parents_of_spans_blocked_by_name: This batch's name/locally blocked + span IDs to their parents. + + Returns: + The merged ``{dropped_span_id -> parent}`` map used for reparenting. """ - if hasattr(span, "_parent"): - try: - span._parent = parent - return - except Exception: - pass + reparent_map: Dict[Any, Any] = {} try: - setattr(span, "parent", parent) + if BLOCKED_LOCAL_PARENT_MAP: + reparent_map.update(BLOCKED_LOCAL_PARENT_MAP) except Exception: - logger.debug("Failed to reparent span %s", getattr(span, "name", ""), exc_info=True) + pass + reparent_map.update(parents_of_root_spans_blocked) + reparent_map.update(parents_of_spans_blocked_by_name) + return reparent_map def shutdown(self) -> None: - """ - Shutdown the exporter. - """ + """Shutdown the wrapped exporter, suppressing shutdown errors.""" try: self._exporter.shutdown() except Exception: pass def force_flush(self, timeout_millis: int = 30000) -> Any: - """ - Force flush the exporter. + """Force flush the wrapped exporter. Args: - timeout_millis: The timeout in milliseconds. + timeout_millis: The flush timeout in milliseconds. Returns: - The result of the force flush operation. + The wrapped exporter's flush result, or ``True`` if it raised. """ try: return self._exporter.force_flush(timeout_millis) diff --git a/netra/exporters/utils.py b/netra/exporters/utils.py index cb444282..33b181a5 100644 --- a/netra/exporters/utils.py +++ b/netra/exporters/utils.py @@ -1,14 +1,25 @@ import logging import threading import time -from typing import Optional, Set +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, cast from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.trace import INVALID_SPAN_ID, SpanContext from netra.config import Config +from netra.processors.root_instrument_filter_processor import ( + ROOT_BLOCK_CANDIDATE_FIELD, + get_root_block_candidates, + root_block_candidates_contains, +) logger = logging.getLogger(__name__) +# Span attributes set by LocalFilteringSpanProcessor and read back here to +# decide per-span local blocking (kept in sync with that processor's keys). +LOCAL_BLOCKED_SPANS_ATTR = "netra.local_blocked_spans" +LOCAL_BLOCKED_FLAG_ATTR = "netra.local_blocked" + _trial_status_lock = threading.Lock() _trial_blocked_at: Optional[float] = None _blocked_trace_ids: Set[str] = set() @@ -123,3 +134,384 @@ def get_trace_id(span: ReadableSpan) -> str: except Exception as e: logger.debug("Error extracting trace ID from span: %s", e) return "" + + +class PatternMatcher: + """Matches span names against a set of block patterns. + + Pattern grammar: + - ``"Foo"`` exact match: ``name == "Foo"``. + - ``"Foo.*"`` prefix match: ``name`` starts with ``"Foo."``. + - ``"*.Query"`` suffix match: ``name`` ends with ``".Query"``. + + A pattern that both starts and ends with ``*`` is treated as an exact match; + empty patterns are ignored. Patterns are classified once at construction so + matching is a constant-time set lookup plus tuple ``startswith``/``endswith``. + """ + + def __init__(self, patterns: Sequence[str]) -> None: + """Classify *patterns* into exact / prefix / suffix buckets once. + + Args: + patterns: The block patterns to compile. Empty patterns are ignored. + """ + exact: Set[str] = set() + prefixes: List[str] = [] + suffixes: List[str] = [] + for pattern in patterns: + if not pattern: + continue + if pattern.endswith("*") and not pattern.startswith("*"): + prefixes.append(pattern[:-1]) + elif pattern.startswith("*") and not pattern.endswith("*"): + suffixes.append(pattern[1:]) + else: + exact.add(pattern) + self._exact = exact + # str.startswith / str.endswith accept a tuple of candidates. + self._prefixes = tuple(prefixes) + self._suffixes = tuple(suffixes) + + def matches(self, name: str) -> bool: + """Check whether *name* matches any configured pattern. + + Args: + name: The span name to test. + + Returns: + ``True`` if *name* matches an exact, prefix, or suffix pattern. + """ + return name in self._exact or name.startswith(self._prefixes) or name.endswith(self._suffixes) + + +def get_span_id(span: ReadableSpan) -> Optional[int]: + """Return *span*'s own span ID. + + Args: + span: The span to inspect. + + Returns: + The integer span ID, or ``None`` if unavailable. + """ + span_context = getattr(span, "context", None) + if span_context is None: + return None + return cast(Optional[int], getattr(span_context, "span_id", None)) + + +def normalize_parent(parent: Any) -> Optional[SpanContext]: + """Normalize a span's parent link, collapsing an invalid parent to ``None``. + + Mirrors ``RootInstrumentFilterProcessor._get_parent_span_context`` so the + in-batch overlay and the registry store parents identically. A ``None`` + result means the span is a trace root. + + Args: + parent: The span's ``parent`` ``SpanContext`` (or ``None``). + + Returns: + The parent ``SpanContext``, or ``None`` when the span is a trace root. + """ + if parent is None: + return None + parent_id = getattr(parent, "span_id", None) + if parent_id is None or parent_id == INVALID_SPAN_ID: + return None + return cast(Optional[SpanContext], parent) + + +def _read_attribute(span: ReadableSpan, key: str) -> Any: + """Read attribute *key* off *span*, tolerating read-only proxies and misses. + + ``ReadableSpan.attributes`` may be a plain mapping or a read-only proxy. + + Args: + span: The span to read the attribute from. + key: The attribute key to look up. + + Returns: + The attribute value, or ``None`` when it is absent or unreadable. + """ + attrs = getattr(span, "attributes", None) + if not attrs: + return None + try: + if hasattr(attrs, "get"): + return attrs.get(key) + return attrs[key] + except Exception: + return None + + +def read_local_block_patterns(span: ReadableSpan) -> List[str]: + """Read the per-span local-block patterns set by ``LocalFilteringSpanProcessor``. + + Args: + span: The span to read patterns from. + + Returns: + The non-empty string patterns, or an empty list if none are set. + """ + value = _read_attribute(span, LOCAL_BLOCKED_SPANS_ATTR) + if isinstance(value, (list, tuple)) and all(isinstance(v, str) for v in value): + return [v for v in value if v] + return [] + + +def has_local_block_flag(span: ReadableSpan) -> bool: + """Check whether the processor explicitly flagged *span* as locally blocked. + + Args: + span: The span to check. + + Returns: + ``True`` if the span carries a truthy local-block flag. + """ + return bool(_read_attribute(span, LOCAL_BLOCKED_FLAG_ATTR)) + + +def is_root_block_candidate(span: ReadableSpan) -> bool: + """Check whether *span* carries the durable root-block-candidate marker. + + The marker is a plain instance attribute (not an OTel span attribute), so it + cannot be evicted by the span attribute limit. + + Args: + span: The span to check. + + Returns: + ``True`` if the span was marked as a root-block candidate at start. + """ + return bool(getattr(span, ROOT_BLOCK_CANDIDATE_FIELD, False)) + + +def find_root_spans_blocked( + spans: Sequence[ReadableSpan], + parents_of_spans_blocked_by_name: Optional[Dict[int, Optional[SpanContext]]] = None, +) -> Tuple[Set[int], Dict[int, Optional[SpanContext]]]: + """Find the root-block candidates in *spans* that must be dropped. + + A candidate is dropped when it is *root-connected*: it is a trace root (no + local parent), or every ancestor between it and the trace root is itself + dropped. Walking up the ancestor chain stops at the first surviving ancestor + — an allowed instrument, a netra/manual span, or a non-instrumentation span — + and never crosses a remote (cross-process) parent link. + + ``parents_of_spans_blocked_by_name`` are spans dropped for *other* reasons (a + global or per-span local name block) that are removed from the exported tree + just like candidates. Folding them in lets the walk "see through" them: a + candidate whose only surviving ancestor was a name-blocked span becomes + root-connected once that span is dropped, and so must be dropped too rather + than promoted to a root it is not allowed to be. + + The candidate set is the registry snapshot overlaid with any span in the + current batch that carries the durable candidacy marker. The overlay makes + the decision robust: a blocked root still in the batch is dropped even if + its registry entry was evicted (TTL/overflow) or cleared, because the marker + travels with the span. Only ancestor chains reachable from this batch are + evaluated, so the cost is proportional to the batch (plus its ancestry) + rather than to the whole registry. + + Args: + spans: The batch of spans being exported. + parents_of_spans_blocked_by_name: ``{span_id -> parent_span_context}`` for + spans dropped by name/local rules, treated as transparent dropped nodes. + + Returns: + ``(root_connected_ids, parents_of_root_spans_blocked)`` where + ``parents_of_root_spans_blocked`` maps each dropped span ID to its parent + ``SpanContext`` (``None`` for a true root) for reparenting. + """ + candidates = _collect_root_block_candidates(spans) + # Only genuine root-block candidates trigger a chain walk; name/local drops + # alone (with no candidate in play) are handled by the caller directly. + if not candidates: + return set(), {} + + if parents_of_spans_blocked_by_name: + for span_id, parent_ctx in parents_of_spans_blocked_by_name.items(): + candidates.setdefault(span_id, parent_ctx) + + root_connected_ids = _find_root_connected_candidates(spans, candidates) + parents_of_root_spans_blocked: Dict[int, Optional[SpanContext]] = { + sid: candidates[sid] for sid in root_connected_ids + } + return root_connected_ids, parents_of_root_spans_blocked + + +def _collect_root_block_candidates(spans: Sequence[ReadableSpan]) -> Dict[int, Optional[SpanContext]]: + """Merge in-batch candidacy markers with the cross-batch registry. + + The in-batch markers let a marked span be recognised even if its registry + entry was evicted or cleared before export. The registry is only consulted + for *cross-batch* ancestry — a batch parent whose candidacy was recorded in + an earlier batch. When no batch parent references such an entry, the full + (locked, O(registry)) snapshot is skipped entirely. + + Args: + spans: The batch of spans being exported. + + Returns: + A ``{span_id -> parent_span_context}`` map of every relevant candidate, + with in-batch markers overriding the registry snapshot. + """ + in_batch: Dict[int, Optional[SpanContext]] = {} + parent_ids: Set[int] = set() + for span in spans: + span_id = get_span_id(span) + if is_root_block_candidate(span) and span_id is not None and span_id != INVALID_SPAN_ID: + in_batch[span_id] = normalize_parent(getattr(span, "parent", None)) + parent = getattr(span, "parent", None) + parent_id = getattr(parent, "span_id", None) if parent is not None else None + if parent_id is not None and parent_id != INVALID_SPAN_ID: + parent_ids.add(parent_id) + + # A registry entry can only be walked if it is the parent of some batch span + # (every in-batch candidate contributes its own parent id here too, so + # multi-hop cross-batch chains are covered). Parents already resolvable + # in-batch never need the registry. + if not root_block_candidates_contains(parent_ids - set(in_batch)): + return in_batch + + candidates = get_root_block_candidates() + candidates.update(in_batch) + return candidates + + +def _find_root_connected_candidates( + spans: Sequence[ReadableSpan], candidates: Dict[int, Optional[SpanContext]] +) -> Set[int]: + """Return the candidate span IDs that are root-connected (and so dropped). + + Traversal is seeded only from batch spans and their parents, so unrelated + registry entries are never walked. + + Args: + spans: The batch of spans being exported. + candidates: The ``{span_id -> parent_span_context}`` candidate map from + :func:`_collect_root_block_candidates`. + + Returns: + The set of candidate span IDs that are root-connected. + """ + memo: Dict[int, bool] = {} + root_connected_ids: Set[int] = set() + + def resolves_to_root(start_id: int) -> bool: + """Resolve whether *start_id* is a root-connected candidate, memoized. + + Walks the candidate ancestry chain iteratively (no recursion, so a + deeply nested trace cannot blow the Python stack). Every node on a + linear candidate chain shares the terminal verdict — the chain is + dropped iff it walks all the way up to a true root — so the resolved + result is written back to the whole walked path in one pass. + + Args: + start_id: The span ID to resolve. + + Returns: + ``True`` if *start_id* is a candidate whose ancestry reaches a root. + """ + walked_ids: List[int] = [] + seen_ids: Set[int] = set() + current_id = start_id + while True: + if current_id in memo: + reaches_root = memo[current_id] + break + if current_id not in candidates: + # Not a candidate -> a surviving ancestor. Stops the walk. + reaches_root = False + break + if current_id in seen_ids: + # Cycle guard: treat as non-dropped to avoid looping forever. + reaches_root = False + break + + # Fresh candidate node: it shares the chain's terminal verdict. + walked_ids.append(current_id) + seen_ids.add(current_id) + + parent_ctx = candidates[current_id] + if parent_ctx is None: + reaches_root = True # candidate is a true trace root + break + if getattr(parent_ctx, "is_remote", False): + reaches_root = False # do not walk across a process boundary + break + parent_id = getattr(parent_ctx, "span_id", None) + if parent_id is None or parent_id == INVALID_SPAN_ID: + reaches_root = True + break + + current_id = parent_id + + for span_id in walked_ids: + memo[span_id] = reaches_root + if reaches_root: + root_connected_ids.add(span_id) + return reaches_root + + for span in spans: + span_id = get_span_id(span) + if span_id is not None and span_id != INVALID_SPAN_ID: + resolves_to_root(span_id) + parent_ctx = getattr(span, "parent", None) + parent_id = getattr(parent_ctx, "span_id", None) if parent_ctx is not None else None + if parent_id is not None and parent_id != INVALID_SPAN_ID: + resolves_to_root(parent_id) + + return root_connected_ids + + +def reparent_spans(spans: Sequence[ReadableSpan], dropped_span_parents: Dict[Any, Any]) -> None: + """Reparent each span past any dropped ancestor onto its first survivor. + + Walks the chain of dropped parents (following ``dropped_span_parents``) until + a surviving parent — or ``None`` (promote to root) — is reached. + + Args: + spans: The surviving spans whose parents may need rewriting. + dropped_span_parents: A ``{dropped_span_id -> parent}`` map of dropped + spans to their own parents. + """ + if not dropped_span_parents: + return + + for span in spans: + parent = getattr(span, "parent", None) + if parent is None: + continue + + new_parent = parent + visited: Set[Any] = set() + changed = False + while new_parent is not None: + parent_span_id = getattr(new_parent, "span_id", None) + if parent_span_id not in dropped_span_parents or parent_span_id in visited: + break + visited.add(parent_span_id) + new_parent = dropped_span_parents[parent_span_id] + changed = True + + if changed: + set_span_parent(span, new_parent) + + +def set_span_parent(span: ReadableSpan, parent: Any) -> None: + """Set *span*'s parent, preferring the private ``_parent`` slot. + + Args: + span: The span to reparent. + parent: The new parent ``SpanContext`` (``None`` to promote to root). + """ + if hasattr(span, "_parent"): + try: + span._parent = parent + return + except Exception: + pass + try: + setattr(span, "parent", parent) + except Exception: + logger.debug("Failed to reparent span %s", getattr(span, "name", ""), exc_info=True) diff --git a/netra/processors/root_instrument_filter_processor.py b/netra/processors/root_instrument_filter_processor.py index a900e3ed..ee1ba94f 100644 --- a/netra/processors/root_instrument_filter_processor.py +++ b/netra/processors/root_instrument_filter_processor.py @@ -2,39 +2,104 @@ import threading import time from collections import OrderedDict -from typing import Optional, Set, cast +from typing import Dict, Optional, Set, cast from opentelemetry import context as otel_context -from opentelemetry import trace from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor -from opentelemetry.trace import INVALID_SPAN_ID +from opentelemetry.trace import INVALID_SPAN_ID, SpanContext logger = logging.getLogger(__name__) -_LOCAL_BLOCKED_ATTR = "netra.local_blocked" _INSTRUMENTATION_PREFIXES = ("opentelemetry.instrumentation.", "netra.instrumentation.") -_MAX_BLOCKED_TRACES = 4096 -_BLOCKED_TRACE_TTL_SECONDS = 600.0 +_MAX_ROOT_CANDIDATES = 4096 +_ROOT_CANDIDATE_TTL_SECONDS = 600.0 + +# Durable per-span marker set on a span the moment it is classified as a +# root-block candidate. It is a plain *instance attribute* — deliberately NOT +# an OTel span attribute — so it travels with the span into its export batch +# without consuming the span's bounded attribute capacity and without being +# evicted by ``OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT`` when later instrumentation adds +# attributes. Being off the attribute map, it is never serialised to the +# backend, so nothing needs to strip it before export. The exporter can still +# recognise a blocked root even if the registry entry was evicted (TTL/overflow) +# or cleared (shutdown) between ``on_start`` and export. +ROOT_BLOCK_CANDIDATE_FIELD = "_netra_root_block_candidate" + +# Process-global registry of "root-block candidate" spans: spans emitted by an +# auto-instrumentation library that is *not* permitted to produce root-level +# spans. Maps ``span_id -> (parent_span_context_or_None, ended_at)`` where +# ``ended_at`` is ``None`` while the span is still active and the monotonic end +# timestamp once it has ended. +# +# The registry lets the ``FilteringSpanExporter`` resolve *cross-batch* +# ancestry — reparenting a child that exports in a later batch than its dropped +# ancestor. It is a supplement, not the source of truth: candidacy is carried +# durably on the span via ``ROOT_BLOCK_CANDIDATE_FIELD`` (see above). The TTL +# clock only starts when a candidate *ends*, so a long-lived root is never +# evicted while still active, and ended entries survive long enough for +# late-exporting children to still find their dropped ancestor's parent. +_root_candidates_lock = threading.Lock() +ROOT_BLOCK_CANDIDATES: "OrderedDict[int, tuple[Optional[SpanContext], Optional[float]]]" = OrderedDict() + + +def get_root_block_candidates() -> Dict[int, Optional[SpanContext]]: + """Return a snapshot ``{span_id -> parent_span_context}`` of the current + root-block candidate registry. + + The snapshot decouples the exporter's read from concurrent writes by + ``RootInstrumentFilterProcessor.on_start``. + + Returns: + A plain dict mapping each candidate span's ID to the ``SpanContext`` of + its parent (``None`` for a candidate that is itself a trace root). + """ + with _root_candidates_lock: + return {span_id: parent_ctx for span_id, (parent_ctx, _ended_at) in ROOT_BLOCK_CANDIDATES.items()} -class RootInstrumentFilterProcessor(SpanProcessor): # type: ignore[misc] - """Block root spans (and their entire subtree) whose instrumentation is - not in the allowed *root_instruments* set. +def root_block_candidates_contains(span_ids: Set[int]) -> bool: + """Return whether any of *span_ids* is currently a recorded candidate. + + A cheap membership probe (set lookups under the lock, no copy) that lets the + exporter skip the full :func:`get_root_block_candidates` snapshot on batches + whose spans reference no cross-batch candidate ancestor. + + Args: + span_ids: Span IDs to probe against the registry. - When an auto-instrumentation root span (e.g. FastAPI, Flask) is not - permitted, this processor: + Returns: + ``True`` if at least one ID is a live candidate. + """ + if not span_ids: + return False + with _root_candidates_lock: + return any(span_id in ROOT_BLOCK_CANDIDATES for span_id in span_ids) - 1. Marks it with ``netra.local_blocked = True``. - 2. Records its **trace_id** in a bounded, TTL-evicted registry. - 3. Marks every subsequent child span that shares the same trace ID. - Tracking by trace ID (rather than individual span IDs) guarantees that - the block propagates correctly even in async frameworks where a parent - span may end before all of its children have started. +class RootInstrumentFilterProcessor(SpanProcessor): # type: ignore[misc] + """Record spans from auto-instrumentation libraries that are not permitted + to produce root-level spans, so the exporter can drop-and-reparent them. + + When an auto-instrumentation span comes from a library outside the allowed + *root_instruments* set, the processor marks it as a **root-block candidate** + with a durable instance-level marker (``ROOT_BLOCK_CANDIDATE_FIELD``) and + also records it, along with its parent ``SpanContext``, in a shared TTL-evicted + registry (``ROOT_BLOCK_CANDIDATES``) used for cross-batch reparenting. + + The actual drop decision is made at export time by + :class:`~netra.exporters.filtering_span_exporter.FilteringSpanExporter`: + a candidate is dropped only when it is *root-connected* — i.e. it is a + trace root, or every ancestor up to the trace root is also a dropped + candidate. Dropped candidates have their children reparented onto the + dropped span's parent (``None`` for a true root, so the child becomes the + new root). This peel repeats recursively: if the promoted child is itself + from a non-root instrumentation it is dropped too, until a survivor is + reached (an allowed instrument, a netra decorator / manual span, or any + non-instrumentation span). Spans created directly through netra decorators or ``Netra.start_span`` - are never filtered — only spans from recognised auto-instrumentation + are never candidates — only spans from recognised auto-instrumentation libraries (scope prefix ``opentelemetry.instrumentation.*`` or ``netra.instrumentation.*``) are subject to the allow-list. @@ -45,45 +110,53 @@ class RootInstrumentFilterProcessor(SpanProcessor): # type: ignore[misc] def __init__(self, allowed_root_instrument_names: Set[str]) -> None: self._allowed: frozenset[str] = frozenset(allowed_root_instrument_names) - self._blocked_trace_ids: OrderedDict[int, float] = OrderedDict() - self._lock = threading.Lock() def on_start( self, span: Span, parent_context: Optional[otel_context.Context] = None, ) -> None: - """Evaluate whether *span* belongs to a blocked trace and mark it - accordingly. + """Record *span* as a root-block candidate when it comes from a + non-allowed auto-instrumentation library. Args: span: The span that is being started. parent_context: The parent context of the span. """ try: - self._process_span_start(span, parent_context) + self._process_span_start(span) except Exception: logger.debug("RootInstrumentFilterProcessor.on_start failed", exc_info=True) def on_end(self, span: ReadableSpan) -> None: - """Prune expired entries from the blocked-trace registry. + """Start the TTL clock for a just-ended candidate, then prune expired entries. - The registry is **not** cleared on a per-span basis — entries - survive until they expire via TTL so that late-starting children - still see their trace as blocked. + Entries are **not** cleared per-span — once a candidate ends its entry + survives for ``_ROOT_CANDIDATE_TTL_SECONDS`` so that children exporting + after their (already-ended) blocked ancestor can still be reparented. + The TTL is measured from the *end* time, so an active long-lived root is + never evicted before it finishes. Args: span: The span that is being ended. """ try: - self._evict_stale_traces() + self._mark_ended(span) + self._evict_stale_candidates() except Exception: pass def shutdown(self) -> None: - """Release all resources held by the processor.""" - with self._lock: - self._blocked_trace_ids.clear() + """No-op — the candidate registry is **not** cleared here. + + This processor is registered before the exporter's span processor, so + clearing the shared registry on shutdown would empty it *before* the + exporter's final flush runs — causing still-buffered blocked root spans + to slip through as exported roots. The registry is process-global and + already bounded (``_MAX_ROOT_CANDIDATES``) and TTL-evicted, so leaving + it intact is safe; clearing it here is not. + """ + return None def force_flush(self, timeout_millis: int = 30000) -> bool: """No-op — this processor does not buffer data. @@ -100,104 +173,115 @@ def force_flush(self, timeout_millis: int = 30000) -> bool: # Core logic # ------------------------------------------------------------------ - def _process_span_start( - self, - span: Span, - parent_context: Optional[otel_context.Context], - ) -> None: - """Decide whether *span* should be blocked. + def _process_span_start(self, span: Span) -> None: + """Record *span* as a candidate if it is an auto-instrumentation span + whose instrument is not in the allowed root set. - A span is classified as a **child** (non-root) when either the - *parent_context* carries a valid span or the span object itself - records a valid ``parent_span_id``. The two-pronged check handles - instrumentations (e.g. HTTPX, OpenAI) that set the parent link on - the span directly without propagating through the active context. + Every qualifying span is recorded regardless of whether it is a root + or a child at start time. A child only *becomes* a root after its + ancestors are dropped, which is resolved at export; recording it now + ensures the exporter can peel it recursively. Args: span: The span that is being started. - parent_context: The parent context of the span. """ - parent_span_id = self._resolve_parent_span_id(parent_context) - if parent_span_id is None or parent_span_id == INVALID_SPAN_ID: - parent_span_id = self._get_parent_span_id_from_span(span) + if not self._is_from_instrumentation_library(span): + return + + instr_name = self._extract_instrumentation_name(span) + if instr_name is None or instr_name in self._allowed: + return + + span_id = self._get_own_span_id(span) + if span_id is None or span_id == INVALID_SPAN_ID: + return - is_child = parent_span_id is not None and parent_span_id != INVALID_SPAN_ID + parent_ctx = self._get_parent_span_context(span) + # Durable marker first: it must survive even if the registry entry is + # later evicted/cleared before this span reaches the exporter. + self._mark_candidate(span) + self._record_candidate(span_id, parent_ctx) - if is_child: - self._maybe_block_child(span) - else: - self._maybe_block_root(span) + @staticmethod + def _mark_candidate(span: Span) -> None: + """Stamp *span* with the durable root-block-candidate marker. - def _maybe_block_child(self, span: Span) -> None: - """Block *span* if its trace is in the blocked registry. + The marker is a plain instance attribute (see ``ROOT_BLOCK_CANDIDATE_FIELD``), + not an OTel span attribute, so it cannot be evicted by the span attribute + limit and is never exported. Args: - span: A child (non-root) span that is being started. + span: The candidate span being started. """ - trace_id = self._get_trace_id(span) - if trace_id is None: - return - with self._lock: - if trace_id in self._blocked_trace_ids: - self._mark_blocked(span) + try: + setattr(span, ROOT_BLOCK_CANDIDATE_FIELD, True) + except Exception: + pass + + @staticmethod + def _record_candidate(span_id: int, parent_ctx: Optional[SpanContext]) -> None: + """Register *span_id* as an **active** root-block candidate with its parent. - def _maybe_block_root(self, span: Span) -> None: - """Block *span* if it is an auto-instrumentation root whose name - is not in the allowed set. + The entry's TTL clock is left unstarted (``ended_at = None``) until the + span ends, so a candidate that stays open longer than the TTL is never + evicted while still active. Args: - span: A root span that is being started. + span_id: The candidate span's own span ID. + parent_ctx: The candidate's parent ``SpanContext`` (``None`` when it + is a trace root). """ - if not self._is_from_instrumentation_library(span): - return + with _root_candidates_lock: + ROOT_BLOCK_CANDIDATES[span_id] = (parent_ctx, None) + ROOT_BLOCK_CANDIDATES.move_to_end(span_id) + while len(ROOT_BLOCK_CANDIDATES) > _MAX_ROOT_CANDIDATES: + ROOT_BLOCK_CANDIDATES.popitem(last=False) - instr_name = self._extract_instrumentation_name(span) - if instr_name is None or instr_name in self._allowed: - return + @staticmethod + def _mark_ended(span: ReadableSpan) -> None: + """Start the TTL clock for *span*'s candidate entry, if it has one. - trace_id = self._get_trace_id(span) - if trace_id is not None: - with self._lock: - self._blocked_trace_ids[trace_id] = time.monotonic() - self._blocked_trace_ids.move_to_end(trace_id) - self._evict_overflow() - self._mark_blocked(span) + Args: + span: The span that is ending. + """ + ctx = getattr(span, "context", None) + span_id = getattr(ctx, "span_id", None) if ctx is not None else None + if span_id is None: + return + with _root_candidates_lock: + entry = ROOT_BLOCK_CANDIDATES.get(span_id) + if entry is not None: + ROOT_BLOCK_CANDIDATES[span_id] = (entry[0], time.monotonic()) @staticmethod - def _resolve_parent_span_id( - parent_context: Optional[otel_context.Context], - ) -> Optional[int]: - """Return the parent span's ``span_id`` from *parent_context*. + def _get_own_span_id(span: Span) -> Optional[int]: + """Return *span*'s own ``span_id``. Args: - parent_context: The context passed to ``on_start``. + span: The span to inspect. Returns: - The parent ``span_id``, or ``None`` if unavailable. + The integer span ID, or ``None`` if unavailable. """ - if parent_context is None: - return None - parent_span = trace.get_current_span(parent_context) - if parent_span is None: - return None - sc = parent_span.get_span_context() - if sc is None: + ctx = getattr(span, "context", None) or getattr(span, "get_span_context", lambda: None)() + if ctx is None: return None - return cast(Optional[int], sc.span_id) + return cast(Optional[int], getattr(ctx, "span_id", None)) @staticmethod - def _get_parent_span_id_from_span(span: Span) -> Optional[int]: - """Extract the parent ``span_id`` from the span's internal state. + def _get_parent_span_context(span: Span) -> Optional[SpanContext]: + """Return the parent ``SpanContext`` recorded on *span*. - The OTel SDK ``Span`` stores the parent ``SpanContext`` directly, - which is authoritative even when the active-context-based - *parent_context* has no current span. + The OTel SDK ``Span`` stores its parent ``SpanContext`` directly, which + is authoritative for reparenting. A ``None`` result means the span is a + trace root (its children should be promoted to roots when it is + dropped). Args: span: The span to inspect. Returns: - The parent ``span_id``, or ``None`` if unavailable. + The parent ``SpanContext``, or ``None`` for a root span. """ parent = getattr(span, "parent", None) if parent is None: @@ -205,23 +289,7 @@ def _get_parent_span_id_from_span(span: Span) -> Optional[int]: parent_id = getattr(parent, "span_id", None) if parent_id is None or parent_id == INVALID_SPAN_ID: return None - return cast(Optional[int], parent_id) - - @staticmethod - def _get_trace_id(span: object) -> Optional[int]: - """Return the ``trace_id`` carried by *span*. - - Args: - span: Any span-like object with a ``context`` or - ``get_span_context`` accessor. - - Returns: - The integer trace ID, or ``None``. - """ - ctx = getattr(span, "context", None) or getattr(span, "get_span_context", lambda: None)() - if ctx is None: - return None - return cast(Optional[int], getattr(ctx, "trace_id", None)) + return cast(Optional[SpanContext], parent) @staticmethod def _is_from_instrumentation_library(span: Span) -> bool: @@ -271,32 +339,18 @@ def _extract_instrumentation_name(span: Span) -> Optional[str]: return base if base else name return name - @staticmethod - def _mark_blocked(span: Span) -> None: - """Set the ``netra.local_blocked`` attribute on *span*. + def _evict_stale_candidates(self) -> None: + """Evict entries whose span ended more than ``_ROOT_CANDIDATE_TTL_SECONDS`` ago. - Args: - span: The span to mark. + Active candidates (``ended_at is None``) are never TTL-evicted. Scanning + stops at the first entry that is still active or not yet stale; a + long-lived active entry may therefore defer reclamation of stale entries + behind it, but memory stays bounded by ``_MAX_ROOT_CANDIDATES``. """ - try: - span.set_attribute(_LOCAL_BLOCKED_ATTR, True) - except Exception: - pass - - def _evict_stale_traces(self) -> None: - """Remove entries older than ``_BLOCKED_TRACE_TTL_SECONDS``.""" - cutoff = time.monotonic() - _BLOCKED_TRACE_TTL_SECONDS - with self._lock: - while self._blocked_trace_ids: - _, ts = next(iter(self._blocked_trace_ids.items())) - if ts > cutoff: + cutoff = time.monotonic() - _ROOT_CANDIDATE_TTL_SECONDS + with _root_candidates_lock: + while ROOT_BLOCK_CANDIDATES: + _, (_, ended_at) = next(iter(ROOT_BLOCK_CANDIDATES.items())) + if ended_at is None or ended_at > cutoff: break - self._blocked_trace_ids.popitem(last=False) - - def _evict_overflow(self) -> None: - """Trim the registry to ``_MAX_BLOCKED_TRACES``. - - Must be called while holding ``self._lock``. - """ - while len(self._blocked_trace_ids) > _MAX_BLOCKED_TRACES: - self._blocked_trace_ids.popitem(last=False) + ROOT_BLOCK_CANDIDATES.popitem(last=False) diff --git a/tests/test_root_instrument_reparenting.py b/tests/test_root_instrument_reparenting.py new file mode 100644 index 00000000..dfb8014e --- /dev/null +++ b/tests/test_root_instrument_reparenting.py @@ -0,0 +1,463 @@ +"""Tests for root-instrument filtering with drop-and-reparent semantics. + +When a root span comes from an instrumentation that is not in +``root_instruments``, only that span is dropped and its children are reparented +onto its parent — recursively — instead of discarding the whole subtree. + +The tests drive the real :class:`RootInstrumentFilterProcessor` (to populate the +shared candidate registry) and the real :class:`FilteringSpanExporter` (to make +the drop + reparent decision), using lightweight span doubles so parent/scope +wiring is explicit. +""" + +import threading +import time + +import pytest + +from netra.exporters.filtering_span_exporter import FilteringSpanExporter +from netra.processors import root_instrument_filter_processor as rifp +from netra.processors.root_instrument_filter_processor import ROOT_BLOCK_CANDIDATE_FIELD, RootInstrumentFilterProcessor + +pytestmark = pytest.mark.unit + + +@pytest.fixture(autouse=True) +def _clear_registry(): + """Isolate the process-global candidate registry between tests.""" + with rifp._root_candidates_lock: + rifp.ROOT_BLOCK_CANDIDATES.clear() + yield + with rifp._root_candidates_lock: + rifp.ROOT_BLOCK_CANDIDATES.clear() + + +class FakeSpanContext: + def __init__(self, span_id: int, trace_id: int = 0xABC, is_remote: bool = False) -> None: + self.span_id = span_id + self.trace_id = trace_id + self.is_remote = is_remote + + +class FakeScope: + def __init__(self, name: str) -> None: + self.name = name + + +class FakeSpan: + """Minimal stand-in exposing the surface the processor/exporter read.""" + + def __init__( + self, + span_id: int, + scope_name: str | None, + parent_ctx: FakeSpanContext | None = None, + name: str = "span", + trace_id: int = 0xABC, + ) -> None: + self.context = FakeSpanContext(span_id, trace_id) + self.instrumentation_scope = FakeScope(scope_name) if scope_name else None + self.parent = parent_ctx + self.name = name + self.attributes: dict = {} + + def get_span_context(self) -> FakeSpanContext: + return self.context + + def set_attribute(self, key: str, value: object) -> None: + self.attributes[key] = value + + +class RecordingExporter: + def __init__(self) -> None: + self.exported: list = [] + + def export(self, spans): + self.exported.extend(spans) + from opentelemetry.sdk.trace.export import SpanExportResult + + return SpanExportResult.SUCCESS + + def shutdown(self) -> None: + pass + + def force_flush(self, timeout_millis: int = 30000) -> bool: + return True + + +def scope(name: str) -> str: + return f"netra.instrumentation.{name}" + + +def make_pipeline(allowed: set[str]): + processor = RootInstrumentFilterProcessor(allowed) + recorder = RecordingExporter() + exporter = FilteringSpanExporter(recorder, patterns=[]) + return processor, exporter, recorder + + +def exported_ids(recorder: RecordingExporter) -> set[int]: + return {s.context.span_id for s in recorder.exported} + + +# --------------------------------------------------------------------------- +# Processor: candidate recording +# --------------------------------------------------------------------------- + + +def test_processor_records_only_non_allowed_instrumentation(): + processor = RootInstrumentFilterProcessor({"openai"}) + + fastapi = FakeSpan(1, scope("fastapi")) + openai = FakeSpan(2, scope("openai")) + manual = FakeSpan(3, "my.app.tracer") + + for span in (fastapi, openai, manual): + processor.on_start(span) + + candidates = rifp.get_root_block_candidates() + assert set(candidates) == {1} # only the non-allowed FastAPI span + + +# --------------------------------------------------------------------------- +# Exporter: drop + reparent +# --------------------------------------------------------------------------- + + +def test_drops_root_and_reparents_child(): + processor, exporter, recorder = make_pipeline({"openai"}) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + openai = FakeSpan(2, scope("openai"), parent_ctx=fastapi.context) + + processor.on_start(fastapi) + processor.on_start(openai) + exporter.export([fastapi, openai]) + + assert exported_ids(recorder) == {2} + assert openai.parent is None # promoted to root + + +def test_recursive_peel_promotes_deepest_allowed_span(): + processor, exporter, recorder = make_pipeline({"openai"}) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + asgi = FakeSpan(2, scope("asgi"), parent_ctx=fastapi.context) + openai = FakeSpan(3, scope("openai"), parent_ctx=asgi.context) + + for span in (fastapi, asgi, openai): + processor.on_start(span) + exporter.export([fastapi, asgi, openai]) + + assert exported_ids(recorder) == {3} + assert openai.parent is None + + +def test_peel_stops_at_allowed_ancestor(): + processor, exporter, recorder = make_pipeline({"openai"}) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + openai = FakeSpan(2, scope("openai"), parent_ctx=fastapi.context) + httpx = FakeSpan(3, scope("httpx"), parent_ctx=openai.context) + + for span in (fastapi, openai, httpx): + processor.on_start(span) + exporter.export([fastapi, openai, httpx]) + + # FastAPI dropped; OpenAI promoted to root; HTTPX kept under OpenAI untouched + assert exported_ids(recorder) == {2, 3} + assert openai.parent is None + assert httpx.parent is openai.context + + +def test_non_root_span_under_surviving_parent_is_not_dropped(): + # A non-allowed instrumentation span that is *not* root-connected stays. + processor, exporter, recorder = make_pipeline({"openai"}) + + manual_root = FakeSpan(1, "my.app.tracer", parent_ctx=None) + fastapi = FakeSpan(2, scope("fastapi"), parent_ctx=manual_root.context) + + processor.on_start(manual_root) + processor.on_start(fastapi) + exporter.export([manual_root, fastapi]) + + assert exported_ids(recorder) == {1, 2} + assert fastapi.parent is manual_root.context + + +def test_whole_branch_with_no_allowed_span_is_dropped(): + processor, exporter, recorder = make_pipeline({"openai"}) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + httpx = FakeSpan(2, scope("httpx"), parent_ctx=fastapi.context) + + processor.on_start(fastapi) + processor.on_start(httpx) + exporter.export([fastapi, httpx]) + + assert exported_ids(recorder) == set() + + +def test_cross_batch_reparenting_via_registry(): + processor, exporter, recorder = make_pipeline({"openai"}) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + openai = FakeSpan(2, scope("openai"), parent_ctx=fastapi.context) + + processor.on_start(fastapi) + processor.on_start(openai) + + # Batch 1: only the (dropped) FastAPI root exports. + exporter.export([fastapi]) + assert exported_ids(recorder) == set() + + # Batch 2: the child exports later and is still reparented to root. + exporter.export([openai]) + assert exported_ids(recorder) == {2} + assert openai.parent is None + + +def test_remote_parent_root_is_not_dropped(): + processor, exporter, recorder = make_pipeline({"openai"}) + + remote_parent = FakeSpanContext(999, is_remote=True) + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=remote_parent) + + processor.on_start(fastapi) + exporter.export([fastapi]) + + # Continues an upstream distributed trace → not peeled. + assert exported_ids(recorder) == {1} + + +def test_allowed_root_instrument_exports_normally(): + processor, exporter, recorder = make_pipeline({"fastapi", "openai"}) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + openai = FakeSpan(2, scope("openai"), parent_ctx=fastapi.context) + + processor.on_start(fastapi) + processor.on_start(openai) + exporter.export([fastapi, openai]) + + assert exported_ids(recorder) == {1, 2} + assert openai.parent is fastapi.context + + +# --------------------------------------------------------------------------- +# Robustness: candidacy must survive registry eviction / clear +# --------------------------------------------------------------------------- + + +def test_evicted_registry_still_drops_in_batch_root(): + # The blocked root leaks unless candidacy is carried durably on the span: + # simulate the registry being evicted (overflow/TTL) between on_start and + # export. The FastAPI root must still be dropped via its span marker. + processor, exporter, recorder = make_pipeline({"openai"}) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + openai = FakeSpan(2, scope("openai"), parent_ctx=fastapi.context) + + processor.on_start(fastapi) + processor.on_start(openai) + + # Registry wiped before export (models 4096-overflow / TTL eviction). + with rifp._root_candidates_lock: + rifp.ROOT_BLOCK_CANDIDATES.clear() + + exporter.export([fastapi, openai]) + + assert exported_ids(recorder) == {2} # FastAPI still dropped + assert openai.parent is None + + +def test_candidate_marker_is_off_the_attribute_map(): + processor, exporter, recorder = make_pipeline({"openai"}) + + # Non-root-connected candidate (under a surviving manual root) is kept. + manual_root = FakeSpan(1, "my.app.tracer", parent_ctx=None) + fastapi = FakeSpan(2, scope("fastapi"), parent_ctx=manual_root.context) + + processor.on_start(manual_root) + processor.on_start(fastapi) + + # Candidacy is carried as a plain instance attribute, never an OTel span + # attribute, so it cannot be evicted by the attribute limit and is never + # exported. + assert getattr(fastapi, ROOT_BLOCK_CANDIDATE_FIELD, False) is True # marked at on_start + assert ROOT_BLOCK_CANDIDATE_FIELD not in fastapi.attributes + assert getattr(manual_root, ROOT_BLOCK_CANDIDATE_FIELD, False) is False # manual span untouched + + exporter.export([manual_root, fastapi]) + + assert exported_ids(recorder) == {1, 2} + # Nothing to strip: the marker was never in the exported attribute map. + assert ROOT_BLOCK_CANDIDATE_FIELD not in fastapi.attributes + + +def test_shutdown_does_not_clear_registry_before_flush(): + processor, exporter, recorder = make_pipeline({"openai"}) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + openai = FakeSpan(2, scope("openai"), parent_ctx=fastapi.context) + + processor.on_start(fastapi) + processor.on_start(openai) + + # Processor shuts down first (registered before the exporter); the registry + # must survive so the exporter's final flush still drops the blocked root. + processor.shutdown() + assert len(rifp.ROOT_BLOCK_CANDIDATES) >= 1 + + exporter.export([fastapi, openai]) + assert exported_ids(recorder) == {2} + assert openai.parent is None + + +def test_unrelated_registry_entries_do_not_affect_batch(): + # A large backlog of candidates from other traces must not be dropped or + # traversed as part of this batch (correctness + bounded work). + processor, exporter, recorder = make_pipeline({"openai"}) + + for other in range(1000, 1500): + processor.on_start(FakeSpan(other, scope("fastapi"), parent_ctx=None, trace_id=other)) + + fastapi = FakeSpan(1, scope("fastapi"), parent_ctx=None) + openai = FakeSpan(2, scope("openai"), parent_ctx=fastapi.context) + processor.on_start(fastapi) + processor.on_start(openai) + + exporter.export([fastapi, openai]) + + # Only spans in this batch are affected; unrelated candidates are neither + # exported nor reparented here. + assert exported_ids(recorder) == {2} + assert openai.parent is None + + +# --------------------------------------------------------------------------- +# Concurrency +# --------------------------------------------------------------------------- + + +@pytest.mark.thread_safety +def test_concurrent_on_start_is_safe_and_bounded(): + processor = RootInstrumentFilterProcessor({"openai"}) + errors: list = [] + + def worker(base: int) -> None: + try: + for i in range(200): + processor.on_start(FakeSpan(base + i, scope("fastapi"))) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(t * 1000,)) for t in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + assert len(rifp.ROOT_BLOCK_CANDIDATES) <= rifp._MAX_ROOT_CANDIDATES + + +# --------------------------------------------------------------------------- +# Interaction between name/local blocking and the root-instrument peel +# --------------------------------------------------------------------------- + + +def test_name_blocked_ancestor_drops_disallowed_candidate_child(): + # A root dropped by a global name pattern must not let a disallowed + # instrumentation child be promoted to root in its place. + processor = RootInstrumentFilterProcessor({"openai"}) + recorder = RecordingExporter() + exporter = FilteringSpanExporter(recorder, patterns=["BlockedRoot"]) + + blocked_root = FakeSpan(1, None, parent_ctx=None, name="BlockedRoot") # name-blocked, not a candidate + fastapi = FakeSpan(2, scope("fastapi"), parent_ctx=blocked_root.context, name="fastapi.request") + + processor.on_start(blocked_root) + processor.on_start(fastapi) + exporter.export([blocked_root, fastapi]) + + # Root dropped by name; FastAPI dropped because it would otherwise become a + # root from an instrument excluded from root_instruments. + assert exported_ids(recorder) == set() + + +def test_name_blocked_ancestor_keeps_allowed_grandchild_as_root(): + # With an allowed span beneath the name-blocked root, it survives and is + # promoted to root; the disallowed intermediate is dropped. + processor = RootInstrumentFilterProcessor({"openai"}) + recorder = RecordingExporter() + exporter = FilteringSpanExporter(recorder, patterns=["BlockedRoot"]) + + blocked_root = FakeSpan(1, None, parent_ctx=None, name="BlockedRoot") + fastapi = FakeSpan(2, scope("fastapi"), parent_ctx=blocked_root.context, name="fastapi.request") + openai = FakeSpan(3, scope("openai"), parent_ctx=fastapi.context, name="openai.chat") + + for span in (blocked_root, fastapi, openai): + processor.on_start(span) + exporter.export([blocked_root, fastapi, openai]) + + assert exported_ids(recorder) == {3} + assert openai.parent is None + + +# --------------------------------------------------------------------------- +# TTL: an active candidate must not be evicted before it ends +# --------------------------------------------------------------------------- + + +def test_active_candidate_not_evicted_by_unrelated_span_end(): + processor = RootInstrumentFilterProcessor({"openai"}) + + # Long-lived blocked root starts and stays active. + root = FakeSpan(1, scope("fastapi"), parent_ctx=None) + processor.on_start(root) + + # An unrelated (allowed) span cycles through start/end many times, each of + # which triggers an eviction pass. The still-active root must survive. + for other in range(2, 12): + unrelated = FakeSpan(other, scope("openai"), parent_ctx=None, trace_id=other) + processor.on_start(unrelated) + processor.on_end(unrelated) + + assert 1 in rifp.ROOT_BLOCK_CANDIDATES # active root retained regardless of elapsed time + + # Once the root ends its TTL clock starts; age it past the TTL and confirm + # it is then reclaimable. + processor.on_end(root) + with rifp._root_candidates_lock: + parent_ctx, _ended = rifp.ROOT_BLOCK_CANDIDATES[1] + rifp.ROOT_BLOCK_CANDIDATES[1] = (parent_ctx, time.monotonic() - rifp._ROOT_CANDIDATE_TTL_SECONDS - 1) + trigger = FakeSpan(99, scope("openai"), parent_ctx=None, trace_id=99) + processor.on_start(trigger) + processor.on_end(trigger) + + assert 1 not in rifp.ROOT_BLOCK_CANDIDATES # ended + stale -> evicted + + +# --------------------------------------------------------------------------- +# Marker durability: real spans under a tight attribute limit +# --------------------------------------------------------------------------- + + +def test_candidate_marker_survives_span_attribute_limit(): + from opentelemetry.sdk.trace import SpanLimits, TracerProvider + + from netra.exporters.utils import is_root_block_candidate + + provider = TracerProvider(span_limits=SpanLimits(max_attributes=1)) + provider.add_span_processor(RootInstrumentFilterProcessor({"openai"})) + tracer = provider.get_tracer("netra.instrumentation.fastapi") + + span = tracer.start_span("fastapi.request") + # Instrumentation records its own attribute; with max_attributes=1 this would + # evict an attribute-based marker. + span.set_attribute("http.method", "GET") + span.end() + + assert is_root_block_candidate(span) is True + assert ROOT_BLOCK_CANDIDATE_FIELD not in dict(span.attributes or {})