Repository: dask/distributed

Issue:
Metrics introduce non-trivial overhead
I recently noticed that our per-task overhead increased to an extend that is not critical but still concerning.

To test this, I ran a very simple workload

```python
    from distributed import Client
    with Client() as c:
        def inc(x):
            return x + 1
        futs = c.map(inc, range(10_000))
```

First run, this shows quite some whitespace

![image](https://github.com/dask/distributed/assets/8629629/91cabe1b-8573-421b-bb86-2d24b11a4290)

but I anticipate that most of this whitespace is due to `worker-saturation` being set and the mapped tasks are actually detected as root tasks. Workers are just starving.

Disabling worker saturation again reveals a bit of a different problem with metrics. Running a py-spy profile indicates that this line is taking up about 5% of CPU time on the main thread (and another ~2-5% on the worker thread)

https://github.com/dask/distributed/blob/e0dbb35e8d4ae12081bb36e51d62008051a9d2c5/distributed/metrics.py#L325

This is not great but the actual problem is that this is also blocking the event loop, increasing all kinds of latencies and overheads as well.

Digging a bit into why this is so expensive, I noticed that the number of callbacks registered here
https://github.com/dask/distributed/blob/e0dbb35e8d4ae12081bb36e51d62008051a9d2c5/distributed/metrics.py#L212-L217

is increasing monotonically for the duration of this compute until it drops at the very end. In this rather trivial example, I get about 1200 callbacks that are evaluated on every `ContextMeter.digest_metric`.

A side-effect of this is that the dashboard much more often shows `deserialize-**` (grey) blocks because [this](https://github.com/dask/distributed/blob/e0dbb35e8d4ae12081bb36e51d62008051a9d2c5/distributed/worker.py#L2205-L2212) typically trivial task deserialization now takes nontrivial time since `_deserialize` is instrumented. (In my example deserializing `inc` supposedly takes 30-120ms)

cc @crusaderky 

Discussion:
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/distributed/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

       20 files  ±    0         20 suites  ±0   10h 39m 44s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "duration of all tests") - 50m 50s
  3 701 tests +    3    3 591 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "passed tests") +    4     106 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "skipped / disabled tests") ±  0  4 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "failed tests")  - 1 
34 974 runs   - 792  33 262 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "passed tests")  - 749  1 708 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "skipped / disabled tests")  - 40  4 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "failed tests")  - 3 

