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