Wrap TaskLocals in a single Arc.

Author: iksteen

src/generic.rs

@@ -91,7 +91,7 @@ where
     R: ContextExt,
 {
     if let Some(locals) = R::get_task_locals() {
-        Ok(locals.event_loop.clone_ref(py).into_bound(py))
+        Ok(locals.0.event_loop.clone_ref(py).into_bound(py))
     } else {
         get_running_loop(py)
     }
@@ -585,7 +585,7 @@ where
 {
     let (cancel_tx, cancel_rx) = oneshot::channel();
 
-    let py_fut = create_future(locals.event_loop.bind(py).clone())?;
+    let py_fut = create_future(locals.0.event_loop.bind(py).clone())?;
     py_fut.call_method1(
         "add_done_callback",
         (PyDoneCallback {
@@ -638,7 +638,7 @@ where
                         get_panic_message(&e.into_panic())
                     );
                     let _ = set_result(
-                        locals.event_loop.bind(py),
+                        locals.0.event_loop.bind(py),
                         future_tx2.bind(py),
                         Err(RustPanic::new_err(panic_message)),
                     )
@@ -990,7 +990,7 @@ where
 {
     let (cancel_tx, cancel_rx) = oneshot::channel();
 
-    let py_fut = create_future(locals.event_loop.clone_ref(py).into_bound(py))?;
+    let py_fut = create_future(locals.0.event_loop.clone_ref(py).into_bound(py))?;
     py_fut.call_method1(
         "add_done_callback",
         (PyDoneCallback {
@@ -1020,7 +1020,7 @@ where
                 }
 
                 let _ = set_result(
-                    locals2.event_loop.bind(py),
+                    locals2.0.event_loop.bind(py),
                     future_tx1.bind(py),
                     result.and_then(|val| val.into_py_any(py)),
                 )
@@ -1043,7 +1043,7 @@ where
                         get_panic_message(&e.into_panic())
                     );
                     let _ = set_result(
-                        locals.event_loop.bind(py),
+                        locals.0.event_loop.bind(py),
                         future_tx2.bind(py),
                         Err(RustPanic::new_err(panic_message)),
                     )

src/lib.rs

@@ -469,22 +469,26 @@ fn copy_context(py: Python) -> PyResult<Bound<PyAny>> {
     contextvars(py)?.call_method0("copy_context")
 }
 
-/// Task-local data to store for Python conversions.
+/// Task-local inner structure.
 #[derive(Debug)]
-pub struct TaskLocals {
+struct TaskLocalsInner {
     /// Track the event loop of the Python task
-    event_loop: Arc<Py<PyAny>>,
+    event_loop: Py<PyAny>,
     /// Track the contextvars of the Python task
-    context: Arc<Py<PyAny>>,
+    context: Py<PyAny>,
 }
 
+/// Task-local data to store for Python conversions.
+#[derive(Debug)]
+pub struct TaskLocals(Arc<TaskLocalsInner>);
+
 impl TaskLocals {
     /// At a minimum, TaskLocals must store the event loop.
     pub fn new(event_loop: Bound<PyAny>) -> Self {
-        Self {
-            context: Arc::new(event_loop.py().None()),
-            event_loop: Arc::new(event_loop.into()),
-        }
+        Self(Arc::new(TaskLocalsInner {
+            context: event_loop.py().None(),
+            event_loop: event_loop.into(),
+        }))
     }
 
     /// Construct TaskLocals with the event loop returned by `get_running_loop`
@@ -494,10 +498,10 @@ impl TaskLocals {
 
     /// Manually provide the contextvars for the current task.
     pub fn with_context(self, context: Bound<PyAny>) -> Self {
-        Self {
-            context: Arc::new(context.into()),
-            ..self
-        }
+        Self(Arc::new(TaskLocalsInner {
+            event_loop: self.0.event_loop.clone_ref(context.py()),
+            context: context.into(),
+        }))
     }
 
     /// Capture the current task's contextvars
@@ -507,12 +511,12 @@ impl TaskLocals {
 
     /// Get a reference to the event loop
     pub fn event_loop<'p>(&self, py: Python<'p>) -> Bound<'p, PyAny> {
-        self.event_loop.clone_ref(py).into_bound(py)
+        self.0.event_loop.clone_ref(py).into_bound(py)
     }
 
     /// Get a reference to the python context
     pub fn context<'p>(&self, py: Python<'p>) -> Bound<'p, PyAny> {
-        self.context.clone_ref(py).into_bound(py)
+        self.0.context.clone_ref(py).into_bound(py)
     }
 
     /// Create a clone of the TaskLocals. No longer uses the runtime, use `clone` instead.
@@ -526,10 +530,7 @@ impl Clone for TaskLocals {
     /// Create a clone of the TaskLocals by incrementing the reference counters of the event loop and
     /// contextvars.
     fn clone(&self) -> Self {
-        Self {
-            event_loop: self.event_loop.clone(),
-            context: self.context.clone(),
-        }
+        Self(self.0.clone())
     }
 }