Merge pull request #25 from RustedBytes/enhance-tasks

fix tasks implementation and reduce monkey patching

Author: egorsmkv

python/rsloop/__init__.py

@@ -11,7 +11,6 @@
 import io as __io
 
 __DLL_DIR_HANDLES: list[object] = []
-__TASK_KWARGS: set[str] | None = None
 
 
 def __configure_windows_dll_search_path() -> None:
@@ -161,26 +160,6 @@ def profile() -> __typing.Iterator[None]:
 _os = __os
 
 
-def __task_kwargs() -> set[str]:
-    global __TASK_KWARGS
-    if __TASK_KWARGS is not None:
-        return __TASK_KWARGS
-
-    import inspect as __inspect
-
-    try:
-        signature = __inspect.signature(__asyncio.Task)
-        __TASK_KWARGS = {
-            name
-            for name, parameter in signature.parameters.items()
-            if parameter.kind is __inspect.Parameter.KEYWORD_ONLY
-        }
-    except (TypeError, ValueError):
-        __TASK_KWARGS = {"loop", "name"}
-
-    return __TASK_KWARGS
-
-
 def __get_asyncgen_state(loop: Loop) -> dict[str, object]:
     state = __ASYNCGEN_STATE.get(loop)
     if state is None:
@@ -286,42 +265,6 @@ def __loop_close(self):
         __LOOP_CONFIG.pop(self, None)
 
 
-def __loop_create_task(self, coro, *, name=None, context=None, **kwargs):
-    task_kwargs_supported = __task_kwargs()
-    factory = self.get_task_factory()
-    if factory is not None:
-        factory_kwargs = dict(kwargs)
-        factory_kwargs["name"] = name
-        if context is not None:
-            factory_kwargs["context"] = context
-        task = factory(self, coro, **factory_kwargs)
-        return task
-
-    task_kwargs = {}
-    extra_kwargs = dict(kwargs)
-
-    if "name" in task_kwargs_supported:
-        task_kwargs["name"] = name
-    if context is not None and "context" in task_kwargs_supported:
-        task_kwargs["context"] = context
-    if "eager_start" in extra_kwargs:
-        eager_start = extra_kwargs.pop("eager_start")
-        if "eager_start" in task_kwargs_supported:
-            task_kwargs["eager_start"] = eager_start
-
-    if extra_kwargs:
-        unexpected = next(iter(extra_kwargs))
-        raise TypeError(
-            f"create_task() got an unexpected keyword argument {unexpected!r}"
-        )
-
-    task = __asyncio.Task(coro, loop=self, **task_kwargs)
-    source_traceback = getattr(task, "_source_traceback", None)
-    if source_traceback:
-        del source_traceback[-1]
-    return task
-
-
 def __cancel_all_tasks(loop: Loop) -> None:
     to_cancel = __asyncio.all_tasks(loop)
     if not to_cancel:
@@ -1597,8 +1540,9 @@ async def __loop_create_connection(
 if Loop.close is __ORIG_CLOSE:
     Loop.close = __loop_close
 
-if Loop.create_task is __ORIG_CREATE_TASK:
-    Loop.create_task = __loop_create_task
+# 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.
 
 
 def __install_ssl_tracking() -> None:

src/loop_core.rs

@@ -282,8 +282,40 @@ impl LoopCore {
             captured,
             context_needs_run,
         ));
