1use std::{
4 cell::Cell,
5 panic::AssertUnwindSafe,
6 rc::Rc,
7 time::{Duration, Instant},
8};
9
10use futures::{
11 channel::oneshot,
12 future::{AbortHandle, Abortable, FutureExt},
13};
14use lenso_app_plan::{PlanResolutionError, ResolvedAppPlan};
15use lenso_kernel::{
16 DriverTask, ExecutionAdapterCatalog, LocalTask, PlanValidationError, RuntimeDriver,
17 ShutdownOutcome, TaskOutcome, TerminalOutcome,
18};
19use tokio::sync::Notify;
20
21mod replicated;
22
23pub use replicated::{
24 CrossLaneRequestCatalog, CrossLaneTransferCatalog, LaneCancellationToken,
25 LaneDiagnosticsSnapshot, LaneInvocationOptions, ReplicatedAppRoute, ReplicatedNativeApp,
26 ReplicatedRunnerError,
27};
28
29#[derive(Clone, Debug)]
31pub struct TokioDriver {
32 started_at: Instant,
33 shutdown_requested: Rc<Cell<bool>>,
34 runtime_event: Rc<Notify>,
35 jitter_state: Rc<Cell<u64>>,
36}
37
38impl TokioDriver {
39 pub fn new() -> Self {
41 Self::with_epoch(Instant::now())
42 }
43
44 pub(crate) fn with_epoch(started_at: Instant) -> Self {
45 Self {
46 started_at,
47 shutdown_requested: Rc::new(Cell::new(false)),
48 runtime_event: Rc::new(Notify::new()),
49 jitter_state: Rc::new(Cell::new(
50 u64::try_from(started_at.elapsed().as_nanos()).unwrap_or(u64::MAX)
51 ^ 0x9e37_79b9_7f4a_7c15,
52 )),
53 }
54 }
55
56 pub fn request_shutdown(&self) {
58 self.shutdown_requested.set(true);
59 self.runtime_event.notify_one();
60 }
61}
62
63impl Default for TokioDriver {
64 fn default() -> Self {
65 Self::new()
66 }
67}
68
69impl RuntimeDriver for TokioDriver {
70 fn now(&self) -> Duration {
71 self.started_at.elapsed()
72 }
73
74 fn sleep_until(&self, deadline: Duration) -> futures::future::LocalBoxFuture<'static, ()> {
75 let target = self.started_at + deadline;
76 Box::pin(async move {
77 tokio::time::sleep_until(tokio::time::Instant::from_std(target)).await;
78 })
79 }
80
81 fn yield_now(&self) -> futures::future::LocalBoxFuture<'static, ()> {
82 Box::pin(tokio::task::yield_now())
83 }
84
85 fn wait_for_runtime_event(
86 &self,
87 deadline: Duration,
88 ) -> futures::future::LocalBoxFuture<'static, ()> {
89 let target = self.started_at + deadline;
90 let runtime_event = Rc::clone(&self.runtime_event);
91 Box::pin(async move {
92 tokio::select! {
93 () = runtime_event.notified() => {}
94 () = tokio::time::sleep_until(tokio::time::Instant::from_std(target)) => {}
95 }
96 })
97 }
98
99 fn jitter(&self, maximum: Duration) -> Duration {
100 if maximum.is_zero() {
101 return Duration::ZERO;
102 }
103 let next = self
104 .jitter_state
105 .get()
106 .wrapping_mul(6_364_136_223_846_793_005)
107 .wrapping_add(1_442_695_040_888_963_407);
108 self.jitter_state.set(next);
109 let maximum_nanos = maximum.as_nanos().min(u128::from(u64::MAX));
110 let jitter_nanos = u128::from(next) % maximum_nanos.saturating_add(1);
111 Duration::from_nanos(u64::try_from(jitter_nanos).unwrap_or(u64::MAX))
112 }
113
114 fn spawn_local(&self, task: LocalTask) -> Result<DriverTask, futures::task::SpawnError> {
115 let (abort, registration) = AbortHandle::new_pair();
116 let (completed, completion) = oneshot::channel();
117 tokio::task::spawn_local(async move {
118 let outcome = match AssertUnwindSafe(Abortable::new(task, registration))
119 .catch_unwind()
120 .await
121 {
122 Ok(Ok(())) => TaskOutcome::Completed,
123 Ok(Err(_)) => TaskOutcome::Cancelled,
124 Err(_) => TaskOutcome::Failed,
125 };
126 let _ = completed.send(outcome);
127 });
128 Ok(DriverTask::new(abort, completion))
129 }
130
131 fn shutdown_requested(&self) -> bool {
132 self.shutdown_requested.get()
133 }
134}
135
136pub async fn run<D: RuntimeDriver>(
138 plan: ResolvedAppPlan,
139 driver: D,
140 adapters: ExecutionAdapterCatalog,
141 shutdown_timeout: Duration,
142) -> Result<TerminalOutcome, PlanValidationError> {
143 if let Err(error) = plan.validate() {
144 return Err(match error {
145 PlanResolutionError::UnsupportedSchemaVersion { expected, actual } => {
146 PlanValidationError::UnsupportedSchemaVersion { expected, actual }
147 }
148 error => PlanValidationError::InvalidResolvedPlan {
149 detail: error.to_string(),
150 },
151 });
152 }
153
154 let app = match lenso_kernel::Kernel::start(plan, driver.clone(), adapters).await {
155 Ok(app) => app,
156 Err(error) => return Ok(TerminalOutcome::StartupFailure { error }),
157 };
158 while !driver.shutdown_requested() && !app.is_failed() {
159 let failure_check = driver.now().saturating_add(Duration::from_millis(10));
160 driver.wait_for_runtime_event(failure_check).await;
161 }
162 if let Some(error) = app.terminal_failure() {
163 return Ok(match app.shutdown(shutdown_timeout).await {
164 ShutdownOutcome::Clean => TerminalOutcome::RuntimeFailure { error },
165 ShutdownOutcome::RuntimeFailure {
166 error: cleanup_error,
167 } => TerminalOutcome::RuntimeFailureDuringShutdown {
168 error,
169 cleanup_error,
170 },
171 ShutdownOutcome::Timeout => {
172 TerminalOutcome::RuntimeFailureWithShutdownTimeout { error }
173 }
174 });
175 }
176 Ok(match app.shutdown(shutdown_timeout).await {
177 ShutdownOutcome::Clean => TerminalOutcome::CleanShutdown,
178 ShutdownOutcome::RuntimeFailure { error } => TerminalOutcome::RuntimeFailure { error },
179 ShutdownOutcome::Timeout => TerminalOutcome::ShutdownTimeout,
180 })
181}