diff --git a/ddtrace/internal/_runtime_id.py b/ddtrace/internal/_runtime_id.py index 52bec1c2747..24885f76408 100644 --- a/ddtrace/internal/_runtime_id.py +++ b/ddtrace/internal/_runtime_id.py @@ -67,6 +67,11 @@ def on_runtime_identity_refresh(cb: t.Callable[[str], None]) -> None: _ON_RUNTIME_IDENTITY_REFRESH.add(cb) +def remove_runtime_identity_refresh(cb: t.Callable[[str], None]) -> None: + """Unregister a callback for explicit runtime identity refreshes.""" + _ON_RUNTIME_IDENTITY_REFRESH.discard(cb) + + def get_runtime_identity_refresh_lock() -> t.ContextManager[None]: """Return the lock that serializes a MicroVM identity refresh with its consumers.""" return t.cast(t.ContextManager[None], _RUNTIME_IDENTITY_REFRESH_LOCK) diff --git a/ddtrace/internal/runtime/__init__.py b/ddtrace/internal/runtime/__init__.py index 59c443bd4d2..10886eef0b0 100644 --- a/ddtrace/internal/runtime/__init__.py +++ b/ddtrace/internal/runtime/__init__.py @@ -8,6 +8,7 @@ from ddtrace.internal._runtime_id import on_runtime_id_change from ddtrace.internal._runtime_id import on_runtime_identity_refresh from ddtrace.internal._runtime_id import refresh_identity +from ddtrace.internal._runtime_id import remove_runtime_identity_refresh from ddtrace.internal.serverless import MICROVM_RUN_HOOK_METHOD from ddtrace.internal.serverless import MICROVM_RUN_HOOK_PATH from ddtrace.internal.serverless import in_aws_lambda_microvm @@ -21,6 +22,7 @@ "get_runtime_propagation_envs", "on_runtime_id_change", "on_runtime_identity_refresh", + "remove_runtime_identity_refresh", "listen_for_identity_refresh_hooks", "maybe_refresh_identity", "refresh_identity", diff --git a/ddtrace/internal/telemetry/dependency.py b/ddtrace/internal/telemetry/dependency.py index fbee35e4ba6..bfed97ccb24 100644 --- a/ddtrace/internal/telemetry/dependency.py +++ b/ddtrace/internal/telemetry/dependency.py @@ -130,6 +130,10 @@ def mark_all_metadata_sent(self) -> None: for m in self.metadata: m._mark_sent() + def reset_for_refresh(self) -> None: + """Mark this dependency for reporting to a new worker.""" + self._initial_report_sent = False + def add_metadata(self, cve_id: str, path: str = "", symbol: str = "", line: int = 0) -> bool: """Add or update reachability metadata for a CVE. diff --git a/ddtrace/internal/telemetry/dependency_tracker.py b/ddtrace/internal/telemetry/dependency_tracker.py index 6520ac1e15e..7eea8b64463 100644 --- a/ddtrace/internal/telemetry/dependency_tracker.py +++ b/ddtrace/internal/telemetry/dependency_tracker.py @@ -57,6 +57,7 @@ class DependencyTracker: def __init__(self) -> None: self._imported_dependencies: dict[str, DependencyEntry] = {} + self._report_all = False self._modules_already_imported: set[str] = set() self._lock = Lock() @@ -77,15 +78,12 @@ def collect_report(self) -> Optional[list[dict[str, Any]]]: new_keys = {_normalize_dep_name(d["name"]) for d in new_deps} self._mark_sent(new_keys) - # Skip the re-report scan when SCA is disabled. - # Without SCA, no entry will ever have unsent metadata, so the - # scan over all _imported_dependencies is pure overhead (~887us - # at 10K deps). Only entries created by the SCA hook or with - # metadata attached can trigger needs_report() after initial send. - if not appsec_telemetry_config.SCA_ENABLED: + # Skip re-report scanning when SCA is disabled, except after identity refresh. + if not appsec_telemetry_config.SCA_ENABLED and not self._report_all: return new_deps if new_deps else None re_report_deps = self._collect_rereports(new_keys) + self._report_all = False all_deps = new_deps + re_report_deps return all_deps if all_deps else None @@ -187,11 +185,19 @@ def enable_sca_metadata(self) -> None: if entry.metadata is None: entry.metadata = [] + def refresh(self) -> None: + """Preserve dependency metadata while scheduling a full report for a new worker.""" + with self._lock: + for dependency in self._imported_dependencies.values(): + dependency.reset_for_refresh() + self._report_all = True + def reset(self) -> None: """Reset all state (used on fork / queue reset).""" with self._lock: self._imported_dependencies = {} self._modules_already_imported = set() + self._report_all = False def update_imported_dependencies( diff --git a/ddtrace/internal/telemetry/noop_writer.py b/ddtrace/internal/telemetry/noop_writer.py index 5ed18472393..e96ac168d05 100644 --- a/ddtrace/internal/telemetry/noop_writer.py +++ b/ddtrace/internal/telemetry/noop_writer.py @@ -108,13 +108,19 @@ def set_test_session_token(self, token: Optional[str]) -> None: def _restart_sequence(self) -> None: pass + def _refresh_runtime_identity(self, _runtime_id: str) -> None: + pass + def _fork_writer(self) -> None: pass def _report_dependencies(self) -> Optional[list[dict[str, Any]]]: return None - def _subscribe_worker_changes(self, callback: Any) -> None: + def _subscribe_worker_changes(self, callback: Any, expected_worker: Any) -> None: + pass + + def _unsubscribe_worker_changes(self, callback: Any) -> None: pass def periodic(self, force_flush: bool = False) -> None: diff --git a/ddtrace/internal/telemetry/writer.py b/ddtrace/internal/telemetry/writer.py index b3df7e67ab1..790959c2ad1 100644 --- a/ddtrace/internal/telemetry/writer.py +++ b/ddtrace/internal/telemetry/writer.py @@ -1,3 +1,4 @@ +from collections import deque import itertools import os import traceback @@ -22,7 +23,10 @@ from .._runtime_id import get_ancestor_runtime_id from .._runtime_id import get_parent_runtime_id from .._runtime_id import get_runtime_id +from .._runtime_id import on_runtime_identity_refresh +from .._runtime_id import remove_runtime_identity_refresh from ..periodic import PeriodicService +from ..serverless import in_aws_lambda_microvm from ..utils.formats import get_test_session_token from ..utils.version import version as tracer_version from .constants import TELEMETRY_APM_PRODUCT @@ -172,8 +176,14 @@ def __init__(self, agentless: Optional[bool] = None) -> None: # metric can't both register it (which would create duplicate native contexts / split the # series). Only taken on a cache miss; the hot add path reads the cache lock-free. self._metric_lock = forksafe.Lock() - # Serializes building and publishing the native worker in enable(). - self._enable_lock = forksafe.Lock() + self._is_microvm = in_aws_lambda_microvm() + # Serialize worker construction/publication and lifecycle transitions. MicroVM identity + # refresh re-enters this lock through enable(), so use an RLock in that environment. + self._enable_lock = forksafe.RLock() if self._is_microvm else forksafe.Lock() + # A MicroVM refresh replaces the worker in-process. Guard the complete worker operation + # (lookup through native call) so refresh cannot stop the old worker between those steps. + # Reuse the lifecycle lock; None preserves lock-free worker access outside MicroVMs. + self._worker_access_lock = self._enable_lock if self._is_microvm else None # Callbacks notified whenever the native worker is replaced or torn down. Handles issued by # a worker die with it, so anything holding one (the trace exporter, for its trace_api.* # health metrics) has to be handed the new one rather than keeping a stale clone. @@ -188,9 +198,17 @@ def __init__(self, agentless: Optional[bool] = None) -> None: self._payload_file_dir: Optional[str] = None self._test_session_token: Optional[str] = get_test_session_token() self.started = False + # Preserve lifecycle state across a failed identity refresh so a retry starts the replacement. + self._identity_refresh_started = False # Product enablement is tracked so the version can be passed alongside each change. self._product_versions: dict[str, str] = {product.value: tracer_version for product in TELEMETRY_APM_PRODUCT} + # Native worker rebuilds lose configuration state. Match libdatadog's bounded store. + self._configurations: deque[tuple[str, Optional[str], str, Optional[str]]] = deque(maxlen=5000) + # Native worker rebuilds lose integration state; keep the latest payload for replay. + self._integrations: dict[str, tuple[Optional[str], bool, Optional[bool], Optional[bool], Optional[str]]] = {} + # Native worker rebuilds lose product-change state; keep the latest status for replay. + self._product_statuses: dict[str, bool] = {} if self._enabled: # Captures unhandled exceptions during application start up @@ -208,6 +226,8 @@ def __init__(self, agentless: Optional[bool] = None) -> None: # is marked abandoned in the child; the replacement worker starts lazily after # all child hooks have completed, without unparking the inherited Tokio runtime. forksafe.register(self._fork_writer) + if self._is_microvm: + on_runtime_identity_refresh(self._refresh_runtime_identity) get_logger("ddtrace").addHandler(DDTelemetryErrorHandler(self)) def _build_worker(self) -> "TelemetryWorker": @@ -274,6 +294,35 @@ def _build_worker(self) -> "TelemetryWorker": install_time=config.INSTALL_TIME, ) + def _replay_worker_state(self, worker: "TelemetryWorker") -> None: + # Rebuilt native workers start with empty stores. Replay each accepted configuration event + # in sequence order so dynamic updates and repeated values retain their original semantics. + if self._is_microvm: + origin_cls = _native_telemetry_enums()["origin"] + for name, value, origin, config_id in self._configurations: + seq_id = next(self._sequence_configurations) + worker.add_configuration( + name, + value, + getattr(origin_cls, origin, origin_cls.unknown), + config_id, + seq_id, + ) + # Replay the latest integration state for the rebuilt worker. + for integration_name, state in self._integrations.items(): + version, patched, compatible, auto_patched, error_msg = state + worker.add_integration( + integration_name, + version, + patched, + compatible, + auto_patched, + error_msg, + ) + # Replay the latest product state for the rebuilt worker. + for product, status in self._product_statuses.items(): + worker.add_product_change(product, status, self._product_versions.get(product, tracer_version)) + def enable(self) -> bool: """ Enable the instrumentation telemetry collection service. If the service has already been @@ -287,6 +336,8 @@ def enable(self) -> bool: with self._enable_lock: # extra check to skip the self._worker check on the hotter path + if not self._enabled: + return False if self._worker is not None: return True # type: ignore[unreachable] @@ -299,22 +350,27 @@ def enable(self) -> bool: self._worker = worker # Adopt the recorders onto the new worker while still holding the lock, so no thread # can observe a published worker whose recorders still point at the previous one. + self._replay_worker_state(worker) _bind_metric_recorders(self, worker) self._notify_worker_changed(worker) - # Every process starts its worker so it heartbeats with its own session id. - # app-started is emitted only by the root process; this is enforced inside the - # worker via emit_app_lifecycle (set in _build_worker), so calling start() in a - # forked child schedules heartbeats without re-emitting app-started. - # The root process defers app-started until startup configuration has been reported - # (products load + report_configuration run after enable()); see app_started(), which is - # invoked once products are loaded. (Forked children never emit app-started, so just start.) - if get_parent_runtime_id() is None: - if not self.started: - self.add_configurations(get_python_config_vars()) - else: - worker.start() - self.started = True + # Every process starts its worker so it heartbeats with its own session id. + # app-started is emitted only by the root process; this is enforced inside the + # worker via emit_app_lifecycle (set in _build_worker), so calling start() in a + # forked child schedules heartbeats without re-emitting app-started. + # The root process defers app-started until startup configuration has been reported + # (products load + report_configuration run after enable()); see app_started(), which is + # invoked once products are loaded. (Forked children never emit app-started, so just start.) + # worker.start() must stay under this lock: it uses the local `worker` reference built + # above, and a concurrent identity refresh (which shares this lock in MicroVMs) could + # otherwise stop/replace self._worker in between, leaving this thread starting an + # already-discarded worker and marking self.started against the wrong generation. + if get_parent_runtime_id() is None: + if not self.started: + self.add_configurations(get_python_config_vars()) + else: + worker.start() + self.started = True # Subscribe before replaying, so an endpoint registered in between is forwarded twice # rather than lost; the native worker dedupes them (ASM API security). @@ -333,6 +389,13 @@ def enable(self) -> bool: def app_started(self) -> None: """Emit the root process's app-started event, exactly once.""" + if self._worker_access_lock: + with self._worker_access_lock: + self._app_started_without_lock() + else: + self._app_started_without_lock() + + def _app_started_without_lock(self) -> None: if self.started: return if not self.enable() or self._worker is None: @@ -358,19 +421,77 @@ def app_started(self) -> None: def _get_shared_worker(self): """Return the native telemetry worker for this process, so the trace exporter can report its ``trace_api.*`` health metrics through the same worker instead of spawning - a second one. + a second one. ``enable()`` is idempotent and always called to ensure existence of the worker. """ + if self._worker_access_lock: + with self._worker_access_lock: + return self._get_shared_worker_without_lock() + return self._get_shared_worker_without_lock() + + def _get_shared_worker_without_lock(self): self.enable() return self._worker + def _stop_worker(self, send_app_closing: bool, reason: str) -> None: + worker = self._worker + if worker is None: + return + self._worker = None + self.started = False + _unbind_metric_recorders(self) + self._notify_worker_changed(None) + try: + # NOTE: send_app_closing is currently ignored by the native worker + # (it always emits app-closing in the origin process); see + # TelemetryWorker.stop in ddtrace/internal/native/_native.pyi. + worker.stop(send_app_closing=send_app_closing) + except Exception: + log.debug("Failed to stop the native telemetry worker %s", reason, exc_info=True) + + def _discard_worker(self, reason: str) -> None: + """Discard a worker without flushing its queued telemetry. + + MicroVM identity refresh requires a native discard operation. Falling back to stop() would + normally flush the old runtime's queue because the current native stop API ignores its + lifecycle argument. Older native dependencies do not expose drop() on TelemetryWorker, + so stopping the worker is the only compatible way to let the application continue serving + requests while the replacement worker is built. + """ + worker = self._worker + if worker is None: + return + discard = getattr(worker, "drop", None) + if discard is None: + log.warning( + "Native TelemetryWorker does not support discard; stopping the worker %s. " + "Upgrade the native ddtrace dependency to avoid flushing stale telemetry.", + reason, + ) + self._stop_worker(False, reason) + return + discard() + self._worker = None + self.started = False + _unbind_metric_recorders(self) + self._notify_worker_changed(None) + def disable(self) -> None: """ Disable the telemetry collection service and drop the existing integrations and events Once disabled, telemetry collection can not be re-enabled. """ + if self._worker_access_lock: + with self._worker_access_lock: + self._disable_without_lock() + else: + self._disable_without_lock() + + def _disable_without_lock(self) -> None: self._enabled = False + if self._is_microvm: + remove_runtime_identity_refresh(self._refresh_runtime_identity) if endpoint_collection.on_endpoint_registered == self._record_endpoint: endpoint_collection.on_endpoint_registered = None if self._deps_collector is not None: @@ -379,21 +500,44 @@ def disable(self) -> None: except Exception: log.debug("Failed to stop the telemetry dependency collector", exc_info=True) self._deps_collector = None - if self._worker is not None: - try: - # NOTE: send_app_closing is currently ignored by the native worker - # (it always emits app-closing in the origin process); see - # TelemetryWorker.stop in ddtrace/internal/native/_native.pyi. - self._worker.stop(send_app_closing=get_parent_runtime_id() is None) - except Exception: - log.debug("Failed to stop the native telemetry worker", exc_info=True) - self._worker = None - self.started = False - _unbind_metric_recorders(self) - self._notify_worker_changed(None) + self._stop_worker(get_parent_runtime_id() is None, "during shutdown") + + def _subscribe_worker_changes( + self, + callback: "Callable[[Optional[TelemetryWorker]], None]", + expected_worker: Optional["TelemetryWorker"], + late_callback: Optional["Callable[[Optional[TelemetryWorker]], None]"] = None, + ) -> None: + if self._worker_access_lock: + with self._worker_access_lock: + self._subscribe_worker_changes_without_lock(callback, expected_worker, late_callback) + else: + self._subscribe_worker_changes_without_lock(callback, expected_worker, late_callback) + + def _unsubscribe_worker_changes(self, callback: "Callable[[Optional[TelemetryWorker]], None]") -> None: + if self._worker_access_lock: + with self._worker_access_lock: + self._unsubscribe_worker_changes_without_lock(callback) + else: + self._unsubscribe_worker_changes_without_lock(callback) + + def _unsubscribe_worker_changes_without_lock(self, callback: "Callable[[Optional[TelemetryWorker]], None]") -> None: + self._worker_subscribers[:] = [ + subscriber for subscriber in self._worker_subscribers if subscriber() not in (None, callback) + ] - def _subscribe_worker_changes(self, callback: "Callable[[Optional[TelemetryWorker]], None]") -> None: + def _subscribe_worker_changes_without_lock( + self, + callback: "Callable[[Optional[TelemetryWorker]], None]", + expected_worker: Optional["TelemetryWorker"], + late_callback: Optional["Callable[[Optional[TelemetryWorker]], None]"] = None, + ) -> None: if any(subscriber() == callback for subscriber in self._worker_subscribers): + if self._is_microvm and self._worker is not expected_worker and late_callback is not None: + try: + late_callback(self._worker) + except Exception: + log.debug("Telemetry worker subscriber failed during registration", exc_info=True) return writer_ref = weakref.ref(self) @@ -408,6 +552,15 @@ def remove_subscriber(subscriber: weakref.WeakMethod) -> None: return self._worker_subscribers.append(weakref.WeakMethod(callback, remove_subscriber)) + # The exporter is initialized with expected_worker before subscribing. A MicroVM can + # replace that worker during runtime-identity refresh between those operations, so + # synchronize a late subscriber with the current worker while holding this lock. + # Avoid invoking the callback again when the worker is unchanged. + if self._is_microvm and self._worker is not expected_worker: + try: + (late_callback or callback)(self._worker) + except Exception: + log.debug("Telemetry worker subscriber failed during registration", exc_info=True) def _notify_worker_changed(self, worker: Optional["TelemetryWorker"]) -> None: for subscriber in list(self._worker_subscribers): @@ -420,6 +573,13 @@ def _notify_worker_changed(self, worker: Optional["TelemetryWorker"]) -> None: log.debug("Telemetry worker subscriber failed", exc_info=True) def enable_agentless_client(self, enabled: bool = True) -> None: + if self._worker_access_lock: + with self._worker_access_lock: + self._enable_agentless_client_without_lock(enabled) + else: + self._enable_agentless_client_without_lock(enabled) + + def _enable_agentless_client_without_lock(self, enabled: bool = True) -> None: if self._agentless == enabled: return @@ -428,20 +588,13 @@ def enable_agentless_client(self, enabled: bool = True) -> None: if enabled and not agentless_config.api_key: log.debug("Cannot switch telemetry to agentless mode: no Datadog API key found") return - # Rebuild the worker against the new endpoint/api_key. It is called early, # before heavy traffic. + # Make sure to restart the worker if it was already running. + if self._worker is not None: - # Make sure to restart the worker if it was already running. was_started = self.started - try: - self._worker.stop(send_app_closing=False) - except Exception: - log.debug("Failed to stop the native telemetry worker during agentless switch", exc_info=True) - self._worker = None - self.started = False - _unbind_metric_recorders(self) - self._notify_worker_changed(None) + self._stop_worker(False, "during agentless switch") self.enable() if was_started: self.app_started() @@ -460,20 +613,29 @@ def add_integration( :param str integration_name: name of patched module :param bool auto_enabled: True if module is enabled in _monkey.PATCH_MODULES """ + if self._worker_access_lock: + with self._worker_access_lock: + self._add_integration_without_lock(integration_name, patched, auto_patched, error_msg, version) + else: + self._add_integration_without_lock(integration_name, patched, auto_patched, error_msg, version) + + def _add_integration_without_lock( + self, + integration_name: str, + patched: bool, + auto_patched: Optional[bool] = None, + error_msg: Optional[str] = None, + version: str = "", + ) -> None: if not self.enable() or self._worker is None: return compatible = None if error_msg is None else (error_msg == "") - self._worker.add_integration( - integration_name, - version or None, - patched, - compatible, - auto_patched, - # Preserve the failure detail so the backend keeps the message/stack for diagnosing - # patch failures; empty means "compatible, no error" -> send null. - error_msg or None, - ) + state = (version or None, patched, compatible, auto_patched, error_msg or None) + self._integrations[integration_name] = state + # Preserve the failure detail so the backend keeps the message/stack for diagnosing + # patch failures; empty means "compatible, no error" -> send null. + self._worker.add_integration(integration_name, *state) def attach_dependency_metadata( self, @@ -508,6 +670,12 @@ def _report_dependencies(self) -> Optional[list]: Returns the reported dependency records or ``None`` when nothing was reported for testing. """ + if self._worker_access_lock: + with self._worker_access_lock: + return self._report_dependencies_without_lock() + return self._report_dependencies_without_lock() + + def _report_dependencies_without_lock(self) -> Optional[list]: if not self._enabled or self._worker is None: return None deps = self._dependency_tracker.collect_report() @@ -534,6 +702,13 @@ def _record_endpoint(self, endpoint: HttpEndPoint) -> None: if not appsec_telemetry_config.ENDPOINT_COLLECTION_ENABLED or not self._enabled: return + if self._worker_access_lock: + with self._worker_access_lock: + self._record_endpoint_without_lock(endpoint) + else: + self._record_endpoint_without_lock(endpoint) + + def _record_endpoint_without_lock(self, endpoint: HttpEndPoint) -> None: worker = self._worker if worker is None: return @@ -556,6 +731,13 @@ def _report_endpoints(self) -> None: if not appsec_telemetry_config.ENDPOINT_COLLECTION_ENABLED or not self._enabled: return + if self._worker_access_lock: + with self._worker_access_lock: + self._report_endpoints_without_lock() + else: + self._report_endpoints_without_lock() + + def _report_endpoints_without_lock(self) -> None: worker = self._worker if worker is None: return @@ -571,10 +753,17 @@ def _report_endpoints(self) -> None: def product_activated(self, product: str, status: bool) -> None: """Updates the product enablement state and emits an app-product-change.""" + if self._worker_access_lock: + with self._worker_access_lock: + self._product_activated_without_lock(product, status) + else: + self._product_activated_without_lock(product, status) + + def _product_activated_without_lock(self, product: str, status: bool) -> None: if not self.enable() or self._worker is None: return - version = self._product_versions.get(product, tracer_version) - self._worker.add_product_change(product, status, version) + self._product_statuses[product] = status + self._worker.add_product_change(product, status, self._product_versions.get(product, tracer_version)) def add_configuration( self, @@ -594,14 +783,30 @@ def add_configuration( # convert unsupported types to strings configuration_value = str(configuration_value) + if self._worker_access_lock: + with self._worker_access_lock: + self._add_configuration_without_lock(configuration_name, configuration_value, origin, config_id) + else: + self._add_configuration_without_lock(configuration_name, configuration_value, origin, config_id) + + def _add_configuration_without_lock( + self, + configuration_name: str, + configuration_value: Any, + origin: str, + config_id: Optional[str], + ) -> None: if not self.enable() or self._worker is None: return seq_id = next(self._sequence_configurations) + serialized_value = _config_value_to_str(configuration_value) + if self._is_microvm: + self._configurations.append((configuration_name, serialized_value, origin, config_id)) origin_cls = _native_telemetry_enums()["origin"] self._worker.add_configuration( configuration_name, - _config_value_to_str(configuration_value), + serialized_value, getattr(origin_cls, origin, origin_cls.unknown), config_id, seq_id, @@ -609,26 +814,33 @@ def add_configuration( def add_configurations(self, configuration_list: list[tuple[str, str, str]]) -> None: """Creates and queues a list of configurations""" + if self._worker_access_lock: + with self._worker_access_lock: + self._add_configurations_without_lock(configuration_list) + else: + self._add_configurations_without_lock(configuration_list) + + def _add_configurations_without_lock(self, configuration_list: list[tuple[str, str, str]]) -> None: if not self.enable() or self._worker is None: return origin_cls = _native_telemetry_enums()["origin"] for name, value, origin in configuration_list: seq_id = next(self._sequence_configurations) + serialized_value = _config_value_to_str(value) + if self._is_microvm: + self._configurations.append((name, serialized_value, origin, None)) self._worker.add_configuration( - name, _config_value_to_str(value), getattr(origin_cls, origin, origin_cls.unknown), None, seq_id + name, serialized_value, getattr(origin_cls, origin, origin_cls.unknown), None, seq_id ) def add_log(self, level, message: str, stack_trace: str = "", tags: Optional[dict] = None) -> None: """ - Queues log. This event is meant to send library logs to Datadog's backend through the Telemetry intake. + Queues log. This event is meant to send library logs to Datadog's telemetry intake. This will make support cycles easier and ensure we know about potentially silent issues in libraries. """ if tags is None: tags = {} - if not self.enable() or self._worker is None: - return - tags_str = None if tags: tags_str = ",".join(["%s:%s" % (k, str(v).lower()) for k, v in tags.items()]) @@ -645,6 +857,17 @@ def add_log(self, level, message: str, stack_trace: str = "", tags: Optional[dic data["stack_trace"] = stack_trace identifier = hash(data) & 0xFFFFFFFFFFFFFFFF + if self._worker_access_lock: + with self._worker_access_lock: + self._add_log_without_lock(level, message, stack_trace, tags_str, identifier) + else: + self._add_log_without_lock(level, message, stack_trace, tags_str, identifier) + + def _add_log_without_lock( + self, level, message: str, stack_trace: str, tags_str: Optional[str], identifier: int + ) -> None: + if not self.enable() or self._worker is None: + return self._worker.add_log( identifier, message, @@ -743,19 +966,28 @@ def _register_metric_context( self._metric_contexts[key] = context return context - # The four ``add_*_metric`` methods inline the hot path (worker fetch + cached-context lookup - # + add_point) rather than delegating to a shared helper: metric points are recorded in tight - # loops, so avoiding the extra Python call frame per point measurably lowers the cost. + # The four ``add_*_metric`` methods keep lock acquisition at the public boundary and the hot + # metric body in a separate ``_without_lock`` helper, so MicroVM refreshes can synchronize + # worker access without duplicating the metric implementation. def add_count_metric( self, namespace: TELEMETRY_NAMESPACE, name: str, value: int = 1, tags: Optional[MetricTagType] = None ) -> None: """Queues count metric""" # Metric recording sits in hot paths (every IAST aspect, every propagation inject), so keep - # both branches lean. ``_worker`` is only ever set while enabled (``disable()`` clears it), - # so the ``_enabled`` test belongs inside this branch: the enabled path stays a single + # both wrapper paths lean. ``_worker`` is only ever set while enabled (``disable()`` clears it), + # so the ``_enabled`` test belongs inside the helper: the enabled path stays a single # attribute load, while the disabled path short-circuits without paying for an ``enable()`` # call frame on every point. + if self._worker_access_lock: + with self._worker_access_lock: + self._add_count_metric_without_lock(namespace, name, value, tags) + else: + self._add_count_metric_without_lock(namespace, name, value, tags) + + def _add_count_metric_without_lock( + self, namespace: TELEMETRY_NAMESPACE, name: str, value: int, tags: Optional[MetricTagType] + ) -> None: worker = self._worker if worker is None: if not self._enabled or not self.enable(): @@ -775,6 +1007,15 @@ def add_gauge_metric( self, namespace: TELEMETRY_NAMESPACE, name: str, value: float, tags: Optional[MetricTagType] = None ) -> None: """Queues gauge metric""" + if self._worker_access_lock: + with self._worker_access_lock: + self._add_gauge_metric_without_lock(namespace, name, value, tags) + else: + self._add_gauge_metric_without_lock(namespace, name, value, tags) + + def _add_gauge_metric_without_lock( + self, namespace: TELEMETRY_NAMESPACE, name: str, value: float, tags: Optional[MetricTagType] + ) -> None: worker = self._worker if worker is None: if not self._enabled or not self.enable(): @@ -794,6 +1035,15 @@ def add_rate_metric( self, namespace: TELEMETRY_NAMESPACE, name: str, value: float, tags: Optional[MetricTagType] = None ) -> None: """Queues rate metric""" + if self._worker_access_lock: + with self._worker_access_lock: + self._add_rate_metric_without_lock(namespace, name, value, tags) + else: + self._add_rate_metric_without_lock(namespace, name, value, tags) + + def _add_rate_metric_without_lock( + self, namespace: TELEMETRY_NAMESPACE, name: str, value: float, tags: Optional[MetricTagType] + ) -> None: worker = self._worker if worker is None: if not self._enabled or not self.enable(): @@ -813,6 +1063,15 @@ def add_distribution_metric( self, namespace: TELEMETRY_NAMESPACE, name: str, value: float, tags: Optional[MetricTagType] = None ) -> None: """Queues distributions metric""" + if self._worker_access_lock: + with self._worker_access_lock: + self._add_distribution_metric_without_lock(namespace, name, value, tags) + else: + self._add_distribution_metric_without_lock(namespace, name, value, tags) + + def _add_distribution_metric_without_lock( + self, namespace: TELEMETRY_NAMESPACE, name: str, value: float, tags: Optional[MetricTagType] + ) -> None: worker = self._worker if worker is None: if not self._enabled or not self.enable(): @@ -839,14 +1098,18 @@ def periodic(self, force_flush: bool = False) -> None: """ # Fallback trigger for the deferred root app-started (e.g. shutdown, CI visibility, tests # that flush without going through product load). No-op once already started. - self.app_started() + if self._worker_access_lock: + with self._worker_access_lock: + self._periodic_without_lock(force_flush) + else: + self._periodic_without_lock(force_flush) + def _periodic_without_lock(self, force_flush: bool = False) -> None: + self.app_started() if self._worker is None: return - if config.DEPENDENCY_COLLECTION: self._report_dependencies() - if force_flush: try: self._worker.flush() @@ -854,6 +1117,13 @@ def periodic(self, force_flush: bool = False) -> None: log.debug("Failed to flush the native telemetry worker", exc_info=True) def app_shutdown(self) -> None: + if self._worker_access_lock: + with self._worker_access_lock: + self._app_shutdown_without_lock() + else: + self._app_shutdown_without_lock() + + def _app_shutdown_without_lock(self) -> None: if self._worker is not None: # The native stop() unconditionally drains the buffer and sends an # app-closing event, so there's no need for an additional flush @@ -868,20 +1138,20 @@ def set_test_session_token(self, token: Optional[str]) -> None: The token is baked into the native worker's endpoint, so the worker is rebuilt to apply it (it is set once per test, before traffic). """ + if self._worker_access_lock: + with self._worker_access_lock: + self._set_test_session_token_without_lock(token) + else: + self._set_test_session_token_without_lock(token) + + def _set_test_session_token_without_lock(self, token: Optional[str]) -> None: self._test_session_token = token or None if not self._enabled: return # Rebuild the worker so the new token takes effect (without the # non-reversible semantics of disable()). if self._worker is not None: - try: - self._worker.stop(send_app_closing=False) - except Exception: - log.debug("Failed to stop the native telemetry worker while setting test token", exc_info=True) - self._worker = None - self.started = False - _unbind_metric_recorders(self) - self._notify_worker_changed(None) + self._stop_worker(False, "while setting test token") self.enable() def set_payload_file_dir(self, output_dir: str) -> None: @@ -893,23 +1163,23 @@ def set_payload_file_dir(self, output_dir: str) -> None: app-started with full content. app-closing is captured when the worker later stops (``app_shutdown``), which still points at the same file:// endpoint. """ + if self._worker_access_lock: + with self._worker_access_lock: + self._set_payload_file_dir_without_lock(output_dir) + else: + self._set_payload_file_dir_without_lock(output_dir) + + def _set_payload_file_dir_without_lock(self, output_dir: str) -> None: self._payload_file_dir = output_dir if not self._enabled: return was_started = self.started if self._worker is not None: - try: - self._worker.stop(send_app_closing=False) - except Exception: - log.debug("Failed to stop the native telemetry worker while enabling payload files", exc_info=True) - self._worker = None - self.started = False - _unbind_metric_recorders(self) - self._notify_worker_changed(None) + self._stop_worker(False, "while enabling payload files") self.enable() - # Re-emit app-started against the file:// worker if it had already started, so the offline - # payload directory captures the lifecycle event rather than nothing. if was_started: + # Re-emit app-started against the file:// worker if it had already started, so the offline + # payload directory captures the lifecycle event rather than nothing. self.app_started() def _restart_sequence(self) -> None: @@ -917,6 +1187,28 @@ def _restart_sequence(self) -> None: # worker owns the message-batch seq_id and resets it when rebuilt. TelemetryWriter._sequence_configurations = itertools.count(1) + def _refresh_runtime_identity(self, _runtime_id: str) -> None: + # Identity refresh is MicroVM-only; non-MicroVM writers have no worker access lock. + if self._worker_access_lock is None: + return + with self._worker_access_lock: + self._refresh_runtime_identity_without_lock() + + def _refresh_runtime_identity_without_lock(self) -> None: + if not self._enabled: + return + was_started = self.started or getattr(self, "_identity_refresh_started", False) + self._identity_refresh_started = was_started + if self._worker is not None: + self._discard_worker("while refreshing runtime identity") + self._dependency_tracker.refresh() + # The old worker is already discarded; do not report refresh success when enable() fails. + if not self.enable(): + raise RuntimeError("Failed to rebuild the telemetry worker during identity refresh") + if was_started: + self.app_started() + self._identity_refresh_started = False + def _fork_writer(self) -> None: # Runs in the child after a Python-managed fork. Drop the inherited worker handle # without shutting it down: the shared runtime is marked abandoned, and rebuilding @@ -929,6 +1221,13 @@ def _fork_writer(self) -> None: # This hook is registered before the tracer's _child_after_fork (TelemetryWriter is # constructed before the tracer), so it always runs before the trace-exporter rebuild # that calls _get_shared_worker() to get a new telemetry client. + if self._worker_access_lock: + with self._worker_access_lock: + self._fork_writer_without_lock() + else: + self._fork_writer_without_lock() + + def _fork_writer_without_lock(self) -> None: self._worker = None self.started = False # Contexts belong to the worker the parent built; the child rebuilds lazily and the diff --git a/ddtrace/internal/writer/writer.py b/ddtrace/internal/writer/writer.py index e01f2aa157d..733158aeb0d 100644 --- a/ddtrace/internal/writer/writer.py +++ b/ddtrace/internal/writer/writer.py @@ -2,6 +2,7 @@ import binascii from collections import defaultdict from collections.abc import Sequence +from functools import partial import gzip import os import socket @@ -74,6 +75,7 @@ if TYPE_CHECKING: # pragma: no cover from ddtrace.internal.http import HTTPConnection # noqa:F401 + from ddtrace.internal.telemetry.writer import TelemetryWriter from ddtrace.vendor.dogstatsd import DogStatsd @@ -896,6 +898,8 @@ def __init__( self._accepting_writes = True self._exporter_dropped = False self._owner_pid = os.getpid() + self._telemetry_writer: Optional[TelemetryWriter] = None + self._telemetry_worker_subscribed = False # Native exporter methods require exclusive access because PyO3 rejects # overlapping mutable borrows. @@ -908,6 +912,7 @@ def __del__(self) -> None: try: if getattr(self, "_owner_pid", None) != os.getpid(): return + self._unsubscribe_telemetry_worker() exporter = getattr(self, "_exporter", None) if exporter is not None and not getattr(self, "_exporter_dropped", False): self._shutdown_exporter(exporter) @@ -999,14 +1004,37 @@ def _create_exporter(self) -> native.TraceExporter: exporter = builder.build(get_native_runtime()) if shared_worker is not None: exporter.set_telemetry_handle(shared_worker) - telemetry_writer._subscribe_worker_changes(self._on_telemetry_worker_changed) + late_callback = ( + partial(self._on_telemetry_worker_changed, exporter=exporter) if telemetry_writer._is_microvm else None + ) + if telemetry_writer._is_microvm: + telemetry_writer._subscribe_worker_changes( + self._on_telemetry_worker_changed, shared_worker, late_callback + ) + else: + telemetry_writer._subscribe_worker_changes(self._on_telemetry_worker_changed, shared_worker) + self._telemetry_writer = telemetry_writer + self._telemetry_worker_subscribed = True return exporter - def _on_telemetry_worker_changed(self, worker: "Optional[native.TelemetryWorker]") -> None: + def _unsubscribe_telemetry_worker(self) -> None: + if not self._telemetry_worker_subscribed: + return + telemetry_writer = self._telemetry_writer + self._telemetry_writer = None + self._telemetry_worker_subscribed = False + if telemetry_writer is not None: + telemetry_writer._unsubscribe_worker_changes(self._on_telemetry_worker_changed) + + def _on_telemetry_worker_changed( + self, + worker: "Optional[native.TelemetryWorker]", + exporter: "Optional[native.TraceExporter]" = None, + ) -> None: """Follow the telemetry writer onto a rebuilt worker (or off a stopped one).""" try: with self._exporter_lock: - self._exporter.set_telemetry_handle(worker) + (exporter if exporter is not None else self._exporter).set_telemetry_handle(worker) except Exception: log.debug("Failed to re-point the trace exporter at the telemetry worker", exc_info=True) @@ -1039,9 +1067,11 @@ def set_test_session_token(self, token: Optional[str]) -> None: def shutdown_exporter(self) -> None: """Tear down the native exporter without going through ``stop()``.""" + self._unsubscribe_telemetry_worker() self._shutdown_exporter(self._exporter) def _drop_exporter(self) -> None: + self._unsubscribe_telemetry_worker() with self._exporter_lock: self._exporter.drop() @@ -1346,10 +1376,12 @@ def _stop_service( def on_shutdown(self): if self._exporter_dropped: + self._unsubscribe_telemetry_worker() return try: self.periodic() finally: + self._unsubscribe_telemetry_worker() self._shutdown_exporter(self._exporter) diff --git a/releasenotes/notes/fix-telemetry-worker-identity-refresh-19fe1775be48b696.yaml b/releasenotes/notes/fix-telemetry-worker-identity-refresh-19fe1775be48b696.yaml new file mode 100644 index 00000000000..b814d6f2584 --- /dev/null +++ b/releasenotes/notes/fix-telemetry-worker-identity-refresh-19fe1775be48b696.yaml @@ -0,0 +1,5 @@ +--- +fixes: + - | + internal: Fixes an issue where telemetry events can be associated with a previous + AWS Lambda MicroVM invocation after the runtime identity is refreshed. diff --git a/tests/appsec/sca/test_telemetry.py b/tests/appsec/sca/test_telemetry.py index 04e243b777f..70a0b409b18 100644 --- a/tests/appsec/sca/test_telemetry.py +++ b/tests/appsec/sca/test_telemetry.py @@ -282,6 +282,7 @@ def _make_writer_and_tracker(sca_enabled=False, deps=None, enabled=True): appsec_telemetry_config.SCA_ENABLED = sca_enabled writer = TelemetryWriter.__new__(TelemetryWriter) writer._metric_lock = MagicMock() + writer._worker_access_lock = MagicMock() writer._enabled = enabled # The native worker is mocked: _report_dependencies() forwards to worker.add_dependency and # returns the reported records, which is what these tests assert on. diff --git a/tests/telemetry/test_dependency.py b/tests/telemetry/test_dependency.py index af409957139..55afaa7dcf0 100644 --- a/tests/telemetry/test_dependency.py +++ b/tests/telemetry/test_dependency.py @@ -452,9 +452,57 @@ def test_report_dependencies_no_rereport_without_new_metadata(self): result = tracker.collect_report() - # No new deps, no new metadata -> None assert result is None + def test_refresh_preserves_and_rereports_dependency_metadata(self): + from unittest.mock import patch + + _, tracker = _make_writer_and_tracker(sca_enabled=True) + entry = DependencyEntry(name="requests", version="2.28.0", metadata=[]) + entry.add_metadata("CVE-1", "requests.sessions", "send", 10) + entry.mark_initial_sent() + entry.mark_all_metadata_sent() + tracker._imported_dependencies["requests"] = entry + + tracker.refresh() + + with ( + patch("ddtrace.internal.telemetry.dependency_tracker.modules") as mock_modules, + patch("ddtrace.internal.telemetry.dependency_tracker.telemetry_config") as mock_config, + ): + mock_config.DEPENDENCY_COLLECTION = True + mock_modules.get_newly_imported_modules.return_value = set() + + result = tracker.collect_report() + + assert result is not None + assert result[0]["name"] == "requests" + assert len(result[0]["metadata"]) == 1 + assert json.loads(result[0]["metadata"][0]["value"])["id"] == "CVE-1" + + def test_refresh_rereports_dependencies_when_sca_disabled(self): + from unittest.mock import patch + + _, tracker = _make_writer_and_tracker(sca_enabled=False) + entry = DependencyEntry(name="requests", version="2.28.0") + entry.mark_initial_sent() + tracker._imported_dependencies["requests"] = entry + + tracker.refresh() + + with ( + patch("ddtrace.internal.telemetry.dependency_tracker.modules") as mock_modules, + patch("ddtrace.internal.telemetry.dependency_tracker.telemetry_config") as mock_config, + ): + mock_config.DEPENDENCY_COLLECTION = True + mock_modules.get_newly_imported_modules.return_value = set() + + result = tracker.collect_report() + second_result = tracker.collect_report() + + assert result == [{"name": "requests", "version": "2.28.0"}] + assert second_result is None + def test_rereport_includes_all_metadata_per_rfc(self): """Re-report includes ALL metadata (sent + unsent) per RFC.""" from unittest.mock import patch diff --git a/tests/telemetry/test_writer.py b/tests/telemetry/test_writer.py index c3769e9663f..fb956758bb9 100644 --- a/tests/telemetry/test_writer.py +++ b/tests/telemetry/test_writer.py @@ -1,6 +1,8 @@ +from collections import deque import os import sys import sysconfig +import threading from typing import Any from typing import Optional from unittest import mock @@ -840,6 +842,459 @@ def test_otel_exporter_otlp_headers_telemetry_omitted(test_agent_session, run_py assert "OTEL_EXPORTER_OTLP_TIMEOUT" in configurations +def test_microvm_identity_refresh_rebuilds_worker_with_new_runtime_id(monkeypatch): + """MicroVM identity refresh must not leave native telemetry on the old runtime ID.""" + from ddtrace.internal import _runtime_id + from ddtrace.internal import runtime + import ddtrace.internal.native as native + from ddtrace.internal.telemetry.writer import TelemetryWriter + + workers = [] + original_runtime_id = runtime.get_runtime_id() + callbacks = set() + monkeypatch.setattr(_runtime_id, "_ON_RUNTIME_IDENTITY_REFRESH", callbacks) + + class FakeTelemetryWorker: + def __init__(self, native_runtime, **kwargs): + self.native_runtime = native_runtime + self.kwargs = kwargs + self.start_calls = 0 + self.drop_calls = 0 + self.stop_calls = [] + self.configurations = [] + self.integrations = [] + self.products = [] + workers.append(self) + + def start(self): + self.start_calls += 1 + + def drop(self): + self.drop_calls += 1 + + def stop(self, send_app_closing=True): + self.stop_calls.append(send_app_closing) + + def add_configuration(self, *args): + self.configurations.append(args) + + def add_integration(self, *args): + self.integrations.append(args) + + def add_product_change(self, *args): + self.products.append(args) + + def __getattr__(self, name): + def _noop(*args, **kwargs): + pass + + return _noop + + with ( + mock.patch.object(telemetry_config, "TELEMETRY_ENABLED", True), + mock.patch.object(telemetry_config, "DEPENDENCY_COLLECTION", False), + mock.patch("ddtrace.internal.telemetry.writer.in_aws_lambda_microvm", return_value=True), + mock.patch.object(native, "TelemetryWorker", FakeTelemetryWorker), + mock.patch("ddtrace.internal.native_runtime.get_native_runtime", return_value=object()), + ): + writer = TelemetryWriter(agentless=False) + monkeypatch.setattr(ddtrace.internal.telemetry, "telemetry_writer", writer) + try: + writer.app_started() + first_worker = workers[-1] + first_runtime_id = first_worker.kwargs["runtime_id"] + writer.add_integration("integration", True, True, "failure", "1.2.3") + writer.add_configuration("test.config", "before-refresh", "env_var", "test-id") + writer.add_configurations([("bulk.config", True, "env_var")]) + writer.product_activated(TELEMETRY_APM_PRODUCT.PROFILER, True) + assert first_worker.start_calls == 1 + runtime.refresh_identity() + + assert first_worker.drop_calls == 1 + assert first_worker.stop_calls == [] + assert workers[-1] is writer._worker + assert workers[-1].kwargs["runtime_id"] == runtime.get_runtime_id() + assert workers[-1].kwargs["session_id"] == runtime.get_runtime_id() + assert workers[-1].kwargs["runtime_id"] != first_runtime_id + assert workers[-1].start_calls == 1 + assert workers[-1].integrations == [("integration", "1.2.3", True, False, True, "failure")] + assert workers[-1].products == [(TELEMETRY_APM_PRODUCT.PROFILER, True, _pep440_to_semver())] + test_configurations = [args for args in workers[-1].configurations if args[0] == "test.config"] + first_test_configurations = [args for args in first_worker.configurations if args[0] == "test.config"] + assert [(args[0], args[1]) for args in test_configurations] == [("test.config", "before-refresh")] + assert test_configurations[0][4] > first_test_configurations[0][4] + bulk_configurations = [args for args in workers[-1].configurations if args[0] == "bulk.config"] + assert [(args[0], args[1]) for args in bulk_configurations] == [("bulk.config", "true")] + finally: + writer.disable() + _runtime_id._RUNTIME_ID = original_runtime_id + assert callbacks == set() + + +def test_microvm_identity_refresh_discard_failure_propagates(): + from ddtrace.internal.telemetry.writer import TelemetryWriter + + worker = mock.Mock() + worker.drop.side_effect = RuntimeError("discard failed") + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._enabled = True + writer._worker = worker + writer._worker_access_lock = threading.RLock() + writer.started = True + writer._dependency_tracker = mock.Mock() + + with pytest.raises(RuntimeError, match="discard failed"): + writer._refresh_runtime_identity("new-runtime-id") + + assert writer._worker is worker + assert writer.started is True + writer._dependency_tracker.refresh.assert_not_called() + + +def test_microvm_identity_refresh_stops_legacy_worker_without_drop(): + from ddtrace.internal.telemetry.writer import TelemetryWriter + + class LegacyTelemetryWorker: + def __init__(self): + self.stop_calls = [] + + def stop(self, send_app_closing): + self.stop_calls.append(send_app_closing) + + worker = LegacyTelemetryWorker() + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._worker = worker + writer.started = True + writer._worker_subscribers = [] + + with mock.patch("ddtrace.internal.telemetry.writer._unbind_metric_recorders"): + writer._discard_worker("while refreshing runtime identity") + + assert worker.stop_calls == [False] + assert writer._worker is None + assert writer.started is False + + +def test_microvm_identity_refresh_retries_failed_rebuild(): + from ddtrace.internal.telemetry.writer import TelemetryWriter + + lock = threading.RLock() + old_worker = mock.Mock() + replacement_worker = mock.Mock() + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._enabled = True + writer._is_microvm = True + writer._worker = old_worker + writer._enable_lock = lock + writer._worker_access_lock = lock + writer._worker_subscribers = [] + writer._metric_contexts = {} + writer.started = True + writer._deps_collector = mock.Mock() + writer._dependency_tracker = mock.Mock() + writer._report_dependencies = mock.Mock() + writer._replay_worker_state = mock.Mock() + writer.add_configurations = mock.Mock() + writer._build_worker = mock.Mock(side_effect=[RuntimeError("rebuild failed"), replacement_worker]) + + with ( + mock.patch("ddtrace.internal.telemetry.writer._bind_metric_recorders"), + mock.patch("ddtrace.internal.telemetry.writer._unbind_metric_recorders"), + ): + with pytest.raises(RuntimeError, match="Failed to rebuild the telemetry worker"): + writer._refresh_runtime_identity("new-runtime-id") + + assert writer._worker is None + assert writer._identity_refresh_started is True + + writer._refresh_runtime_identity("new-runtime-id") + + assert writer._worker is replacement_worker + replacement_worker.start.assert_called_once_with() + assert writer._identity_refresh_started is False + assert writer._build_worker.call_count == 2 + + +def test_microvm_identity_refresh_emits_new_runtime_id(telemetry_writer, test_agent_session, monkeypatch): + """Telemetry emitted after identity refresh uses the refreshed runtime ID.""" + from ddtrace.internal import _runtime_id + from ddtrace.internal import runtime + + original_runtime_id = runtime.get_runtime_id() + callbacks = set() + monkeypatch.setattr(_runtime_id, "_ON_RUNTIME_IDENTITY_REFRESH", callbacks) + if not hasattr(telemetry_writer._worker, "drop"): + pytest.skip("native telemetry worker does not provide discard support") + + telemetry_writer._is_microvm = True + telemetry_writer._worker_access_lock = telemetry_writer._enable_lock + runtime.on_runtime_identity_refresh(telemetry_writer._refresh_runtime_identity) + + try: + telemetry_writer.periodic(force_flush=True) + test_agent_session.clear() + + runtime.refresh_identity() + telemetry_writer.periodic(force_flush=True) + + app_started_events = test_agent_session.get_events("app-started") + assert len(app_started_events) == 1 + assert app_started_events[0]["runtime_id"] == runtime.get_runtime_id() + finally: + telemetry_writer.disable() + _runtime_id._RUNTIME_ID = original_runtime_id + + +def test_non_microvm_writer_does_not_register_identity_refresh(monkeypatch): + """Identity refresh callbacks are only registered for MicroVM writers.""" + from ddtrace.internal import _runtime_id + from ddtrace.internal.telemetry.writer import TelemetryWriter + + callbacks = set() + monkeypatch.setattr(_runtime_id, "_ON_RUNTIME_IDENTITY_REFRESH", callbacks) + + with ( + mock.patch.object(telemetry_config, "TELEMETRY_ENABLED", True), + mock.patch.object(TelemetryWriter, "enable", return_value=True), + mock.patch.object(TelemetryWriter, "install_excepthook"), + mock.patch("ddtrace.internal.telemetry.writer.in_aws_lambda_microvm", return_value=False), + mock.patch("ddtrace.internal.telemetry.writer.atexit.register"), + mock.patch("ddtrace.internal.telemetry.writer.forksafe.register"), + ): + writer = TelemetryWriter(agentless=False) + + assert writer._worker_access_lock is None + assert callbacks == set() + + with mock.patch.object(writer, "_refresh_runtime_identity_without_lock") as refresh: + writer._refresh_runtime_identity("new-runtime-id") + refresh.assert_not_called() + + +@pytest.mark.parametrize("is_microvm", [True, False]) +def test_worker_subscription_initial_sync_is_microvm_only(is_microvm): + class Subscriber: + def __init__(self): + self.workers = [] + + def on_worker_changed(self, worker): + self.workers.append(worker) + + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._is_microvm = is_microvm + writer._worker_access_lock = threading.RLock() if is_microvm else None + writer._worker_subscribers = [] + first_worker = object() + second_worker = object() + writer._worker = first_worker + subscriber = Subscriber() + + writer._subscribe_worker_changes(subscriber.on_worker_changed, first_worker) + writer._notify_worker_changed(second_worker) + + assert subscriber.workers == [second_worker] + + +def test_microvm_worker_subscription_syncs_replacement(): + class Subscriber: + def __init__(self): + self.workers = [] + + def on_worker_changed(self, worker): + self.workers.append(worker) + + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._is_microvm = True + writer._worker_access_lock = threading.RLock() + writer._worker_subscribers = [] + first_worker = object() + second_worker = object() + writer._worker = second_worker + subscriber = Subscriber() + + writer._subscribe_worker_changes(subscriber.on_worker_changed, first_worker) + + assert subscriber.workers == [second_worker] + + +def test_microvm_worker_subscription_can_sync_a_local_target(): + class Subscriber: + def on_worker_changed(self, worker): + pass + + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._is_microvm = True + writer._worker_access_lock = threading.RLock() + writer._worker_subscribers = [] + first_worker = object() + second_worker = object() + writer._worker = second_worker + subscriber = Subscriber() + late_callback = mock.Mock() + + writer._subscribe_worker_changes(subscriber.on_worker_changed, first_worker, late_callback) + + late_callback.assert_called_once_with(second_worker) + + +def test_non_microvm_worker_subscription_does_not_use_late_callback(): + class Subscriber: + def on_worker_changed(self, worker): + pass + + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._is_microvm = False + writer._worker_access_lock = None + writer._worker_subscribers = [] + writer._worker = object() + + subscriber = Subscriber() + late_callback = mock.Mock() + writer._subscribe_worker_changes(subscriber.on_worker_changed, object(), late_callback) + + late_callback.assert_not_called() + + +@pytest.mark.parametrize("is_microvm", [True, False]) +def test_configuration_replay_store_is_microvm_only(is_microvm): + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._is_microvm = is_microvm + writer._worker = mock.Mock() + writer._configurations = deque(maxlen=5000) + writer.enable = mock.Mock(return_value=True) + + with mock.patch( + "ddtrace.internal.telemetry.writer._native_telemetry_enums", + return_value={"origin": mock.Mock(env_var="env_var", unknown="unknown")}, + ): + writer._add_configuration_without_lock("single.config", True, "env_var", None) + writer._add_configurations_without_lock([("bulk.config", False, "env_var")]) + + assert writer._worker.add_configuration.call_count == 2 + if is_microvm: + assert list(writer._configurations) == [ + ("single.config", "true", "env_var", None), + ("bulk.config", "false", "env_var", None), + ] + else: + assert not writer._configurations + + +def test_worker_subscription_can_be_removed(): + class Subscriber: + def __init__(self): + self.workers = [] + + def on_worker_changed(self, worker): + self.workers.append(worker) + + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._is_microvm = True + writer._worker_access_lock = threading.RLock() + writer._worker_subscribers = [] + worker = object() + writer._worker = worker + subscriber = Subscriber() + + writer._subscribe_worker_changes(subscriber.on_worker_changed, worker) + writer._unsubscribe_worker_changes(subscriber.on_worker_changed) + writer._notify_worker_changed(object()) + + assert subscriber.workers == [] + + +def test_microvm_metric_recording_serializes_identity_refresh(): + """A MicroVM refresh waits for an in-flight metric write on the old worker.""" + from ddtrace.internal.telemetry.writer import TelemetryWriter + + class TrackingLock: + def __init__(self): + self._lock = threading.RLock() + self._attempts = 0 + self._attempts_lock = threading.Lock() + self.metric_acquired = threading.Event() + self.refresh_attempted = threading.Event() + + def __enter__(self): + with self._attempts_lock: + self._attempts += 1 + if self._attempts == 1: + self.metric_acquired.set() + elif self._attempts == 2: + self.refresh_attempted.set() + self._lock.acquire() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self._lock.release() + + class BlockingWorker: + def __init__(self): + self.point_entered = threading.Event() + self.release_point = threading.Event() + self.stopped = False + + def add_point(self, context, value): + self.point_entered.set() + assert self.release_point.wait(5) + assert not self.stopped + + def drop(self): + self.stopped = True + + def stop(self, send_app_closing): + self.stopped = True + + lock = TrackingLock() + worker = BlockingWorker() + writer = TelemetryWriter.__new__(TelemetryWriter) + writer._is_microvm = True + writer._enable_lock = lock + writer._worker_access_lock = lock + writer._enabled = True + writer._worker = worker + writer._metric_contexts = {} + writer._metric_lock = threading.Lock() + writer._worker_subscribers = [] + writer.started = False + writer._dependency_tracker = mock.Mock() + writer.enable = lambda: True + writer.app_started = lambda: None + + metric_errors = [] + refresh_errors = [] + + def record_metric(): + try: + writer.add_count_metric("tracers", "test.metric") + except BaseException as error: + metric_errors.append(error) + + def refresh_identity(): + try: + writer._refresh_runtime_identity("new-runtime-id") + except BaseException as error: + refresh_errors.append(error) + + with mock.patch("ddtrace.internal.telemetry.writer.register_metric_context", return_value=object()): + metric_thread = threading.Thread(target=record_metric) + metric_thread.start() + assert lock.metric_acquired.wait(5) + assert worker.point_entered.wait(5) + + refresh_thread = threading.Thread(target=refresh_identity) + refresh_thread.start() + assert lock.refresh_attempted.wait(5) + + worker.release_point.set() + metric_thread.join(5) + refresh_thread.join(5) + + assert metric_errors == [] + assert refresh_errors == [] + assert worker.stopped is True + + def test_dd_api_key_app_key_telemetry_omitted(telemetry_writer, test_agent_session): """DD_API_KEY and DD_APP_KEY values are excluded from configuration telemetry. diff --git a/tests/tracer/test_writer.py b/tests/tracer/test_writer.py index 2c9e3571fcb..cb6184b266a 100644 --- a/tests/tracer/test_writer.py +++ b/tests/tracer/test_writer.py @@ -1362,6 +1362,47 @@ def test_writer_telemetry_enabled_on_linux( mock_builder.set_restart_after_fork.assert_called_once_with(False) +def test_microvm_exporter_late_subscription_updates_local_exporter(): + import ddtrace.internal.writer.writer as writer_module + + first_worker = object() + replacement_worker = object() + exporter = mock.Mock() + + class Builder: + def __getattr__(self, name): + if name == "build": + return lambda runtime: exporter + return lambda *args, **kwargs: self + + class TelemetryWriter: + _is_microvm = True + + def _get_shared_worker(self): + return first_worker + + def _subscribe_worker_changes(self, callback, expected_worker, late_callback): + assert expected_worker is first_worker + late_callback(replacement_worker) + + def _unsubscribe_worker_changes(self, callback): + pass + + telemetry_writer = TelemetryWriter() + with ( + override_global_config(dict(_telemetry_enabled=True, _health_metrics_enabled=False)), + mock.patch.object(writer_module, "telemetry_writer", telemetry_writer), + mock.patch("ddtrace.internal.telemetry.telemetry_writer", telemetry_writer), + mock.patch.object(writer_module, "_build_base_exporter_builder", return_value=Builder()), + mock.patch.object(writer_module, "get_native_runtime", return_value=object()), + ): + writer = NativeWriter("http://localhost:8126/v0.5/traces", sync_mode=True) + + assert writer._exporter is exporter + assert exporter.set_telemetry_handle.call_args_list == [mock.call(first_worker), mock.call(replacement_worker)] + writer.shutdown_exporter() + + @pytest.mark.subprocess(err=None, env={"DD_APPSEC_ENABLED": "false"}) def test_writer_telemetry_platform_mock_does_not_rebuild_exporter_on_import_cold(): """Import-cold: override_global_config must not construct NativeWriter before the test writer."""