xBis7 commented on code in PR #63932:
URL: https://github.com/apache/airflow/pull/63932#discussion_r3148600576
##########
shared/observability/src/airflow_shared/observability/metrics/stats.py:
##########
@@ -62,72 +72,239 @@ 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,
+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, _instance_pid, _export_legacy_names
+ _factory = factory
+ _backend = None
+ _instance_pid = None
+ _export_legacy_names = export_legacy_names
+
+
+def _get_backend() -> StatsLogger | NoStatsLogger:
+ """Return the current backend, re-initializing if the process has been
forked."""
+ global _backend, _instance_pid
+ current_pid = os.getpid()
+
+ if _backend is not None and _instance_pid != current_pid:
+ log.info(
+ "Stats backend was created in PID %s but accessed in PID %s.
Re-initializing.",
+ _instance_pid,
+ current_pid,
+ )
+ _backend = None
+ _instance_pid = None
+
+ if _backend is None:
+ factory = _factory if _factory is not None else NoStatsLogger
+ try:
+ _backend = factory()
+ _instance_pid = current_pid
+ except (socket.gaierror, ImportError) as e:
+ log.error("Could not configure StatsClient: %s, using
NoStatsLogger instead.", e)
+ _backend = NoStatsLogger()
+ _instance_pid = current_pid
+
+ 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]
+ extra_tags = {k: v for k, v in (tags or {}).items() if k not in
name_var_set}
+ return formatted_name, extra_tags
+
+
+def _defined(**kwargs: Any) -> dict[str, Any]:
+ """Return only kwargs with a meaningful (non-None, non-empty) value."""
+ result = {}
+ for k, v in kwargs.items():
+ if v is None:
+ continue
+ try:
+ if len(v) == 0:
+ continue
+ except TypeError:
+ pass # Numerics and bools don't have len, always include
+ result[k] = v
+ return result
+
+
+def incr(
+ stat: str,
+ count: int | None = None,
+ rate: int | float | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+) -> None:
+ """Increment a counter metric."""
+ legacy_name, extra_tags = _get_legacy_stat_name_and_tags(stat, tags)
+ if legacy_name is not None:
+ _get_backend().incr(legacy_name, **_defined(count=count, rate=rate,
tags=extra_tags))
+ _get_backend().incr(stat, **_defined(count=count, rate=rate, tags=tags))
+
+
+def decr(
+ stat: str,
+ count: int | None = None,
+ rate: int | float | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+) -> None:
+ """Decrement a counter metric."""
+ legacy_name, extra_tags = _get_legacy_stat_name_and_tags(stat, tags)
+ if legacy_name is not None:
+ _get_backend().decr(legacy_name, **_defined(count=count, rate=rate,
tags=extra_tags))
+ _get_backend().decr(stat, **_defined(count=count, rate=rate, tags=tags))
+
+
+def gauge(
+ stat: str,
+ value: float,
+ rate: int | float | None = None,
+ delta: bool | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+) -> None:
+ """Set a gauge metric."""
+ legacy_name, extra_tags = _get_legacy_stat_name_and_tags(stat, tags)
+ if legacy_name is not None:
+ _get_backend().gauge(legacy_name, value, **_defined(rate=rate,
delta=delta, tags=extra_tags))
+ _get_backend().gauge(stat, value, **_defined(rate=rate, delta=delta,
tags=tags))
+
+
+def timing(
+ stat: str,
+ dt: DeltaType,
+ *,
+ tags: dict[str, Any] | None = None,
+) -> None:
+ """Record a timing metric."""
+ legacy_name, extra_tags = _get_legacy_stat_name_and_tags(stat, tags)
+ if legacy_name is not None:
+ _get_backend().timing(legacy_name, dt, **_defined(tags=extra_tags))
+ _get_backend().timing(stat, dt, **_defined(tags=tags))
+
+
+def timer(
+ stat: str | None = None,
+ tags: dict[str, Any] | None = None,
+ **kwargs,
+) -> Timer:
+ """
+ Context manager that times a block and emits a timer metric.
+
+ When ``stat`` is None, returns the backend timer directly for stopwatch
+ use (duration is available but no metric is emitted).
+
+ When legacy name export is enabled and the metric has a legacy name in the
+ registry, both the legacy metric name and the regular name are timed
+ simultaneously via an ExitStack. Otherwise, the backend timer is returned
+ directly, preserving the ability to access ``timer.duration`` on the
returned object.
+ """
+ if stat is None:
+ return _get_backend().timer()
+
+ regular_kw: dict[str, Any] = {**kwargs}
+ if tags:
+ regular_kw["tags"] = tags
+
+ legacy_name, extra_tags = _get_legacy_stat_name_and_tags(stat, tags)
+ if legacy_name is not None:
+ legacy_kw: dict[str, Any] = {**kwargs}
+ if extra_tags:
+ legacy_kw["tags"] = extra_tags
+
+ stack = ExitStack()
+ stack.enter_context(
+ cast(
+ "AbstractContextManager[Any]",
+ _get_backend().timer(legacy_name, **legacy_kw),
)
- # Setting the instance to None, will force re-initialization.
- cls.instance = None
- cls._instance_pid = None
+ )
+ real_timer = stack.enter_context(_get_backend().timer(stat,
**regular_kw))
+ return _ExitStackTimer(stack, real_timer)
Review Comment:
Fixed.
--
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]
