ashb commented on code in PR #63932:
URL: https://github.com/apache/airflow/pull/63932#discussion_r3137885938
##########
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:
`_ExitStackTimer` does not support the stopwatch protocol
`timer()` declares `-> Timer` as its return type, and
`Timer`/`TimerProtocol` defines a stopwatch-style API: `.start()`,
`.stop(send=True|False)`, `.duration`. The scheduler uses this pattern today:
```python
timer = Stats.timer("scheduler.critical_section_duration")
timer.start()
try:
# ...
timer.stop(send=True)
except OperationalError:
timer.stop(send=False)
```
When `timer()` returns an `_ExitStackTimer` (i.e. when the metric has a
legacy name in the registry — the `if legacy_name is not None:` branch here),
this stopwatch API is broken:
- `.stop(send=True)` — records `duration` and stops the *regular* backend
timer, but the `ExitStack` is never exited, so the legacy timer is never
stopped/sent. Resource leak + missing legacy metric.
- `.stop(send=False)` — records `duration` only. Neither timer is properly
cleaned up.
This is not a live bug today because the metrics that use stopwatch style
(`scheduler.critical_section_duration`,
`scheduler.critical_section_query_duration`) both have `legacy_name: "-"` in
the YAML, so `timer()` returns a plain backend timer, not `_ExitStackTimer`.
But if anyone later adds a legacy name to a metric that uses stopwatch style,
it will silently break.
Additionally, there is a double-start issue when `_ExitStackTimer` is used
as a context manager. `timer()` calls `stack.enter_context()` on both backend
timers, which calls `__enter__()` → `start()` on each — both are now running.
Then when the caller does `with _exit_stack_timer as t:`, `__enter__()`
(inherited from `Timer`) calls `self.start()` → `self.real_timer.start()`,
which starts the regular backend timer a **second** time (overwriting its
`_start_time` and re-starting the underlying statsd timer). The legacy timer on
the stack is *not* re-started, so the two timers measure from different start
points (legacy from `timer()` call time, regular from `with` entry time). The
delta is microseconds in practice, but it is an asymmetry that could confuse
someone reading the code.
Suggested fix: don't eagerly start the timers in `timer()`. Instead, have
`_ExitStackTimer.__enter__` create and start both timers, and store them for
`__exit__` to clean up. This would fix both the stopwatch protocol and the
double-start. Alternatively, if stopwatch mode genuinely is not needed for
dual-emit metrics, document that clearly and narrow the return type.
##########
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
Review Comment:
The PID check on every call to `_get_backend()` could be replaced by
`os.register_at_fork(after_in_child=...)` that resets `_backend = None`. Then
`_get_backend()` simplifies to just `if _backend is None` — no PID tracking
needed.
This also removes the need to call `_get_backend()` multiple times per
emission (currently happens in `incr`/`decr`/`gauge`/`timing`/`timer` when a
legacy name exists).
--
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]
