Skip to main content

wingfoil/
lib.rs

1mod proxy_stream;
2#[cfg(feature = "aeron")]
3mod py_aeron;
4mod py_csv;
5mod py_element;
6#[cfg(feature = "etcd")]
7mod py_etcd;
8mod py_fix;
9mod py_fluvio;
10#[cfg(feature = "iceoryx2")]
11mod py_iceoryx2;
12mod py_kafka;
13mod py_kdb;
14mod py_latency;
15mod py_otlp;
16mod py_prometheus;
17mod py_stream;
18mod py_web;
19mod py_zmq;
20mod types;
21
22use ::wingfoil::{Dep, Node, NodeOperators};
23use py_element::*;
24use py_stream::*;
25use types::ToPyResult;
26
27use pyo3::prelude::*;
28use std::rc::Rc;
29use std::time::Duration;
30
31#[pyclass(unsendable, name = "Node")]
32#[derive(Clone)]
33pub(crate) struct PyNode(Rc<dyn Node>);
34
35impl PyNode {
36    pub(crate) fn new(node: Rc<dyn Node>) -> Self {
37        Self(node)
38    }
39}
40
41#[pymethods]
42impl PyNode {
43    /// Counts how many times upstream node has ticked.
44    fn count(&self) -> PyStream {
45        self.0.count().as_py_stream()
46    }
47
48    #[pyo3(signature = (realtime, start=None, duration=None, cycles=None))]
49    fn run(
50        &self,
51        py: Python<'_>,
52        realtime: bool,
53        start: Option<Py<PyAny>>,
54        duration: Option<Py<PyAny>>,
55        cycles: Option<u32>,
56    ) -> PyResult<()> {
57        let (run_mode, run_for) =
58            types::parse_run_args(py, realtime, start, duration, cycles).to_pyresult()?;
59
60        // Convert fat pointer to (addr, vtable) pair which is Send+Sync
61        let node_ptr = Rc::as_ptr(&self.0);
62        let (addr, vtable): (usize, usize) = unsafe { std::mem::transmute(node_ptr) };
63
64        // Release GIL during the run to allow async tasks to acquire it
65        // SAFETY: The Rc is kept alive by self for the duration of this call
66        let result = py.detach(move || {
67            // Reconstruct the fat pointer from (addr, vtable)
68            let node_ptr: *const dyn Node = unsafe { std::mem::transmute((addr, vtable)) };
69            // Temporarily reconstruct the Rc without taking ownership
70            let node = unsafe { Rc::from_raw(node_ptr) };
71            let result = ::wingfoil::NodeOperators::run(&node, run_mode, run_for);
72            std::mem::forget(node); // Don't drop the Rc (self.0 still owns it)
73            result
74        });
75        result.to_pyresult()?;
76        Ok(())
77    }
78}
79
80/// A node that ticks at the specified period
81#[pyfunction]
82fn ticker(seconds: f64) -> PyNode {
83    let ticker = ::wingfoil::ticker(Duration::from_secs_f64(seconds));
84    PyNode::new(ticker)
85}
86
87/// A stream that ticks once, on first engine cycle
88#[pyfunction]
89fn constant(val: Py<PyAny>) -> PyStream {
90    let strm = ::wingfoil::constant(PyElement::new(val));
91    PyStream(strm)
92}
93
94/// maps steams a amd b into a new stream using func (e.g lambda a, b: a + b)
95#[pyfunction]
96fn bimap(a: Py<PyAny>, b: Py<PyAny>, func: Py<PyAny>) -> PyResult<PyStream> {
97    Python::attach(|py| {
98        let a = a
99            .as_ref()
100            .extract::<PyRef<PyStream>>(py)
101            .map_err(|_| types::py_type_error("bimap: first argument must be a Stream"))
102            .to_pyresult()?
103            .inner_stream();
104        let b = b
105            .as_ref()
106            .extract::<PyRef<PyStream>>(py)
107            .map_err(|_| types::py_type_error("bimap: second argument must be a Stream"))
108            .to_pyresult()?
109            .inner_stream();
110        let stream = ::wingfoil::try_bimap(
111            Dep::Active(a),
112            Dep::Active(b),
113            move |a: PyElement, b: PyElement| {
114                Python::attach(|py: Python<'_>| {
115                    let res = func
116                        .call1(py, (a.value(), b.value()))
117                        .map_err(py_callback_error)?;
118                    Ok(PyElement::new(res))
119                })
120            },
121        );
122        Ok(PyStream(stream))
123    })
124}
125
126#[pyclass(unsendable, name = "Graph")]
127#[derive(Clone)]
128pub(crate) struct PyGraph(Vec<Rc<dyn Node>>);
129
130#[pymethods]
131impl PyGraph {
132    #[new]
133    fn new(nodes: Vec<Py<PyAny>>) -> PyResult<Self> {
134        Python::attach(|py| {
135            let mut roots: Vec<Rc<dyn Node>> = Vec::new();
136            for obj in nodes {
137                if let Ok(stream) = obj.extract::<PyRef<PyStream>>(py) {
138                    roots.push(stream.0.clone().as_node());
139                } else if let Ok(node) = obj.extract::<PyRef<PyNode>>(py) {
140                    roots.push(node.0.clone());
141                } else {
142                    return Err(pyo3::exceptions::PyTypeError::new_err(
143                        "Graph components must be Stream or Node",
144                    ));
145                }
146            }
147            Ok(PyGraph(roots))
148        })
149    }
150
151    #[pyo3(signature = (realtime, start=None, duration=None, cycles=None))]
152    fn run(
153        &self,
154        py: Python<'_>,
155        realtime: bool,
156        start: Option<Py<PyAny>>,
157        duration: Option<Py<PyAny>>,
158        cycles: Option<u32>,
159    ) -> PyResult<()> {
160        let (run_mode, run_for) =
161            types::parse_run_args(py, realtime, start, duration, cycles).to_pyresult()?;
162
163        let mut ptrs: Vec<(usize, usize)> = Vec::with_capacity(self.0.len());
164        for node in &self.0 {
165            let node_ptr = Rc::as_ptr(node);
166            let (addr, vtable): (usize, usize) = unsafe { std::mem::transmute(node_ptr) };
167            ptrs.push((addr, vtable));
168        }
169
170        let result = py.detach(move || {
171            let mut roots: Vec<Rc<dyn Node>> = Vec::with_capacity(ptrs.len());
172            for (addr, vtable) in ptrs {
173                let node_ptr: *const dyn Node = unsafe { std::mem::transmute((addr, vtable)) };
174                let node = unsafe { Rc::from_raw(node_ptr) };
175                roots.push(node.clone());
176                std::mem::forget(node);
177            }
178
179            let mut graph = ::wingfoil::Graph::new(roots, run_mode, run_for);
180            graph.run()
181        });
182        result.to_pyresult()?;
183        Ok(())
184    }
185}
186
187/// Wingfoil is a blazingly fast, highly scalable stream processing
188/// framework designed for latency-critical use cases such as electronic
189/// trading and real-time AI systems
190#[pymodule]
191fn _wingfoil(module: &Bound<'_, PyModule>) -> PyResult<()> {
192    let env = env_logger::Env::default().default_filter_or("info");
193    env_logger::Builder::from_env(env).init();
194    module.add_function(wrap_pyfunction!(ticker, module)?)?;
195    module.add_function(wrap_pyfunction!(constant, module)?)?;
196    module.add_function(wrap_pyfunction!(bimap, module)?)?;
197    module.add_function(wrap_pyfunction!(py_csv::py_csv_read, module)?)?;
198    #[cfg(feature = "etcd")]
199    module.add_function(wrap_pyfunction!(py_etcd::py_etcd_sub, module)?)?;
200    module.add_function(wrap_pyfunction!(py_fluvio::py_fluvio_sub, module)?)?;
201    module.add_function(wrap_pyfunction!(py_kafka::py_kafka_sub, module)?)?;
202    module.add_function(wrap_pyfunction!(py_kdb::py_kdb_read, module)?)?;
203    module.add_function(wrap_pyfunction!(py_kdb::py_kdb_write, module)?)?;
204    module.add_function(wrap_pyfunction!(py_zmq::py_zmq_sub, module)?)?;
205    #[cfg(feature = "iceoryx2")]
206    module.add_function(wrap_pyfunction!(py_iceoryx2::py_iceoryx2_sub, module)?)?;
207    #[cfg(feature = "aeron")]
208    module.add_function(wrap_pyfunction!(py_aeron::py_aeron_sub, module)?)?;
209    #[cfg(feature = "etcd")]
210    module.add_function(wrap_pyfunction!(py_zmq::py_zmq_sub_etcd, module)?)?;
211    module.add_function(wrap_pyfunction!(py_fix::py_fix_connect, module)?)?;
212    module.add_function(wrap_pyfunction!(py_fix::py_fix_connect_tls, module)?)?;
213    module.add_function(wrap_pyfunction!(py_fix::py_fix_accept, module)?)?;
214    module.add_class::<PyNode>()?;
215    module.add_class::<PyStream>()?;
216    module.add_class::<PyGraph>()?;
217    #[cfg(feature = "iceoryx2")]
218    module.add_class::<py_iceoryx2::PyIceoryx2ServiceVariant>()?;
219    #[cfg(feature = "iceoryx2")]
220    module.add_class::<py_iceoryx2::PyIceoryx2Mode>()?;
221    #[cfg(feature = "aeron")]
222    module.add_class::<py_aeron::PyAeronMode>()?;
223    module.add_class::<py_prometheus::PyPrometheusExporter>()?;
224    module.add_class::<py_latency::PyLatency>()?;
225    module.add_class::<py_latency::PyTracedBytes>()?;
226    module.add_class::<py_web::PyWebServer>()?;
227    module.add("__version__", env!("CARGO_PKG_VERSION"))?;
228    Ok(())
229}