Merge pull request #26 from RustedBytes/reduce-py

move more to rust side

Author: egorsmkv

python/rsloop/__init__.py

@@ -7,8 +7,15 @@
 import ssl as __ssl
 import sys as __sys
 import typing as __typing
-import asyncio.base_events as __asyncio_base_events
-import io as __io
+
+from ._loop import PyLoop as Loop
+from ._loop import __version__
+from ._loop import new_event_loop
+from ._loop import start_server as __start_server
+from ._loop import open_connection as __open_connection
+from ._loop import profiler_running
+from ._loop import start_profiler
+from ._loop import stop_profiler
 
 __DLL_DIR_HANDLES: list[object] = []
 
@@ -49,14 +56,6 @@ def __configure_windows_dll_search_path() -> None:
 
 __configure_windows_dll_search_path()
 
-from ._loop import PyLoop as Loop
-from ._loop import __version__
-from ._loop import open_connection as __open_connection
-from ._loop import profiler_running as __profiler_running
-from ._loop import start_profiler as __start_profiler
-from ._loop import start_server as __start_server
-from ._loop import stop_profiler as __stop_profiler
-
 __all__: tuple[str, ...] = (
     "Loop",
     "__version__",
@@ -102,24 +101,6 @@ def __set_event_loop(loop: Loop | None) -> None:
     __asyncio.events.set_event_loop = __set_event_loop
 
 
-def new_event_loop() -> Loop:
-    return Loop()
-
-
-def profiler_running() -> bool:
-    return __profiler_running()
-
-
-def start_profiler() -> None:
-    """Start a Tracy profiling session."""
-    __start_profiler()
-
-
-def stop_profiler() -> None:
-    """Stop the active Tracy profiling session."""
-    __stop_profiler()
-
-
 @__contextlib.contextmanager
 def profile() -> __typing.Iterator[None]:
     """Context manager wrapper around ``start_profiler()`` / ``stop_profiler()``."""
@@ -141,130 +122,17 @@ def profile() -> __typing.Iterator[None]:
 __ORIG_SOCK_RECVFROM_INTO = getattr(Loop, "sock_recvfrom_into", None)
 __ORIG_SOCK_SENDTO = getattr(Loop, "sock_sendto", None)
 __ORIG_SOCK_SENDFILE = getattr(Loop, "sock_sendfile", None)
-__ORIG_RUN_FOREVER = Loop.run_forever
-__ORIG_RUN_UNTIL_COMPLETE = Loop.run_until_complete
-__ORIG_SHUTDOWN_ASYNCGENS = Loop.shutdown_asyncgens
-__ORIG_CLOSE = Loop.close
-__ORIG_CREATE_TASK = Loop.create_task
 __USE_FAST_STREAMS = __os.environ.get("RSLOOP_USE_FAST_STREAMS", "1") != "0"
-__ASYNCGEN_STATE: dict[Loop, dict[str, object]] = {}
-__LOOP_CONFIG: dict[Loop, dict[str, object]] = {}
 
 if __USE_FAST_STREAMS and __asyncio.open_connection is __ORIG_OPEN_CONNECTION:
     __asyncio.open_connection = __open_connection
 if __USE_FAST_STREAMS and __asyncio.start_server is __ORIG_START_SERVER:
     __asyncio.start_server = __start_server
 
 _asyncio = __asyncio
-_io = __io
 _os = __os
 
 
-def __get_asyncgen_state(loop: Loop) -> dict[str, object]:
-    state = __ASYNCGEN_STATE.get(loop)
-    if state is None:
-        state = {
-            "active": set(),
-            "shutdown_called": False,
-            "old_hooks": None,
-        }
-        __ASYNCGEN_STATE[loop] = state
-    return state
-
-
-def __get_loop_config(loop: Loop) -> dict[str, object]:
-    config = __LOOP_CONFIG.get(loop)
-    if config is None:
-        config = {
-            "slow_callback_duration": 0.1,
-        }
-        __LOOP_CONFIG[loop] = config
-    return config
-
-
-@__contextlib.contextmanager
-def __asyncgen_hooks_installed(loop: Loop) -> __typing.Iterator[None]:
-    state = __get_asyncgen_state(loop)
-    old_hooks = __sys.get_asyncgen_hooks()
-    state["old_hooks"] = old_hooks
-    __sys.set_asyncgen_hooks(
-        firstiter=lambda agen: __asyncgen_firstiter_hook(loop, agen),
-        finalizer=lambda agen: __asyncgen_finalizer_hook(loop, agen),
-    )
-    try:
-        yield None
-    finally:
-        saved_hooks = state.get("old_hooks")
-        if saved_hooks is not None:
-            __sys.set_asyncgen_hooks(*saved_hooks)
-            state["old_hooks"] = None
-
-
-def __asyncgen_firstiter_hook(loop: Loop, agen) -> None:
-    state = __get_asyncgen_state(loop)
-    if state["shutdown_called"]:
-        import warnings as __warnings
-
-        __warnings.warn(
-            f"asynchronous generator {agen!r} was scheduled after "
-            f"loop.shutdown_asyncgens() call",
-            ResourceWarning,
-            source=loop,
-        )
-    state["active"].add(agen)
-
-
-def __asyncgen_finalizer_hook(loop: Loop, agen) -> None:
-    state = __get_asyncgen_state(loop)
-    state["active"].discard(agen)
-    if not loop.is_closed():
-        loop.call_soon_threadsafe(loop.create_task, agen.aclose())
-
-
-async def __loop_shutdown_asyncgens(self):
-    state = __get_asyncgen_state(self)
-    state["shutdown_called"] = True
-
-    if not state["active"]:
-        return
-
-    closing_agens = list(state["active"])
-    state["active"].clear()
-
-    results = await __asyncio.gather(
-        *(agen.aclose() for agen in closing_agens),
-        return_exceptions=True,
-    )
-
-    for result, agen in zip(results, closing_agens):
-        if isinstance(result, Exception):
-            self.call_exception_handler(
-                {
-                    "message": f"an error occurred during closing of asynchronous generator {agen!r}",
-                    "exception": result,
-                    "asyncgen": agen,
-                }
-            )
-
-
-def __loop_run_forever(self):
-    with __asyncgen_hooks_installed(self):
-        return __ORIG_RUN_FOREVER(self)
-
-
-def __loop_run_until_complete(self, future):
-    with __asyncgen_hooks_installed(self):
-        return __ORIG_RUN_UNTIL_COMPLETE(self, future)
-
-
-def __loop_close(self):
-    try:
-        return __ORIG_CLOSE(self)
-    finally:
-        __ASYNCGEN_STATE.pop(self, None)
-        __LOOP_CONFIG.pop(self, None)
-
-
 def __cancel_all_tasks(loop: Loop) -> None:
     to_cancel = __asyncio.all_tasks(loop)
     if not to_cancel:
@@ -574,14 +442,6 @@ async def __loop_sock_sendfile(
             file.seek(offset + total_sent)
 
 
-def __get_slow_callback_duration(self):
-    return __get_loop_config(self)["slow_callback_duration"]
-
-
-def __set_slow_callback_duration(self, value):
-    __get_loop_config(self)["slow_callback_duration"] = float(value)
-
-
 async def __loop_create_datagram_endpoint(
     self,
     protocol_factory,
@@ -1522,24 +1382,6 @@ async def __loop_create_connection(
 if __ORIG_SOCK_SENDFILE is None:
     Loop.sock_sendfile = __loop_sock_sendfile
 
-if not hasattr(Loop, "slow_callback_duration"):
-    Loop.slow_callback_duration = property(
-        __get_slow_callback_duration,
-        __set_slow_callback_duration,
-    )
-
-if Loop.run_forever is __ORIG_RUN_FOREVER:
-    Loop.run_forever = __loop_run_forever
-
-if Loop.run_until_complete is __ORIG_RUN_UNTIL_COMPLETE:
-    Loop.run_until_complete = __loop_run_until_complete
-
-if Loop.shutdown_asyncgens is __ORIG_SHUTDOWN_ASYNCGENS:
-    Loop.shutdown_asyncgens = __loop_shutdown_asyncgens
-
-if Loop.close is __ORIG_CLOSE:
-    Loop.close = __loop_close
-
 # Keep the Rust implementation on the hot path. It already handles task
 # factories and keyword forwarding, while the Python wrapper adds measurable
 # overhead in task-heavy workloads.

src/lib.rs

@@ -21,7 +21,10 @@ pub use fast_streams::{open_connection, start_server, PyFastStreamReader, PyFast
 pub use loop_core::{LoopCommand, LoopCore};
 pub use process_transport::{PyProcessPipeTransport, PyProcessTransport};
 pub use profiler::{profiler_running, start_profiler, stop_profiler};
-pub use python_api::{future_done_stop, new_event_loop, signal_bridge, PyLoop};
+pub use python_api::{
+    asyncgen_finalizer_hook, asyncgen_firstiter_hook, future_done_stop, new_event_loop,
+    signal_bridge, PyLoop,
+};
 pub use stream_transport::{PyServer, PyStreamTransport};
 
 use pyo3::prelude::*;
@@ -39,6 +42,8 @@ fn _loop(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyFastStreamReader>()?;
     m.add_class::<PyFastStreamWriter>()?;
     m.add_function(wrap_pyfunction!(new_event_loop, m)?)?;
+    m.add_function(wrap_pyfunction!(asyncgen_firstiter_hook, m)?)?;
+    m.add_function(wrap_pyfunction!(asyncgen_finalizer_hook, m)?)?;
     m.add_function(wrap_pyfunction!(future_done_stop, m)?)?;
     m.add_function(wrap_pyfunction!(open_connection, m)?)?;
     m.add_function(wrap_pyfunction!(profiler_running, m)?)?;
@@ -47,5 +52,13 @@ fn _loop(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(start_server, m)?)?;
     m.add_function(wrap_pyfunction!(stop_profiler, m)?)?;
     m.add("_future_done_stop", m.getattr("future_done_stop")?)?;
+    m.add(
+        "_asyncgen_firstiter_hook",
+        m.getattr("asyncgen_firstiter_hook")?,
+    )?;
+    m.add(
+        "_asyncgen_finalizer_hook",
+        m.getattr("asyncgen_finalizer_hook")?,
+    )?;
     Ok(())
 }

src/loop_core.rs

@@ -18,7 +18,7 @@ use crate::runtime::run_runtime_thread;
 use crate::stream_transport::{ReaderTarget, ServerCore, ServerListener, StreamTransportCore};
 use crossbeam_channel::Sender;
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyTuple};
+use pyo3::types::{PyDict, PySet, PyTuple};
 
 const READY_DRAIN_SLICE: usize = 64;
 
@@ -135,7 +135,9 @@ pub struct LoopState {
     pub closed: bool,
     pub running: bool,
     pub stopping: bool,
+    pub slow_callback_duration: f64,
     pub asyncgens_shutdown_called: bool,
+    pub active_asyncgens: Option<Py<PySet>>,
     pub executor_shutdown_called: bool,
     pub signal_handlers: HashMap<i32, SignalHandlerTemplate>,
     pub previous_signal_handlers: HashMap<i32, Py<PyAny>>,
@@ -152,7 +154,9 @@ impl LoopState {
             closed: false,
             running: false,
             stopping: false,
+            slow_callback_duration: 0.1,
             asyncgens_shutdown_called: false,
+            active_asyncgens: None,
             executor_shutdown_called: false,
             signal_handlers: HashMap::new(),
             previous_signal_handlers: HashMap::new(),