xBis7 commented on code in PR #63932:
URL: https://github.com/apache/airflow/pull/63932#discussion_r3120195163
##########
shared/observability/src/airflow_shared/observability/metrics/stats.py:
##########
@@ -62,71 +71,235 @@ 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
-
- 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
+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
- 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
+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()
- type.__setattr__(cls, "instance", instance)
+ 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
- return getattr(instance, name)
+ 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
- def initialize(cls, *, factory: Callable) -> None:
- type.__setattr__(cls, "instance", None)
- type.__setattr__(cls, "factory", factory)
+ return _backend
- @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 _get_legacy_stat(stat: str, variables: dict[str, Any]) -> str | None:
+ """
+ Look up and format the legacy name for a metric from the registry.
-if TYPE_CHECKING:
- Stats: StatsLogger
-else:
+ Returns the formatted legacy name, or None if the metric has no legacy
name.
+ Raises ValueError if the metric is not in the registry or required
variables are missing.
+ """
+ stat_from_registry = _registry.get(name=stat)
+
+ if not stat_from_registry:
+ raise ValueError(
+ f"Metric '{stat}' not found in the registry. Add the metric to the
YAML file before using it."
+ )
+
+ legacy_name = stat_from_registry.get("legacy_name", "-")
+
+ if legacy_name == "-":
+ return None
+
+ required_vars = stat_from_registry.get("name_variables", [])
+ missing_vars = set(required_vars) - set(variables.keys())
+ 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 variables: {sorted(variables.keys())}. "
+ f"Provide all required variables."
+ )
+
+ legacy_vars = {k: variables[k] for k in required_vars if k in variables}
+ return legacy_name.format(**legacy_vars)
+
+
+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 _merged(
+ tags: dict[str, Any] | None,
+ legacy_name_tags: dict[str, Any] | None,
+) -> dict[str, Any] | None:
+ """Merge tags and legacy_name_tags into a single dict, returning None if
both are empty."""
+ if not tags and not legacy_name_tags:
+ return None
+ merged: dict[str, Any] = {}
+ if tags:
+ merged.update(tags)
+ if legacy_name_tags:
+ merged.update(legacy_name_tags)
+ return merged or None
+
+
+def _emit_legacy(stat: str, legacy_name_tags: dict[str, Any]) -> str:
+ """Return the formatted legacy stat name, raising ValueError if it has
none."""
+ legacy_stat = _get_legacy_stat(stat, legacy_name_tags)
+ if legacy_stat is None:
+ raise ValueError(f"Stat '{stat}' doesn't have a legacy name registered
in the YAML file.")
+ return legacy_stat
+
+
+def incr(
+ stat: str,
+ count: int | None = None,
+ rate: int | float | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+) -> None:
+ """Increment a counter metric."""
+ if _export_legacy_names and legacy_name_tags is not None:
+ _get_backend().incr(
+ _emit_legacy(stat, legacy_name_tags), **_defined(count=count,
rate=rate, tags=tags)
+ )
+ _get_backend().incr(stat, **_defined(count=count, rate=rate,
tags=_merged(tags, legacy_name_tags)))
+
+
+def decr(
+ stat: str,
+ count: int | None = None,
+ rate: int | float | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+) -> None:
+ """Decrement a counter metric."""
+ if _export_legacy_names and legacy_name_tags is not None:
+ _get_backend().decr(
+ _emit_legacy(stat, legacy_name_tags), **_defined(count=count,
rate=rate, tags=tags)
+ )
+ _get_backend().decr(stat, **_defined(count=count, rate=rate,
tags=_merged(tags, legacy_name_tags)))
+
+
+def gauge(
+ stat: str,
+ value: float,
+ rate: int | float | None = None,
+ delta: bool | None = None,
+ *,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+) -> None:
+ """Set a gauge metric."""
+ if _export_legacy_names and legacy_name_tags is not None:
+ _get_backend().gauge(
+ _emit_legacy(stat, legacy_name_tags), value, **_defined(rate=rate,
delta=delta, tags=tags)
+ )
+ _get_backend().gauge(
+ stat, value, **_defined(rate=rate, delta=delta, tags=_merged(tags,
legacy_name_tags))
+ )
+
+
+def timing(
+ stat: str,
+ dt: DeltaType,
+ *,
+ tags: dict[str, Any] | None = None,
+ legacy_name_tags: dict[str, Any] | None = None,
+) -> None:
+ """Record a timing metric."""
+ if _export_legacy_names and legacy_name_tags is not None:
+ _get_backend().timing(_emit_legacy(stat, legacy_name_tags), dt,
**_defined(tags=tags))
+ _get_backend().timing(stat, dt, **_defined(tags=_merged(tags,
legacy_name_tags)))
+
+
+def timer(
+ stat: str | None = None,
+ tags: dict[str, Any] | None = None,
+ legacy_name_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_tags`` is provided and legacy name export is enabled,
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}
+ merged = _merged(tags, legacy_name_tags)
+ if merged:
+ regular_kw["tags"] = merged
+
+ if _export_legacy_names and legacy_name_tags is not None:
+ legacy_kw: dict[str, Any] = {**kwargs}
+ if tags is not None:
+ legacy_kw["tags"] = tags
+
+ stack = ExitStack()
+ stack.enter_context(
+ cast(
+ "AbstractContextManager[Any]",
+ _get_backend().timer(_emit_legacy(stat, legacy_name_tags),
**legacy_kw),
+ )
+ )
+ stack.enter_context(_get_backend().timer(stat, **regular_kw))
+ return cast("Timer", stack)
Review Comment:
Fixed. I went with the wrapper approach.
--
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]
