Wrap task locals loop and context in Arc.

Author: iksteen

src/async_std.rs

@@ -94,7 +94,7 @@ impl ContextExt for AsyncStdRuntime {
             .try_with(|c| {
                 c.borrow()
                     .as_ref()
-                    .map(|locals| Python::attach(|py| locals.clone_ref(py)))
+                    .map(|locals| locals.clone_ref())
             })
             .unwrap_or_default()
     }

src/generic.rs

@@ -91,7 +91,7 @@ where
     R: ContextExt,
 {
     if let Some(locals) = R::get_task_locals() {
-        Ok(locals.event_loop.into_bound(py))
+        Ok(locals.event_loop.clone_ref(py).into_bound(py))
     } else {
         get_running_loop(py)
     }
@@ -597,11 +597,11 @@ where
     let future_tx2 = future_tx1.clone_ref(py);
 
     R::spawn(async move {
-        let locals2 = Python::attach(|py| locals.clone_ref(py));
+        let locals2 = locals.clone_ref();
 
         if let Err(e) = R::spawn(async move {
             let result = R::scope(
-                Python::attach(|py| locals2.clone_ref(py)),
+                locals2.clone_ref(),
                 Cancellable::new_with_cancel_rx(fut, cancel_rx),
             )
             .await;
@@ -1002,11 +1002,11 @@ where
     let future_tx2 = future_tx1.clone_ref(py);
 
     R::spawn_local(async move {
-        let locals2 = Python::attach(|py| locals.clone_ref(py));
+        let locals2 = locals.clone_ref();
 
         if let Err(e) = R::spawn_local(async move {
             let result = R::scope_local(
-                Python::attach(|py| locals2.clone_ref(py)),
+                locals2.clone_ref(),
                 Cancellable::new_with_cancel_rx(fut, cancel_rx),
             )
             .await;
@@ -1510,7 +1510,7 @@ impl SenderGlue {
             self.tx
                 .lock()
                 .unwrap()
-                .send(py, self.locals.clone_ref(py), item)
+                .send(py, self.locals.clone_ref(), item)
         })
     }
     pub fn close(&mut self) -> PyResult<()> {

src/lib.rs

@@ -395,6 +395,7 @@ pub mod doc_test {
 }
 
 use std::future::Future;
+use std::sync::Arc;
 
 use futures::channel::oneshot;
 use pyo3::{call::PyCallArgs, prelude::*, sync::PyOnceLock, types::PyDict};
@@ -472,17 +473,17 @@ fn copy_context(py: Python) -> PyResult<Bound<PyAny>> {
 #[derive(Debug)]
 pub struct TaskLocals {
     /// Track the event loop of the Python task
-    event_loop: Py<PyAny>,
+    event_loop: Arc<Py<PyAny>>,
     /// Track the contextvars of the Python task
-    context: Py<PyAny>,
+    context: Arc<Py<PyAny>>,
 }
 
 impl TaskLocals {
     /// At a minimum, TaskLocals must store the event loop.
     pub fn new(event_loop: Bound<PyAny>) -> Self {
         Self {
-            context: event_loop.py().None(),
-            event_loop: event_loop.into(),
+            context: Arc::new(event_loop.py().None()),
+            event_loop: Arc::new(event_loop.into()),
         }
     }
 
@@ -494,7 +495,7 @@ impl TaskLocals {
     /// Manually provide the contextvars for the current task.
     pub fn with_context(self, context: Bound<PyAny>) -> Self {
         Self {
-            context: context.into(),
+            context: Arc::new(context.into()),
             ..self
         }
     }
@@ -516,10 +517,10 @@ impl TaskLocals {
 
     /// Create a clone of the TaskLocals by incrementing the reference counters of the event loop and
     /// contextvars.
-    pub fn clone_ref(&self, py: Python<'_>) -> Self {
+    pub fn clone_ref(&self) -> Self {
         Self {
-            event_loop: self.event_loop.clone_ref(py),
-            context: self.context.clone_ref(py),
+            event_loop: self.event_loop.clone(),
+            context: self.context.clone(),
         }
     }
 }

src/tokio.rs

@@ -111,7 +111,7 @@ impl ContextExt for TokioRuntime {
         TASK_LOCALS
             .try_with(|c| {
                 c.get()
-                    .map(|locals| Python::attach(|py| locals.clone_ref(py)))
+                    .map(|locals| locals.clone_ref())
             })
             .unwrap_or_default()
     }