Unwind some more future/couritine issues

Author: dwoz

salt/minion.py

@@ -3114,8 +3114,9 @@ def cleanup_subprocesses(self):
         # Add an extra fallback in case a forked process leaks through
         multiprocessing.active_children()
         self.subprocess_list.cleanup()
-        if self.schedule:
-            self.schedule.cleanup_subprocesses()
+        schedule = getattr(self, "schedule", None)
+        if schedule:
+            schedule.cleanup_subprocesses()
 
     def _setup_core(self):
         """
@@ -3405,7 +3406,11 @@ def destroy(self):
             self.req_channel.close()
         if hasattr(self, "periodic_callbacks"):
             for cb in self.periodic_callbacks.values():
-                cb.stop()
+                if hasattr(cb, "stop"):
+                    cb.stop()
+                elif asyncio.isfuture(cb) or isinstance(cb, asyncio.Task):
+                    cb.cancel()
+            self.periodic_callbacks.clear()
 
     # pylint: disable=W1701
     def __del__(self):

salt/transport/tcp.py

@@ -1513,7 +1513,12 @@ async def publish(
         """
         if not self.pub_sock:
             self.connect()
-        self.pub_sock.send(payload)
+        log.debug("TCP PublishServer publishing payload (%d bytes)", len(payload))
+        result = self.pub_sock.send(payload)
+        if asyncio.iscoroutine(result) or isinstance(result, asyncio.Future):
+            await result
+        else:
+            log.debug("TCP PublishServer publish returned %r", result)
 
     def close(self):
         self._closing = True
@@ -1607,15 +1612,23 @@ def __init__(self, host, port, path, io_loop=None):
         self.stream = None
         self.unpacker = salt.utils.msgpack.Unpacker(raw=False)
         self._connecting_future = None
+        self._asyncio_loop = None
+        if hasattr(self.io_loop, "asyncio_loop"):
+            self._asyncio_loop = self.io_loop.asyncio_loop
 
     def connected(self):
         return self.stream is not None and not self.stream.closed()
 
-    def connect(self, callback=None, timeout=None):
+    def _legacy_connect(self, callback=None, timeout=None):
         """
         Connect to the IPC socket
         """
         if self._connecting_future is not None and not self._connecting_future.done():
+            log.debug(
+                "%s connect reuse pending future (closing=%s)",
+                self.__class__.__name__,
+                self._closing,
+            )
             future = self._connecting_future
         else:
             if self._connecting_future is not None:
@@ -1624,6 +1637,12 @@ def connect(self, callback=None, timeout=None):
             future = tornado.concurrent.Future()
             self._connecting_future = future
             # self._connect(timeout)
+            log.debug(
+                "%s connect spawning _connect (closing=%s, timeout=%s)",
+                self.__class__.__name__,
+                self._closing,
+                timeout,
+            )
             self.io_loop.spawn_callback(self._connect, timeout)
 
         if callback is not None:
@@ -1636,6 +1655,34 @@ def handle_future(future):
 
         return future
 
