Repository: dask/distributed

Issue:
Flaky tests: `coroutine 'InProc.write' was never awaited`
Catch-all issue for flaky tests with tracebacks like
```python
    def unraisable_exception_runtest_hook() -> Generator[None, None, None]:
        with catch_unraisable_exception() as cm:
            yield
            if cm.unraisable:
                if cm.unraisable.err_msg is not None:
                    err_msg = cm.unraisable.err_msg
                else:
                    err_msg = "Exception ignored in"
                msg = f"{err_msg}: {cm.unraisable.object!r}\n\n"
                msg += "".join(
                    traceback.format_exception(
                        cm.unraisable.exc_type,
                        cm.unraisable.exc_value,
                        cm.unraisable.exc_traceback,
                    )
                )
>               warnings.warn(pytest.PytestUnraisableExceptionWarning(msg))
E               pytest.PytestUnraisableExceptionWarning: Exception ignored in: <coroutine object InProc.write at 0x7fd864183ec0>
E               
E               Traceback (most recent call last):
E                 File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.8/warnings.py", line 506, in _warn_unawaited_coroutine
E                   warn(msg, category=RuntimeWarning, stacklevel=2, source=coro)
E               RuntimeWarning: coroutine 'InProc.write' was never awaited
```

I would imagine there's some core problem that could be fixed to make all these tests work.

Some known failures (feel free to add more in comments):
* https://github.com/dask/distributed/issues/6507
* https://github.com/dask/distributed/issues/6506
* https://github.com/dask/distributed/issues/6451
* https://github.com/dask/distributed/issues/6549 (this one is interesting because it's with a TCP comm, which indicates the problem isn't with InProc but rather with comm handling)
* https://github.com/dask/distributed/issues/6769

@graingert do you have any ideas? Weren't you looking at this at some point?


Discussion:
maybe this? https://github.com/python/cpython/pull/31003
ah no it's this https://github.com/dask/distributed/blob/96c13a5bc26e67d18cee256c2a110102cef7eac6/distributed/batched.py#L94-L96
xref https://github.com/dask/distributed/issues/6389 
Would that be resolved just by switching to asyncio? I don't actually see what's causing it to not be awaited. Is it if the `_background_send` coroutine gets cancelled while it's awaiting `self.comm.write`? AFAIU `_background_send` actually can't be cancelled since we don't have a handle on it.
> Would that be resolved just by switching to asyncio

no it's switching to `await ` that fixes it here - because it moves the yield point after the Coroutine has been awaited

> Is it if the _background_send coroutine gets cancelled while it's awaiting self.comm.write

No the issue is the generator getting garbage collected and raising the GeneratorExit out of the yield here rather than the Future that write goes on to `yield from`
Ah, so just switching to `async` instead of `@gen.coroutine`.
Is there any known issue that would stop us from switching it?
XREF: #6389 
@graingert and I will write up a ticket or comment on what approach we're going to take here. Refactoring BatchedSend to use asyncio right now may not help/be hard, because we don't have something to hang the long-lived task off of yet. We'll probably just ignore the warning for now.

Resolving patch:
diff --git a/distributed/batched.py b/distributed/batched.py
--- a/distributed/batched.py
+++ b/distributed/batched.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import logging
 from collections import deque
 from typing import Any
@@ -94,9 +95,24 @@ def _background_send(self):
             self.batch_count += 1
             self.next_deadline = time() + self.interval
             try:
-                nbytes = yield self.comm.write(
-                    payload, serializers=self.serializers, on_error="raise"
-                )
+                # NOTE: Since `BatchedSend` doesn't have a handle on the running
+                # `_background_send` coroutine, the only thing with a reference to this
+                # coroutine is the event loop itself. If the event loop stops while
+                # we're waiting on a `write`, the `_background_send` coroutine object
+                # may be garbage collected. If that happens, the `yield coro` will raise
+                # `GeneratorExit`. But because this is an old-school `gen.coroutine`,
+                # and we're using `yield` and not `await`, the `write` coroutine object
+                # will not actually have been awaited, and it will remain sitting around
+                # for someone to retrieve it. At interpreter exit, this will warn
+                # sommething like `RuntimeWarning: coroutine 'TCP.write' was never
+                # awaited`. By using the `closing` contextmanager, the `write` coroutine
+                # object is always cleaned up, even if `yield` raises `GeneratorExit`.
+                with contextlib.closing(
+                    self.comm.write(
+                        payload, serializers=self.serializers, on_error="raise"
+                    )
+                ) as coro:
+                    nbytes = yield coro
                 if nbytes < 1e6:
                     self.recent_message_log.append(payload)
                 else:
@@ -160,9 +176,13 @@ def close(self, timeout=None):
             try:
                 if self.buffer:
                     self.buffer, payload = [], self.buffer
-                    yield self.comm.write(
-                        payload, serializers=self.serializers, on_error="raise"
-                    )
+                    # See note in `_background_send` for explanation of `closing`.
+                    with contextlib.closing(
+                        self.comm.write(
+                            payload, serializers=self.serializers, on_error="raise"
+                        )
+                    ) as coro:
+                        yield coro
             except CommClosedError:
                 pass
             yield self.comm.close()


Test patch: