lenso_runner/replicated/
transfer.rs

1use std::{
2    any::Any, collections::BTreeMap, fmt, marker::PhantomData, rc::Rc, sync::Arc, time::Instant,
3};
4
5use futures::future::LocalBoxFuture;
6use lenso_app_plan::ResolvedAppPlan;
7use lenso_kernel::{
8    CancellationToken, InvocationContext, NativeRequestEndpoint, RequestCapability, RuntimeFailure,
9};
10use tokio::sync::watch;
11
12use super::{LaneRoute, LaneTask};
13
14trait RequestTransferFactory: fmt::Debug + Send + Sync {
15    fn endpoint(&self, provider_lane: LaneRoute, epoch: Instant) -> Rc<dyn NativeRequestEndpoint>;
16}
17
18struct TypedRequestTransferFactory<C: RequestCapability> {
19    operations: &'static [&'static str],
20    capability: PhantomData<fn() -> C>,
21}
22
23impl<C: RequestCapability> fmt::Debug for TypedRequestTransferFactory<C> {
24    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
25        formatter
26            .debug_struct("TypedRequestTransferFactory")
27            .field("capability", &C::ID)
28            .finish_non_exhaustive()
29    }
30}
31
32impl<C> RequestTransferFactory for TypedRequestTransferFactory<C>
33where
34    C: RequestCapability,
35    C::Request: Send,
36    C::Response: Send,
37    C::DomainError: Send,
38{
39    fn endpoint(&self, provider_lane: LaneRoute, epoch: Instant) -> Rc<dyn NativeRequestEndpoint> {
40        Rc::new(CrossLaneRequestEndpoint::<C> {
41            operations: self.operations,
42            provider_lane,
43            epoch,
44            capability: PhantomData,
45        })
46    }
47}
48
49/// Native request types registered for zero-serialization cross-lane transfer.
50#[derive(Clone, Debug, Default)]
51pub struct CrossLaneRequestCatalog {
52    factories: BTreeMap<&'static str, Arc<dyn RequestTransferFactory>>,
53}
54
55impl CrossLaneRequestCatalog {
56    /// Creates an empty catalog.
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    /// Registers one generated request Capability whose values are `Send`.
62    #[must_use]
63    pub fn with_request<C>(mut self, operations: &'static [&'static str]) -> Self
64    where
65        C: RequestCapability,
66        C::Request: Send,
67        C::Response: Send,
68        C::DomainError: Send,
69    {
70        self.factories.insert(
71            C::ID,
72            Arc::new(TypedRequestTransferFactory::<C> {
73                operations,
74                capability: PhantomData,
75            }),
76        );
77        self
78    }
79
80    pub(super) fn contains(&self, capability_id: &str) -> bool {
81        self.factories.contains_key(capability_id)
82    }
83
84    pub(super) fn validate_plan(
85        &self,
86        plan: &ResolvedAppPlan,
87    ) -> Result<(), super::ReplicatedRunnerError> {
88        for binding in plan.capability_bindings() {
89            let consumer = plan
90                .module_instance(binding.consumer_instance())
91                .expect("validated binding consumer should exist");
92            let provider = plan
93                .module_instance(binding.provider_instance())
94                .expect("validated binding provider should exist");
95            let endpoint = provider
96                .provided_capabilities()
97                .iter()
98                .find(|endpoint| endpoint.capability_id() == binding.capability_id())
99                .expect("validated provider endpoint should exist");
100            if consumer.execution_lane() != provider.execution_lane()
101                && !endpoint.request_operations().is_empty()
102                && !self.contains(binding.capability_id())
103            {
104                return Err(
105                    super::ReplicatedRunnerError::MissingCrossLaneRequestTransfer {
106                        capability: binding.capability_id().to_owned(),
107                    },
108                );
109            }
110        }
111        Ok(())
112    }
113
114    pub(super) fn endpoint(
115        &self,
116        capability_id: &str,
117        provider_lane: LaneRoute,
118        epoch: Instant,
119    ) -> Option<Rc<dyn NativeRequestEndpoint>> {
120        self.factories
121            .get(capability_id)
122            .map(|factory| factory.endpoint(provider_lane, epoch))
123    }
124}
125
126struct CrossLaneRequestEndpoint<C: RequestCapability> {
127    operations: &'static [&'static str],
128    provider_lane: LaneRoute,
129    epoch: Instant,
130    capability: PhantomData<fn() -> C>,
131}
132
133impl<C: RequestCapability> fmt::Debug for CrossLaneRequestEndpoint<C> {
134    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
135        formatter
136            .debug_struct("CrossLaneRequestEndpoint")
137            .field("capability", &C::ID)
138            .finish_non_exhaustive()
139    }
140}
141
142impl<C> NativeRequestEndpoint for CrossLaneRequestEndpoint<C>
143where
144    C: RequestCapability,
145    C::Request: Send,
146    C::Response: Send,
147    C::DomainError: Send,
148{
149    fn capability_id(&self) -> &'static str {
150        C::ID
151    }
152
153    fn descriptor_version(&self) -> &'static str {
154        C::DESCRIPTOR_VERSION
155    }
156
157    fn operations(&self) -> &'static [&'static str] {
158        self.operations
159    }
160
161    fn invoke(
162        &self,
163        operation: &str,
164        request: Box<dyn Any>,
165        context: InvocationContext,
166    ) -> LocalBoxFuture<'static, Result<Result<Box<dyn Any>, Box<dyn Any>>, RuntimeFailure>> {
167        let Ok(request) = request.downcast::<C::Request>() else {
168            return Box::pin(futures::future::ready(Err(
169                RuntimeFailure::ProtocolViolation { capability: C::ID },
170            )));
171        };
172        let provider_lane = self.provider_lane.clone();
173        let operation = operation.to_owned();
174        let epoch = self.epoch;
175        Box::pin(async move {
176            let provider_lane = provider_lane.upgrade().ok_or_else(lane_unavailable::<C>)?;
177            let caller_instance = context
178                .caller_instance()
179                .ok_or_else(|| RuntimeFailure::InvalidResolvedPlan {
180                    detail: format!("cross-lane invocation of `{}` has no planned caller", C::ID),
181                })?
182                .to_owned();
183            let cancellation = context.cancellation();
184            let deadline = context.deadline();
185            let request_id = context.request_id();
186            let (transferred, cancellation_signal) =
187                TransferredInvocationContext::capture(&context);
188            let (completed, completion) = futures::channel::oneshot::channel();
189            let command: LaneTask = Box::new(move |lane| {
190                Box::pin(async move {
191                    let local_cancellation = CancellationToken::new();
192                    let (context, mut transferred_cancellation) =
193                        transferred.restore(local_cancellation.clone());
194                    if *transferred_cancellation.borrow() {
195                        local_cancellation.cancel();
196                    }
197                    let handle = match lane.request_handle::<C>(&caller_instance) {
198                        Ok(handle) => handle,
199                        Err(error) => {
200                            let _ = completed.send(Err(error));
201                            return;
202                        }
203                    };
204                    let invocation = handle.invoke_with_context(&operation, context, *request);
205                    tokio::pin!(invocation);
206                    let result = if local_cancellation.is_cancelled() {
207                        invocation.await
208                    } else {
209                        tokio::select! {
210                            result = &mut invocation => result,
211                            () = wait_for_cancellation(&mut transferred_cancellation) => {
212                                local_cancellation.cancel();
213                                invocation.await
214                            }
215                        }
216                    };
217                    let _ = completed.send(result);
218                })
219            });
220
221            let send = provider_lane.send(command);
222            tokio::pin!(send);
223            let cancelled = cancellation.cancelled();
224            tokio::pin!(cancelled);
225            let result = match deadline {
226                Some(deadline) => {
227                    let sleep = tokio::time::sleep_until((epoch + deadline).into());
228                    tokio::pin!(sleep);
229                    tokio::select! {
230                        result = &mut send => result.map_err(|_| lane_unavailable::<C>())?,
231                        () = &mut cancelled => {
232                            cancellation_signal.send_replace(true);
233                            return Err(RuntimeFailure::Cancelled { request_id });
234                        }
235                        () = &mut sleep => {
236                            cancellation_signal.send_replace(true);
237                            return Err(RuntimeFailure::DeadlineExceeded { request_id });
238                        }
239                    }
240                    tokio::select! {
241                        biased;
242                        result = completion => result.map_err(|_| lane_unavailable::<C>())?,
243                        () = &mut cancelled => {
244                            cancellation_signal.send_replace(true);
245                            return Err(RuntimeFailure::Cancelled { request_id });
246                        }
247                        () = &mut sleep => {
248                            cancellation_signal.send_replace(true);
249                            return Err(RuntimeFailure::DeadlineExceeded { request_id });
250                        }
251                    }
252                }
253                None => {
254                    tokio::select! {
255                        result = &mut send => result.map_err(|_| lane_unavailable::<C>())?,
256                        () = &mut cancelled => {
257                            cancellation_signal.send_replace(true);
258                            return Err(RuntimeFailure::Cancelled { request_id });
259                        }
260                    }
261                    tokio::select! {
262                        biased;
263                        result = completion => result.map_err(|_| lane_unavailable::<C>())?,
264                        () = &mut cancelled => {
265                            cancellation_signal.send_replace(true);
266                            return Err(RuntimeFailure::Cancelled { request_id });
267                        }
268                    }
269                }
270            };
271            result.map(|domain| {
272                domain
273                    .map(|response| Box::new(response) as Box<dyn Any>)
274                    .map_err(|error| Box::new(error) as Box<dyn Any>)
275            })
276        })
277    }
278}
279
280fn lane_unavailable<C: RequestCapability>() -> RuntimeFailure {
281    RuntimeFailure::Internal {
282        detail: format!("provider lane for `{}` is unavailable", C::ID),
283    }
284}
285
286#[derive(Debug)]
287struct TransferredInvocationContext {
288    request_id: u64,
289    deadline: Option<std::time::Duration>,
290    cancellation: watch::Receiver<bool>,
291    extensions: Vec<lenso_kernel::InvocationExtension>,
292    sealed_extensions: Vec<lenso_kernel::SealedInvocationExtension>,
293}
294
295impl TransferredInvocationContext {
296    fn capture(context: &InvocationContext) -> (Self, watch::Sender<bool>) {
297        let (cancellation_signal, cancellation) = watch::channel(context.is_cancelled());
298        (
299            Self {
300                request_id: context.request_id(),
301                deadline: context.deadline(),
302                cancellation,
303                extensions: context.extensions().cloned().collect(),
304                sealed_extensions: context.sealed_extensions().cloned().collect(),
305            },
306            cancellation_signal,
307        )
308    }
309
310    fn restore(
311        self,
312        cancellation: CancellationToken,
313    ) -> (InvocationContext, watch::Receiver<bool>) {
314        let mut context = InvocationContext::new(self.request_id, self.deadline, cancellation);
315        for extension in self.extensions {
316            context = context
317                .with_extension(extension.key(), extension.value().to_vec())
318                .expect("captured ordinary Invocation Context extension remains valid");
319        }
320        for extension in self.sealed_extensions {
321            context = context
322                .with_sealed_extension(extension)
323                .expect("captured sealed Invocation Context extension remains valid");
324        }
325        (context, self.cancellation)
326    }
327}
328
329async fn wait_for_cancellation(cancellation: &mut watch::Receiver<bool>) {
330    loop {
331        if *cancellation.borrow_and_update() {
332            return;
333        }
334        if cancellation.changed().await.is_err() {
335            return;
336        }
337    }
338}
lenso_runner/replicated/transfer.rs

lenso_runner/replicated/
transfer.rs
