Skip to main content

wingfoil/
lib.rs

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