-        self.send_command(LoopCommand::ScheduleReady(Arc::clone(&ready)))
-            .map_err(|err| pyo3::exceptions::PyRuntimeError::new_err(err.to_string()))?;
+
+        match self.try_enqueue_local_ready(ReadyItem::Callback(Arc::clone(&ready))) {
+            Ok(()) => return Ok(ready),
+            Err(ReadyItem::Callback(callback)) => {
+                if self
+                    .try_enqueue_active_ready(ReadyItem::Callback(callback))
+                    .is_ok()
+                {
+                    return Ok(ready);
+                }
+            }
+            Err(
+                ReadyItem::Stop
+                | ReadyItem::FutureSetResult { .. }
+                | ReadyItem::FutureSetException { .. }
+                | ReadyItem::StreamTransportRead(_)
+                | ReadyItem::ProcessTransport(_),
+            ) => unreachable!("schedule_callback only enqueues callback ready items"),
+        }
+
+        self.command_tx
+            .send(LoopCommand::ScheduleReady(Arc::clone(&ready)))
+            .map_err(|_| {
+                pyo3::exceptions::PyRuntimeError::new_err(LoopCoreError::ChannelClosed.to_string())
+            })?;
+        #[cfg(target_os = "linux")]
+        if let Some(waker) = self
+            .runtime_waker
+            .lock()
+            .expect("poisoned runtime waker")
+            .as_ref()
+        {
+            waker.wake_by_ref();
+        }
         Ok(ready)
     }
 