+    async def connect(self, callback=None, timeout=None):
+        if self._asyncio_loop is None:
+            # Fall back to the legacy tornado.Future based implementation
+            future = self._legacy_connect(callback=callback, timeout=timeout)
+            loop = asyncio.get_running_loop()
+            await salt.utils.asynchronous._ensure_task(loop, future)
+            return True
+
+        if self._connecting_future is None or self._connecting_future.done():
+            self._connecting_future = self._asyncio_loop.create_future()
+
+            async def runner():
+                try:
+                    await self._connect(timeout=timeout)
+                except Exception as exc:  # pylint: disable=broad-except
+                    if not self._connecting_future.done():
+                        self._connecting_future.set_exception(exc)
+                else:
+                    if not self._connecting_future.done():
+                        self._connecting_future.set_result(True)
+
+            asyncio.ensure_future(runner(), loop=self._asyncio_loop)
+
+        result = await self._connecting_future
+        if callback is not None:
+            self.io_loop.spawn_callback(callback, result)
+        return result
+
     async def _connect(self, timeout=None):
         """
         Connect to a running IPCServer
@@ -1652,9 +1699,19 @@ async def _connect(self, timeout=None):
         self.stream = None
         if timeout is not None:
             timeout_at = time.monotonic() + timeout
+            log.debug(
+                "%s will timeout connection at %.3f (in %.3fs)",
+                self.__class__.__name__,
+                timeout_at,
+                timeout,
+            )
 
         while True:
             if self._closing:
+                log.debug(
+                    "%s connect aborting due to closing flag",
+                    self.__class__.__name__,
+                )
                 break
 
             if self.stream is None:
@@ -1664,6 +1721,12 @@ async def _connect(self, timeout=None):
                 )
             try:
                 await self.stream.connect(sock_addr)
+                log.debug(
+                    "%s connected to %s (closing=%s)",
+                    self.__class__.__name__,
+                    sock_addr,
+                    self._closing,
+                )
                 self._connecting_future.set_result(True)
                 break
             except Exception as e:  # pylint: disable=broad-except
@@ -1674,6 +1737,12 @@ async def _connect(self, timeout=None):
                     if self.stream is not None:
                         self.stream.close()
                         self.stream = None
+                    log.debug(
+                        "%s connection to %s failed: %s",
+                        self.__class__.__name__,
+                        sock_addr,
+                        e,
+                    )
                     self._connecting_future.set_exception(e)
                     break
 
@@ -1728,7 +1797,14 @@ async def send(self, msg, timeout=None, tries=None):
         if not self.connected():
             await self.connect()
         pack = salt.transport.frame.frame_msg_ipc(msg, raw_body=True)
+        log.debug(
+            "%s sending %d bytes (closing=%s)",
+            self.__class__.__name__,
+            len(pack),
+            self._closing,
+        )
         await self.stream.write(pack)
+        log.debug("%s send complete", self.__class__.__name__)
 
 
 class RequestClient(salt.transport.base.RequestClient):

salt/transport/zeromq.py

@@ -575,7 +575,8 @@ async def request_handler(self):
                 continue
             except Exception as exc:  # pylint: disable=broad-except
                 log.error(
-                    "Exception in request handler",
+                    "Exception in request handler: %r",
+                    exc,
                     exc_info_on_loglevel=logging.DEBUG,
                 )
                 continue

salt/utils/asynchronous.py

@@ -19,9 +19,16 @@
 from typing import Any
 
 try:  # Optional dependency during the transition away from tornado.
+    import tornado  # type: ignore
+    import tornado.concurrent  # type: ignore
     import tornado.ioloop  # type: ignore
+    from tornado.platform.asyncio import to_asyncio_future  # type: ignore
 except ImportError:  # pragma: no cover - tornado optional
     tornado = None  # type: ignore
+    to_asyncio_future = None  # type: ignore
+    _TORNADO_FUTURE_TYPES: tuple[type[Any], ...] = ()
+else:
+    _TORNADO_FUTURE_TYPES = (tornado.concurrent.Future,)  # type: ignore[attr-defined]
 
 log = logging.getLogger(__name__)
 
@@ -32,8 +39,94 @@ def _ensure_task(loop: asyncio.AbstractEventLoop, result: Any) -> Any:
     return it unchanged.
     """
 
