Repository: dask/distributed

Issue:
Prometheus: "other" tasks count is confusing
In Prometheus, we export the number of tasks on the worker using

https://github.com/dask/distributed/blob/f83025935383d13033fe0dcd4af2ee689078cb40/distributed/worker_state_machine.py#L3299-L3324

Erred tasks can be an important indicator of something very wrong going on. I just went through a real use case where "other" was very large, and I don't know for sure what I'm looking at - huge amount of erred tasks, released tasks, a bug in the code which hides a third state?

Solution 1: add a new set to the WorkerStateMachine, containing all erred tasks, just for the purpose of this counting
Solution 2: rewrite the task count as an inc/dec counter whenever there is a transition
Solution 3: https://github.com/dask/distributed/pull/7411 (which is the same as 2, but natively with prometheus_client)
Better instrumentation for `Worker.gather_dep`
[Task queuing](https://github.com/dask/distributed/issues/7213) has been proven to significantly improve performance by reducing [root task overproduction](https://github.com/dask/distributed/issues/5555)

In recent benchmarks and tests I noticed that one major source for root task overproduction is not necessarily that reducers are not assigned fast enough to the workers but that the workers are unable to run these tasks since they need to fetch dependencies first. If average root task runtime is much smaller than it takes to fetch dependencies, this can cause workers to run many data producers before it has the possibility to run a reducer.

Right now, we're almost blind to this situation but could be exposing much better metrics on the dashboard (or Prometheus).

Specifically, I'm interested in

- How much time do tasks spend in the ready queue before they are worked on? Can we calculate averages on TaskGroups/Prefix? TaskGroups per Worker?
- How much time do tasks spend in any state, e.g. fetch. In general, how long are wait times in our queues?
- How long do gather_dep requests typically last, broken down per TaskGroup/Prefix?
- How much time on the gather_dep request are we spending on
  - Connection establishment (e.g. connection pool empty, remote event loop is blocked, handshake takes a while)
  - Data (de-)serialization
  - Spill-to-disk


Ideally, I would love to get data for a Task X with dependencies deps that tells me

```
X spent 10s in ready queue
-> 8s spent fetching data
  -> 1s connection
  -> 2s spill-to-disk
  -> 2s (de-)serialization
  -> 2s network transfer
  -> 1s idle / event loop busy
-> 2s spent waiting for open slot on the ThreadPool
```

Some of this information is already available, other information we still need to collect. I don't think we have anything that can break it up this way and/or group by TaskGroups or individual tasks.

I think this kind of visibility would help us significantly with making decisions about optimizations, e.g. should we prioritize STA? Should we focus on getting a sendfile implementation up and running? Do connection attempts take way too long because event loops are blocked?


Discussion:

Resolving patch:
diff --git a/distributed/http/worker/prometheus/core.py b/distributed/http/worker/prometheus/core.py
--- a/distributed/http/worker/prometheus/core.py
+++ b/distributed/http/worker/prometheus/core.py
@@ -39,13 +39,15 @@ def collect(self) -> Iterator[Metric]:
             "Number of tasks at worker.",
             labels=["state"],
         )
-        for k, n in ws.task_counts.items():
-            if k == "memory" and hasattr(self.server.data, "slow"):
+        for state, n in ws.task_counter.current_count(by_prefix=False).items():
+            if state == "memory" and hasattr(self.server.data, "slow"):
                 n_spilled = len(self.server.data.slow)
-                tasks.add_metric(["memory"], n - n_spilled)
-                tasks.add_metric(["disk"], n_spilled)
+                if n - n_spilled > 0:
+                    tasks.add_metric(["memory"], n - n_spilled)
+                if n_spilled > 0:
+                    tasks.add_metric(["disk"], n_spilled)
             else:
