use std::{
any::Any,
collections::BTreeMap,
fmt,
marker::PhantomData,
rc::Rc,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::Instant,
};
use futures::future::LocalBoxFuture;
use lenso_app_plan::ResolvedAppPlan;
use lenso_kernel::{
CancellationToken, InvocationContext, NativeRequestEndpoint, RequestCapability, RuntimeFailure,
};
use super::{LaneCommand, LaneRoute};
trait RequestTransferFactory: fmt::Debug + Send + Sync {
fn endpoint(&self, provider_lane: LaneRoute, epoch: Instant) -> Rc<dyn NativeRequestEndpoint>;
}
struct TypedRequestTransferFactory<C: RequestCapability> {
operations: &'static [&'static str],
capability: PhantomData<fn() -> C>,
}
impl<C: RequestCapability> fmt::Debug for TypedRequestTransferFactory<C> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("TypedRequestTransferFactory")
.field("capability", &C::ID)
.finish_non_exhaustive()
}
}
impl<C> RequestTransferFactory for TypedRequestTransferFactory<C>
where
C: RequestCapability,
C::Request: Send,
C::Response: Send,
C::DomainError: Send,
{
fn endpoint(&self, provider_lane: LaneRoute, epoch: Instant) -> Rc<dyn NativeRequestEndpoint> {
Rc::new(CrossLaneRequestEndpoint::<C> {
operations: self.operations,
provider_lane,
epoch,
capability: PhantomData,
})
}
}
#[derive(Clone, Debug, Default)]
pub struct CrossLaneRequestCatalog {
factories: BTreeMap<&'static str, Arc<dyn RequestTransferFactory>>,
}
impl CrossLaneRequestCatalog {
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_request<C>(mut self, operations: &'static [&'static str]) -> Self
where
C: RequestCapability,
C::Request: Send,
C::Response: Send,
C::DomainError: Send,
{
self.factories.insert(
C::ID,
Arc::new(TypedRequestTransferFactory::<C> {
operations,
capability: PhantomData,
}),
);
self
}
pub(super) fn contains(&self, capability_id: &str) -> bool {
self.factories.contains_key(capability_id)
}
pub(super) fn validate_plan(
&self,
plan: &ResolvedAppPlan,
) -> Result<(), super::ReplicatedRunnerError> {
for binding in plan.capability_bindings() {
let consumer = plan
.module_instance(binding.consumer_instance())
.expect("validated binding consumer should exist");
let provider = plan
.module_instance(binding.provider_instance())
.expect("validated binding provider should exist");
let endpoint = provider
.provided_capabilities()
.iter()
.find(|endpoint| endpoint.capability_id() == binding.capability_id())
.expect("validated provider endpoint should exist");
if consumer.execution_lane() != provider.execution_lane()
&& !endpoint.request_operations().is_empty()
&& !self.contains(binding.capability_id())
{
return Err(
super::ReplicatedRunnerError::MissingCrossLaneRequestTransfer {
capability: binding.capability_id().to_owned(),
},
);
}
}
Ok(())
}
pub(super) fn endpoint(
&self,
capability_id: &str,
provider_lane: LaneRoute,
epoch: Instant,
) -> Option<Rc<dyn NativeRequestEndpoint>> {
self.factories
.get(capability_id)
.map(|factory| factory.endpoint(provider_lane, epoch))
}
}
struct CrossLaneRequestEndpoint<C: RequestCapability> {
operations: &'static [&'static str],
provider_lane: LaneRoute,
epoch: Instant,
capability: PhantomData<fn() -> C>,
}
impl<C: RequestCapability> fmt::Debug for CrossLaneRequestEndpoint<C> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CrossLaneRequestEndpoint")
.field("capability", &C::ID)
.finish_non_exhaustive()
}
}
impl<C> NativeRequestEndpoint for CrossLaneRequestEndpoint<C>
where
C: RequestCapability,
C::Request: Send,
C::Response: Send,
C::DomainError: Send,
{
fn capability_id(&self) -> &'static str {
C::ID
}
fn descriptor_version(&self) -> &'static str {
C::DESCRIPTOR_VERSION
}
fn operations(&self) -> &'static [&'static str] {
self.operations
}
fn invoke(
&self,
operation: &str,
request: Box<dyn Any>,
context: InvocationContext,
) -> LocalBoxFuture<'static, Result<Result<Box<dyn Any>, Box<dyn Any>>, RuntimeFailure>> {
let Ok(request) = request.downcast::<C::Request>() else {
return Box::pin(futures::future::ready(Err(
RuntimeFailure::ProtocolViolation { capability: C::ID },
)));
};
let provider_lane = self.provider_lane.clone();
let operation = operation.to_owned();
let epoch = self.epoch;
Box::pin(async move {
let provider_lane = provider_lane.upgrade().ok_or_else(lane_unavailable::<C>)?;
let caller_instance = context
.caller_instance()
.ok_or_else(|| RuntimeFailure::InvalidResolvedPlan {
detail: format!("cross-lane invocation of `{}` has no planned caller", C::ID),
})?
.to_owned();
let cancellation = context.cancellation();
let deadline = context.deadline();
let transferred = TransferredInvocationContext::capture(&context);
let source_completed = Arc::clone(&transferred.source_completed);
let (completed, completion) = futures::channel::oneshot::channel();
let command = LaneCommand::Run(Box::new(move |app| {
Box::pin(async move {
let local_cancellation = CancellationToken::new();
let watcher_token = local_cancellation.clone();
let cancelled = Arc::clone(&transferred.cancelled);
let watcher_completed = Arc::new(AtomicBool::new(false));
let watcher_done = Arc::clone(&watcher_completed);
tokio::task::spawn_local(async move {
while !cancelled.load(Ordering::Acquire)
&& !watcher_done.load(Ordering::Acquire)
{
tokio::task::yield_now().await;
}
if cancelled.load(Ordering::Acquire) {
watcher_token.cancel();
}
});
let context = transferred.restore(local_cancellation);
let result = app
.invoke_with_context::<C>(&caller_instance, &operation, context, *request)
.await;
watcher_completed.store(true, Ordering::Release);
let _ = completed.send(result);
})
}));
let send = provider_lane.send(command);
tokio::pin!(send);
match deadline {
Some(deadline) => {
let sleep = tokio::time::sleep_until((epoch + deadline).into());
tokio::pin!(sleep);
tokio::select! {
result = &mut send => result.map_err(|_| {
source_completed.store(true, Ordering::Release);
lane_unavailable::<C>()
})?,
_ = cancellation.cancelled() => {
source_completed.store(true, Ordering::Release);
return Err(RuntimeFailure::Cancelled { request_id: context.request_id() });
}
_ = &mut sleep => {
source_completed.store(true, Ordering::Release);
return Err(RuntimeFailure::DeadlineExceeded { request_id: context.request_id() });
}
}
}
None => {
tokio::select! {
result = &mut send => result.map_err(|_| {
source_completed.store(true, Ordering::Release);
lane_unavailable::<C>()
})?,
_ = cancellation.cancelled() => {
source_completed.store(true, Ordering::Release);
return Err(RuntimeFailure::Cancelled { request_id: context.request_id() });
}
}
}
}
let result = completion.await.map_err(|_| lane_unavailable::<C>());
source_completed.store(true, Ordering::Release);
result?.map(|domain| {
domain
.map(|response| Box::new(response) as Box<dyn Any>)
.map_err(|error| Box::new(error) as Box<dyn Any>)
})
})
}
}
fn lane_unavailable<C: RequestCapability>() -> RuntimeFailure {
RuntimeFailure::Internal {
detail: format!("provider lane for `{}` is unavailable", C::ID),
}
}
#[derive(Debug)]
struct TransferredInvocationContext {
request_id: u64,
deadline: Option<std::time::Duration>,
cancelled: Arc<AtomicBool>,
source_completed: Arc<AtomicBool>,
extensions: Vec<lenso_kernel::InvocationExtension>,
sealed_extensions: Vec<lenso_kernel::SealedInvocationExtension>,
}
impl TransferredInvocationContext {
fn capture(context: &InvocationContext) -> Self {
let cancelled = Arc::new(AtomicBool::new(context.is_cancelled()));
let cancellation = context.cancellation();
let signal = Arc::clone(&cancelled);
let source_completed = Arc::new(AtomicBool::new(false));
let watcher_completed = Arc::clone(&source_completed);
tokio::task::spawn_local(async move {
while !cancellation.is_cancelled() && !watcher_completed.load(Ordering::Acquire) {
tokio::task::yield_now().await;
}
if cancellation.is_cancelled() {
signal.store(true, Ordering::Release);
}
});
Self {
request_id: context.request_id(),
deadline: context.deadline(),
cancelled,
source_completed,
extensions: context.extensions().cloned().collect(),
sealed_extensions: context.sealed_extensions().cloned().collect(),
}
}
fn restore(&self, cancellation: CancellationToken) -> InvocationContext {
let mut context = InvocationContext::new(self.request_id, self.deadline, cancellation);
for extension in &self.extensions {
context = context
.with_extension(extension.key(), extension.value().to_vec())
.expect("captured ordinary Invocation Context extension remains valid");
}
for extension in &self.sealed_extensions {
context = context
.with_sealed_extension(extension.clone())
.expect("captured sealed Invocation Context extension remains valid");
}
context
}
}