# Source code for distributed.metrics

    from __future__ import annotations

    import asyncio
    import collections
    import sys
    import threading
    import time as timemod
    from collections.abc import Callable, Hashable, Iterator
    from contextlib import contextmanager
    from contextvars import ContextVar
    from dataclasses import dataclass
    from functools import wraps
    from math import nan
    from typing import Literal

    import psutil

    from distributed.compatibility import WINDOWS

    _empty_namedtuple = collections.namedtuple("_empty_namedtuple", ())

    def _psutil_caller(method_name, default=_empty_namedtuple):
        """
        Return a function calling the given psutil *method_name*,
        or returning *default* if psutil fails.
        """
        meth = getattr(psutil, method_name)

        @wraps(meth)
        def wrapper():  # pragma: no cover
            try:
                return meth()
            except RuntimeError:
                # This can happen on some systems (e.g. no physical disk in worker)
                return default()

        return wrapper

    disk_io_counters = _psutil_caller("disk_io_counters")

    net_io_counters = _psutil_caller("net_io_counters")

    class _WindowsTime:
        """Combine time.time() or time.monotonic() with time.perf_counter() to get an
        absolute clock with fine resolution.
        """

        base_timer: Callable[[], float]
        delta: float
        previous: float | None
        next_resync: float
        resync_every: float

        def __init__(
            self, base: Callable[[], float], is_monotonic: bool, resync_every: float = 600.0
        ):
            self.base_timer = base
            self.previous = float("-inf") if is_monotonic else None
            self.next_resync = float("-inf")
            self.resync_every = resync_every

        def time(self) -> float:
            cur = timemod.perf_counter()
            if cur > self.next_resync:
                self.resync()
                self.next_resync = cur + self.resync_every
            cur += self.delta
            if self.previous is not None:
                # Monotonic timer
                if cur <= self.previous:
                    cur = self.previous + 1e-9
                self.previous = cur
            return cur

        def resync(self) -> None:
            _time = self.base_timer
            _perf_counter = timemod.perf_counter
            min_samples = 5
            while True:
                times = [(_time(), _perf_counter()) for _ in range(min_samples * 2)]
                abs_times = collections.Counter(t[0] for t in times)
                first, nfirst = abs_times.most_common()[0]
                if nfirst < min_samples:
                    # System too noisy? Start again
                    continue

                perf_times = [t[1] for t in times if t[0] == first][:-1]
                assert len(perf_times) >= min_samples - 1, perf_times
                self.delta = first - sum(perf_times) / len(perf_times)
                break

    # A high-resolution wall clock timer measuring the seconds since Unix epoch
    if WINDOWS and sys.version_info < (3, 13):
        time = _WindowsTime(timemod.time, is_monotonic=False).time
        monotonic = _WindowsTime(timemod.monotonic, is_monotonic=True).time
    else:
        # Under modern Unixes, time.time() and time.monotonic() should be good enough
        time = timemod.time
        monotonic = timemod.monotonic

    process_time = timemod.process_time

    # Get a per-thread CPU timer function if possible, otherwise
    # use a per-process CPU timer function.
    try:
        # thread_time is not supported on all platforms
        thread_time = timemod.thread_time
    except (AttributeError, OSError):  # pragma: no cover
        thread_time = process_time

    @dataclass
    class MeterOutput:
        start: float
        stop: float
        delta: float
        __slots__ = tuple(__annotations__)

    @contextmanager
    def meter(
        func: Callable[[], float] = timemod.perf_counter,
        floor: float | Literal[False] = 0.0,
    ) -> Iterator[MeterOutput]:
        """Convenience context manager which calls func() before and after the wrapped
        code and calculates the delta.

        Parameters
        ----------
        label: str
            label to pass to the callback
        func: callable
            function to call before and after, which must return a number.
            Besides time, it could return e.g. cumulative network traffic or disk usage.
            Default: :func:`timemod.perf_counter`
        floor: float or False, optional
            Floor the delta to the given value (default: 0). This is useful for strictly
            cumulative functions that can occasionally glitch and go backwards.
            Set to False to disable.

        Yields
        ------
        :class:`MeterOutput` where the ``start`` attribute is populated straight away, while
        ``stop`` and ``delta`` are nan until context exit.
        """
        out = MeterOutput(func(), nan, nan)
        try:
            yield out
        finally:
            out.stop = func()
            out.delta = out.stop - out.start
            if floor is not False:
                out.delta = max(floor, out.delta)

    class ContextMeter:
        """Context-based general purpose meter.

        Usage
        -----
        1. In high level code, call :meth:`add_callback` to install a hook that defines an
           activity
        2. In low level code, typically many stack levels below, log quantitative events
           (e.g. elapsed time, transferred bytes, etc.) so that they will be attributed to
           the high-level code calling it, either with :meth:`meter`,
           :meth:`meter_function`, or :meth:`digest_metric`.

        Examples
        --------
        In the code that e.g. sends a Python object from A to B over the network:
        >>> from distributed.metrics import context_meter
        >>> with context_meter.add_callback(partial(print, "A->B comms:")):
        ...     await send_over_the_network(obj)

        In the serialization utilities, called many stack levels below:
        >>> with context_meter.meter("dumps"):
        ...     pik = pickle.dumps(obj)
        >>> with context_meter.meter("compress"):
        ...     pik = lz4.compress(pik)

        And finally, elsewhere, deep into the TCP stack:
        >>> with context_meter.meter("network-write"):
        ...     await comm.write(frames)

        When you call the top-level code, you'll get::
          A->B comms: dumps 0.012 seconds
          A->B comms: compress 0.034 seconds
          A->B comms: network-write 0.567 seconds
        """

        _callbacks: ContextVar[dict[Hashable, Callable[[Hashable, float, str], None]]]

        def __init__(self):
            self._callbacks = ContextVar(
                f"MetricHook<{id(self)}>._callbacks", default={}  # noqa: B039
            )

        def __reduce__(self):
            assert self is context_meter, "Found copy of singleton"
            return self._unpickle_singleton, ()

        @staticmethod
        def _unpickle_singleton():
            return context_meter

        @contextmanager
        def add_callback(
            self,
            callback: Callable[[Hashable, float, str], None],
            *,
            key: Hashable | None = None,
            allow_offload: bool = False,
        ) -> Iterator[None]:
            """Add a callback when entering the context and remove it when exiting it.
            The callback must accept the same parameters as :meth:`digest_metric`.

            Parameters
            ----------
            callback: Callable
                ``f(label, value, unit)`` to be executed
            key: Hashable, optional
                Unique key for the callback. If two nested calls to ``add_callback`` use the
                same key, suppress the outermost callback.
            allow_offload: bool, optional
                If set to True, this context must be executed inside a running asyncio
                event loop. If a call to :meth:`digest_metric` is performed from a different
                thread, e.g. from inside :func:`distributed.utils.offload`, ensure that
                the callback is executed in the event loop's thread instead.
            """
            if allow_offload:
                loop = asyncio.get_running_loop()
                tid = threading.get_ident()

                def safe_cb(label: Hashable, value: float, unit: str, /) -> None:
                    if threading.get_ident() == tid:
                        callback(label, value, unit)
                    else:  # We're inside offload()
                        loop.call_soon_threadsafe(callback, label, value, unit)

            else:
                safe_cb = callback

            if key is None:
                key = object()
            cbs = self._callbacks.get()
            cbs = cbs.copy()
            cbs[key] = safe_cb
            tok = self._callbacks.set(cbs)
            try:
                yield
            finally:
                tok.var.reset(tok)

        @contextmanager
        def clear_callbacks(self) -> Iterator[None]:
            """Do not trigger any callbacks set outside of this context"""
            tok = self._callbacks.set({})
            try:
                yield
            finally:
                tok.var.reset(tok)

        def digest_metric(self, label: Hashable, value: float, unit: str) -> None:
            """Invoke the currently set context callbacks for an arbitrary quantitative
            metric.
            """
            cbs = self._callbacks.get()
            for cb in cbs.values():
                cb(label, value, unit)

        @contextmanager
        def meter(
            self,
            label: Hashable,
            unit: str = "seconds",
            func: Callable[[], float] = timemod.perf_counter,
            floor: float | Literal[False] = 0.0,
        ) -> Iterator[MeterOutput]:
            """Convenience context manager or decorator which calls func() before and after
            the wrapped code, calculates the delta, and finally calls :meth:`digest_metric`.

            If unit=='seconds', it also subtracts any other calls to :meth:`meter` or
            :meth:`digest_metric` with the same unit performed within the context, so that
            the total is strictly additive.

            Parameters
            ----------
            label: Hashable
                label to pass to the callback
            unit: str, optional
                unit to pass to the callback. Default: seconds
            func: callable
                see :func:`meter`
            floor: bool, optional
                see :func:`meter`

            Yields
            ------
            :class:`MeterOutput` where the ``start`` attribute is populated straight away,
            while ``stop`` and ``delta`` are nan until context exit. In case of multiple
            nested calls to :meth:`meter`, then delta (for seconds only) is reduced by the
            inner metrics, to a minimum of ``floor``.
            """
            if unit != "seconds":
                try:
                    with meter(func, floor=floor) as m:
                        yield m
                finally:
                    self.digest_metric(label, m.delta, unit)
                return

            # If unit=="seconds", subtract time metered from the sub-contexts
            offsets = []

            def callback(label2: Hashable, value2: float, unit2: str) -> None:
                if unit2 == unit:
                    # This must be threadsafe to support callbacks invoked from
                    # distributed.utils.offload; '+=' on a float would not be threadsafe!
                    offsets.append(value2)

            try:
                with self.add_callback(callback), meter(func, floor=False) as m:
                    yield m
            finally:
                delta = m.delta - sum(offsets)
                if floor is not False:
                    delta = max(floor, delta)
                m.delta = delta
                self.digest_metric(label, delta, unit)

    context_meter = ContextMeter()

    class DelayedMetricsLedger:
        """Add-on to :class:`ContextMeter` that helps in the case where:

        - The code to be metered is not easily expressed as a self-contained code block
          e.g. you want to measure latency in the asyncio event loop before and after
          running a task
        - You want to alter the metrics depending on how the code ends; e.g. you want to
          post them differently in case of failure.

        Examples
        --------
        >>> ledger = DelayedMetricsLedger()  # Metering starts here
        >>> async def wrapper():
        ...     with ledger.record():
        ...         return await metered_function()
        >>> task = asyncio.create_task(wrapper())
        >>> # (later, elsewhere)
        >>> try:
        ...     await task
        ...     coarse_time = False
        ... except Exception:
        ...     coarse_time = "failed"
        ...     raise
        ... finally:
        ...     # Metering stops here
        ...     for label, value, unit in ledger.finalize(coarse_time):
        ...         # actually log metrics
        """

        func: Callable[[], float]
        start: float
        metrics: list[tuple[Hashable, float, str]]  # (label, value, unit)

        def __init__(self, func: Callable[[], float] = timemod.perf_counter):
            self.func = func
            self.start = func()
            self.metrics = []

        def _callback(self, label: Hashable, value: float, unit: str) -> None:
            self.metrics.append((label, value, unit))

        @contextmanager
        def record(self, *, key: Hashable | None = None) -> Iterator[None]:
            """Ingest metrics logged with :meth:`ContextMeter.digest_metric` or
            :meth:`ContextMeter.meter` and temporarily store them in :ivar:`metrics`.

            Parameters
            ----------
            key: Hashable, optional
                See :meth:`ContextMeter.add_callback`
            """
            with context_meter.add_callback(self._callback, key=key):
                yield

        def finalize(
            self,
            coarse_time: str | Literal[False] = False,
            floor: float | Literal[False] = 0.0,
        ) -> Iterator[tuple[Hashable, float, str]]:
            """The metered code is terminated, and we now know how to log it.

            Parameters
            ----------
            coarse_time: str | False, optional
                False
                    Yield all acquired metrics, plus an extra time metric, labelled "other",
                    which is the time between creating the DelayedMetricsLedger and
                    calling this method, minus any time logged in the metrics.
                label
                    Yield all acquired non-time metrics.
                    Yield a single metric, labelled <coarse_time>, which is the time
                    between creating the DelayedMetricsLedger and calling this method.
            floor: float | False, optional
                Floor either the "other" or the <coarse_time> metric to this value
                 (default: 0). Set to False to disable.
            """
            stop = self.func()
            delta = stop - self.start

            for label, value, unit in self.metrics:
                if unit != "seconds" or not coarse_time:
                    yield label, value, unit
                if unit == "seconds" and not coarse_time:
                    delta -= value

            if floor is not False:
                delta = max(floor, delta)
            yield coarse_time or "other", delta, "seconds"