-    if asyncio.iscoroutine(result) or isinstance(result, asyncio.Future):
-        return loop.create_task(result)
+    if asyncio.iscoroutine(result):
+        return asyncio.ensure_future(result, loop=loop)
+
+    if isinstance(result, asyncio.Future):
+        try:
+            future_loop = result.get_loop()
+        except AttributeError:
+            future_loop = loop
+        if future_loop is loop:
+            return asyncio.ensure_future(result, loop=loop)
+
+        proxy = loop.create_future()
+
+        def _relay(src_future: asyncio.Future):
+            if proxy.done():
+                return
+            if src_future.cancelled():
+                loop.call_soon_threadsafe(proxy.cancel)
+                return
+            exc = src_future.exception()
+            if exc is not None:
+                loop.call_soon_threadsafe(proxy.set_exception, exc)
+                return
+            loop.call_soon_threadsafe(proxy.set_result, src_future.result())
+
+        result.add_done_callback(_relay)
+        return proxy
+
+    if _TORNADO_FUTURE_TYPES and isinstance(
+        result, _TORNADO_FUTURE_TYPES  # type: ignore[arg-type]
+    ):
+        converted = None
+        if to_asyncio_future is not None:
+            try:
+                converted = to_asyncio_future(result)
+            except TypeError:
+                converted = None
+            except Exception:  # pylint: disable=broad-except
+                log.exception(
+                    "Failed to convert future %r to asyncio future using tornado helper",
+                    result,
+                )
+        if converted is not None:
+            try:
+                return asyncio.ensure_future(converted, loop=loop)
+            except TypeError:
+                converted = None
+        if converted is None:
+            converted = loop.create_future()
+
+            def _relay(src_future):
+                if converted.done():
+                    return
+                if src_future.cancelled():
+                    loop.call_soon_threadsafe(converted.cancel)
+                    return
+                exc = src_future.exception()
+                if exc is not None:
+                    loop.call_soon_threadsafe(converted.set_exception, exc)
+                    return
+                loop.call_soon_threadsafe(converted.set_result, src_future.result())
+
+            result.add_done_callback(_relay)
+        return asyncio.ensure_future(converted, loop=loop)
+
+    if hasattr(result, "add_done_callback"):
+        try:
+            wrapped = asyncio.wrap_future(result, loop=loop)
+            return asyncio.ensure_future(wrapped, loop=loop)
+        except TypeError:
+            bridge = loop.create_future()
+
+            def _relay(src_future):
+                if bridge.done():
+                    return
+                if getattr(src_future, "cancelled", lambda: False)():
+                    loop.call_soon_threadsafe(bridge.cancel)
+                    return
+                exc = getattr(src_future, "exception", lambda: None)()
+                if exc is not None:
+                    loop.call_soon_threadsafe(bridge.set_exception, exc)
+                    return
+                result_value = getattr(src_future, "result", lambda: None)()
+                loop.call_soon_threadsafe(bridge.set_result, result_value)
+
+            result.add_done_callback(_relay)
+            return bridge
+
     return result
 
 
@@ -103,9 +196,23 @@ def remove_timeout(self, handle: asyncio.TimerHandle):
         handle.cancel()
 
     def add_future(self, future: Awaitable, callback: Callable[[asyncio.Future], Any]):
-        fut = asyncio.ensure_future(future, loop=self._loop)
-        fut.add_done_callback(lambda done: self._loop.call_soon(callback, done))
-        return fut
+        scheduled = _ensure_task(self._loop, future)
+        if isinstance(scheduled, asyncio.Future):
+            scheduled.add_done_callback(
+                lambda done: self._loop.call_soon(callback, done)
+            )
+            return scheduled
+        if hasattr(scheduled, "add_done_callback"):
+
+            def _relay(done):
+                try:
+                    loop = self._loop
+                    loop.call_soon_threadsafe(callback, done)
+                except RuntimeError:
+                    loop.call_soon(callback, done)
+
+            scheduled.add_done_callback(_relay)
+        return scheduled
 
     def create_task(self, coro: Awaitable):
         return self._loop.create_task(coro)
@@ -120,16 +227,17 @@ def run_sync(
         **kwargs,
     ):
         result = func(*args, **kwargs)
-        if asyncio.iscoroutine(result) or isinstance(result, asyncio.Future):
+        scheduled = _ensure_task(self._loop, result)
+        if asyncio.iscoroutine(scheduled) or isinstance(scheduled, asyncio.Future):
             if self._loop.is_running():
                 raise RuntimeError("Cannot run_sync on a running event loop")
             policy = asyncio.get_event_loop_policy()
             try:
                 policy.set_event_loop(self._loop)
-                return self._loop.run_until_complete(result)
+                return self._loop.run_until_complete(scheduled)
             finally:
                 policy.set_event_loop(None)
-        return result
+        return scheduled
 
     def start(self):
         policy = asyncio.get_event_loop_policy()