jscheffl commented on code in PR #63932:
URL: https://github.com/apache/airflow/pull/63932#discussion_r3156985864
##########
shared/observability/src/airflow_shared/observability/metrics/stats.py:
##########
@@ -62,72 +78,246 @@ def normalize_name_for_stats(name: str, log_warning: bool
= True) -> str:
return normalized
-class _Stats(type):
- factory: Callable[[], StatsLogger | NoStatsLogger] | None = None
- instance: StatsLogger | NoStatsLogger | None = None
- _instance_pid: int | None = None
-
- def __getattr__(cls, name: str) -> str:
- factory = type.__getattribute__(cls, "factory")
- instance = type.__getattribute__(cls, "instance")
-
- # When using OpenTelemetry, some subprocesses are short-lived and
- # often exit before flushing any metrics.
- #
- # The solution is to register a hook that performs a force flush at
exit.
- # The atexit hook is registered when initializing the instance.
- #
- # The instance gets initialized once per process. In case a process is
forked, then
- # the new subprocess, will inherit the already initialized instance of
the parent process.
- #
- # Store the instance pid so that it can be compared with the current
pid
- # to decide whether to initialize the instance again or not.
- #
- # So far, all forks are resetting their state to remove anything
inherited by the parent.
- # But in the future that might not always be true.
- current_pid = os.getpid()
- if cls.instance and cls._instance_pid != current_pid:
- log.info(
- "Stats instance was created in PID %s but accessed in PID %s.
Re-initializing.",
- cls._instance_pid,
- current_pid,
- )
- # Setting the instance to None, will force re-initialization.
- cls.instance = None
- cls._instance_pid = None
-
- if instance is None:
- if factory is None:
- factory = NoStatsLogger
- type.__setattr__(cls, "factory", factory)
-
- try:
- instance = factory()
- cls._instance_pid = current_pid
- except (socket.gaierror, ImportError) as e:
- log.error("Could not configure StatsClient: %s, using
NoStatsLogger instead.", e)
- instance = NoStatsLogger()
- cls._instance_pid = current_pid
-
- type.__setattr__(cls, "instance", instance)
-
- return getattr(instance, name)
-
- def initialize(cls, *, factory: Callable) -> None:
- type.__setattr__(cls, "instance", None)
- type.__setattr__(cls, "factory", factory)
-
- @classmethod
- def get_constant_tags(cls, *, tags_in_string: str | None) -> list[str]:
- """Get constant DataDog tags to add to all stats."""
- if not tags_in_string:
- return []
- return tags_in_string.split(",")
+def initialize(
+ *,
+ factory: Callable[[], StatsLogger | NoStatsLogger],
+ export_legacy_names: bool,
+) -> None:
+ """Initialize the stats module with a backend factory and legacy name
configuration."""
+ global _factory, _backend, _export_legacy_names
+ _factory = factory
+ _backend = None
+ _export_legacy_names = export_legacy_names
-if TYPE_CHECKING:
- Stats: StatsLogger
-else:
+def _get_backend() -> StatsLogger | NoStatsLogger:
+ """Return the current backend, creating it on first access."""
+ global _backend
+
+ if _backend is None:
+ factory = _factory if _factory is not None else NoStatsLogger
+ try:
+ _backend = factory()
+ except (socket.gaierror, ImportError) as e:
+ log.error("Could not configure StatsClient: %s, using
NoStatsLogger instead.", e)
+ _backend = NoStatsLogger()
+
+ return _backend
+
+
+def _get_registry() -> MetricsRegistry:
+ """Initialize the registry on first use to avoid import-time file I/O."""
+ global _registry
+ if _registry is None:
+ _registry = MetricsRegistry()
+ return _registry
+
+
+def _get_legacy_stat_name_and_tags(
+ stat: str, tags: dict[str, Any] | None
+) -> tuple[str | None, dict[str, Any]]:
+ """
+ Look up and format the legacy name for a metric from the registry.
+
+ Returns (formatted_name, extra_tags) where formatted_name is None when
there
+ is no legacy stat to emit, and extra_tags are the tags not consumed as name
+ variables. Raises ValueError if required name variables are missing from
tags.
+ """
+ _none: tuple[None, dict[str, Any]] = None, {}
+
+ # If the config flag is enabled/disabled.
+ if not _export_legacy_names:
+ return _none
+
+ stat_from_registry = _get_registry().get(name=stat)
+
+ # If the provided stat exists in the registry.
+ if not stat_from_registry:
+ return _none
+
+ legacy_name = stat_from_registry.get("legacy_name", "-")
+
+ # If the registry stat has a legacy name.
+ if legacy_name == "-":
+ return _none
+
+ required_vars = stat_from_registry.get("name_variables", [])
+ provided_vars = set(tags.keys()) if tags else set()
+ missing_vars = set(required_vars) - provided_vars
+ # If there are specified variables in the YAML file that haven't been
provided in the tags param.
+ if missing_vars:
+ raise ValueError(
+ f"Missing required variables for metric '{stat}':
{sorted(missing_vars)}. "
+ f"Required variables found in the registry: {required_vars}. "
+ f"Provided tags: {sorted(provided_vars)}. "
+ f"Provide all required variables as tags."
+ )
+
+ # 'required_vars' are the ones found in the registry.
+ # If there is a variable that exists in the tags and not in the registry,
+ # then it's extra, and it will be set as a tag for the legacy stat.
+ name_var_set = set(required_vars)
+ formatted_name = legacy_name.format(**{k: tags[k] for k in required_vars})
# type: ignore[index]
Review Comment:
Cool, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
