replace gen.coroutine usage with async def

distributed/batched.py

@@ -72,11 +72,10 @@ def __repr__(self):
 
     __str__ = __repr__
 
-    @gen.coroutine
-    def _background_send(self):
+    async def _background_send(self):
         while not self.please_stop:
             try:
-                yield self.waker.wait(self.next_deadline)
+                await self.waker.wait(self.next_deadline)
                 self.waker.clear()
             except gen.TimeoutError:
                 pass
@@ -91,7 +90,7 @@ def _background_send(self):
             self.batch_count += 1
             self.next_deadline = time() + self.interval
             try:
-                nbytes = yield self.comm.write(
+                nbytes = await self.comm.write(
                     payload, serializers=self.serializers, on_error="raise"
                 )
                 if nbytes < 1e6:
@@ -142,8 +141,8 @@ def send(self, *msgs: dict) -> None:
         if self.next_deadline is None:
             self.waker.set()
 
-    @gen.coroutine
-    def close(self, timeout=None):
+
+    async def close(self, timeout=None):
         """Flush existing messages and then close comm
 
         If set, raises `tornado.util.TimeoutError` after a timeout.
@@ -152,17 +151,17 @@ def close(self, timeout=None):
             return
         self.please_stop = True
         self.waker.set()
-        yield self.stopped.wait(timeout=timeout)
+        await self.stopped.wait(timeout=timeout)
         if not self.comm.closed():
             try:
                 if self.buffer:
                     self.buffer, payload = [], self.buffer
-                    yield self.comm.write(
+                    await self.comm.write(
                         payload, serializers=self.serializers, on_error="raise"
                     )
             except CommClosedError:
                 pass
-            yield self.comm.close()
+            await self.comm.close()
 
     def abort(self):
         if self.comm is None:

distributed/comm/tcp.py

@@ -12,8 +12,6 @@
 from ssl import SSLCertVerificationError, SSLError
 from typing import Any, ClassVar
 
-from tornado import gen
-
 try:
     import ssl
 except ImportError:
@@ -327,8 +325,7 @@ async def write(self, msg, serializers=None, on_error="message"):
 
         return frames_nbytes_total
 
-    @gen.coroutine
-    def close(self):
+    async def close(self):
         # We use gen.coroutine here rather than async def to avoid errors like
         # Task was destroyed but it is pending!
         # Triggered by distributed.deploy.tests.test_local::test_silent_startup
@@ -338,7 +335,7 @@ def close(self):
             try:
                 # Flush the stream's write buffer by waiting for a last write.
                 if stream.writing():
-                    yield stream.write(b"")
+                    await stream.write(b"")
                 stream.socket.shutdown(socket.SHUT_RDWR)
             except OSError:
                 pass