# Source code for distributed.spans

    from __future__ import annotations

    import copy
    import uuid
    import weakref
    from collections import defaultdict
    from collections.abc import Hashable, Iterable, Iterator, Mapping
    from contextlib import contextmanager
    from itertools import islice
    from typing import TYPE_CHECKING, Any, TypedDict

    import dask.config
    from dask.typing import Key

    from distributed.collections import sum_mappings
    from distributed.itertools import ffill
    from distributed.metrics import time

    if TYPE_CHECKING:
        # Needed to avoid Sphinx WARNING: more than one target found for cross-reference
        # 'TaskState' and 'WorkerState'"
        # https://github.com/agronholm/sphinx-autodoc-typehints#dealing-with-circular-imports
        from distributed import Scheduler, Worker
        from distributed import scheduler as scheduler_module
        from distributed.client import SourceCode
        from distributed.scheduler import TaskGroup, TaskStateState

    CONTEXTS_WITH_SPAN_ID = ("execute", "p2p")

    class SpanMetadata(TypedDict):
        collections: list[dict]

    @contextmanager
    def span(*tags: str) -> Iterator[str]:
        """Tag group of tasks to be part of a certain group, called a span.

        This context manager can be nested, thus creating sub-spans. If you close and
        re-open a span context manager with the same tag, you'll end up with two separate
        spans.

        Every cluster defines a global "default" span when no span has been defined by the
        client; the default span is automatically closed and reopened when all tasks
        associated to it have been completed; in other words the cluster is idle save for
        tasks that are explicitly annotated by a span. Note that, in some edge cases, you
        may end up with overlapping default spans, e.g. if a worker crashes and all unique
        tasks that were in memory on it need to be recomputed.

        You may capture the ID of a span on the client to match it with the
        :class:`~distributed.spans.Span` objects the scheduler:

        >>> client = Client()
        >>> with span("my workflow") as span_id:
        ...     client.submit(lambda: "Hello world!").result()
        >>> client.cluster.scheduler.extensions["spans"].spans[span_id]
        Span<name=('my workflow',), id=5dc9b908-116b-49a5-b0d7-5a681f49a111>

        Notes
        -----
        You may retrieve the current span with ``dask.get_annotations().get("span")``.
        You can do so in the client code as well as from inside a task.
        """
        if not tags:
            raise ValueError("Must specify at least one span tag")

        annotation = dask.get_annotations().get("span")
        prev_tags = annotation["name"] if annotation else ()
        # You must specify the full history of IDs, not just the parent, because
        # otherwise you would not be able to uniquely identify grandparents when
        # they have no tasks of their own.
        prev_ids = annotation["ids"] if annotation else ()
        ids = tuple(str(uuid.uuid4()) for _ in tags)
        with dask.annotate(span={"name": prev_tags + tags, "ids": prev_ids + ids}):
            yield ids[-1]

    class Span:
        #: (<tag>, <tag>, ...)
        #: Matches ``TaskState.annotations["span"]["name"]``, both on the scheduler and the
        #: worker.
        name: tuple[str, ...]

        #: Unique ID, generated by :func:`~distributed.span` and
        #: taken from ``TaskState.annotations["span"]["id"][-1]``.
        #: Matches ``distributed.scheduler.TaskState.group.span_id``
        #: and ``distributed.worker_state_machine.TaskState.span_id``.
        id: str

        _parent: weakref.ref[Span] | None

        #: Direct children of this span, sorted by creation time
        children: list[Span]

        #: Task groups *directly* belonging to this span.
        #:
        #: See Also
        #: --------
        #  traverse_groups
        #:
        #: Notes
        #: -----
        #: TaskGroups are forgotten by the Scheduler when the last task is forgotten, but
        #: remain referenced here indefinitely. If a user calls compute() twice on the same
        #: collection, you'll have more than one group with the same tg.name in this set!
        #: For the same reason, while the same TaskGroup object is guaranteed to be attached
        #: to exactly one Span, you may have different TaskGroups with the same key attached
        #: to different Spans.
        groups: set[TaskGroup]

        #: Time when the span first appeared on the scheduler.
        #: The same property on parent spans is always less than or equal to this.
        #:
        #: See Also
        #: --------
        #: start
        #: stop
        enqueued: float

        #: Source code snippets, if it was sent by the client.
        #: We're using a dict without values as an insertion-sorted set.
        _code: dict[tuple[SourceCode, ...], None]
        _metadata: SpanMetadata | None

        _cumulative_worker_metrics: defaultdict[tuple[Hashable, ...], float]

        #: reference to SchedulerState.total_nthreads_history
        _total_nthreads_history: list[tuple[float, int]]
        #: Length of total_nthreads_history when this span was enqueued
        _total_nthreads_offset: int

        # Support for weakrefs to a class with __slots__
        __weakref__: Any

        __slots__ = tuple(__annotations__)
        _metadata_seen: set[int] = set()

        def __init__(
            self,
            name: tuple[str, ...],
            id_: str,
            parent: Span | None,
            total_nthreads_history: list[tuple[float, int]],
        ):
            self.name = name
            self.id = id_
            self._parent = weakref.ref(parent) if parent is not None else None
            self.enqueued = time()
            self.children = []
            self.groups = set()
            self._code = {}
            self._metadata = None

            # Don't cast int metrics to float
            self._cumulative_worker_metrics = defaultdict(int)

            assert len(total_nthreads_history) > 0
            self._total_nthreads_history = total_nthreads_history
            self._total_nthreads_offset = len(total_nthreads_history) - 1

        def __repr__(self) -> str:
            return f"Span<name={self.name}, id={self.id}>"

        @property
        def parent(self) -> Span | None:
            if self._parent:
                out = self._parent()
                assert out
                return out
            return None

        def add_metadata(self, metadata: SpanMetadata) -> None:
            """Add metadata to the span, e.g. code snippets"""
            id_ = id(metadata)
            if id_ in self._metadata_seen:
                return
            self._metadata_seen.add(id_)
            if self._metadata is None:
                self._metadata = copy.deepcopy(metadata)
            else:
                self._metadata["collections"].extend(metadata["collections"])

        @property
        def annotation(self) -> dict[str, tuple[str, ...]] | None:
            """Rebuild the dask graph annotation which contains the full id history

            Note that this may not match the original annotation in case of TaskGroup
            collision.
            """
            if self.name == ("default",):
                return None
            ids = []
            node: Span | None = self
            while node:
                ids.append(node.id)
                node = node.parent
            return {"name": self.name, "ids": tuple(reversed(ids))}

        def traverse_spans(self) -> Iterator[Span]:
            """Top-down recursion of all spans belonging to this branch off span tree,
            including self
            """
            yield self
            for child in self.children:
                yield from child.traverse_spans()

        def traverse_groups(self) -> Iterator[TaskGroup]:
            """All TaskGroups belonging to this branch of span tree"""
            for span in self.traverse_spans():
                yield from span.groups

        @property
        def start(self) -> float:
            """Earliest time when a task belonging to this span tree started computing;
            0 if no task has *finished* computing yet.

            Notes
            -----
            This is not updated until at least one task has *finished* computing.
            It could move backwards as tasks complete.

            See Also
            --------
            enqueued
            stop
            distributed.scheduler.TaskGroup.start
            """
            out = min(
                (tg.start for tg in self.traverse_groups() if tg.start != 0.0),
                default=0.0,
            )
            if out:
                # absorb small errors in worker delay calculation
                out = max(out, self.enqueued)
            return out

        @property
        def stop(self) -> float:
            """When this span tree finished computing, or current timestamp if it didn't
            finish yet.

            Notes
            -----
            This differs from ``TaskGroup.stop`` when there aren't unfinished tasks; is also
            will never be zero.

            See Also
            --------
            enqueued
            start
            done
            distributed.scheduler.TaskGroup.stop
            """
            if self.done:
                out = max(tg.stop for tg in self.traverse_groups())
            else:
                out = time()
            # absorb small errors in worker delay calculation, as well as in time() not
            # being perfectly monotonic
            return max(out, self.enqueued)

        @property
        def metadata(self) -> SpanMetadata | None:
            return self._metadata

        @property
        def states(self) -> dict[TaskStateState, int]:
            """The number of tasks currently in each state in this span tree;
            e.g. ``{"memory": 10, "processing": 3, "released": 4, ...}``.

            See Also
            --------
            distributed.scheduler.TaskGroup.states
            """
            return sum_mappings(tg.states for tg in self.traverse_groups())

        @property
        def done(self) -> bool:
            """Return True if all tasks in this span tree are completed; False otherwise.

            Notes
            -----
            This property may transition from True to False, e.g. when a new sub-span is
            added or when a worker that contained the only replica of a task in memory
            crashes and the task need to be recomputed.

            See Also
            --------
            distributed.scheduler.TaskGroup.done
            """
            return all(tg.done for tg in self.traverse_groups())

        @property
        def all_durations(self) -> dict[str, float]:
            """Cumulative duration of all completed actions in this span tree, by action

            See Also
            --------
            duration
            distributed.scheduler.TaskGroup.all_durations
            """
            return sum_mappings(tg.all_durations for tg in self.traverse_groups())

        @property
        def duration(self) -> float:
            """The total amount of time spent on all tasks in this span tree

            See Also
            --------
            all_durations
            distributed.scheduler.TaskGroup.duration
            """
            return sum(tg.duration for tg in self.traverse_groups())

        @property
        def nbytes_total(self) -> int:
            """The total number of bytes that this span tree has produced

            See Also
            --------
            distributed.scheduler.TaskGroup.nbytes_total
            """
            return sum(tg.nbytes_total for tg in self.traverse_groups())

        @property
        def code(self) -> list[tuple[SourceCode, ...]]:
            """Code snippets, sent by the client on compute(), persist(), and submit().

            Only populated if ``distributed.diagnostics.computations.nframes`` is non-zero.
            """
            # Deduplicate, but preserve order
            return list(
                dict.fromkeys(sc for child in self.traverse_spans() for sc in child._code)
            )

        @property
        def cumulative_worker_metrics(self) -> dict[tuple[Hashable, ...], float]:
            """Replica of ``Worker.digests_total`` and
            ``Scheduler.cumulative_worker_metrics``, but only for the metrics that can be
            attributed to the current span tree. The span id has been removed from the key.

            At the moment of writing, all keys are
            ``("execute", <task prefix>, <activity>, <unit>)``
            or
            ``("p2p", <where>, <activity>, <unit>)``
            but more may be added in the future with a different format; please test e.g.
            for ``k[0] == "execute"``.
            """
            out = sum_mappings(
                child._cumulative_worker_metrics for child in self.traverse_spans()
            )
            known_seconds = sum(
                v for k, v in out.items() if k[0] == "execute" and k[-1] == "seconds"
            )
            # Besides rounding errors, you may get negative unknown seconds if a user
            # manually invokes `context_meter.digest_metric`.
            unknown_seconds = max(0.0, self.active_cpu_seconds - known_seconds)

            out["execute", "N/A", "idle or other spans", "seconds"] = unknown_seconds
            return out

        @staticmethod
        def merge(*items: Span) -> Span:
            """Merge multiple spans into a synthetic one.
            The input spans must not be related with each other.
            """
            if not items:
                raise ValueError("Nothing to merge")
            out = Span(
                name=("(merged)",),
                id_="(merged)",
                parent=None,
                total_nthreads_history=items[0]._total_nthreads_history,
            )
            out._total_nthreads_offset = min(
                child._total_nthreads_offset for child in items
            )
            out.children.extend(items)
            out.enqueued = min(child.enqueued for child in items)
            return out

        def _nthreads_timeseries(self) -> Iterator[tuple[float, int]]:
            """Yield (timestamp, number of threads across the cluster), forward-fill"""
            stop = self.stop if self.done else 0
            for t, n in islice(
                self._total_nthreads_history, self._total_nthreads_offset, None
            ):
                if stop and t >= stop:
                    break
                yield max(self.enqueued, t), n

        def _active_timeseries(self) -> Iterator[tuple[float, bool]]:
            """If this span is the output of :meth:`merge`, yield
            (timestamp, True if at least one input span is active), forward-fill.
            """
            if self.id != "(merged)":
                yield self.enqueued, True
                yield self.stop, False
                return

            events = []
            for child in self.children:
                events += [(child.enqueued, 1), (child.stop, -1)]
            # enqueued <= stop by construction.
            # Occasionally, enqueued == stop, e.g. when the clock is adjusted backwards.
            # Prevent negative n_active when this happens.
            events.sort(key=lambda el: el[0])

            n_active = 0
            for t, delta in events:
                if not n_active:
                    assert delta == 1
                    yield t, True
                n_active += delta
                if not n_active:
                    yield t, False

        @property
        def nthreads_intervals(self) -> list[tuple[float, float, int]]:
            """
            Returns
            -------
            List of tuples:

            - begin timestamp
            - end timestamp
            - Scheduler.total_nthreads during this interval

            When the Span is the output of :meth:`merge`, the intervals may not be
            contiguous.

            See Also
            --------
            enqueued
            stop
            active_cpu_seconds
            distributed.scheduler.SchedulerState.total_nthreads
            """
            nthreads_t, nthreads_count = zip(*self._nthreads_timeseries())
            is_active_t, is_active_flag = zip(*self._active_timeseries())
            t_interp = sorted({*nthreads_t, *is_active_t})
            nthreads_count_interp = ffill(t_interp, nthreads_t, nthreads_count, left=0)
            is_active_flag_interp = ffill(t_interp, is_active_t, is_active_flag, left=False)
            return [
                (t0, t1, n)
                for t0, t1, n, active in zip(
                    t_interp, t_interp[1:], nthreads_count_interp, is_active_flag_interp
                )
                if active
            ]

        @property
        def active_cpu_seconds(self) -> float:
            """Return number of CPU seconds that were made available on the cluster while
            this Span was running; in other words
            ``(Span.stop - Span.enqueued) * Scheduler.total_nthreads``.

            This accounts for workers joining and leaving the cluster while this Span was
            active. If this Span is the output of :meth:`merge`, do not count gaps between
            input spans.

            See Also
            --------
            enqueued
            stop
            nthreads_intervals
            distributed.scheduler.SchedulerState.total_nthreads
            """
            return sum((t1 - t0) * nthreads for t0, t1, nthreads in self.nthreads_intervals)

    class SpansSchedulerExtension:
        """Scheduler extension for spans support"""

        scheduler: Scheduler

        #: All Span objects by id
        spans: dict[str, Span]

        #: Only the spans that don't have any parents, sorted by creation time.
        #: This is a convenience helper structure to speed up searches.
        root_spans: list[Span]

        #: All spans, keyed by their full name and sorted by creation time.
        #: This is a convenience helper structure to speed up searches.
        spans_search_by_name: defaultdict[tuple[str, ...], list[Span]]

        #: All spans, keyed by the individual tags that make up their name and sorted by
        #: creation time.
        #: This is a convenience helper structure to speed up searches.
        #:
        #: See Also
        #: --------
        #: find_by_tags
        #: merge_by_tags
        spans_search_by_tag: defaultdict[str, list[Span]]

        def __init__(self, scheduler: Scheduler):
            self.scheduler = scheduler
            self.spans = {}
            self.root_spans = []
            self.spans_search_by_name = defaultdict(list)
            self.spans_search_by_tag = defaultdict(list)

        def observe_tasks(
            self,
            tss: Iterable[scheduler_module.TaskState],
            code: tuple[SourceCode, ...],
            span_metadata: SpanMetadata,
        ) -> dict[Key, dict]:
            """Acknowledge the existence of runnable tasks on the scheduler. These may
            either be new tasks, tasks that were previously unrunnable, or tasks that were
            already fed into this method already.

            Attach newly observed tasks to either the desired span or to ("default", ).
            Update TaskGroup.span_id and wipe TaskState.annotations["span"].

            Returns
            -------
            Updated 'span' annotations: {key: {"name": (..., ...), "ids": (..., ...)}}
            """
            out = {}
            default_span = None

            for ts in tss:
                if ts.annotations is None:
                    ts.annotations = dict()
                # You may have different tasks belonging to the same TaskGroup but to
                # different spans. If that happens, arbitrarily force everything onto the
                # span of the earliest encountered TaskGroup.
                tg = ts.group
                if tg.span_id:
                    span = self.spans[tg.span_id]
                else:
                    ann = ts.annotations.get("span")
                    if ann:
                        span = self._ensure_span(ann["name"], ann["ids"])
                    else:
                        if not default_span:
                            default_span = self._ensure_default_span()
                        span = default_span

                    tg.span_id = span.id
                    span.groups.add(tg)

                if code:
                    span._code[code] = None
                if span_metadata:
                    span.add_metadata(span_metadata)

                # The span may be completely different from the one referenced by the
                # annotation, due to the TaskGroup collision issue explained above.
                if ann := span.annotation:
                    ts.annotations["span"] = out[ts.key] = ann
                else:
                    ts.annotations.pop("span", None)

            return out

        def _ensure_default_span(self) -> Span:
            """Return the currently active default span, or create one if the previous one
            terminated. In other words, do not reuse the previous default span if all tasks
            that were not explicitly annotated with :func:`spans` on the client side are
            finished.
            """
            defaults = self.spans_search_by_name["default",]
            if defaults and not defaults[-1].done:
                return defaults[-1]
            return self._ensure_span(("default",), (str(uuid.uuid4()),))

        def _ensure_span(self, name: tuple[str, ...], ids: tuple[str, ...]) -> Span:
            """Create Span if it doesn't exist and return it"""
            try:
                return self.spans[ids[-1]]
            except KeyError:
                pass

            assert len(name) == len(ids)
            assert len(name) > 0

            parent = None
            for i in range(1, len(name)):
                parent = self._ensure_span(name[:i], ids[:i])

            span = Span(
                name=name,
                id_=ids[-1],
                parent=parent,
                total_nthreads_history=self.scheduler.total_nthreads_history,
            )
            self.spans[span.id] = span
            self.spans_search_by_name[name].append(span)
            for tag in name:
                self.spans_search_by_tag[tag].append(span)
            if parent:
                parent.children.append(span)
            else:
                self.root_spans.append(span)

            return span

        def find_by_tags(self, *tags: str) -> Iterator[Span]:
            """Yield all spans that contain any of the given tags.
            When a tag is shared both by a span and its (grand)children, only return the
            parent.
            """
            by_level = defaultdict(list)
            for tag in tags:
                for sp in self.spans_search_by_tag[tag]:
                    by_level[len(sp.name)].append(sp)

            seen = set()
            for _, level in sorted(by_level.items()):
                seen.update(level)
                for sp in level:
                    if sp.parent not in seen:
                        yield sp

        def merge_all(self) -> Span:
            """Return a synthetic Span which is the sum of all spans"""
            return Span.merge(*self.root_spans)

        def merge_by_tags(self, *tags: str) -> Span:
            """Return a synthetic Span which is the sum of all spans containing the given
            tags
            """
            return Span.merge(*self.find_by_tags(*tags))

        def heartbeat(
            self, ws: scheduler_module.WorkerState, data: dict[tuple[Hashable, ...], float]
        ) -> None:
            """Triggered by :meth:`SpansWorkerExtension.heartbeat`.

            Populate :meth:`Span.cumulative_worker_metrics` with data from the worker.

            See Also
            --------
            SpansWorkerExtension.heartbeat
            Span.cumulative_worker_metrics
            """
            for (context, span_id, *other), v in data.items():
                assert isinstance(span_id, str)
                span = self.spans[span_id]
                span._cumulative_worker_metrics[(context, *other)] += v

    class SpansWorkerExtension:
        """Worker extension for spans support"""

        worker: Worker
        digests_total_since_heartbeat: dict[tuple[Hashable, ...], float]

        def __init__(self, worker: Worker):
            self.worker = worker
            self.digests_total_since_heartbeat = {}

        def collect_digests(
            self, digests_total_since_heartbeat: Mapping[Hashable, float]
        ) -> None:
            # Note: this method may be called spuriously by Worker._register_with_scheduler,
            # but when it does it's guaranteed not to find any metrics
            assert not self.digests_total_since_heartbeat
            self.digests_total_since_heartbeat = {
                k: v
                for k, v in digests_total_since_heartbeat.items()
                if isinstance(k, tuple) and k[0] in CONTEXTS_WITH_SPAN_ID
            }

        def heartbeat(self) -> dict[tuple[Hashable, ...], float]:
            """Apportion the metrics that do have a span to the Spans on the scheduler

            Returns
            -------
            ``{(context, span_id, prefix, activity, unit): value}}``

            See Also
            --------
            SpansSchedulerExtension.heartbeat
            Span.cumulative_worker_metrics
            distributed.worker.Worker.get_metrics
            """
            out = self.digests_total_since_heartbeat
            self.digests_total_since_heartbeat = {}
            return out