For more details on these failures, see [this check](https://github.com/dask/distributed/runs/14769143610).

Results for commit 6a595fe4. ± Comparison against base commit 9b9f948b.

[test-results]:data:application/gzip;base64,H4sIAAYkpGQC/02NywrDIBBFfyW47sL3oz8T1ESQJrEYXZX+eychsW4Gzrlw5oNCXOYdPQeKHwPaayx/mmq2JaYNmGmmOShYy7EzhcmN4169P5wwnXvFNziCZVPBxgVUy4xzzimDOX7lup1dbhS/sHUZlbS5q6uwvlXXPbnP+rSusQAgaYURYebMBUktFswIK7niysAljkxBW+cDQd8fDnzowxQBAAA=


Patch:
diff --git a/distributed/dashboard/components/scheduler.py b/distributed/dashboard/components/scheduler.py
index a274c7ec..c80758cc 100644
--- a/distributed/dashboard/components/scheduler.py
+++ b/distributed/dashboard/components/scheduler.py
@@ -3392,6 +3392,7 @@ class FinePerformanceMetrics(DashboardComponent):
     substantial_change: bool
     visible_functions: list[str]
     visible_activities: list[str]
+    stacked_chart_visible_activities: list[str]
     function_selector: MultiChoice
     span_tag_selector: MultiChoice
     unit_selector: Select
@@ -3409,6 +3410,7 @@ class FinePerformanceMetrics(DashboardComponent):
         self.substantial_change = False
         self.visible_functions = []
         self.visible_activities = []
+        self.stacked_chart_visible_activities = []
         self.task_exec_by_activity_chart = None
         self.task_exec_by_prefix_chart = None
         self.get_data_by_activity_chart = None
@@ -3651,7 +3653,6 @@ class FinePerformanceMetrics(DashboardComponent):
             self.visible_functions = sorted(visible_functions)
 
         if visible_activities != set(self.visible_activities):
-            self.substantial_change = True
             self.visible_activities = sorted(visible_activities)
 
         (
@@ -3668,18 +3669,29 @@ class FinePerformanceMetrics(DashboardComponent):
         --------
         _build_pie_chart
         """
-        total_value = sum(data.values())
-        percent_k = 100.0 / total_value if total_value else 0.0
-        angle_k = 2.0 * math.pi / total_value if total_value else 0.0
         activities = self.visible_activities
+        # Generate palette based on all visible activities to make sure that
+        # colors match between the plots
+        palette = self._get_palette()
         values = [data[activity] for activity in activities]
-        total_text = self._format(sum(values))
+
+        # Sort by values from largest to smallest
+        # Hide activities that are missing in the current plot from the legend
+        idx = [i for i, v in sorted(enumerate(values), key=lambda el: -el[1]) if v]
+        activities = [activities[i] for i in idx]
+        palette = [palette[i] for i in idx]
+        values = [values[i] for i in idx]
+
+        total_value = sum(values)
+        total_text = self._format(total_value)
+        percent_k = 100.0 / total_value if total_value else 0.0
+        angle_k = 2.0 * math.pi / total_value if total_value else 0.0
         return {
             "activity": activities,
             "value": values,
             "text": [self._format(v) + f" ({v * percent_k:.0f}%)" for v in values],
             "angle": [v * angle_k for v in values],
-            "color": self._get_palette(),
+            "color": palette,
             "total_text": [total_text] * len(values),
         }
 
@@ -3735,13 +3747,21 @@ class FinePerformanceMetrics(DashboardComponent):
             "____total_text": [self._format(v) for v in func_totals],
         }
 
+        sc_visible_activities = []
         for activity in self.visible_activities:
             values = [data[function, activity] for function in self.visible_functions]
+            if not any(values):
+                continue
+            sc_visible_activities.append(activity)
             out[activity] = values
             out[f"__{activity}_text"] = [
                 self._format(v) + f" ({v * perc_ki:.0f}%)"
                 for v, perc_ki in zip(values, perc_k)
             ]
+        if sc_visible_activities != self.stacked_chart_visible_activities:
+            self.substantial_change = True
+            self.stacked_chart_visible_activities[:] = sc_visible_activities
+
         return out, max(func_totals, default=0.0)
 
     def _build_task_execution_by_prefix_chart(self) -> figure:
@@ -3779,13 +3799,20 @@ class FinePerformanceMetrics(DashboardComponent):
         barchart = self.task_exec_by_prefix_chart
         assert barchart is not None
         barchart.x_range = FactorRange(*self.visible_functions)
+
+        palette = [
+            p
+            for p, a in zip(self._get_palette(), self.visible_activities)
+            if a in self.stacked_chart_visible_activities
+        ]
+        assert len(palette) == len(self.stacked_chart_visible_activities)
         renderers = barchart.vbar_stack(
-            self.visible_activities,
+            self.stacked_chart_visible_activities,
             x="__functions",
             width=0.9,
             source=self.task_exec_by_prefix_src,
-            color=self._get_palette(),
-            legend_label=self.visible_activities,
+            color=palette,
+            legend_label=self.stacked_chart_visible_activities,
         )
 
         # Create or refresh hovertools on top of base tools
diff --git a/distributed/metrics.py b/distributed/metrics.py
index 2783e5b3..30a95552 100755
--- a/distributed/metrics.py
+++ b/distributed/metrics.py
@@ -189,10 +189,10 @@ class ContextMeter:
       A->B comms: network-write 0.567 seconds
     """
 
-    _callbacks: ContextVar[list[Callable[[Hashable, float, str], None]]]
+    _callbacks: ContextVar[dict[Hashable, Callable[[Hashable, float, str], None]]]
 
     def __init__(self):
-        self._callbacks = ContextVar(f"MetricHook<{id(self)}>._callbacks", default=[])
+        self._callbacks = ContextVar(f"MetricHook<{id(self)}>._callbacks", default={})
 
     def __reduce__(self):
         assert self is context_meter, "Found copy of singleton"
@@ -204,13 +204,28 @@ class ContextMeter:
 
     @contextmanager
     def add_callback(
-        self, callback: Callable[[Hashable, float, str], None]
+        self,
+        callback: Callable[[Hashable, float, str], None],
+        *,
+        key: Hashable | None = None,
     ) -> 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.
         """
+        if key is None:
+            key = object()
         cbs = self._callbacks.get()
-        tok = self._callbacks.set(cbs + [callback])
+        cbs = cbs.copy()
+        cbs[key] = callback
+        tok = self._callbacks.set(cbs)
         try:
             yield
         finally:
@@ -221,7 +236,7 @@ class ContextMeter:
         metric.
         """
         cbs = self._callbacks.get()
-        for cb in cbs:
+        for cb in cbs.values():
             cb(label, value, unit)
 
     @contextmanager
@@ -234,9 +249,10 @@ class ContextMeter:
     ) -> 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`.
-        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.
+
+        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
         ----------
@@ -256,10 +272,19 @@ class ContextMeter:
         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 == "seconds":
+            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)
@@ -316,14 +341,20 @@ class DelayedMetricsLedger:
         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) -> Iterator[None]:
+    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(
-            lambda label, value, unit: self.metrics.append((label, value, unit))
-        ):
+        with context_meter.add_callback(self._callback, key=key):
             yield
 
     def finalize(
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 567662c0..40655c6f 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -4884,6 +4884,7 @@ class Scheduler(SchedulerState, ServerNode):
         exception=None,
         stimulus_id=None,
         traceback=None,
+        run_id=None,
         **kwargs,
     ):
         """Mark that a task has erred on a particular worker"""
@@ -4893,6 +4894,11 @@ class Scheduler(SchedulerState, ServerNode):
         if ts is None or ts.state != "processing":
             return {}, {}, {}
 
+        if ts.run_id != run_id:
+            if ts.processing_on and ts.processing_on.address == worker:
+                return self._transition(key, "released", stimulus_id)
+            return {}, {}, {}
+
         if ts.retries > 0:
             ts.retries -= 1
             return self._transition(key, "waiting", stimulus_id)
diff --git a/distributed/worker.py b/distributed/worker.py
index 7d89998c..d588c74a 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -2242,6 +2242,7 @@ class Worker(BaseWorker, ServerNode):
             return ExecuteFailureEvent.from_exception(
                 exc,
                 key=key,
+                run_id=run_id,
                 stimulus_id=f"run-spec-deserialize-failed-{time()}",
             )
 
@@ -2366,6 +2367,7 @@ class Worker(BaseWorker, ServerNode):
             return ExecuteFailureEvent.from_exception(
                 result,
                 key=key,
+                run_id=run_id,
                 start=result["start"],
                 stop=result["stop"],
                 stimulus_id=f"task-erred-{time()}",
@@ -2376,6 +2378,7 @@ class Worker(BaseWorker, ServerNode):
             return ExecuteFailureEvent.from_exception(
                 exc,
                 key=key,
+                run_id=run_id,
                 stimulus_id=f"execute-unknown-error-{time()}",
             )
 
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
index 504eddfe..db32765d 100644
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -482,6 +482,7 @@ class TaskErredMsg(SendMessageToScheduler):
     op = "task-erred"
 
     key: str
+    run_id: int
     exception: Serialize
     traceback: Serialize | None
     exception_text: str
@@ -497,11 +498,12 @@ class TaskErredMsg(SendMessageToScheduler):
 
     @staticmethod
     def from_task(
-        ts: TaskState, stimulus_id: str, thread: int | None = None
+        ts: TaskState, run_id: int, stimulus_id: str, thread: int | None = None
     ) -> TaskErredMsg:
         assert ts.exception
         return TaskErredMsg(
             key=ts.key,
+            run_id=run_id,
             exception=ts.exception,
             traceback=ts.traceback,
             exception_text=ts.exception_text,
@@ -903,6 +905,7 @@ class ExecuteSuccessEvent(ExecuteDoneEvent):
 
 @dataclass
 class ExecuteFailureEvent(ExecuteDoneEvent):
+    run_id: int  # FIXME: Utilize the run ID in all ExecuteDoneEvents
     start: float | None
     stop: float | None
     exception: Serialize
@@ -921,6 +924,7 @@ class ExecuteFailureEvent(ExecuteDoneEvent):
         err_or_msg: BaseException | ErrorMessage,
         *,
         key: str,
+        run_id: int,
         start: float | None = None,
         stop: float | None = None,
         stimulus_id: str,
@@ -932,6 +936,7 @@ class ExecuteFailureEvent(ExecuteDoneEvent):
 
         return cls(
             key=key,
+            run_id=run_id,
             start=start,
             stop=stop,
             exception=msg["exception"],
@@ -945,6 +950,7 @@ class ExecuteFailureEvent(ExecuteDoneEvent):
     def dummy(
         key: str,
         *,
+        run_id: int = 1,
         stimulus_id: str,
     ) -> ExecuteFailureEvent:
         """Build a dummy event, with most attributes set to a reasonable default.
@@ -952,6 +958,7 @@ class ExecuteFailureEvent(ExecuteDoneEvent):
         """
         return ExecuteFailureEvent(
             key=key,
+            run_id=run_id,
             start=None,
             stop=None,
             exception=Serialize(None),
@@ -2025,6 +2032,7 @@ class WorkerState:
         traceback: Serialize | None,
         exception_text: str,
         traceback_text: str,
+        run_id: int,
         *,
         stimulus_id: str,
     ) -> RecsInstrs:
@@ -2035,6 +2043,7 @@ class WorkerState:
         ts.state = "error"
         smsg = TaskErredMsg.from_task(
             ts,
+            run_id=run_id,
             stimulus_id=stimulus_id,
             thread=self.threads.get(ts.key),
         )
@@ -2048,6 +2057,7 @@ class WorkerState:
         traceback: Serialize | None,
         exception_text: str,
         traceback_text: str,
+        run_id: int,
         *,
         stimulus_id: str,
     ) -> RecsInstrs:
@@ -2438,7 +2448,7 @@ class WorkerState:
                 # Third-party MutableMappings (dask-cuda etc.) may have other use cases
                 # for this.
                 msg = error_message(e)
-                return {ts: tuple(msg.values())}, []
+                return {ts: tuple(msg.values()) + (run_id,)}, []
 
             stop = time()
             if stop - start > 0.005:
@@ -2859,7 +2869,9 @@ class WorkerState:
                 )
             )
         elif ts.state == "error":
-            instructions.append(TaskErredMsg.from_task(ts, stimulus_id=ev.stimulus_id))
+            instructions.append(
+                TaskErredMsg.from_task(ts, run_id=ev.run_id, stimulus_id=ev.stimulus_id)
+            )
         elif ts.state in {
             "released",
             "fetch",
@@ -3042,6 +3054,7 @@ class WorkerState:
                 ev.traceback,
                 ev.exception_text,
                 ev.traceback_text,
+                ts.run_id,
             )
             for ts in self._gather_dep_done_common(ev)
         }
@@ -3179,6 +3192,7 @@ class WorkerState:
             ev.traceback,
             ev.exception_text,
             ev.traceback_text,
+            ev.run_id,
         )
         return recs, instr
 
@@ -3621,7 +3635,7 @@ class BaseWorker(abc.ABC):
 
         @wraps(func)
         async def wrapper() -> StateMachineEvent:
-            with ledger.record():
+            with ledger.record(key="async-instruction"):
                 return await func(*args, **kwargs)
 
         task = asyncio.create_task(wrapper(), name=task_name)
@@ -3650,8 +3664,11 @@ class BaseWorker(abc.ABC):
             logger.exception("async instruction handlers should never raise!")
             raise
 
-        with ledger.record():
-            # Capture metric events in _transition_to_memory()
+        # Capture metric events in _transition_to_memory()
+        # As this may trigger calls to _start_async_instruction for more tasks,
+        # make sure we don't endlessly pile up context_meter callbacks by specifying
+        # the same key as in _start_async_instruction.
+        with ledger.record(key="async-instruction"):
             self.handle_stimulus(stim)
 
         self._finalize_metrics(stim, ledger, span_id)


Test patch:
diff --git a/distributed/tests/test_cancelled_state.py b/distributed/tests/test_cancelled_state.py
index 7f95ab36..1609972f 100644
--- a/distributed/tests/test_cancelled_state.py
+++ b/distributed/tests/test_cancelled_state.py
@@ -14,6 +14,7 @@ from distributed.utils_test import (
     _LockedCommPool,
     assert_story,
     async_poll_for,
+    freeze_batched_send,
     gen_cluster,
     inc,
     lock_inc,
@@ -488,6 +489,8 @@ async def test_resumed_cancelled_handle_compute(
         await lock_compute.release()
         await exit_compute.wait()
 
+        await async_poll_for(lambda: f3.key not in b.state.tasks, timeout=5)
+
     f1 = c.submit(inc, 1, key="f1", workers=[a.address])
     f2 = c.submit(inc, f1, key="f2", workers=[a.address])
     f3 = c.submit(inc, f2, key="f3", workers=[b.address])
@@ -569,8 +572,7 @@ async def test_resumed_cancelled_handle_compute(
         )
 
     elif wait_for_processing and raise_error:
-        with pytest.raises(RuntimeError, match="test error"):
-            await f3
+        assert await f4 == 4 + 2
 
         assert_story(
             b.state.story(f3.key),
@@ -581,20 +583,31 @@ async def test_resumed_cancelled_handle_compute(
                 (f3.key, "resumed", "released", "cancelled", {}),
                 (f3.key, "cancelled", "waiting", "executing", {}),
                 (f3.key, "executing", "error", "error", {}),
-                # FIXME: (distributed#7489)
+                (
+                    f3.key,
+                    "error",
+                    "released",
+                    "released",
+                    {f2.key: "released", f3.key: "forgotten"},
+                ),
+                (f3.key, "released", "forgotten", "forgotten", {f2.key: "forgotten"}),
+                (f3.key, "ready", "executing", "executing", {}),
+                (f3.key, "executing", "memory", "memory", {}),
             ],
         )
     else:
         assert False, "unreachable"
 
 
+@pytest.mark.parametrize("raise_error", [True, False])
 @gen_cluster(client=True)
-async def test_cancelled_handle_compute(c, s, a, b):
+async def test_cancelled_handle_compute(c, s, a, b, raise_error):
     """
     Given the history of a task
     executing -> cancelled
 
-    A handle_compute should properly restore executing.
+    A handle_compute should cause the result of the cancelled task to be rejected
+    by the scheduler and the task to be re-run.
 
     See Also
     --------
@@ -611,6 +624,8 @@ async def test_cancelled_handle_compute(c, s, a, b):
     def block(x, lock, enter_event, exit_event):
         enter_event.set()
         with lock:
+            if raise_error:
+                raise RuntimeError("test error")
             return x + 1
 
     f1 = c.submit(inc, 1, key="f1", workers=[a.address])
@@ -650,22 +665,151 @@ async def test_cancelled_handle_compute(c, s, a, b):
 
     assert await f4 == 4 + 2
 
-    story = b.state.story(f3.key)
+    if raise_error:
+        assert_story(
+            b.state.story(f3.key),
+            expect=[
+                (f3.key, "ready", "executing", "executing", {}),
+                (f3.key, "executing", "released", "cancelled", {}),
+                (f3.key, "cancelled", "waiting", "executing", {}),
+                (f3.key, "executing", "error", "error", {}),
+                (
+                    f3.key,
+                    "error",
+                    "released",
+                    "released",
+                    {f2.key: "released", f3.key: "forgotten"},
+                ),
+                (f3.key, "released", "forgotten", "forgotten", {f2.key: "forgotten"}),
+                (f3.key, "ready", "executing", "executing", {}),
+                (f3.key, "executing", "memory", "memory", {}),
+            ],
+        )
+    else:
+        assert_story(
+            b.state.story(f3.key),
+            expect=[
+                (f3.key, "ready", "executing", "executing", {}),
+                (f3.key, "executing", "released", "cancelled", {}),
+                (f3.key, "cancelled", "waiting", "executing", {}),
+                (f3.key, "executing", "memory", "memory", {}),
+                (
+                    f3.key,
+                    "memory",
+                    "released",
+                    "released",
+                    {f2.key: "released", f3.key: "forgotten"},
+                ),
+                (f3.key, "released", "forgotten", "forgotten", {f2.key: "forgotten"}),
+                (f3.key, "ready", "executing", "executing", {}),
+                (f3.key, "executing", "memory", "memory", {}),
+            ],
+        )
+
+
+@gen_cluster(client=True)
+async def test_cancelled_task_error_rejected(c, s, a, b):
+    """
+    Given the history of a task
+    executing -> cancelled
+
+    An error in the cancelled task is rejected by the scheduler and superseded
+    by a more recent run on another worker.
+
+    """
+    # This test is heavily using set_restrictions to simulate certain scheduler
+    # decisions of placing keys
+
+    lock_erring = Lock()
+    enter_compute_erring = Event()
+    exit_compute_erring = Event()
+    lock_successful = Lock()
+    enter_compute_successful = Event()
+    exit_compute_successful = Event()
+
+    await lock_erring.acquire()
+    await lock_successful.acquire()
+
+    def block(x, lock, enter_event, exit_event, raise_error):
+        enter_event.set()
+        try:
+            with lock:
+                if raise_error:
+                    raise RuntimeError("test_error")
+                return x + 1
+        finally:
+            exit_event.set()
+
+    f1 = c.submit(inc, 1, key="f1", workers=[a.address])
+    f2 = c.submit(inc, f1, key="f2", workers=[a.address])
+    f3 = c.submit(
+        block,
+        f2,
+        lock=lock_erring,
+        enter_event=enter_compute_erring,
+        exit_event=exit_compute_erring,
+        raise_error=True,
+        key="f3",
+        workers=[b.address],
+    )
+
+    f4 = c.submit(sum, [f1, f3], key="f4", workers=[b.address])
+
+    await enter_compute_erring.wait()
+
+    async def release_all_futures():
+        futs = [f1, f2, f3, f4]
+        for fut in futs:
+            fut.release()
+
+        while any(fut.key in s.tasks for fut in futs):
+            await asyncio.sleep(0.05)
+
+    with freeze_batched_send(s.stream_comms[b.address]):
+        await release_all_futures()
+
+        f1 = c.submit(inc, 1, key="f1", workers=[a.address])
+        f2 = c.submit(inc, f1, key="f2", workers=[a.address])
+        f3 = c.submit(
+            block,
+            f2,
+            lock=lock_successful,
+            enter_event=enter_compute_successful,
+            exit_event=exit_compute_successful,
+            raise_error=False,
+            key="f3",
+            workers=[a.address],
+        )
+        f4 = c.submit(sum, [f1, f3], key="f4", workers=[b.address])
+
+        await wait_for_state(f3.key, "processing", s)
+        await enter_compute_successful.wait()
+
+        await lock_erring.release()
+        await wait_for_state(f3.key, "error", b)
+
+    await lock_successful.release()
+    assert await f4 == 4 + 2
+
     assert_story(
         b.state.story(f3.key),
         expect=[
             (f3.key, "ready", "executing", "executing", {}),
-            (f3.key, "executing", "released", "cancelled", {}),
-            (f3.key, "cancelled", "waiting", "executing", {}),
-            (f3.key, "executing", "memory", "memory", {}),
+            (f3.key, "executing", "error", "error", {}),
             (
                 f3.key,
-                "memory",
+                "error",
                 "released",
                 "released",
-                {f2.key: "released", f3.key: "forgotten"},
+                {f3.key: "forgotten"},
             ),
-            (f3.key, "released", "forgotten", "forgotten", {f2.key: "forgotten"}),
+            (f3.key, "released", "forgotten", "forgotten", {}),
+        ],
+    )
+
+    assert_story(
+        a.state.story(f3.key),
+        expect=[
             (f3.key, "ready", "executing", "executing", {}),
             (f3.key, "executing", "memory", "memory", {}),
         ],
@@ -787,7 +931,7 @@ def test_workerstate_executing_failure_to_fetch(ws_with_running_task):
     - executing -> long-running -> cancelled -> resumed(fetch)
 
     The task execution later terminates with a failure.
-    This is an edge case interaction between work stealing and a task that does not
+    This is an edge case interaction involving task cancellation and a task that does not
     deterministically succeed or fail when run multiple times or on different workers.
 
     Test that the task is fetched from the other worker. This is to avoid having to deal
diff --git a/distributed/tests/test_metrics.py b/distributed/tests/test_metrics.py
index 3f6edab2..cbb9e2a7 100644
--- a/distributed/tests/test_metrics.py
+++ b/distributed/tests/test_metrics.py
@@ -77,25 +77,34 @@ def test_meter_floor(kwargs, delta):
 
 
 def test_context_meter():
-    it = iter([123, 124])
+    it = iter([123, 124, 125, 126])
     cbs = []
 
     with metrics.context_meter.add_callback(lambda l, v, u: cbs.append((l, v, u))):
-        with metrics.context_meter.meter("m1", func=lambda: next(it)) as m:
-            assert m.start == 123
-            assert math.isnan(m.stop)
-            assert math.isnan(m.delta)
+        with metrics.context_meter.meter("m1", func=lambda: next(it)) as m1:
+            assert m1.start == 123
+            assert math.isnan(m1.stop)
+            assert math.isnan(m1.delta)
+        with metrics.context_meter.meter("m2", func=lambda: next(it), unit="foo") as m2:
+            assert m2.start == 125
+            assert math.isnan(m2.stop)
+            assert math.isnan(m2.delta)
+
         metrics.context_meter.digest_metric("m1", 2, "seconds")
         metrics.context_meter.digest_metric("m1", 1, "foo")
 
     # Not recorded - out of context
     metrics.context_meter.digest_metric("m1", 123, "foo")
 
-    assert m.start == 123
-    assert m.stop == 124
-    assert m.delta == 1
+    assert m1.start == 123
+    assert m1.stop == 124
+    assert m1.delta == 1
+    assert m2.start == 125
+    assert m2.stop == 126
+    assert m2.delta == 1
     assert cbs == [
         ("m1", 1, "seconds"),
+        ("m2", 1, "foo"),
         ("m1", 2, "seconds"),
         ("m1", 1, "foo"),
     ]
@@ -199,3 +208,43 @@ def test_delayed_metrics_ledger():
         ("foo", 10, "bytes"),
         ("other", 20, "seconds"),
     ]
+
+
+def test_context_meter_keyed():
+    cbs = []
+
+    def cb(tag, key):
+        return metrics.context_meter.add_callback(
+            lambda l, v, u: cbs.append((tag, l)), key=key
+        )
+
+    with cb("x", key="x"), cb("y", key="y"):
+        metrics.context_meter.digest_metric("l1", 1, "u")
+        with cb("z", key="x"):
+            metrics.context_meter.digest_metric("l2", 2, "u")
+        metrics.context_meter.digest_metric("l3", 3, "u")
+
+    assert cbs == [
+        ("x", "l1"),
+        ("y", "l1"),
+        ("z", "l2"),
+        ("y", "l2"),
+        ("x", "l3"),
+        ("y", "l3"),
+    ]
+
+
+def test_delayed_metrics_ledger_keyed():
+    l1 = metrics.DelayedMetricsLedger()
+    l2 = metrics.DelayedMetricsLedger()
+    l3 = metrics.DelayedMetricsLedger()
+
+    with l1.record(key="x"), l2.record(key="y"):
+        metrics.context_meter.digest_metric("l1", 1, "u")
+        with l3.record(key="x"):
+            metrics.context_meter.digest_metric("l2", 2, "u")
+        metrics.context_meter.digest_metric("l3", 3, "u")
+
+    assert l1.metrics == [("l1", 1, "u"), ("l3", 3, "u")]
+    assert l2.metrics == [("l1", 1, "u"), ("l2", 2, "u"), ("l3", 3, "u")]
+    assert l3.metrics == [("l2", 2, "u")]
diff --git a/distributed/tests/test_worker_metrics.py b/distributed/tests/test_worker_metrics.py
index 419408cb..7efd261f 100644
--- a/distributed/tests/test_worker_metrics.py
+++ b/distributed/tests/test_worker_metrics.py
@@ -618,3 +618,22 @@ async def test_new_metrics_during_heartbeat(c, s, a):
     assert a.digests_total["execute", span.id, "x", "test", "test"] == n
     assert s.cumulative_worker_metrics["execute", "x", "test", "test"] == n
     assert span.cumulative_worker_metrics["execute", "x", "test", "test"] == n
+
+
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1)],
+    config={"distributed.scheduler.worker-saturation": float("inf")},
+)
+async def test_delayed_ledger_is_not_reentrant(c, s, a):
+    """https://github.com/dask/distributed/issues/7949
+
+    Test that, when there's a long chain of task done -> task start events,
+    the callbacks added by the delayed ledger don't pile up on top of each other.
+    """
+
+    def f(_):
+        return len(context_meter._callbacks.get())
+
+    out = await c.gather(c.map(f, range(1000)))
+    assert max(out) < 10
diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
index b2817363..1c87d1d3 100644
--- a/distributed/tests/test_worker_state_machine.py
+++ b/distributed/tests/test_worker_state_machine.py
@@ -531,6 +531,7 @@ def test_executefailure_to_dict():
     ev = ExecuteFailureEvent(
         stimulus_id="test",
         key="x",
+        run_id=1,
         start=123.4,
         stop=456.7,
         exception=Serialize(ValueError("foo")),
@@ -546,6 +547,7 @@ def test_executefailure_to_dict():
         "stimulus_id": "test",
         "handled": 11.22,
         "key": "x",
+        "run_id": 1,
         "start": 123.4,
         "stop": 456.7,
         "exception": "<Serialize: foo>",
@@ -571,6 +573,7 @@ def test_executefailure_dummy():
     ev = ExecuteFailureEvent.dummy("x", stimulus_id="s")
     assert ev == ExecuteFailureEvent(
         key="x",
+        run_id=1,
         start=None,
         stop=None,
         exception=Serialize(None),