Repository: dask/distributed

Issue:
Resumed tasks may not deallocate (or over-deallocate) resources
### Use case
In jupyter notebook, a user starts a long-ish task. Then, they realise that it should have resources, so they resubmit it before the original task terminated.
Or the other way around, they submit a task with resources, realise they should have no or different resources, and resubmit.

### Expected behaviour
After the original task terminates, `available_resources == total_resources`.

### Actual behaviour
`available_resources` are released according to the second invocation of `client.submit`. 
This could keep them occupied forever or send them above `total_resources`.

### Reproducer
given `available_resources = {"R": 1}, total_resources={"R": 1}`:
```python
f = client.submit(sleep, 2)
sleep(1)
f = client.submit(sleep, 2, resources={"R": 1})
```
Now we have `available_resources = {"R": 2}, total_resources={"R": 1}`.


Discussion:

Resolving patch:
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
@@ -2679,9 +2679,18 @@ def _handle_compute_task(self, ev: ComputeTaskEvent) -> RecsInstrs:
             ts.traceback_text = ""
             ts.priority = priority
             ts.duration = ev.duration
-            ts.resource_restrictions = ev.resource_restrictions
             ts.annotations = ev.annotations
 
+            # If we receive ComputeTaskEvent twice for the same task, resources may have
+            # changed, but the task is still running. Preserve the previous resource
+            # restrictions so that they can be properly released when it eventually
+            # completes.
+            if not (
+                ts.state in ("cancelled", "resumed")
+                and ts.previous in ("executing", "long-running")
+            ):
+                ts.resource_restrictions = ev.resource_restrictions
+
             if self.validate:
                 assert ev.who_has.keys() == ev.nbytes.keys()
                 for dep_workers in ev.who_has.values():


Test patch:
diff --git a/distributed/tests/test_resources.py b/distributed/tests/test_resources.py
--- a/distributed/tests/test_resources.py
+++ b/distributed/tests/test_resources.py
@@ -14,6 +14,7 @@
 from distributed.worker_state_machine import (
     ComputeTaskEvent,
     Execute,
+    ExecuteFailureEvent,
     ExecuteSuccessEvent,
     FreeKeysEvent,
     TaskFinishedMsg,
@@ -508,3 +509,64 @@ async def test_resources_from_python_override_config(c, s, a, b):
     info = c.scheduler_info()
     for worker in [a, b]:
         assert info["workers"][worker.address]["resources"] == {"my_resources": 10}
+
+
+@pytest.mark.parametrize("done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent])
+def test_cancelled_with_resources(ws_with_running_task, done_ev_cls):
+    """Test transition loop of a task with resources:
+
+    executing -> cancelled -> released
+    """
+    ws = ws_with_running_task
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(FreeKeysEvent(keys=["x"], stimulus_id="s1"))
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="s2"))
+    assert ws.available_resources == {"R": 1}
+    assert "x" not in ws.tasks
+
+
+@pytest.mark.parametrize("done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent])
+def test_resumed_with_resources(ws_with_running_task, done_ev_cls):
+    """Test transition loop of a task with resources:
+
+    executing -> cancelled -> resumed(fetch) -> (complete execution)
+    """
+    ws = ws_with_running_task
+    ws2 = "127.0.0.1:2"
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(FreeKeysEvent(keys=["x"], stimulus_id="s1"))
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s2")
+    )
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="s3"))
+    assert ws.available_resources == {"R": 1}
+
+
+@pytest.mark.parametrize("done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent])
+def test_resumed_with_different_resources(ws_with_running_task, done_ev_cls):
+    """A task with resources is cancelled and then resumed to the same state, but with
+    different resources. This is actually possible in case of manual cancellation from
+    the client, followed by resubmit.
+    """
+    ws = ws_with_running_task
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(FreeKeysEvent(keys=["x"], stimulus_id="s1"))
+    assert ws.available_resources == {"R": 0}
+
+    instructions = ws.handle_stimulus(
+        ComputeTaskEvent.dummy("x", stimulus_id="s2", resource_restrictions={"R": 0.4})
+    )
+    assert not instructions
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(done_ev_cls.dummy(key="x", stimulus_id="s3"))
+    assert ws.available_resources == {"R": 1}