nautilus_common/python/
wrappers.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Registry of the strong references which keep registered components' Python wrappers alive.
17//!
18//! Component inners hold only a weak reference to their Python wrapper, so something must own the
19//! wrapper for as long as the component stays registered. This registry is that owner, and it is
20//! thread-local for the same reason the component and actor registries are: those hold
21//! `Rc<UnsafeCell<..>>` and cannot leave their thread, so a wrapper registry with wider visibility
22//! would desync from the registrations it shadows.
23
24use std::{cell::RefCell, rc::Rc};
25
26use ahash::AHashMap;
27use nautilus_core::python::to_pyruntime_err;
28use nautilus_model::identifiers::ComponentId;
29use pyo3::prelude::*;
30
31use super::msgbus::PyMessageBusScope;
32
33thread_local! {
34    static PYTHON_WRAPPERS: RefCell<AHashMap<ComponentId, RegisteredWrapper>> =
35        RefCell::new(AHashMap::new());
36}
37
38/// Retains the strong reference which keeps the Python wrapper for `component_id` alive.
39pub fn retain_python_wrapper(
40    component_id: ComponentId,
41    wrapper: Py<PyAny>,
42    message_bus: Rc<PyMessageBusScope>,
43) {
44    let displaced = PYTHON_WRAPPERS.with_borrow_mut(|wrappers| {
45        wrappers.insert(
46            component_id,
47            RegisteredWrapper {
48                wrapper,
49                message_bus,
50            },
51        )
52    });
53
54    if displaced.is_some() {
55        log::warn!("Replaced the retained Python wrapper for {component_id}");
56    }
57
58    // Dropping a wrapper can run Python finalization which re-enters Rust, so the value leaves the
59    // registry before the borrow ends
60    drop(displaced);
61}
62
63/// Releases the strong reference retained for `component_id`.
64pub fn release_python_wrapper(component_id: ComponentId) {
65    let released = PYTHON_WRAPPERS.with_borrow_mut(|wrappers| wrappers.remove(&component_id));
66
67    drop(released);
68}
69
70/// Returns the Python wrapper retained for `component_id`, or `None` when nothing is retained.
71#[must_use]
72pub fn get_python_wrapper(component_id: ComponentId) -> Option<Py<PyAny>> {
73    Python::attach(|py| {
74        PYTHON_WRAPPERS.with_borrow(|wrappers| {
75            wrappers
76                .get(&component_id)
77                .map(|registered| registered.wrapper.clone_ref(py))
78        })
79    })
80}
81
82/// Returns message state owned by the receiver's registered runtime thread.
83///
84/// Inspect the unborrowed Python receiver: PyO3's unsendable borrow check panics on a foreign thread.
85/// The state is shared separately so lifecycle callbacks do not re-borrow their component.
86///
87/// # Errors
88///
89/// Returns a runtime error when this thread does not retain the wrapper.
90pub fn get_python_message_bus(wrapper: &Bound<'_, PyAny>) -> PyResult<Rc<PyMessageBusScope>> {
91    PYTHON_WRAPPERS.with_borrow(|wrappers| {
92        wrappers
93            .values()
94            .find(|registered| registered.wrapper.as_ptr() == wrapper.as_ptr())
95            .map(|registered| Rc::clone(&registered.message_bus))
96            .ok_or_else(|| {
97                to_pyruntime_err("Component must be registered on the calling runtime thread")
98            })
99    })
100}
101
102struct RegisteredWrapper {
103    wrapper: Py<PyAny>,
104    message_bus: Rc<PyMessageBusScope>,
105}
106
107#[cfg(test)]
108mod tests {
109    use pyo3::{ffi::c_str, types::PyModule, wrap_pyfunction};
110    use rstest::rstest;
111
112    use super::*;
113
114    /// Lets a finalizing Python wrapper re-enter the registry.
115    #[pyfunction]
116    fn wrapper_is_retained(component_id: &str) -> bool {
117        get_python_wrapper(ComponentId::from(component_id)).is_some()
118    }
119
120    #[rstest]
121    fn test_wrapper_finalization_re_enters_an_unborrowed_registry() {
122        Python::initialize();
123
124        Python::attach(|py| {
125            let module = PyModule::new(py, "test_wrapper_finalization").unwrap();
126            module
127                .add_function(wrap_pyfunction!(wrapper_is_retained, &module).unwrap())
128                .unwrap();
129
130            let code = c_str!(
131                r#"
132OBSERVED = []
133
134
135class Finalizing:
136    def __del__(self):
137        OBSERVED.append(wrapper_is_retained("Finalizing-Component"))
138"#
139            );
140            py.run(code, Some(&module.dict()), None).unwrap();
141
142            let finalizing = module.getattr("Finalizing").unwrap();
143            let component_id = ComponentId::from("Finalizing-Component");
144
145            retain_python_wrapper(
146                component_id,
147                finalizing.call0().unwrap().unbind(),
148                Rc::default(),
149            );
150
151            // The registry holds the only reference to each wrapper, so both the displaced wrapper
152            // and the released one run `__del__` while being dropped
153            retain_python_wrapper(
154                component_id,
155                finalizing.call0().unwrap().unbind(),
156                Rc::default(),
157            );
158            release_python_wrapper(component_id);
159
160            let observed = module
161                .getattr("OBSERVED")
162                .unwrap()
163                .extract::<Vec<bool>>()
164                .unwrap();
165
166            assert_eq!(observed, vec![true, false]);
167        });
168    }
169}
nautilus_common/python/wrappers.rs

nautilus_common/python/
wrappers.rs
