Repository: dask/distributed

Issue:
SerializeableLock fails when used with resources
For some reason the lock appears to be locked elsewhere.  This only happens when I have two threads and two processes, restricting usage with resources:

```python
from dask.distributed import Client
import time
from dask.utils import SerializableLock


def dummy(i: int) -> int:
    return i


def constrained(lock: SerializableLock, dummy: int):
    assert not lock.locked()
    with lock:
        time.sleep(0.005)
    assert not lock.locked()


N_TASKS = 1000
lock = SerializableLock()
futures = []

if __name__ == "__main__":
    with Client(resources={"foo": 1}, n_workers=2, threads_per_worker=2) as client:
        for i in range(N_TASKS):
            x = client.submit(dummy, i)
            futures.append(client.submit(constrained, lock, x, resources={"foo": 1}))
        client.gather(futures)
```

Things that might be wrong:

1. Resources are bad?  Should check to see what other tasks are running when this failure occurs
2. Accidentally shared state between the two processes?  Should try running this maybe with `dask-worker` processes created separately


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
@@ -1746,8 +1746,6 @@ def _transition_released_waiting(
 
         if ts.waiting_for_data:
             self.waiting_for_data_count += 1
-        elif ts.resource_restrictions:
-            recommendations[ts] = "constrained"
         else:
             recommendations[ts] = "ready"
 
@@ -1831,6 +1829,9 @@ def _transition_waiting_ready(
                 assert dep.key in self.data or dep.key in self.actors
                 assert dep.state == "memory"
 
+        if ts.resource_restrictions:
+            return {ts: "constrained"}, []
+
         ts.state = "ready"
         assert ts.priority is not None
         heapq.heappush(self.ready, (ts.priority, ts.key))
@@ -1961,6 +1962,11 @@ def _transition_from_resumed(
             assert next_state in {"waiting", "fetch"}, next_state
             assert ts._previous in {"executing", "long-running", "flight"}, ts._previous
 
+            if ts._previous in ("executing", "long-running"):
+                self._release_resources(ts)
+                self.executing.discard(ts)
+                self.long_running.discard(ts)
+
             if next_state != finish:
                 recs, instructions = self._transition_generic_released(
                     ts, stimulus_id=stimulus_id
@@ -3202,6 +3208,22 @@ def validate_state(self) -> None:
         for tss in self.data_needed.values():
             assert len({ts.key for ts in tss}) == len(tss)
 
+        self._validate_resources()
+
+    def _validate_resources(self) -> None:
+        """Assert that available_resources + resources held by tasks = total_resources"""
+        assert self.total_resources.keys() == self.available_resources.keys()
+        total = self.total_resources.copy()
+        for k, v in self.available_resources.items():
+            assert v > -1e-9, self.available_resources
+            total[k] -= v
+        for ts in self.all_running_tasks:
+            for k, v in ts.resource_restrictions.items():
+                assert v >= 0, (ts, ts.resource_restrictions)
+                total[k] -= v
+
+        assert all((abs(v) < 1e-9) for v in total.values()), total
+
 
 class BaseWorker(abc.ABC):
     """Wrapper around the :class:`WorkerState` that implements instructions handling.


Test patch:
diff --git a/distributed/diagnostics/tests/test_worker_plugin.py b/distributed/diagnostics/tests/test_worker_plugin.py
--- a/distributed/diagnostics/tests/test_worker_plugin.py
+++ b/distributed/diagnostics/tests/test_worker_plugin.py
@@ -139,6 +139,7 @@ def failing(x):
 async def test_superseding_task_transitions_called(c, s, w):
     expected_notifications = [
         {"key": "task", "start": "released", "finish": "waiting"},
+        {"key": "task", "start": "waiting", "finish": "ready"},
         {"key": "task", "start": "waiting", "finish": "constrained"},
         {"key": "task", "start": "constrained", "finish": "executing"},
         {"key": "task", "start": "executing", "finish": "memory"},
diff --git a/distributed/tests/test_utils_test.py b/distributed/tests/test_utils_test.py
--- a/distributed/tests/test_utils_test.py
+++ b/distributed/tests/test_utils_test.py
@@ -975,4 +975,5 @@ def test_ws_with_running_task(ws_with_running_task):
     ts = ws.tasks["x"]
     assert ts.resource_restrictions == {"R": 1}
     assert ws.available_resources == {"R": 0}
+    assert ws.total_resources == {"R": 1}
     assert ts.state in ("executing", "long-running")
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
@@ -25,11 +25,14 @@
 )
 from distributed.worker_state_machine import (
     AcquireReplicasEvent,
+    AddKeysMsg,
     ComputeTaskEvent,
+    Execute,
     ExecuteFailureEvent,
     ExecuteSuccessEvent,
     FreeKeysEvent,
     GatherDep,
+    GatherDepSuccessEvent,
     Instruction,
     PauseEvent,
     RecommendationsConflict,
@@ -37,6 +40,7 @@
     ReleaseWorkerDataMsg,
     RescheduleEvent,
     RescheduleMsg,
+    SecedeEvent,
     SerializedTask,
     StateMachineEvent,
     TaskState,
@@ -1053,6 +1057,99 @@ def test_gather_priority(ws):
     ]
 
 
+@pytest.mark.parametrize("state", ["executing", "long-running"])
+def test_task_acquires_resources(ws, state):
+    ws.available_resources = {"R": 1}
+    ws.total_resources = {"R": 1}
+
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            key="x", resource_restrictions={"R": 1}, stimulus_id="compute"
+        )
+    )
+    if state == "long-running":
+        ws.handle_stimulus(
+            SecedeEvent(key="x", compute_duration=1.0, stimulus_id="secede")
+        )
+    assert ws.tasks["x"].state == state
+    assert ws.available_resources == {"R": 0}
+
+
+@pytest.mark.parametrize(
+    "done_ev_cls,done_status",
+    [(ExecuteSuccessEvent, "memory"), (ExecuteFailureEvent, "error")],
+)
+def test_task_releases_resources(ws_with_running_task, done_ev_cls, done_status):
+    ws = ws_with_running_task
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="success"))
+    assert ws.tasks["x"].state == done_status
+    assert ws.available_resources == {"R": 1}
+
+
+def test_task_with_dependencies_acquires_resources(ws):
+    ws.available_resources = {"R": 1}
+    ws.total_resources = {"R": 1}
+    ws2 = "127.0.0.1:2"
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "y", who_has={"x": [ws2]}, resource_restrictions={"R": 1}, stimulus_id="s1"
+        )
+    )
+    assert ws.tasks["x"].state == "flight"
+    assert ws.tasks["y"].state == "waiting"
+    assert ws.available_resources == {"R": 1}
+
+    instructions = ws.handle_stimulus(
+        GatherDepSuccessEvent(
+            worker=ws2, data={"x": 123}, total_nbytes=8, stimulus_id="s2"
+        )
+    )
+    assert instructions == [
+        AddKeysMsg(keys=["x"], stimulus_id="s2"),
+        Execute(key="y", stimulus_id="s2"),
+    ]
+    assert ws.tasks["y"].state == "executing"
+    assert ws.available_resources == {"R": 0}
+
+
+@pytest.mark.parametrize(
+    "done_ev_cls,done_status",
+    [
+        (ExecuteSuccessEvent, "memory"),
+        pytest.param(
+            ExecuteFailureEvent,
+            "error",
+            marks=pytest.mark.xfail(
+                reason="distributed#6682,distributed#6689,distributed#6693"
+            ),
+        ),
+    ],
+)
+def test_resumed_task_releases_resources(
+    ws_with_running_task, done_ev_cls, done_status
+):
+    ws = ws_with_running_task
+    assert ws.available_resources == {"R": 0}
+    ws2 = "127.0.0.1:2"
+
+    ws.handle_stimulus(FreeKeysEvent("cancel", ["x"]))
+    assert ws.tasks["x"].state == "cancelled"
+    assert ws.available_resources == {"R": 0}
+
+    instructions = ws.handle_stimulus(
+        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="compute")
+    )
+    assert not instructions
+    assert ws.tasks["x"].state == "resumed"
+    assert ws.available_resources == {"R": 0}
+
+    ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="s2"))
+    assert ws.tasks["x"].state == done_status
+    assert ws.available_resources == {"R": 1}
+
+
 @gen_cluster()
 async def test_clean_log(s, a, b):
     """Test that brand new workers start with a clean log"""
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -2462,7 +2462,6 @@ def ws_with_running_task(ws, request):
             SecedeEvent(key="x", compute_duration=1.0, stimulus_id="secede")
         )
     assert ws.tasks["x"].state == request.param
-    assert ws.available_resources == {"R": 0}
     yield ws