Ensure uniqueness of steal stimulus ID

distributed/stealing.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import logging
+import uuid
 from collections import defaultdict, deque
 from math import log2
 from time import time
@@ -233,7 +234,9 @@ def move_task_request(self, ts, victim, thief) -> str:
         try:
             if ts in self.in_flight:
                 return "in-flight"
-            stimulus_id = f"steal-{time()}"
+            # Stimulus IDs are used to verify the response, see
+            # `move_task_confirm`. Therefore, this must be truly unique.
+            stimulus_id = f"steal-{uuid.uuid4().hex}"
 
             key = ts.key
             self.remove_key_from_stealable(ts)
@@ -291,7 +294,7 @@ async def move_task_confirm(self, *, key, state, stimulus_id, worker=None):
                 self.in_flight[ts] = d
                 return
         except KeyError:
-            self.log(("already-aborted", key, state, stimulus_id))
+            self.log(("already-aborted", key, state, worker, stimulus_id))
             return
 
         thief = d["thief"]

distributed/tests/test_steal.py

@@ -1193,3 +1193,40 @@ async def test_correct_bad_time_estimate(c, s, *workers):
     assert any(s.tasks[f.key] in steal.key_stealable for f in futures)
     await wait(futures)
     assert all(w.data for w in workers), [sorted(w.data) for w in workers]
+
+
+@gen_cluster(client=True)
+async def test_steal_stimulus_id_unique(c, s, *workers):
+    steal = s.extensions["stealing"]
+    num_futs = 1_000
+    from distributed import Lock
+
+    lock = Lock()
+    async with lock:
+
+        def blocked(x, lock):
+            lock.acquire()
+
+        # Setup all tasks on worker 0 such that victim/thief relation is the
+        # same for all tasks.
+        futures = c.map(
+            blocked, range(num_futs), lock=lock, workers=[workers[0].address]
+        )
+        # Ensure all tasks are assigned to the worker since otherwise the
+        # move_task_request fails.
+        while len(workers[0].tasks) != num_futs:
+            await asyncio.sleep(0.1)
+        tasks = [s.tasks[f.key] for f in futures]
+        w0 = s.workers[workers[0].address]
+        w1 = s.workers[workers[1].address]
+        # Generating the move task requests as fast as possible increases the
+        # chance of duplicates if the uniqueness is not guaranteed.
+        for ts in tasks:
+            steal.move_task_request(ts, w0, w1)
+        stimulus_ids = set()
+        # Values stored in in_flight are used for response verification.
+        # Therefore all stimulus IDs are stored here and must be unique
+        for dct in steal.in_flight.values():
+            stimulus_ids.add(dct["stimulus_id"])
+        assert len(stimulus_ids) == num_futs
+        await c.cancel(futures)