Merge branch 'srothh/transport-class-hierarchy' into srothh/async-task-identifier

Author: srothh

sentry_sdk/api.py

@@ -226,6 +226,14 @@ def flush(
     return get_client().flush(timeout=timeout, callback=callback)
 
 
+@clientmethod
+async def flush_async(
+    timeout: Optional[float] = None,
+    callback: Optional[Callable[[int, float], None]] = None,
+) -> None:
+    return await get_client().flush_async(timeout=timeout, callback=callback)
+
+
 def start_span(**kwargs: Any) -> Span:
     """
     Start and return a span.

sentry_sdk/client.py

@@ -38,7 +38,6 @@
 from sentry_sdk.integrations.dedupe import DedupeIntegration
 from sentry_sdk.sessions import SessionFlusher
 from sentry_sdk.envelope import Envelope
-from sentry_sdk.utils import mark_sentry_task_internal
 
 from sentry_sdk.profiler.continuous_profiler import setup_continuous_profiler
 from sentry_sdk.profiler.transaction_profiler import (
@@ -216,6 +215,12 @@ def close(self, *args: Any, **kwargs: Any) -> None:
     def flush(self, *args: Any, **kwargs: Any) -> None:
         return None
 
+    async def close_async(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+    async def flush_async(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
     def __enter__(self) -> BaseClient:
         return self
 
@@ -919,14 +924,6 @@ def get_integration(
 
         return self.integrations.get(integration_name)
 
-    def _close_transport(self) -> Optional[asyncio.Task[None]]:
-        """Close transport and return cleanup task if any."""
-        if self.transport is not None:
-            cleanup_task = self.transport.kill()  # type: ignore
-            self.transport = None
-            return cleanup_task
-        return None
-
     def _close_components(self) -> None:
         """Kill all client components in the correct order."""
         self.session_flusher.kill()
@@ -935,14 +932,7 @@ def _close_components(self) -> None:
         if self.monitor:
             self.monitor.kill()
 
-    async def _close_components_async(self) -> None:
-        """Async version of _close_components that properly awaits transport cleanup."""
-        self._close_components()
-        cleanup_task = self._close_transport()
-        if cleanup_task is not None:
-            await cleanup_task
-
-    def close(  # type: ignore[override]
+    def close(
         self,
         timeout: Optional[float] = None,
         callback: Optional[Callable[[int, float], None]] = None,
@@ -951,40 +941,45 @@ def close(  # type: ignore[override]
         Close the client and shut down the transport. Arguments have the same
         semantics as :py:meth:`Client.flush`. When using the async transport, close needs to be awaited to block.
         """
-
-        async def _flush_and_close(
-            timeout: Optional[float], callback: Optional[Callable[[int, float], None]]
-        ) -> None:
-
-            await self._flush_async(timeout=timeout, callback=callback)
-            await self._close_components_async()
-
         if self.transport is not None:
             if isinstance(self.transport, AsyncHttpTransport) and hasattr(
                 self.transport, "loop"
             ):
+                logger.debug(
+                    "close() used with AsyncHttpTransport, aborting. Please use close_async() instead."
+                )
+                return
+            self.flush(timeout=timeout, callback=callback)
+            self._close_components()
+            self.transport.kill()
+            self.transport = None
 
-                try:
-                    with mark_sentry_task_internal():
-                        flush_task = self.transport.loop.create_task(
-                            _flush_and_close(timeout, callback)
-                        )
-                except RuntimeError:
-                    # Shutdown the components anyway
-                    self._close_components()
-                    self._close_transport()
-                    logger.warning("Event loop not running, aborting close.")
-                    return None
-                # Enforce flush before shutdown
-                return flush_task
-            else:
-                self.flush(timeout=timeout, callback=callback)
-                self._close_components()
-                self._close_transport()
-
-        return None
+    async def close_async(
+        self,
+        timeout: Optional[float] = None,
+        callback: Optional[Callable[[int, float], None]] = None,
+    ) -> None:
+        """
+        Asynchronously close the client and shut down the transport. Arguments have the same
+        semantics as :py:meth:`Client.flush_async`.
+        """
+        if self.transport is not None:
+            if not (
+                isinstance(self.transport, AsyncHttpTransport)
+                and hasattr(self.transport, "loop")
+            ):
+                logger.debug(
+                    "close_async() used with non-async transport, aborting. Please use close() instead."
+                )
+                return
+            await self.flush_async(timeout=timeout, callback=callback)
+            self._close_components()
+            kill_task = self.transport.kill()  # type: ignore
+            if kill_task is not None:
+                await kill_task
+            self.transport = None
 
-    def flush(  # type: ignore[override]
+    def flush(
         self,
         timeout: Optional[float] = None,
         callback: Optional[Callable[[int, float], None]] = None,
@@ -1000,38 +995,40 @@ def flush(  # type: ignore[override]
             if isinstance(self.transport, AsyncHttpTransport) and hasattr(
                 self.transport, "loop"
             ):
-                try:
-                    with mark_sentry_task_internal():
-                        return self.transport.loop.create_task(
-                            self._flush_async(timeout, callback)
-                        )
-                except RuntimeError:
-                    logger.warning("Event loop not running, aborting flush.")
-                    return None
-            else:
-                self._flush_sync(timeout, callback)
-        return None
-
-    def _flush_sync(
-        self, timeout: Optional[float], callback: Optional[Callable[[int, float], None]]
-    ) -> None:
-        """Synchronous flush implementation."""
-        if timeout is None:
-            timeout = self.options["shutdown_timeout"]
+                logger.debug(
+                    "flush() used with AsyncHttpTransport, aborting. Please use flush_async() instead."
+                )
+                return
+            if timeout is None:
+                timeout = self.options["shutdown_timeout"]
+            self._flush_components()
 
-        self._flush_components()
-        if self.transport is not None:
             self.transport.flush(timeout=timeout, callback=callback)
 
-    async def _flush_async(
-        self, timeout: Optional[float], callback: Optional[Callable[[int, float], None]]
+    async def flush_async(
+        self,
+        timeout: Optional[float] = None,
+        callback: Optional[Callable[[int, float], None]] = None,
     ) -> None:
-        """Asynchronous flush implementation."""
-        if timeout is None:
-            timeout = self.options["shutdown_timeout"]
+        """
+        Asynchronously wait for the current events to be sent.
 
-        self._flush_components()
+        :param timeout: Wait for at most `timeout` seconds. If no `timeout` is provided, the `shutdown_timeout` option value is used.
+
+        :param callback: Is invoked with the number of pending events and the configured timeout.
+        """
         if self.transport is not None:
+            if not (
+                isinstance(self.transport, AsyncHttpTransport)
+                and hasattr(self.transport, "loop")
+            ):
+                logger.debug(
+                    "flush_async() used with non-async transport, aborting. Please use flush() instead."
+                )
+                return
+            if timeout is None:
+                timeout = self.options["shutdown_timeout"]
+            self._flush_components()
             flush_task = self.transport.flush(timeout=timeout, callback=callback)  # type: ignore
             if flush_task is not None:
                 await flush_task

sentry_sdk/integrations/asyncio.py

@@ -58,9 +58,7 @@ async def _flush() -> None:
             if not isinstance(client.transport, AsyncHttpTransport):
                 return
 
-            task = client.close()  # type: ignore
-            if task is not None:
-                await task
+            await client.close_async()
         except Exception:
             logger.warning("Sentry flush failed during loop shutdown", exc_info=True)
 

tests/test_client.py

@@ -1732,7 +1732,7 @@ async def test_async_proxy(monkeypatch, testcase):
             )
             assert proxy_headers == testcase["arg_proxy_headers"]
 
-    await client.close()
+    await client.close_async()
 
 
 @pytest.mark.parametrize(
@@ -1820,4 +1820,4 @@ async def test_async_socks_proxy(testcase):
         f"but got {str(type(client.transport._pool))}"
     )
 
-    await client.close()
+    await client.close_async()

tests/test_transport.py

@@ -204,9 +204,9 @@ async def test_transport_works_async(
     capture_message("löl")
 
     if client_flush_method == "close":
-        await client.close(timeout=2.0)
+        await client.close_async(timeout=2.0)
     if client_flush_method == "flush":
-        await client.flush(timeout=2.0)
+        await client.flush_async(timeout=2.0)
 
     out, err = capsys.readouterr()
     assert not err and not out
@@ -230,7 +230,7 @@ async def test_transport_works_async(
     # Therefore, we need to explicitly close the client to clean up the worker task
     assert any("Sending envelope" in record.msg for record in caplog.records) == debug
     if client_flush_method == "flush":
-        await client.close(timeout=2.0)
+        await client.close_async(timeout=2.0)
 
 
 @pytest.mark.parametrize(
@@ -833,7 +833,7 @@ def background_thread_work():
     thread.join()
     assert not exception_from_thread
     assert captured_from_thread
-    await client.close(timeout=2.0)
+    await client.close_async(timeout=2.0)
     assert capturing_server.captured
 
 
@@ -860,7 +860,7 @@ async def test_async_transport_event_loop_closed_scenario(
                 )
 
     client.transport.loop = original_loop
-    await client.close(timeout=2.0)
+    await client.close_async(timeout=2.0)
 
 
 @pytest.mark.asyncio
@@ -882,7 +882,7 @@ async def send_message(i):
 
     tasks = [send_message(i) for i in range(num_messages)]
     await asyncio.gather(*tasks)
-    await client.close(timeout=2.0)
+    await client.close_async(timeout=2.0)
     assert len(capturing_server.captured) == num_messages
 
 
@@ -916,7 +916,7 @@ async def send_message(i):
     await asyncio.sleep(0.1)
     # New request should be dropped due to rate limiting
     assert len(capturing_server.captured) == 0
-    await client.close(timeout=2.0)
+    await client.close_async(timeout=2.0)
 
 
 @pytest.mark.asyncio
@@ -938,4 +938,4 @@ async def test_async_two_way_ssl_authentication():
     options = client.transport._get_pool_options()
     assert options["ssl_context"] is not None
 
-    await client.close()
+    await client.close_async()