@@ -372,19 +404,23 @@ impl LoopCore {
             let mut processed_since_refill = 0_usize;
             loop {
                 if ready_batch.is_empty() || processed_since_refill >= READY_DRAIN_SLICE {
-                    let mut pending = pending_ready.lock().expect("poisoned pending ready queue");
-                    if !pending.is_empty() {
-                        if ready_batch.is_empty() {
-                            std::mem::swap(&mut ready_batch, &mut *pending);
-                        } else {
-                            pending.append(&mut ready_batch);
-                            std::mem::swap(&mut ready_batch, &mut *pending);
+                    let should_check_pending =
+                        ready_batch.is_empty() || wake_pending.load(Ordering::Acquire);
+                    if should_check_pending {
+                        let mut pending =
+                            pending_ready.lock().expect("poisoned pending ready queue");
+                        if !pending.is_empty() {
+                            if ready_batch.is_empty() {
+                                std::mem::swap(&mut ready_batch, &mut *pending);
+                            } else {
+                                pending.append(&mut ready_batch);
+                                std::mem::swap(&mut ready_batch, &mut *pending);
+                            }
+                        }
+                        if pending.is_empty() {
+                            wake_pending.store(false, Ordering::Release);
                         }
                     }
-                    if pending.is_empty() {
-                        wake_pending.store(false, Ordering::Release);
-                    }
-                    drop(pending);
 
                     // Prioritize cross-thread wakeups such as signals and transport
                     // connection_lost notifications so they cannot be starved by a

src/python_api.rs

@@ -43,6 +43,7 @@ use crate::tls::{client_tls_settings, server_tls_settings};
 static ASYNCIO_TASK_CLS: OnceLock<Py<PyAny>> = OnceLock::new();
 static ASYNCIO_FUTURE_CLS: OnceLock<Py<PyAny>> = OnceLock::new();
 static ASYNCIO_GET_RUNNING_LOOP_FN: OnceLock<Py<PyAny>> = OnceLock::new();
+static ASYNCIO_TASK_KWARG_SUPPORT: OnceLock<TaskKwargSupport> = OnceLock::new();
 #[cfg(any(Py_3_12, all(Py_3_11, not(Py_LIMITED_API))))]
 static ASYNCIO_FUTURE_LOOP_KWNAMES: OnceLock<Py<PyTuple>> = OnceLock::new();
 #[cfg(any(Py_3_12, all(Py_3_11, not(Py_LIMITED_API))))]
@@ -56,6 +57,12 @@ static ASYNCIO_TASK_LOOP_NAME_CONTEXT_KWNAMES: OnceLock<Py<PyTuple>> = OnceLock:
 
 type ResolvedStreamAddrinfo = (i32, i32, i32, Py<PyAny>);
 
+struct TaskKwargSupport {
+    name: bool,
+    context: bool,
+    eager_start: bool,
+}
+
 struct TcpServerSocketOptions {
     family: i32,
     flags: i32,
@@ -151,6 +158,52 @@ fn asyncio_get_running_loop_fn(py: Python<'_>) -> PyResult<&Py<PyAny>> {
     Ok(ASYNCIO_GET_RUNNING_LOOP_FN.get_or_init(|| loaded))
 }
 
+fn asyncio_task_kwarg_support(py: Python<'_>) -> PyResult<&'static TaskKwargSupport> {
+    if let Some(cached) = ASYNCIO_TASK_KWARG_SUPPORT.get() {
+        return Ok(cached);
+    }
+
+    let inspect = py.import("inspect")?;
+    let signature = match inspect
+        .getattr("signature")?
+        .call1((asyncio_task_cls(py)?.clone_ref(py),))
+    {
+        Ok(signature) => signature,
+        Err(_) => {
+            return Ok(ASYNCIO_TASK_KWARG_SUPPORT.get_or_init(|| TaskKwargSupport {
+                name: true,
+                context: false,
+                eager_start: false,
+            }));
+        }
+    };
+    let parameters = signature.getattr("parameters")?;
+    let keyword_only = inspect.getattr("Parameter")?.getattr("KEYWORD_ONLY")?;
+    let mut support = TaskKwargSupport {
+        name: false,
+        context: false,
+        eager_start: false,
+    };
+
+    for kwarg_name in ["name", "context", "eager_start"] {
+        let parameter = match parameters.get_item(kwarg_name) {
+            Ok(parameter) => parameter,
+            Err(_) => continue,
+        };
+        if !parameter.getattr("kind")?.eq(&keyword_only)? {
+            continue;
+        }
+        match kwarg_name {
+            "name" => support.name = true,
+            "context" => support.context = true,
+            "eager_start" => support.eager_start = true,
+            _ => {}
+        }
+    }
+
+    Ok(ASYNCIO_TASK_KWARG_SUPPORT.get_or_init(|| support))
+}
+
 #[inline]
 fn call_callable_noargs(py: Python<'_>, callable: &Py<PyAny>) -> PyResult<Py<PyAny>> {
     unsafe { Bound::from_owned_ptr_or_err(py, ffi::compat::PyObject_CallNoArgs(callable.as_ptr())) }
@@ -330,6 +383,22 @@ fn create_asyncio_task_with_kwargs(
     asyncio_task_cls(py)?.call(py, (coro,), Some(&task_kwargs))
 }
 
+fn trim_task_source_traceback(py: Python<'_>, task: &Py<PyAny>) -> PyResult<()> {
+    let Ok(source_traceback) = task.getattr(py, "_source_traceback") else {
+        return Ok(());
+    };
+    if source_traceback.is_none(py) {
+        return Ok(());
+    }
+
+    let source_traceback = source_traceback.bind(py);
+    if source_traceback.len()? == 0 {
+        return Ok(());
+    }
+
+    source_traceback.del_item(source_traceback.len()? - 1)
+}
+
 fn call_protocol_factory(
     py: Python<'_>,
     loop_obj: &Py<PyAny>,
@@ -1298,22 +1367,14 @@ impl PyLoop {
     ) -> PyResult<Py<PyAny>> {
         let core = Arc::clone(&slf.borrow(py).core);
         let loop_obj = Self::as_py_any(py, &slf);
-        let task_kwargs = PyDict::new(py);
-        if let Some(kwargs_in) = kwargs.as_ref() {
-            for (key, value) in kwargs_in.bind(py).iter() {
-                task_kwargs.set_item(key, value)?;
-            }
-        }
-        if let Some(name) = name.as_ref() {
-            task_kwargs.set_item("name", name)?;
-        }
-        if let Some(context) = context.as_ref() {
-            task_kwargs.set_item("context", context)?;
-        }
-        if let Some(eager_start) = eager_start {
-            task_kwargs.set_item("eager_start", eager_start)?;
-        }
-        let has_kwargs = !task_kwargs.is_empty();
+        let task_kwarg_support = asyncio_task_kwarg_support(py)?;
+        let extra_kwargs = kwargs
+            .as_ref()
+            .is_some_and(|kwargs| !kwargs.bind(py).is_empty());
+        let has_kwargs = extra_kwargs
+            || name.is_some()
+            || (context.is_some() && task_kwarg_support.context)
+            || (eager_start.is_some() && task_kwarg_support.eager_start);
 
         if !core.has_task_factory() && !has_kwargs && core.on_runtime_thread() {
             return create_asyncio_task_for_running_loop(py, coro);
@@ -1329,23 +1390,88 @@ impl PyLoop {
         } else {
             None
         };
+
+        if task_factory.is_none() && extra_kwargs {
+            let unexpected = kwargs
+                .as_ref()
+                .and_then(|kwargs| kwargs.bind(py).iter().next().map(|(key, _)| key))
+                .expect("non-empty kwargs when extra_kwargs is true");
+            let unexpected = unexpected.repr()?.extract::<String>()?;
+            return Err(PyTypeError::new_err(format!(
+                "create_task() got an unexpected keyword argument {unexpected}"
+            )));
+        }
+
+        let task_kwargs = if has_kwargs || task_factory.is_some() {
+            let task_kwargs = PyDict::new(py);
+            if let Some(kwargs_in) = kwargs.as_ref() {
+                for (key, value) in kwargs_in.bind(py).iter() {
+                    task_kwargs.set_item(key, value)?;
+                }
+            }
+            if task_factory.is_some() {
+                let factory_name = name
+                    .as_ref()
+                    .map(|name| name.clone_ref(py))
+                    .unwrap_or_else(|| py.None());
+                task_kwargs.set_item("name", factory_name)?;
+            } else if task_kwarg_support.name {
+                if let Some(name) = name.as_ref() {
+                    task_kwargs.set_item("name", name)?;
+                }
+            }
+            if let Some(context) = context.as_ref() {
+                if task_factory.is_some() || task_kwarg_support.context {
+                    task_kwargs.set_item("context", context)?;
+                }
+            }
+            if let Some(eager_start) = eager_start {
+                if task_factory.is_some() || task_kwarg_support.eager_start {
+                    task_kwargs.set_item("eager_start", eager_start)?;
+                }
+            }
+            Some(task_kwargs)
+        } else {
+            None
+        };
+
         if let Some(factory) = task_factory {
-            let created = factory.call(py, (loop_obj.clone_ref(py), coro), Some(&task_kwargs))?;
+            let created = factory.call(py, (loop_obj.clone_ref(py), coro), task_kwargs.as_ref())?;
             return Ok(created);
         }
 
+        let trim_source_traceback = core.get_debug();
         if is_current_running_loop(py, &loop_obj)? {
-            if !has_kwargs {
-                return create_asyncio_task_for_running_loop(py, coro);
+            let created = if !has_kwargs {
+                create_asyncio_task_for_running_loop(py, coro)?
+            } else {
+                create_asyncio_task_with_kwargs(
+                    py,
+                    None,
+                    coro,
+                    task_kwargs.as_ref().expect("task kwargs"),
+                )?
+            };
+            if trim_source_traceback {
+                trim_task_source_traceback(py, &created)?;
             }
-            return create_asyncio_task_with_kwargs(py, None, coro, &task_kwargs);
+            return Ok(created);
         }
 
-        if !has_kwargs {
-            return create_asyncio_task_for_loop(py, &loop_obj, coro, name, context);
+        let created = if !has_kwargs {
+            create_asyncio_task_for_loop(py, &loop_obj, coro, name, context)?
+        } else {
+            create_asyncio_task_with_kwargs(
+                py,
+                Some(&loop_obj),
+                coro,
+                task_kwargs.as_ref().expect("task kwargs"),
+            )?
+        };
+        if trim_source_traceback {
+            trim_task_source_traceback(py, &created)?;
         }
-
-        create_asyncio_task_with_kwargs(py, Some(&loop_obj), coro, &task_kwargs)
+        Ok(created)
     }
 
     fn set_task_factory(&self, factory: Option<Py<PyAny>>) {