Support a Unix Domain Socket for the StatsD and Datadog metrics clients (#70517)

Airflow's metrics clients could only be pointed at a host and port, so
deployments that expose the daemon over a Unix Domain Socket (for example the
Datadog Kubernetes operator, which mounts a socket and sets DD_DOGSTATSD_URL)
had no way to use it from Airflow.

Add a ``[metrics] statsd_socket_path`` option, honored by both backends: the
DataDog client is built with ``socket_path`` and the StatsD client uses
``statsd.UnixSocketStatsClient``. ``statsd_host`` and ``statsd_port`` are now
nullable, and only explicitly-configured values are passed to the client; when
nothing is set the DataDog client falls back to its own environment variables
(DD_AGENT_HOST / DD_DOGSTATSD_URL). ``statsd_socket_path`` takes precedence over
``statsd_host`` / ``statsd_port`` when both are set — the DataDog client is
given all of them and prefers the socket, the StatsD client is built on the
socket — and a warning is logged in that case.

Author: sortega

airflow-core/docs/administration-and-deployment/logging-monitoring/metrics.rst

@@ -42,6 +42,36 @@ then add the following lines to your configuration file e.g. ``airflow.cfg``
     statsd_port = 8125
     statsd_prefix = airflow
 
+Sending metrics over a Unix Domain Socket
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Metrics can be sent over a Unix Domain Socket instead of UDP by setting
+``statsd_socket_path``:
+
+.. code-block:: ini
+
+    [metrics]
+    statsd_on = True
+    statsd_socket_path = /var/run/statsd/statsd.sock
+    statsd_prefix = airflow
+
+When ``statsd_socket_path`` is set, ``statsd_host``, ``statsd_port``, and
+``statsd_ipv6`` are ignored.
+
+The standard StatsD backend uses a stream Unix socket. When the Datadog backend
+is enabled, both stream and datagram Unix sockets are supported:
+
+.. code-block:: ini
+
+    [metrics]
+    statsd_datadog_enabled = True
+    statsd_socket_path = /var/run/datadog/dsd.socket
+    statsd_prefix = airflow
+
+For maximum compatibility, configure a plain filesystem path. The Datadog
+backend additionally accepts ``unix://``, ``unixgram://``, and
+``unixstream://`` URLs.
+
 If you want to use a custom StatsD client instead of the default one provided by Airflow,
 the following key must be added to the configuration file alongside the module path of your
 custom StatsD client. This module must be available on your :envvar:`PYTHONPATH`.
@@ -51,6 +81,10 @@ custom StatsD client. This module must be available on your :envvar:`PYTHONPATH`
     [metrics]
     statsd_custom_client_path = x.y.customclient
 
+When ``statsd_socket_path`` is configured, a custom client must inherit from
+``statsd.UnixSocketStatsClient``. Otherwise, it must inherit from
+``statsd.StatsClient``.
+
 See :doc:`../modules_management` for details on how Python and Airflow manage modules.
 
 .. note::

airflow-core/src/airflow/config_templates/config.yml

@@ -1206,7 +1206,8 @@ metrics:
       default: "False"
     statsd_host:
       description: |
-        Specifies the host address where the StatsD daemon (or server) is running
+        Specifies the host address where the StatsD daemon (or server) is running.
+        Ignored when ``statsd_socket_path`` is set.
       version_added: 2.0.0
       type: string
       example: ~
@@ -1220,11 +1221,25 @@ metrics:
       default: "False"
     statsd_port:
       description: |
-        Specifies the port on which the StatsD daemon (or server) is listening to
+        Specifies the port on which the StatsD daemon (or server) is listening to.
+        Ignored when ``statsd_socket_path`` is set.
       version_added: 2.0.0
       type: integer
       example: ~
       default: "8125"
+    statsd_socket_path:
+      description: |
+        Path to a Unix Domain Socket used to submit metrics, instead of ``statsd_host``
+        and ``statsd_port``. Applies to both the StatsD and Datadog backends. Takes
+        precedence over ``statsd_host`` / ``statsd_port`` when set. For the Datadog
+        backend (``statsd_datadog_enabled = True``) a plain path or a ``unix://`` /
+        ``unixgram://`` / ``unixstream://`` URL is accepted; for the StatsD backend a
+        plain filesystem path is expected. Note that ``statsd_ipv6`` has no effect when
+        this option is set.
+      version_added: 3.4.0
+      type: string
+      example: "/var/run/statsd/statsd.sock"
+      default: ~
     statsd_prefix:
       description: |
         Defines the namespace for all metrics sent from Airflow to StatsD

airflow-core/src/airflow/observability/metrics/datadog_logger.py

@@ -30,6 +30,7 @@ def get_dogstatsd_logger() -> SafeDogStatsdLogger:
         tags_in_string=conf.get("metrics", "statsd_datadog_tags"),
         host=conf.get("metrics", "statsd_host"),
         port=conf.getint("metrics", "statsd_port"),
+        socket_path=conf.get("metrics", "statsd_socket_path", fallback=None) or None,
         namespace=conf.get("metrics", "statsd_prefix"),
         datadog_metrics_tags=conf.getboolean("metrics", "statsd_datadog_metrics_tags", fallback=True),
         statsd_disabled_tags=conf.get("metrics", "statsd_disabled_tags", fallback=None),

airflow-core/src/airflow/observability/metrics/statsd_logger.py

@@ -23,40 +23,48 @@
 from airflow._shared.observability.metrics import statsd_logger
 from airflow.configuration import conf
 
+log = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from airflow._shared.observability.metrics.statsd_logger import SafeStatsdLogger
 
-log = logging.getLogger(__name__)
-
 
 def get_statsd_logger() -> SafeStatsdLogger:
-    stats_class = conf.getimport("metrics", "statsd_custom_client_path", fallback=None)
+    # Local import to avoid requiring statsd when other backends are used (e.g. Datadog)
+    from statsd import StatsClient, UnixSocketStatsClient
 
-    # no need to check for the scheduler/statsd_on -> this method is only called when it is set
-    # and previously it would crash with None is callable if it was called without it.
-    from statsd import StatsClient
+    socket_path = conf.get("metrics", "statsd_socket_path", fallback=None) or None
+    custom_class = conf.getimport("metrics", "statsd_custom_client_path", fallback=None)
 
-    if stats_class:
-        if not issubclass(stats_class, StatsClient):
+    if custom_class is not None:
+        if socket_path is not None:
+            if not issubclass(custom_class, UnixSocketStatsClient):
+                raise AirflowConfigException(
+                    "Your custom StatsD client must extend the statsd.UnixSocketStatsClient "
+                    "when using a socket path in order to ensure backwards compatibility."
+                )
+        elif not issubclass(custom_class, StatsClient):
             raise AirflowConfigException(
-                "Your custom StatsD client must extend the statsd.StatsClient in order to ensure "
-                "backwards compatibility."
+                "Your custom StatsD client must extend the statsd.StatsClient in order "
+                "to ensure backwards compatibility."
             )
         log.info("Successfully loaded custom StatsD client")
 
-    else:
-        stats_class = StatsClient
-
-    return statsd_logger.get_statsd_logger(
-        stats_class=stats_class,
+    common_kwargs = {
+        "stats_class": custom_class,
+        "prefix": conf.get("metrics", "statsd_prefix"),
+        "influxdb_tags_enabled": conf.getboolean("metrics", "statsd_influxdb_enabled", fallback=False),
+        "statsd_disabled_tags": conf.get("metrics", "statsd_disabled_tags", fallback=None),
+        "metrics_allow_list": conf.get("metrics", "metrics_allow_list", fallback=None),
+        "metrics_block_list": conf.get("metrics", "metrics_block_list", fallback=None),
+        "stat_name_handler": conf.getimport("metrics", "stat_name_handler"),
+        "statsd_influxdb_enabled": conf.getboolean("metrics", "statsd_influxdb_enabled", fallback=False),
+    }
+    if socket_path is not None:
+        return statsd_logger.get_socket_statsd_logger(socket_path=socket_path, **common_kwargs)
+    return statsd_logger.get_udp_statsd_logger(
         host=conf.get("metrics", "statsd_host"),
         port=conf.getint("metrics", "statsd_port"),
-        prefix=conf.get("metrics", "statsd_prefix"),
         ipv6=conf.getboolean("metrics", "statsd_ipv6", fallback=False),
-        influxdb_tags_enabled=conf.getboolean("metrics", "statsd_influxdb_enabled", fallback=False),
-        statsd_disabled_tags=conf.get("metrics", "statsd_disabled_tags", fallback=None),
-        metrics_allow_list=conf.get("metrics", "metrics_allow_list", fallback=None),
-        metrics_block_list=conf.get("metrics", "metrics_block_list", fallback=None),
-        stat_name_handler=conf.getimport("metrics", "stat_name_handler"),
-        statsd_influxdb_enabled=conf.getboolean("metrics", "statsd_influxdb_enabled", fallback=False),
+        **common_kwargs,
     )

airflow-core/tests/unit/observability/metrics/test_stats.py

@@ -39,23 +39,35 @@ def __init__(self, host=None, port=None, prefix=None):
 
 
 class TestStats:
-    def test_load_invalid_custom_stats_client(self):
-        with conf_vars(
-            {
-                ("metrics", "statsd_on"): "True",
-                ("metrics", "statsd_custom_client_path"): f"{__name__}.InvalidCustomStatsd",
-            }
-        ):
+    @pytest.mark.parametrize(
+        ("extra_conf", "expected_message"),
+        [
+            (
+                {},
+                "Your custom StatsD client must extend the statsd.StatsClient in order "
+                "to ensure backwards compatibility.",
+            ),
+            (
+                {("metrics", "statsd_socket_path"): "/var/run/datadog/dsd.socket"},
+                "Your custom StatsD client must extend the statsd.UnixSocketStatsClient "
+                "when using a socket path in order to ensure backwards compatibility.",
+            ),
+        ],
+        ids=["udp", "socket"],
+    )
+    def test_load_invalid_custom_stats_client(self, extra_conf, expected_message):
+        conf = {
+            ("metrics", "statsd_on"): "True",
+            ("metrics", "statsd_custom_client_path"): f"{__name__}.InvalidCustomStatsd",
+            **extra_conf,
+        }
+        with conf_vars(conf):
             importlib.reload(airflow._shared.observability.metrics.stats)
             airflow.observability.stats.initialize(
                 factory=stats_utils.get_stats_factory(), export_legacy_names=True
             )
-            error_message = re.escape(
-                "Your custom StatsD client must extend the statsd."
-                "StatsClient in order to ensure backwards compatibility."
-            )
             # we assert for Exception here instead of AirflowConfigException to not import from shared configuration
-            with pytest.raises(Exception, match=error_message):
+            with pytest.raises(Exception, match=re.escape(expected_message)):
                 airflow.observability.stats.incr("empty_key")
         importlib.reload(airflow._shared.observability.metrics.stats)
 

docs/spelling_wordlist.txt

@@ -396,6 +396,8 @@ dataframes
 DataFusion
 Datafusion
 datafusion
+Datagram
+datagram
 Datalake
 datamodel
 Datanodes

shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py

@@ -148,6 +148,7 @@ def get_dogstatsd_logger(
     tags_in_string: str | None = None,
     host: str | None = None,
     port: int | None = None,
+    socket_path: str | None = None,
     namespace: str | None = None,
     datadog_metrics_tags: bool = True,
     statsd_disabled_tags: str | None = None,
@@ -156,7 +157,7 @@ def get_dogstatsd_logger(
     stat_name_handler: Callable[[str], str] | None = None,
     statsd_influxdb_enabled: bool = False,
 ) -> SafeDogStatsdLogger:
-    """Get DataDog StatsD logger."""
+    """Get Datadog StatsD logger."""
     from datadog import DogStatsd
 
     dogstatsd_kwargs: dict[str, Any] = {
@@ -166,6 +167,8 @@ def get_dogstatsd_logger(
         dogstatsd_kwargs["host"] = host
     if port is not None:
         dogstatsd_kwargs["port"] = port
+    if socket_path is not None:
+        dogstatsd_kwargs["socket_path"] = socket_path
     if namespace is not None:
         dogstatsd_kwargs["namespace"] = namespace
     dogstatsd = DogStatsd(**dogstatsd_kwargs)

shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py

@@ -22,6 +22,8 @@
 from functools import wraps
 from typing import TYPE_CHECKING, TypeVar, cast
 
+from statsd import StatsClient, UnixSocketStatsClient
+
 from .protocols import Timer
 from .validators import (
     PatternAllowListValidator,
@@ -31,8 +33,6 @@
 )
 
 if TYPE_CHECKING:
-    from statsd import StatsClient
-
     from .protocols import DeltaType
     from .validators import ListValidator
 
@@ -156,11 +156,32 @@ def timer(
         return Timer()
 
 
-def get_statsd_logger(
+def _make_safe_statsd_logger(
+    statsd_client: StatsClient,
+    influxdb_tags_enabled: bool,
+    statsd_disabled_tags: str | None,
+    metrics_allow_list: str | None,
+    metrics_block_list: str | None,
+    stat_name_handler: Callable[[str], str] | None,
+    statsd_influxdb_enabled: bool,
+) -> SafeStatsdLogger:
+    metric_tags_validator = PatternBlockListValidator(statsd_disabled_tags)
+    validator = get_validator(metrics_allow_list, metrics_block_list)
+    return SafeStatsdLogger(
+        statsd_client,
+        validator,
+        influxdb_tags_enabled,
+        metric_tags_validator,
+        stat_name_handler,
+        statsd_influxdb_enabled,
+    )
+
+
+def get_udp_statsd_logger(
     *,
-    stats_class: type[StatsClient],
-    host: str | None = None,
-    port: int | None = None,
+    stats_class: type[StatsClient] | None = None,
+    host: str,
+    port: int,
     prefix: str | None = None,
     ipv6: bool = False,
     influxdb_tags_enabled: bool = False,
@@ -170,21 +191,41 @@ def get_statsd_logger(
     stat_name_handler: Callable[[str], str] | None = None,
     statsd_influxdb_enabled: bool = False,
 ) -> SafeStatsdLogger:
-    """Return logger for StatsD."""
-    statsd = stats_class(
-        host=host,
-        port=port,
-        prefix=prefix,
-        ipv6=ipv6,
+    """Return a UDP StatsD logger."""
+    resolved: type[StatsClient] = stats_class if stats_class is not None else StatsClient
+    statsd_client = resolved(host=host, port=port, prefix=prefix, ipv6=ipv6)
+    return _make_safe_statsd_logger(
+        statsd_client,
+        influxdb_tags_enabled,
+        statsd_disabled_tags,
+        metrics_allow_list,
+        metrics_block_list,
+        stat_name_handler,
+        statsd_influxdb_enabled,
     )
 
-    metric_tags_validator = PatternBlockListValidator(statsd_disabled_tags)
-    validator = get_validator(metrics_allow_list, metrics_block_list)
-    return SafeStatsdLogger(
-        statsd,
-        validator,
+
+def get_socket_statsd_logger(
+    *,
+    stats_class: type[UnixSocketStatsClient] | None = None,
+    socket_path: str,
+    prefix: str | None = None,
+    influxdb_tags_enabled: bool = False,
+    statsd_disabled_tags: str | None = None,
+    metrics_allow_list: str | None = None,
+    metrics_block_list: str | None = None,
+    stat_name_handler: Callable[[str], str] | None = None,
+    statsd_influxdb_enabled: bool = False,
+) -> SafeStatsdLogger:
+    """Return a Unix Domain Socket StatsD logger. Note: statsd_ipv6 has no effect in socket mode."""
+    resolved: type[UnixSocketStatsClient] = stats_class if stats_class is not None else UnixSocketStatsClient
+    statsd_client = resolved(socket_path=socket_path, prefix=prefix)
+    return _make_safe_statsd_logger(
+        statsd_client,
         influxdb_tags_enabled,
-        metric_tags_validator,
+        statsd_disabled_tags,
+        metrics_allow_list,
+        metrics_block_list,
         stat_name_handler,
         statsd_influxdb_enabled,
     )