Provide stack for suspended coro in test timeout

distributed/tests/test_utils_test.py

@@ -373,3 +373,21 @@ async def ping_pong():
     assert await write_queue.get() == (b.address, {"op": "ping", "reply": True})
     write_event.set()
     assert await fut == "pong"
+
+
+def test_provide_stack_on_timeout():
+    sleep_time = 30
+
+    async def inner_test(c, s, a, b):
+        await asyncio.sleep(sleep_time)
+
+    test = gen_cluster(client=True, timeout=0.5)(inner_test)
+
+    start = time()
+    with pytest.raises(asyncio.TimeoutError) as exc:
+        test()
+    end = time()
+    assert "inner_test" in str(exc)
+    assert "await asyncio.sleep(sleep_time)" in str(exc)
+    # ensure the task was properly
+    assert end - start < sleep_time / 2

distributed/utils_test.py

@@ -965,10 +965,22 @@ async def coro():
                             args = [c] + args
                         try:
                             future = func(*args, *outer_args, **kwargs)
-                            future = asyncio.wait_for(future, timeout)
+                            task = asyncio.create_task(future)
+
+                            future = asyncio.wait_for(asyncio.shield(task), timeout)
                             result = await future
                             if s.validate:
                                 s.validate_state()
+                        except asyncio.TimeoutError:
+                            assert task
+                            buffer = io.StringIO()
+                            # This stack indicates where the coro/test is
+                            # suspended
+                            task.print_stack(file=buffer)
+                            task.cancel()
+                            while not task.cancelled():
+                                await asyncio.sleep(0.01)
+                            raise TimeoutError(f"Test timeout.\n{buffer.getvalue()}")
                         finally:
                             if client and c.status not in ("closing", "closed"):
                                 await c._close(fast=s.status == Status.closed)