-                tasks.add_metric([k], n)
+                tasks.add_metric([state], n)
         yield tasks
 
         yield GaugeMetricFamily(
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1012,7 +1012,7 @@ async def get_metrics(self) -> dict:
             spilled_memory, spilled_disk = 0, 0
 
         out = dict(
-            task_counts=self.state.task_counts,
+            task_counts=self.state.task_counter.current_count(by_prefix=False),
             bandwidth={
                 "total": self.bandwidth,
                 "workers": dict(self.bandwidth_workers),
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -10,7 +10,7 @@
 import sys
 import warnings
 import weakref
-from collections import defaultdict, deque
+from collections import Counter, defaultdict, deque
 from collections.abc import (
     Callable,
     Collection,
@@ -18,6 +18,7 @@
     Iterator,
     Mapping,
     MutableMapping,
+    Set,
 )
 from copy import copy
 from dataclasses import dataclass, field
@@ -37,13 +38,13 @@
 from tlz import peekn
 
 import dask
-from dask.utils import parse_bytes, typename
+from dask.utils import key_split, parse_bytes, typename
 
 from distributed._stories import worker_story
 from distributed.collections import HeapSet
 from distributed.comm import get_address_host
 from distributed.core import ErrorMessage, error_message
-from distributed.metrics import time
+from distributed.metrics import monotonic, time
 from distributed.protocol import pickle
 from distributed.protocol.serialize import Serialize
 from distributed.sizeof import safe_sizeof as sizeof
@@ -80,16 +81,16 @@
 ]
 
 # TaskState.state subsets
-PROCESSING: set[TaskStateState] = {
+PROCESSING: Set[TaskStateState] = {
     "waiting",
     "ready",
     "constrained",
     "executing",
     "long-running",
 }
-READY: set[TaskStateState] = {"ready", "constrained"}
+READY: Set[TaskStateState] = {"ready", "constrained"}
 # Valid states for a task that is found in TaskState.waiting_for_data
-WAITING_FOR_DATA: set[TaskStateState] = {
+WAITING_FOR_DATA: Set[TaskStateState] = {
     "constrained",
     "executing",
     "fetch",
@@ -216,10 +217,10 @@ def _default_data_size() -> int:
 
 # Note: can't specify __slots__ manually to enable slots in Python <3.10 in a @dataclass
 # that defines any default values
-dc_slots = {"slots": True} if sys.version_info >= (3, 10) else {}
+DC_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {}
 
 
-@dataclass(repr=False, eq=False, **dc_slots)
+@dataclass(repr=False, eq=False, **DC_SLOTS)
 class TaskState:
     """Holds volatile state relating to an individual Dask task.
 
@@ -229,6 +230,8 @@ class TaskState:
 
     #: Task key. Mandatory.
     key: str
+    #: Task prefix (leftmost part of the key)
+    prefix: str = field(init=False)
     #: Task run ID.
     run_id: int = RUN_ID_SENTINEL
     #: A named tuple containing the ``function``, ``args``, ``kwargs`` and ``task``
@@ -301,6 +304,7 @@ class TaskState:
 
     def __post_init__(self) -> None:
         TaskState._instances.add(self)
+        self.prefix = key_split(self.key)
 
     def __repr__(self) -> str:
         if self.state == "cancelled":
@@ -350,8 +354,8 @@ def _to_dict_no_nest(self, *, exclude: Container[str] = ()) -> dict:
         there's an acyclic dependency chain of ~200+ tasks.
         """
         out = recursive_to_dict(self, exclude=exclude, members=True)
-        # Remove all Nones and empty containers
-        return {k: v for k, v in out.items() if v}
+        # Remove all Nones, empty containers, and derived attributes
+        return {k: v for k, v in out.items() if v and k != "prefix"}
 
 
 @dataclass
@@ -1244,6 +1248,10 @@ class WorkerState:
     #: Typically disabled in production.
     validate: bool
 
+    #: Current number of tasks and cumulative elapsed time in each state,
+    #: both broken down by :attr:`prefix`
+    task_counter: TaskCounter
+
     #: Total number of state transitions so far.
     #: See also :attr:`log` and :attr:`transition_counter_max`.
     transition_counter: int
@@ -1319,6 +1327,7 @@ def __init__(
         self.transfer_message_bytes_limit = transfer_message_bytes_limit
         self.log = deque(maxlen=100_000)
         self.stimulus_log = deque(maxlen=10_000)
+        self.task_counter = TaskCounter()
         self.transition_counter = 0
         self.transition_counter_max = transition_counter_max
         self.transfer_incoming_bytes_limit = transfer_incoming_bytes_limit
@@ -1413,6 +1422,7 @@ def _ensure_task_exists(
             logger.debug("Data task %s already known (stimulus_id=%s)", ts, stimulus_id)
         except KeyError:
             self.tasks[key] = ts = TaskState(key)
+            self.task_counter.new_task(ts)
         if not ts.priority:
             assert priority
             ts.priority = priority
@@ -2678,11 +2688,13 @@ def _transitions(self, recommendations: Recs, *, stimulus_id: str) -> Instructio
         """
         instructions = []
         tasks = set()
+        initial_states: dict[TaskState, TaskStateState] = {}
 
         def process_recs(recs: Recs) -> None:
             while recs:
                 ts, finish = recs.popitem()
                 tasks.add(ts)
+                initial_states.setdefault(ts, ts.state)
                 a_recs, a_instructions = self._transition(
                     ts, finish, stimulus_id=stimulus_id
                 )
@@ -2702,6 +2714,8 @@ def process_recs(recs: Recs) -> None:
         instructions += a_instructions
         process_recs(a_recs)
 
+        self.task_counter.transitions(initial_states)
+
         if self.validate:
             # Full state validation is very expensive
             for ts in tasks:
@@ -2725,6 +2739,7 @@ def _handle_update_data(self, ev: UpdateDataEvent) -> RecsInstrs:
                 ts = self.tasks[key]
             except KeyError:
                 self.tasks[key] = ts = TaskState(key)
+                self.task_counter.new_task(ts)
 
             recommendations[ts] = ("memory", value, RUN_ID_SENTINEL)
             self.log.append(
@@ -2822,6 +2837,8 @@ def _handle_compute_task(self, ev: ComputeTaskEvent) -> RecsInstrs:
             )
         except KeyError:
             self.tasks[ev.key] = ts = TaskState(ev.key)
+            self.task_counter.new_task(ts)
+
         self.log.append((ev.key, "compute-task", ts.state, ev.stimulus_id, time()))
         ts.run_id = ev.run_id
         recommendations: Recs = {}
@@ -3248,37 +3265,12 @@ def _to_dict(self, *, exclude: Container[str] = ()) -> dict:
             "stimulus_log": self.stimulus_log,
             "transition_counter": self.transition_counter,
             "tasks": self.tasks,
+            "task_counts": dict(self.task_counter.current_count()),
+            "task_cumulative_elapsed": dict(self.task_counter.cumulative_elapsed()),
         }
         info = {k: v for k, v in info.items() if k not in exclude}
         return recursive_to_dict(info, exclude=exclude)
 
-    @property
-    def task_counts(self) -> dict[TaskStateState | Literal["other"], int]:
-        # Actors can be in any state other than {fetch, flight, missing}
-        n_actors_in_memory = sum(
-            self.tasks[key].state == "memory" for key in self.actors
-        )
-
-        out: dict[TaskStateState | Literal["other"], int] = {
-            # Key measure for occupancy.
-            # Also includes cancelled(executing) and resumed(executing->fetch)
-            "executing": len(self.executing),
-            # Also includes cancelled(long-running) and resumed(long-running->fetch)
-            "long-running": len(self.long_running),
-            "memory": len(self.data) + n_actors_in_memory,
-            "ready": len(self.ready),
-            "constrained": len(self.constrained),
-            "waiting": len(self.waiting),
-            "fetch": self.fetch_count,
-            "missing": len(self.missing_dep_flight),
-            # Also includes cancelled(flight) and resumed(flight->waiting)
-            "flight": len(self.in_flight_tasks),
-        }
-        # released | error
-        out["other"] = other = len(self.tasks) - sum(out.values())
-        assert other >= 0
-        return out
-
     ##############
     # Validation #
     ##############
@@ -3557,6 +3549,14 @@ def validate_state(self) -> None:
         for ts in self.tasks.values():
             self.validate_task(ts)
 
+        expect_state_count = Counter(
+            (ts.prefix, ts.state) for ts in self.tasks.values()
+        )
+        assert self.task_counter.current_count() == expect_state_count, (
+            self.task_counter.current_count(),
+            expect_state_count,
+        )
+
         if self.transition_counter_max:
             assert self.transition_counter < self.transition_counter_max
 
@@ -3718,6 +3718,116 @@ def digest_metric(self, name: str, value: float) -> None:
         """Log an arbitrary numerical metric"""
 
 
+class TaskCounter:
+    # When the monotonic timer was last sampled
+    _previous_ts: float | None
+    # Current task counts, per prefix, per state, on the worker
+    _current_count: Counter[tuple[str, TaskStateState]]
+    # Tasks that have just been inserted in the WorkerState and will be immediately
+    # transitioned to a new state
+    _new_tasks: set[TaskState]
+    # Ever-increasing cumulative task runtimes, per prefix, including tasks that have
+    # left a state or even that don't exist anymore, updated as of the latest call
+    # to cumulative_elapsed() or transitions()
+    _cumulative_elapsed: defaultdict[tuple[str, TaskStateState], float]
+
+    __slots__ = tuple(__annotations__)
+
+    def __init__(self) -> None:
+        self._previous_ts = None
+        self._current_count = Counter()
+        self._new_tasks = set()
+        self._cumulative_elapsed = defaultdict(float)
+
+    def current_count(self, by_prefix: bool = True) -> Counter:
+        """Return current count of tasks.
+
+        Parameters
+        ----------
+        by_prefix: bool, optional
+            True (default)
+                Return counter of (task prefix, task state) -> count
+            False
+                Return counter of task state -> count
+        """
+        if by_prefix:
+            return self._current_count
+
+        out: Counter[TaskStateState] = Counter()
+        for (_, state), n in self._current_count.items():
+            out[state] += n
+        return out
+
+    def cumulative_elapsed(self, by_prefix: bool = True) -> Mapping[Any, float]:
+        """Ever-increasing cumulative task runtimes, including tasks that have left a
+        state or even that don't exist anymore, updated as of the moment when this
+        method is called.
+
+        Parameters
+        ----------
+        by_prefix: bool, optional
+            True (default)
+                Return mapping of (task prefix, task state) -> seconds
+            False
+                Return mapping of task state -> seconds
+        """
+        if self._current_count:
+            assert self._previous_ts is not None
+            now = monotonic()
+            elapsed = now - self._previous_ts
+            self._previous_ts = now
+            for k, n_tasks in self._current_count.items():
+                self._cumulative_elapsed[k] += elapsed * n_tasks
+
+        if by_prefix:
+            return self._cumulative_elapsed
+
+        out: defaultdict[TaskStateState, float] = defaultdict(float)
+        for (_, state), n in self._cumulative_elapsed.items():
+            out[state] += n
+        return out
+
+    def new_task(self, ts: TaskState) -> None:
+        """A new task has just been created and will be immediately fed into the
+        recommendations for transitions
+        """
+        self._new_tasks.add(ts)
+
+    def transitions(self, prev_states: dict[TaskState, TaskStateState]) -> None:
+        """Tasks have just transitioned to a new state"""
+        if not prev_states and not self._new_tasks:
+            return
+
+        now = monotonic()
+        if self._current_count:
+            assert self._previous_ts is not None
+            elapsed = now - self._previous_ts
+            for k, n_tasks in self._current_count.items():
+                self._cumulative_elapsed[k] += elapsed * n_tasks
+        self._previous_ts = now
+
+        for ts, prev_state in prev_states.items():
+            if ts.state != "forgotten":
+                self._current_count[ts.prefix, ts.state] += 1
+
+            if ts in self._new_tasks:
+                self._new_tasks.discard(ts)
+            else:
+                dec_count = self._current_count[ts.prefix, prev_state] - 1
+                if dec_count > 0:
+                    self._current_count[ts.prefix, prev_state] = dec_count
+                else:
+                    assert dec_count == 0
+                    del self._current_count[ts.prefix, prev_state]
+
+        for ts in self._new_tasks:
+            # This happens exclusively on a transition from cancelled(flight) to
+            # resumed(flight->waiting) of a task with dependencies; the dependencies
+            # will remain in released state and never transition to anything else.
+            self._current_count[ts.prefix, ts.state] += 1
+        self._new_tasks.clear()
+
+
 class DeprecatedWorkerStateAttribute:
     name: str
     target: str | None


Test patch:
diff --git a/distributed/http/tests/test_core.py b/distributed/http/tests/test_core.py
--- a/distributed/http/tests/test_core.py
+++ b/distributed/http/tests/test_core.py
@@ -23,8 +23,9 @@ async def test_prometheus_api_doc(c, s, a):
     """
     pytest.importorskip("prometheus_client")
 
-    # Some metrics only appear after a task is executed
-    await c.submit(inc, 1)
+    # Some metrics only appear if there are tasks on the cluster
+    fut = c.submit(inc, 1)
+    await fut
     # Semaphore metrics only appear after semaphores are used
     sem = await Semaphore()
     await sem.acquire()
diff --git a/distributed/http/worker/tests/test_worker_http.py b/distributed/http/worker/tests/test_worker_http.py
--- a/distributed/http/worker/tests/test_worker_http.py
+++ b/distributed/http/worker/tests/test_worker_http.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import json
 from unittest import mock
 
@@ -10,6 +9,7 @@
 from distributed import Event, Worker, wait
 from distributed.sizeof import sizeof
 from distributed.utils_test import (
+    async_wait_for,
     fetch_metrics,
     fetch_metrics_body,
     fetch_metrics_sample_names,
@@ -21,6 +21,10 @@
 async def test_prometheus(c, s, a):
     pytest.importorskip("prometheus_client")
 
+    # We need *some* tasks or dask_worker_tasks won't appear
+    fut = c.submit(lambda: 1)
+    await wait(fut)
+
     active_metrics = await fetch_metrics_sample_names(
         a.http_server.port, prefix="dask_worker_"
     )
@@ -88,22 +92,7 @@ async def test_metrics_when_prometheus_client_not_installed(
 async def test_prometheus_collect_task_states(c, s, a):
     pytest.importorskip("prometheus_client")
 
-    async def assert_metrics(**kwargs):
-        expect = {
-            "constrained": 0,
-            "executing": 0,
-            "fetch": 0,
-            "flight": 0,
-            "long-running": 0,
-            "memory": 0,
-            "disk": 0,
-            "missing": 0,
-            "other": 0,
-            "ready": 0,
-            "waiting": 0,
-        }
-        expect.update(kwargs)
-
+    async def assert_metrics(**expect):
         families = await fetch_metrics(a.http_server.port, prefix="dask_worker_")
         actual = {
             sample.labels["state"]: sample.value
@@ -117,24 +106,29 @@ async def assert_metrics(**kwargs):
     ev = Event()
 
     # submit a task which should show up in the prometheus scraping
-    future = c.submit(ev.wait)
-    while not a.state.executing:
-        await asyncio.sleep(0.001)
+    fut1 = c.submit(ev.wait)
+    await async_wait_for(lambda: a.state.executing, timeout=5)
 
     await assert_metrics(executing=1)
 
     await ev.set()
-    await c.gather(future)
+    await wait(fut1)
 
     await assert_metrics(memory=1)
-    a.data.evict()
-    await assert_metrics(disk=1)
 
-    future.release()
+    fut2 = c.submit(lambda: 1)
+    await wait(fut2)
+    await assert_metrics(memory=2)
+
+    a.data.evict()
+    await assert_metrics(memory=1, disk=1)
+    a.data.evict()
+    await assert_metrics(disk=2)
 
-    while future.key in a.state.tasks:
-        await asyncio.sleep(0.001)
+    fut1.release()
+    fut2.release()
 
+    await async_wait_for(lambda: not a.state.tasks, timeout=5)
     await assert_metrics()
 
 
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -3817,14 +3817,9 @@ async def test_transition_counter_max_worker(c, s, a):
     # This is set by @gen_cluster; it's False in production
     assert s.transition_counter_max > 0
     a.state.transition_counter_max = 1
+    fut = c.submit(inc, 2)
     with captured_logger("distributed.worker") as logger:
-        fut = c.submit(inc, 2)
-        while True:
-            try:
-                a.validate_state()
-            except AssertionError:
-                break
-            await asyncio.sleep(0.01)
+        await async_wait_for(lambda: a.state.transition_counter > 0, timeout=5)
 
     assert "TransitionCounterMaxExceeded" in logger.getvalue()
     # Worker state is corrupted. Avoid test failure on gen_cluster teardown.
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -3428,6 +3428,8 @@ async def test_Worker__to_dict(c, s, a):
         "transition_counter",
         "tasks",
         "data_needed",
+        "task_counts",
+        "task_cumulative_elapsed",
     }
     assert d["tasks"]["x"]["key"] == "x"
     assert d["data"] == {"x": None}
diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
--- a/distributed/tests/test_worker_state_machine.py
+++ b/distributed/tests/test_worker_state_machine.py
@@ -5,6 +5,7 @@
 import pickle
 from collections import defaultdict
 from collections.abc import Iterator
+from time import sleep
 
 import pytest
 from tlz import first
@@ -13,6 +14,7 @@
 
 import distributed.profile as profile
 from distributed import Nanny, Worker, wait
+from distributed.compatibility import MACOS, WINDOWS
 from distributed.protocol.serialize import Serialize
 from distributed.scheduler import TaskState as SchedulerTaskState
 from distributed.utils import recursive_to_dict
@@ -20,6 +22,7 @@
     NO_AMM,
     _LockedCommPool,
     assert_story,
+    async_wait_for,
     freeze_data_fetching,
     gen_cluster,
     inc,
@@ -205,8 +208,18 @@ def test_WorkerState__to_dict(ws):
                 "state": "memory",
             },
         },
+        "task_counts": {"['x', 'flight']": 1, "['y', 'memory']": 1},
+        "task_cumulative_elapsed": {
+            "['x', 'flight']": "SNIP",
+            "['y', 'memory']": "SNIP",
+        },
         "transition_counter": 3,
     }
+
+    # timings data (a few microseconds each)
+    for k in actual["task_cumulative_elapsed"]:
+        actual["task_cumulative_elapsed"][k] = "SNIP"
+
     assert actual == expect
 
 
@@ -649,16 +662,16 @@ async def test_lose_replica_during_fetch(c, s, w1, w2, w3, as_deps):
 
         assert len(s.tasks["x"].who_has) == 2
         await w2.close()
-        while len(s.tasks["x"].who_has) > 1:
-            await asyncio.sleep(0.01)
+        await async_wait_for(lambda: len(s.tasks["x"].who_has) == 1, timeout=5)
 
         if as_deps:
             y2 = c.submit(inc, x, key="y2", workers=[w1.address])
         else:
             s.request_acquire_replicas(w1.address, ["x"], stimulus_id="test")
 
-        while w1.state.tasks["x"].who_has != {w3.address}:
-            await asyncio.sleep(0.01)
+        await async_wait_for(
+            lambda: w1.state.tasks["x"].who_has == {w3.address}, timeout=5
+        )
 
     await wait_for_state("x", "memory", w1)
     assert_story(
@@ -1209,7 +1222,7 @@ def test_resumed_task_releases_resources(ws_with_running_task, done_ev_cls):
     assert ws.available_resources == {"R": 0}
     ws2 = "127.0.0.1:2"
 
-    ws.handle_stimulus(FreeKeysEvent("cancel", ["x"]))
+    ws.handle_stimulus(FreeKeysEvent(keys=["x"], stimulus_id="cancel"))
     assert ws.tasks["x"].state == "cancelled"
     assert ws.available_resources == {"R": 0}
 
@@ -1681,47 +1694,63 @@ def test_fetch_count(ws):
     assert len(ws.missing_dep_flight) == 1
 
 
-def test_task_counts(ws):
-    assert ws.task_counts == {
-        "constrained": 0,
-        "executing": 0,
-        "fetch": 0,
-        "flight": 0,
-        "long-running": 0,
-        "memory": 0,
-        "missing": 0,
-        "other": 0,
-        "ready": 0,
-        "waiting": 0,
-    }
-
+def test_task_counter(ws):
+    ws2 = "127.0.0.1:2"
+    ws3 = "127.0.0.1:3"
+    for by_prefix in (False, True):
+        assert ws.task_counter.current_count(by_prefix=by_prefix) == {}
+        assert ws.task_counter.cumulative_elapsed(by_prefix=by_prefix) == {}
 
-def test_task_counts_with_actors(ws):
-    ws.handle_stimulus(ComputeTaskEvent.dummy("x", actor=True, stimulus_id="s1"))
-    assert ws.actors == {"x": None}
-    assert ws.task_counts == {
-        "constrained": 0,
-        "executing": 1,
-        "fetch": 0,
-        "flight": 0,
-        "long-running": 0,
-        "memory": 0,
-        "missing": 0,
-        "other": 0,
-        "ready": 0,
-        "waiting": 0,
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "('y-123', 7)", who_has={"('x-456', 8)": [ws2]}, stimulus_id="s1"
+        ),
+        AcquireReplicasEvent(
+            who_has={"('x-789', 0)": [ws3], "z": [ws3]},
+            nbytes={"('x-789', 0)": 1, "z": 1},
+            stimulus_id="s2",
+        ),
+    )
+    assert ws.task_counter.current_count() == {
+        ("x", "flight"): 2,
+        ("y", "waiting"): 1,
+        ("z", "flight"): 1,
     }
-    ws.handle_stimulus(ExecuteSuccessEvent.dummy("x", value=123, stimulus_id="s2"))
-    assert ws.actors == {"x": 123}
-    assert ws.task_counts == {
-        "constrained": 0,
-        "executing": 0,
-        "fetch": 0,
-        "flight": 0,
-        "long-running": 0,
-        "memory": 1,
-        "missing": 0,
-        "other": 0,
-        "ready": 0,
-        "waiting": 0,
+    assert ws.task_counter.current_count(by_prefix=False) == {"waiting": 1, "flight": 3}
+
+    def assert_time(actual, expect):
+        # timer accuracy in Windows can be very poor;
+        # see awful hack in distributed.metrics
+        margin_lo = 0.099 if WINDOWS else 0
+        # sleep() has been observed to have up to 450ms lag on MacOSX GitHub CI
+        margin_hi = 0.6 if MACOS else 0.1
+        assert expect - margin_lo <= actual < expect + margin_hi
+
+    sleep(0.1)
+    elapsed = ws.task_counter.cumulative_elapsed()
+    # Transitory states are not recorded
+    assert len(elapsed) == 3
+    assert_time(elapsed["x", "flight"], 0.2)
+    assert_time(elapsed["y", "waiting"], 0.1)
+    assert_time(elapsed["z", "flight"], 0.1)
+
+    elapsed = ws.task_counter.cumulative_elapsed(by_prefix=False)
+    assert len(elapsed) == 2
+    assert_time(elapsed["flight"], 0.3)
+    assert_time(elapsed["waiting"], 0.1)
+
+    # Forgotten keys disappear from current_count() and stop accruing time in
+    # cumulative_elapsed()
+    ws.handle_stimulus(FreeKeysEvent(keys=["('y-123', 7)"], stimulus_id="s3"))
+    assert ws.task_counter.current_count() == {
+        ("x", "cancelled"): 1,
+        ("x", "flight"): 1,
+        ("z", "flight"): 1,
     }
+    sleep(0.15)
+    elapsed = ws.task_counter.cumulative_elapsed()
+    assert len(elapsed) == 4
+    assert_time(elapsed["x", "flight"], 0.35)
+    assert_time(elapsed["x", "cancelled"], 0.15)
+    assert_time(elapsed["y", "waiting"], 0.1)
+    assert_time(elapsed["z", "flight"], 0.25)