Repository: dask/distributed

Issue:
AMM can leave a forgotten task forever in missing state
1. The scheduler sends acquire-replicas to the worker
2. Before the task can transition to flight, the task is forgotten on the scheduler. This in turn triggers a release-key command from the scheduler to all workers _which hold the task in memory_. This is unlike in a compute-task request, where the workers that are fetching the key are tracked on the scheduler side and receive a release-key command too.
3. the task eventually transitions fetch->flight->missing
4. the task remains stuck forever in missing state on the worker that was fetching it. The worker will pester the scheduler every second with find_missing.

Reproducer:
```python


@gen_cluster(client=True)
async def test_forget_acquire_replicas(c, s, a, b):
    """
    1. The scheduler sends acquire-replicas to the worker
    2. Before the task can transition to flight, the task is forgotten on the scheduler
       and on the peer workers *holding the replicas*. 
       This is unlike in a compute-task command, where the workers that are fetching
       the key are tracked on the scheduler side and receive a release-key command too.
    3. the task eventually transitions fetch->flight->missing
    4. Test that the task is eventually forgotten everywhere.
    """
    x = c.submit(inc, 2, key="x", workers=[a.address])
    await x
    with freeze_data_fetching(b, jump_start=True):
        s.request_acquire_replicas(b.address, ["x"], stimulus_id="one")
        await wait_for_state("x", "fetch", b)
        x.release()
        while "x" in s.tasks or "x" in a.tasks:
            await asyncio.sleep(0.01)

    while "x" in b.tasks:
        await asyncio.sleep(0.01)
```
The above test times out on the last line.

# Proposed design
At the moment, the scheduler silently ignores missing keys in the `request-refresh-who-has` message.
The scheduler should instead respond stating "I don't know about this key, you should forget about it too". This would trigger the key to be forgotten on the worker too.

# Blockers
- #6342
- #6348

CC @fjetter 



Discussion:
Same problem for tasks in flight:
```python

@gen_cluster(client=True, nthreads=[("", 1)], timeout=3)
async def test_forget_acquire_replicas_flight(c, s, a):
    """If a dependency fetch finishes on a worker after the scheduler already released
    everything, the worker might be stuck with a redundant replica which is never
    cleaned up.
    """
    async with BlockedGatherDep(s.address) as b:
        x = c.submit(inc, 1, key="x", workers=[a.address])
        await x

        s.request_acquire_replicas(b.address, ["x"], stimulus_id="test")
        await b.in_gather_dep.wait()
        assert b.tasks["x"].state == "flight"

        x.release()
        while "x" in s.tasks:
            await asyncio.sleep(0.01)

        b.block_gather_dep.set()
        while b.tasks:
            await asyncio.sleep(0.01)
```

XREF test_forget_data_not_supposed_to_have for tasks acquired through compute-task
> "I don't know about this key, you should forget about it too". This would trigger the key to be forgotten on the worker too.

I remember having similar logic on the scheduler before with "I don't know about this task, please forget it" in various circumstances. This led to deadlocks due to race conditions

Wouldn't a better fix be for the scheduler be instead to remember who it asked to acquire a replica?
Everything goes through the same batched comms, sequentially.
I can't think of any way to create a race condition, where the key is recreated on the scheduler and the scheduler sends an event to the worker about it before the deletion event is processed, without involving RPC.

