if there's an exception in the Client async context manager body then…

… close fast (#6920)

Co-authored-by: fjetter <[email protected]>

distributed/client.py

@@ -1438,7 +1438,12 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
-        await self._close()
+        await self._close(
+            # if we're handling an exception, we assume that it's more
+            # important to deliver that exception than shutdown gracefully.
+            fast=exc_type
+            is not None
+        )
 
     def __exit__(self, exc_type, exc_value, traceback):
         self.close()
@@ -1568,7 +1573,15 @@ def _handle_error(self, exception=None):
         logger.exception(exception)
 
     async def _close(self, fast=False):
-        """Send close signal and wait until scheduler completes"""
+        """
+        Send close signal and wait until scheduler completes
+
+        If fast is True, the client will close forcefully, by cancelling tasks
+        the background _handle_report_task.
+        """
+        # TODO: aclose more forcefully by aborting the RPC and cancelling all
+        # background tasks.
+        # see https://trio.readthedocs.io/en/stable/reference-io.html#trio.aclose_forcefully
         if self.status == "closed":
             return
 

distributed/tests/test_client.py

@@ -28,6 +28,7 @@
 from threading import Semaphore
 from time import sleep
 from typing import Any
+from unittest import mock
 
 import psutil
 import pytest
@@ -7775,6 +7776,24 @@ def __init__(self, *args, **kwargs):
     ]
 
 
+@gen_cluster(client=False, nthreads=[])
+async def test_fast_close_on_aexit_failure(s):
+    class MyException(Exception):
+        pass
+
+    c = Client(s.address, asynchronous=True)
+    with mock.patch.object(c, "_close", wraps=c._close) as _close_proxy:
+        with pytest.raises(MyException):
+            async with c:
+                start = time()
+                raise MyException
+        stop = time()
+
+    assert _close_proxy.mock_calls == [mock.call(fast=True)]
+    assert c.status == "closed"
+    assert (stop - start) < 2
+
+
 @gen_cluster(client=True, nthreads=[])
 async def test_wait_for_workers_no_default(c, s):
     with pytest.warns(