lenso_runner/replicated/
transfer.rs

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

lenso_runner/replicated/
transfer.rs