The RPC commands that can cause a race condition are:
- scatter (it should *definitely* be redesigned on top of batched comms; should I open an issue about it?)
- rebalance (remediation in #4906)
- replicate (remediation in #6578)

[EDIT] see epic #6604

> Wouldn't a better fix be for the scheduler be instead to remember who it asked to acquire a replica?

It's possible, but also unnecessarily complicated IMHO.
Related: #5759 
@fjetter are you happy for me to proceed with my design?
Yes, sounds good. I just wanted to highlight the possibility. if batchedsend ordering protects us, that's good news 👍 

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -4775,13 +4775,30 @@ async def handle_request_refresh_who_has(
         who_has for some keys. Not to be confused with scheduler.who_has, which is a
         synchronous RPC request from a Client.
         """
-        self.stream_comms[worker].send(
-            {
-                "op": "refresh-who-has",
-                "who_has": self.get_who_has(keys),
-                "stimulus_id": stimulus_id,
-            },
-        )
+        who_has = {}
+        free_keys = []
+        for key in keys:
+            if key in self.tasks:
+                who_has[key] = [ws.address for ws in self.tasks[key].who_has]
+            else:
+                free_keys.append(key)
+
+        if who_has:
+            self.stream_comms[worker].send(
+                {
+                    "op": "refresh-who-has",
+                    "who_has": who_has,
+                    "stimulus_id": stimulus_id,
+                }
+            )
+        if free_keys:
+            self.stream_comms[worker].send(
+                {
+                    "op": "free-keys",
+                    "keys": free_keys,
+                    "stimulus_id": stimulus_id,
+                }
+            )
 
     async def handle_worker(self, comm=None, worker=None, stimulus_id=None):
         """


Test patch:
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
@@ -57,6 +57,7 @@
     dec,
     div,
     freeze_batched_send,
+    freeze_data_fetching,
     gen_cluster,
     gen_test,
     inc,
@@ -2785,6 +2786,35 @@ async def test_acquire_replicas_already_in_flight(c, s, a):
         )
 
 
+@pytest.mark.slow
+@gen_cluster(client=True)
+async def test_forget_acquire_replicas(c, s, a, b):
+    """
+    1. The scheduler sends acquire-replicas to the worker
+    2. Before the task can reach the worker, the task is forgotten on the scheduler
+       and on the peer workers holding the replicas.
+       This is unlike in a compute-task command, where the workers that are fetching
+       the key are tracked on the scheduler side and receive a release-key command too.
+    3. The task eventually transitions to missing
+    4. At the next run of find_missing, the worker sends a request-refresh-who-has
+       message to the scheduler
+    5. The scheduler responds with free-keys
+    6. The task is forgotten everywhere.
+    """
+    x = c.submit(inc, 2, key="x", workers=[a.address])
+    await x
+    with freeze_data_fetching(b, jump_start=True):
+        s.request_acquire_replicas(b.address, ["x"], stimulus_id="test")
+        await wait_for_state("x", "fetch", b)
+        x.release()
+        while "x" in s.tasks or "x" in a.state.tasks:
+            await asyncio.sleep(0.01)
+
+    while "x" in b.state.tasks:
+        await asyncio.sleep(0.01)
+    assert "x" not in s.tasks
+
+
 @gen_cluster(client=True)
 async def test_remove_replicas_simple(c, s, a, b):
     futs = c.map(inc, range(10), workers=[a.address])
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
@@ -521,7 +521,7 @@ async def test_lose_replica_during_fetch(c, s, w1, w2, w3, as_deps):
 
 
 @gen_cluster(client=True, nthreads=[("", 1)] * 2)
-async def test_fetch_to_missing(c, s, a, b):
+async def test_fetch_to_missing_on_busy(c, s, a, b):
     """
     1. task x is a dependency of y
     2. scheduler calls handle_compute("y", who_has={"x": [b]}) on a
@@ -531,7 +531,12 @@ async def test_fetch_to_missing(c, s, a, b):
     6. the scheduler responds {"x": []}, because w1 in the meantime has lost the key.
     7. x is transitioned fetch -> missing
     """
-    x = await c.scatter({"x": 1}, workers=[b.address])
+    # Note: submit and scatter are different. If you lose all workers holding the
+    # replicas of a scattered key, the scheduler forgets the task, which in turn would
+    # trigger a free-keys response to request-refresh-who-has.
+    x = c.submit(inc, 1, key="x", workers=[b.address])
+    await x
+
     b.total_in_connections = 0
     # Crucially, unlike with `c.submit(inc, x, workers=[a.address])`, the scheduler
     # doesn't keep track of acquire-replicas requests, so it won't proactively inform a