dataflow_rs/function/
layer.rs

1//! Call layers: around-call middleware for host calls, and the rest of the
2//! chain each one is handed.
3
4use super::erased::BoundCall;
5use super::{Ctx, Outcome};
6use crate::base::Result;
7use std::fmt;
8use std::pin::Pin;
9use std::sync::Arc;
10
11/// A boxed, `Send` future: what a [`CallLayer`] returns.
12pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
13
14/// Around-call middleware for host calls.
15///
16/// A layer on the [`Registry`](super::Registry) wraps every host call of
17/// every run. The built-ins run inline and are never wrapped. A layer calls
18/// [`Next::run`] to go on, or answers without calling it.
19/// [`testing::call`](super::testing::call) runs one call inside the layers
20/// it is given:
21///
22/// ```
23/// # tokio::runtime::Builder::new_current_thread().build().unwrap().block_on(async {
24/// use dataflow_rs::function::{BoxFuture, CallLayer, Next, testing};
25/// use dataflow_rs::{Ctx, Function, Outcome, Registry, Result};
26/// use serde_json::json;
27///
28/// /// Turns any `5xx` into a `503`.
29/// struct Unavailable;
30///
31/// impl CallLayer for Unavailable {
32///     fn call<'a>(
33///         &'a self,
34///         ctx: &'a mut Ctx<'_>,
35///         next: Next<'a>,
36///     ) -> BoxFuture<'a, Result<Outcome>> {
37///         Box::pin(async move {
38///             Ok(match next.run(ctx).await? {
39///                 Outcome::Status(500..) => Outcome::Status(503),
40///                 outcome => outcome,
41///             })
42///         })
43///     }
44/// }
45///
46/// let mut registry = Registry::default();
47/// registry.layer(Unavailable);
48///
49/// /// Fails every call with a `500`.
50/// struct Down;
51///
52/// impl Function for Down {
53///     type Config = ();
54///
55///     async fn call(&self, _ctx: &mut Ctx<'_>, _cfg: &()) -> Result<Outcome> {
56///         Ok(Outcome::Status(500))
57///     }
58/// }
59///
60/// let called = testing::call(Down, json!({})).layer(Unavailable).await;
61/// assert_eq!(called.outcome.unwrap(), Outcome::Status(503));
62/// # });
63/// ```
64pub trait CallLayer: Send + Sync + 'static {
65    /// Wraps one call. `next` is the rest of the chain.
66    fn call<'a>(&'a self, ctx: &'a mut Ctx<'_>, next: Next<'a>) -> BoxFuture<'a, Result<Outcome>>;
67}
68
69/// The rest of a call's chain: the run's layers, then the registry's, then
70/// the function. Each list runs outermost first.
71pub struct Next<'a> {
72    run: &'a [Arc<dyn CallLayer>],
73    registry: &'a [Arc<dyn CallLayer>],
74    call: &'a dyn BoundCall,
75}
76
77impl fmt::Debug for Next<'_> {
78    /// How many layers of each list are left; not the layers or the call.
79    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80        f.debug_struct("Next")
81            .field("run", &self.run.len())
82            .field("registry", &self.registry.len())
83            .finish_non_exhaustive()
84    }
85}
86
87impl<'a> Next<'a> {
88    /// The chain of `run`'s layers, then `registry`'s, around `call`.
89    pub(crate) fn new(
90        run: &'a [Arc<dyn CallLayer>],
91        registry: &'a [Arc<dyn CallLayer>],
92        call: &'a dyn BoundCall,
93    ) -> Self {
94        Self {
95            run,
96            registry,
97            call,
98        }
99    }
100
101    /// Runs the rest of the chain: the next layer, or the function when no
102    /// layer is left.
103    pub fn run<'b>(self, ctx: &'b mut Ctx<'_>) -> BoxFuture<'b, Result<Outcome>>
104    where
105        'a: 'b,
106    {
107        if let Some((layer, run)) = self.run.split_first() {
108            layer.call(ctx, Next { run, ..self })
109        } else if let Some((layer, registry)) = self.registry.split_first() {
110            layer.call(ctx, Next { registry, ..self })
111        } else {
112            self.call.call(ctx)
113        }
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use crate::base::{ProgramId, TaskIx, WorkflowIx};
121    use crate::expr::{CompileCx, ExprEngine};
122    use crate::function::ctx::testing::with_ctx;
123    use crate::function::registry::Entry;
124    use crate::function::{CallInfo, Function, Prepare, Registry};
125    use serde_json::json;
126    use std::sync::Mutex;
127    use std::task::{Context, Poll, Waker};
128
129    type Log = Arc<Mutex<Vec<&'static str>>>;
130
131    /// Polls once: every future in these tests is ready at once.
132    fn ready<F: Future>(future: F) -> F::Output {
133        let mut future = std::pin::pin!(future);
134        match future
135            .as_mut()
136            .poll(&mut Context::from_waker(Waker::noop()))
137        {
138            Poll::Ready(out) => out,
139            Poll::Pending => panic!("a future in these tests is never pending"),
140        }
141    }
142
143    fn info() -> CallInfo<'static> {
144        CallInfo {
145            program: ProgramId::next(),
146            workflow: WorkflowIx::new(0),
147            task: TaskIx::new(1),
148            function: "probe",
149            workflow_id: "w",
150            task_id: "t",
151            channel: None,
152            routing_bucket: None,
153            loop_counter: None,
154            element_index: None,
155        }
156    }
157
158    /// Logs its call, and answers 201.
159    struct Probe(Log);
160
161    impl Function for Probe {
162        type Config = ();
163        async fn call(&self, ctx: &mut Ctx<'_>, _cfg: &()) -> Result<Outcome> {
164            assert_eq!(ctx.info().task_id, "t");
165            self.0.lock().unwrap().push("function");
166            Ok(Outcome::Status(201))
167        }
168    }
169
170    /// Logs its name around the rest of the chain.
171    struct Named(&'static str, Log);
172
173    impl CallLayer for Named {
174        fn call<'a>(
175            &'a self,
176            ctx: &'a mut Ctx<'_>,
177            next: Next<'a>,
178        ) -> BoxFuture<'a, Result<Outcome>> {
179            Box::pin(async move {
180                self.1.lock().unwrap().push(self.0);
181                let outcome = next.run(ctx).await;
182                self.1.lock().unwrap().push(self.0);
183                outcome
184            })
185        }
186    }
187
188    /// Answers without calling the rest of the chain.
189    struct Answer(Log);
190
191    impl CallLayer for Answer {
192        fn call<'a>(
193            &'a self,
194            _ctx: &'a mut Ctx<'_>,
195            _next: Next<'a>,
196        ) -> BoxFuture<'a, Result<Outcome>> {
197            self.0.lock().unwrap().push("answer");
198            Box::pin(std::future::ready(Ok(Outcome::Skip)))
199        }
200    }
201
202    /// The probe's call site, bound through the registry as the compiler
203    /// binds one.
204    fn bound(log: &Log) -> (Registry, Box<dyn BoundCall>) {
205        let mut registry = Registry::default();
206        registry.register("probe", Probe(Arc::clone(log)));
207        registry.layer(Named("registry-outer", Arc::clone(log)));
208        registry.layer(Named("registry-inner", Arc::clone(log)));
209        let Some(Entry::Host(probe)) = registry.entry("probe") else {
210            panic!("probe is registered");
211        };
212        let engine = ExprEngine::default();
213        let mut issues = Vec::new();
214        let mut cx = CompileCx::new(&engine, &mut issues);
215        let spec = probe.spec().clone();
216        let call = probe
217            .compile(&json!({}), &mut cx, &Prepare::new("probe", "w", "t", &spec))
218            .expect("an empty config compiles");
219        drop(cx);
220        assert!(issues.is_empty(), "{issues:?}");
221        (registry, call)
222    }
223
224    #[test]
225    fn the_run_layers_wrap_the_registry_layers_which_wrap_the_function() {
226        let log = Log::default();
227        let (registry, call) = bound(&log);
228        let run: Vec<Arc<dyn CallLayer>> = vec![
229            Arc::new(Named("run-outer", Arc::clone(&log))),
230            Arc::new(Named("run-inner", Arc::clone(&log))),
231        ];
232        let next = Next::new(&run, registry.layers(), &*call);
233        let outcome = with_ctx(&json!({}), info(), None, |mut ctx| {
234            ready(next.run(&mut ctx))
235        });
236        assert_eq!(outcome.unwrap(), Outcome::Status(201));
237        assert_eq!(
238            *log.lock().unwrap(),
239            [
240                "run-outer",
241                "run-inner",
242                "registry-outer",
243                "registry-inner",
244                "function",
245                "registry-inner",
246                "registry-outer",
247                "run-inner",
248                "run-outer",
249            ]
250        );
251    }
252
253    #[test]
254    fn with_no_layers_the_function_is_called() {
255        let log = Log::default();
256        let (_, call) = bound(&log);
257        let next = Next::new(&[], &[], &*call);
258        let outcome = with_ctx(&json!({}), info(), None, |mut ctx| {
259            ready(next.run(&mut ctx))
260        });
261        assert_eq!(outcome.unwrap(), Outcome::Status(201));
262        assert_eq!(*log.lock().unwrap(), ["function"]);
263    }
264
265    #[test]
266    fn a_layer_that_answers_runs_nothing_inside_it() {
267        let log = Log::default();
268        let (registry, call) = bound(&log);
269        let run: Vec<Arc<dyn CallLayer>> = vec![Arc::new(Answer(Arc::clone(&log)))];
270        let next = Next::new(&run, registry.layers(), &*call);
271        let outcome = with_ctx(&json!({}), info(), None, |mut ctx| {
272            ready(next.run(&mut ctx))
273        });
274        assert_eq!(outcome.unwrap(), Outcome::Skip);
275        assert_eq!(*log.lock().unwrap(), ["answer"]);
276    }
277}
dataflow_rs/function/layer.rs

dataflow_rs/function/
layer.rs
