1use std::collections::HashMap;
30use std::sync::{Arc, Mutex};
31use std::time::{Duration, SystemTime};
32
33use lightshuttle_manifest::{InterpolationContext, Interpolator};
34use tokio::sync::{broadcast, watch};
35use tracing::{Instrument, debug, info, info_span, instrument, warn};
36
37const EVENT_CHANNEL_CAPACITY: usize = 256;
41
42use crate::error::RuntimeError;
43use crate::lifecycle::error::LifecycleError;
44use crate::lifecycle::plan::LifecyclePlan;
45use crate::lifecycle::status::{LifecycleEvent, NodeStatus};
46use crate::runtime::{ContainerId, ContainerRuntime};
47use lightshuttle_spec::{ContainerSpec, ResourceOutputs};
48
49const DEFAULT_HEALTHCHECK_TIMEOUT: Duration = Duration::from_secs(60);
52
53#[derive(Clone)]
55struct NodeHandle {
56 status_tx: Arc<watch::Sender<NodeStatus>>,
57 status_rx: watch::Receiver<NodeStatus>,
58 outputs_tx: Arc<watch::Sender<Option<ResourceOutputs>>>,
59 outputs_rx: watch::Receiver<Option<ResourceOutputs>>,
60 container_id: Arc<Mutex<Option<ContainerId>>>,
61 started_at: Arc<Mutex<Option<SystemTime>>>,
62}
63
64pub(super) struct NodeSnapshot {
67 pub(super) status: NodeStatus,
69 pub(super) started_at: Option<SystemTime>,
71 pub(super) container_id: Option<ContainerId>,
73}
74
75pub struct LifecycleManager<R: ContainerRuntime + 'static> {
109 plan: Arc<LifecyclePlan>,
110 runtime: Arc<R>,
111 nodes: HashMap<String, NodeHandle>,
112 event_tx: broadcast::Sender<LifecycleEvent>,
113 extra_env: Arc<HashMap<String, String>>,
114}
115
116impl<R: ContainerRuntime + 'static> LifecycleManager<R> {
117 #[must_use]
121 pub fn new(plan: LifecyclePlan, runtime: R) -> (Self, broadcast::Receiver<LifecycleEvent>) {
122 let (event_tx, event_rx) = broadcast::channel(EVENT_CHANNEL_CAPACITY);
123 let mut nodes: HashMap<String, NodeHandle> = HashMap::new();
124 for node in plan.nodes() {
125 let (status_tx, status_rx) = watch::channel(NodeStatus::Pending);
126 let (outputs_tx, outputs_rx) = watch::channel(None);
127 nodes.insert(
128 node.name.clone(),
129 NodeHandle {
130 status_tx: Arc::new(status_tx),
131 status_rx,
132 outputs_tx: Arc::new(outputs_tx),
133 outputs_rx,
134 container_id: Arc::new(Mutex::new(None)),
135 started_at: Arc::new(Mutex::new(None)),
136 },
137 );
138 }
139 let manager = Self {
140 plan: Arc::new(plan),
141 runtime: Arc::new(runtime),
142 nodes,
143 event_tx,
144 extra_env: Arc::new(HashMap::new()),
145 };
146 (manager, event_rx)
147 }
148
149 #[must_use]
159 pub fn with_env(mut self, env: HashMap<String, String>) -> Self {
160 self.extra_env = Arc::new(env);
161 self
162 }
163
164 pub fn check_required_env(&self) -> Result<(), LifecycleError> {
178 let report = self.plan.env_report(&self.extra_env);
179 if report.has_missing() {
180 Err(LifecycleError::MissingEnvVars {
181 names: report.missing(),
182 })
183 } else {
184 Ok(())
185 }
186 }
187
188 pub async fn start_all(&self) -> Result<(), LifecycleError> {
206 let mut handles: Vec<tokio::task::JoinHandle<Result<(), LifecycleError>>> =
207 Vec::with_capacity(self.plan.nodes().len());
208
209 for node in self.plan.nodes() {
210 let mut dep_status_rxs: HashMap<String, watch::Receiver<NodeStatus>> = HashMap::new();
211 let mut dep_outputs_rxs: HashMap<String, watch::Receiver<Option<ResourceOutputs>>> =
212 HashMap::new();
213 for dep in &node.depends_on {
214 let handle = self
215 .nodes
216 .get(dep)
217 .ok_or_else(|| LifecycleError::ResourceNotFound(dep.clone()))?;
218 dep_status_rxs.insert(dep.clone(), handle.status_rx.clone());
219 dep_outputs_rxs.insert(dep.clone(), handle.outputs_rx.clone());
220 }
221
222 let node_handle = self.nodes[&node.name].clone();
223 let spec = node.spec.clone();
224 let own_outputs = node.outputs.clone();
225 let name = node.name.clone();
226 let runtime = Arc::clone(&self.runtime);
227 let event_tx = self.event_tx.clone();
228 let extra_env = Arc::clone(&self.extra_env);
229
230 let task = tokio::spawn(async move {
231 start_one(
232 name,
233 spec,
234 own_outputs,
235 runtime,
236 node_handle,
237 dep_status_rxs,
238 dep_outputs_rxs,
239 event_tx,
240 extra_env,
241 )
242 .await
243 });
244 handles.push(task);
245 }
246
247 let mut first_error: Option<LifecycleError> = None;
248 for handle in handles {
249 match handle.await {
250 Ok(Ok(())) => {}
251 Ok(Err(err)) => {
252 if first_error.is_none() {
253 first_error = Some(err);
254 }
255 }
256 Err(join_err) => {
257 if first_error.is_none() {
258 first_error = Some(LifecycleError::Start {
259 resource: "<panicked task>".to_owned(),
260 source: RuntimeError::InvalidSpec(join_err.to_string()),
261 });
262 }
263 }
264 }
265 }
266
267 if let Some(err) = first_error {
268 warn!(error = %err, "start_all failed; rolling back");
269 let _ = self.stop_all(Duration::from_secs(10)).await;
270 return Err(err);
271 }
272
273 let _ = self.event_tx.send(LifecycleEvent::StackStarted);
274 info!(
275 "stack started: {} resource(s) healthy",
276 self.plan.nodes().len()
277 );
278 Ok(())
279 }
280
281 #[instrument(skip_all, fields(resources = self.plan.nodes().len()))]
294 pub async fn stop_all(&self, grace: Duration) -> Result<(), LifecycleError> {
295 let _ = self.event_tx.send(LifecycleEvent::StackStopping);
296
297 let mut errors: Vec<(String, RuntimeError)> = Vec::new();
298 for node in self.plan.nodes().iter().rev() {
299 let Some(handle) = self.nodes.get(&node.name) else {
300 continue;
301 };
302 let id = {
303 let guard = handle
304 .container_id
305 .lock()
306 .expect("container_id mutex poisoned");
307 guard.clone()
308 };
309 let Some(id) = id else { continue };
310 let stop_span = info_span!("stop", resource = %node.name);
311 match self.runtime.stop(&id, grace).instrument(stop_span).await {
312 Ok(()) => {
313 let _ = handle.status_tx.send(NodeStatus::Stopped);
314 let _ = self.event_tx.send(LifecycleEvent::ResourceStopped {
315 name: node.name.clone(),
316 });
317 }
318 Err(e) => errors.push((node.name.clone(), e)),
319 }
320 }
321
322 let _ = self.event_tx.send(LifecycleEvent::StackStopped);
323
324 if let Some(project) = self.plan.nodes().first().map(|n| n.spec.project.as_str()) {
329 if let Err(e) = self.runtime.teardown_project_network(project).await {
330 warn!(error = %e, "could not remove project network");
331 }
332 }
333
334 if let Some((resource, source)) = errors.into_iter().next() {
335 return Err(LifecycleError::Stop { resource, source });
336 }
337 Ok(())
338 }
339
340 pub async fn run_until_signal(&self, grace: Duration) -> Result<(), LifecycleError> {
376 self.start_all().await?;
377 wait_for_shutdown_signal().await;
378 self.stop_all(grace).await
379 }
380
381 #[instrument(skip(self), fields(resource = %resource))]
401 pub async fn restart_one(&self, resource: &str) -> Result<(), LifecycleError> {
402 let node = self
403 .plan
404 .nodes()
405 .iter()
406 .find(|n| n.name == resource)
407 .ok_or_else(|| LifecycleError::ResourceNotFound(resource.to_owned()))?;
408 let handle = self
409 .nodes
410 .get(resource)
411 .ok_or_else(|| LifecycleError::ResourceNotFound(resource.to_owned()))?;
412
413 let id = {
415 let guard = handle
416 .container_id
417 .lock()
418 .expect("container_id mutex poisoned");
419 guard.clone()
420 };
421 if let Some(id) = id {
422 self.runtime
423 .stop(&id, Duration::from_secs(10))
424 .await
425 .map_err(|source| LifecycleError::Stop {
426 resource: resource.to_owned(),
427 source,
428 })?;
429 *handle
430 .container_id
431 .lock()
432 .expect("container_id mutex poisoned") = None;
433 *handle.started_at.lock().expect("started_at mutex poisoned") = None;
434 let _ = handle.status_tx.send(NodeStatus::Stopped);
435 let _ = self.event_tx.send(LifecycleEvent::ResourceStopped {
436 name: resource.to_owned(),
437 });
438 }
439
440 let _ = handle.status_tx.send(NodeStatus::Pending);
442
443 let mut dep_status_rxs: HashMap<String, watch::Receiver<NodeStatus>> = HashMap::new();
446 let mut dep_outputs_rxs: HashMap<String, watch::Receiver<Option<ResourceOutputs>>> =
447 HashMap::new();
448 for dep in &node.depends_on {
449 let dep_handle = self
450 .nodes
451 .get(dep)
452 .ok_or_else(|| LifecycleError::ResourceNotFound(dep.clone()))?;
453 dep_status_rxs.insert(dep.clone(), dep_handle.status_rx.clone());
454 dep_outputs_rxs.insert(dep.clone(), dep_handle.outputs_rx.clone());
455 }
456
457 start_one(
458 resource.to_owned(),
459 node.spec.clone(),
460 node.outputs.clone(),
461 Arc::clone(&self.runtime),
462 handle.clone(),
463 dep_status_rxs,
464 dep_outputs_rxs,
465 self.event_tx.clone(),
466 Arc::clone(&self.extra_env),
467 )
468 .await
469 }
470
471 #[must_use]
477 pub fn subscribe_events(&self) -> broadcast::Receiver<LifecycleEvent> {
478 self.event_tx.subscribe()
479 }
480
481 pub(super) fn plan_arc(&self) -> &Arc<LifecyclePlan> {
483 &self.plan
484 }
485
486 pub(super) fn runtime_arc(&self) -> &Arc<R> {
488 &self.runtime
489 }
490
491 pub(super) fn snapshot(&self, name: &str) -> Option<NodeSnapshot> {
494 let handle = self.nodes.get(name)?;
495 let status = handle.status_rx.borrow().clone();
496 let started_at = *handle.started_at.lock().expect("started_at mutex poisoned");
497 let container_id = handle
498 .container_id
499 .lock()
500 .expect("container_id mutex poisoned")
501 .clone();
502 Some(NodeSnapshot {
503 status,
504 started_at,
505 container_id,
506 })
507 }
508}
509
510#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
511#[instrument(name = "start", skip_all, fields(resource = %name))]
512async fn start_one<R: ContainerRuntime + 'static>(
513 name: String,
514 spec: ContainerSpec,
515 own_outputs: ResourceOutputs,
516 runtime: Arc<R>,
517 handle: NodeHandle,
518 dep_status_rxs: HashMap<String, watch::Receiver<NodeStatus>>,
519 mut dep_outputs_rxs: HashMap<String, watch::Receiver<Option<ResourceOutputs>>>,
520 event_tx: broadcast::Sender<LifecycleEvent>,
521 extra_env: Arc<HashMap<String, String>>,
522) -> Result<(), LifecycleError> {
523 for (dep_name, mut rx) in dep_status_rxs {
525 loop {
526 let status = rx.borrow_and_update().clone();
527 if status.is_ready() {
528 debug!(node = %name, dep = %dep_name, "dependency ready");
529 break;
530 }
531 if let NodeStatus::Failed { reason } = status {
532 let _ = handle.status_tx.send(NodeStatus::Failed {
533 reason: format!("dependency `{dep_name}` failed: {reason}"),
534 });
535 return Err(LifecycleError::DependencyFailed {
536 resource: name,
537 dependency: dep_name,
538 reason,
539 });
540 }
541 if rx.changed().await.is_err() {
542 let reason = format!("dependency `{dep_name}` watch channel closed");
543 let _ = handle.status_tx.send(NodeStatus::Failed {
544 reason: reason.clone(),
545 });
546 return Err(LifecycleError::DependencyFailed {
547 resource: name,
548 dependency: dep_name,
549 reason,
550 });
551 }
552 }
553 }
554
555 let mut dep_outputs: HashMap<String, ResourceOutputs> = HashMap::new();
557 for (dep_name, rx) in &mut dep_outputs_rxs {
558 loop {
559 if let Some(out) = rx.borrow_and_update().clone() {
560 dep_outputs.insert(dep_name.clone(), out);
561 break;
562 }
563 if rx.changed().await.is_err() {
564 let reason = format!("dependency `{dep_name}` outputs channel closed");
565 let _ = handle.status_tx.send(NodeStatus::Failed {
566 reason: reason.clone(),
567 });
568 return Err(LifecycleError::DependencyFailed {
569 resource: name,
570 dependency: dep_name.clone(),
571 reason,
572 });
573 }
574 }
575 }
576
577 let resolved_spec = match interpolate_and_inject(spec, &dep_outputs, &extra_env) {
579 Ok(s) => s,
580 Err(reason) => {
581 let _ = handle.status_tx.send(NodeStatus::Failed {
582 reason: reason.clone(),
583 });
584 return Err(LifecycleError::Start {
585 resource: name,
586 source: RuntimeError::InvalidSpec(reason),
587 });
588 }
589 };
590
591 let _ = handle.status_tx.send(NodeStatus::Starting);
594 if let Err(source) = runtime.remove(&resolved_spec.name).await {
595 let _ = handle.status_tx.send(NodeStatus::Failed {
596 reason: source.to_string(),
597 });
598 let _ = event_tx.send(LifecycleEvent::ResourceFailed {
599 name: name.clone(),
600 error: source.to_string(),
601 });
602 return Err(LifecycleError::Start {
603 resource: name,
604 source,
605 });
606 }
607
608 let id = match runtime.start(&resolved_spec).await {
610 Ok(id) => id,
611 Err(source) => {
612 let _ = handle.status_tx.send(NodeStatus::Failed {
613 reason: source.to_string(),
614 });
615 let _ = event_tx.send(LifecycleEvent::ResourceFailed {
616 name: name.clone(),
617 error: source.to_string(),
618 });
619 return Err(LifecycleError::Start {
620 resource: name,
621 source,
622 });
623 }
624 };
625
626 {
627 let mut guard = handle
628 .container_id
629 .lock()
630 .expect("container_id mutex poisoned");
631 *guard = Some(id.clone());
632 }
633 {
634 let mut guard = handle.started_at.lock().expect("started_at mutex poisoned");
635 *guard = Some(SystemTime::now());
636 }
637 let _ = handle.status_tx.send(NodeStatus::Running);
638 let _ = event_tx.send(LifecycleEvent::ResourceStarted {
639 name: name.clone(),
640 container_id: id.to_string(),
641 });
642
643 let wait_span = info_span!("wait_healthy", resource = %name);
645 match runtime
646 .wait_healthy(&id, DEFAULT_HEALTHCHECK_TIMEOUT)
647 .instrument(wait_span)
648 .await
649 {
650 Ok(()) => {
651 let _ = handle.outputs_tx.send(Some(own_outputs));
652 let _ = handle.status_tx.send(NodeStatus::Healthy);
653 let _ = event_tx.send(LifecycleEvent::ResourceHealthy { name: name.clone() });
654 Ok(())
655 }
656 Err(RuntimeError::Timeout { .. }) => {
657 let reason = format!("healthcheck timed out after {DEFAULT_HEALTHCHECK_TIMEOUT:?}");
658 let _ = handle.status_tx.send(NodeStatus::Failed {
659 reason: reason.clone(),
660 });
661 let _ = event_tx.send(LifecycleEvent::ResourceFailed {
662 name: name.clone(),
663 error: reason,
664 });
665 Err(LifecycleError::HealthcheckTimeout {
666 resource: name,
667 timeout: DEFAULT_HEALTHCHECK_TIMEOUT,
668 })
669 }
670 Err(source) => {
671 let _ = handle.status_tx.send(NodeStatus::Failed {
672 reason: source.to_string(),
673 });
674 let _ = event_tx.send(LifecycleEvent::ResourceFailed {
675 name: name.clone(),
676 error: source.to_string(),
677 });
678 Err(LifecycleError::Start {
679 resource: name,
680 source,
681 })
682 }
683 }
684}
685
686fn interpolate_and_inject(
693 mut spec: ContainerSpec,
694 dep_outputs: &HashMap<String, ResourceOutputs>,
695 extra_env: &HashMap<String, String>,
696) -> std::result::Result<ContainerSpec, String> {
697 let mut ctx = InterpolationContext::from_env()
698 .with_env(extra_env.iter().map(|(k, v)| (k.clone(), v.clone())));
699 for (name, outputs) in dep_outputs {
700 ctx = ctx.with_resource(name.clone(), outputs.clone());
701 }
702 let interpolator = Interpolator::new(&ctx);
703
704 let mut resolved_env = std::collections::HashMap::with_capacity(spec.env.len());
706 for (k, v) in spec.env.drain() {
707 let resolved = interpolator.resolve(&v).map_err(|e| e.to_string())?;
708 resolved_env.insert(k, resolved);
709 }
710
711 for (dep_name, outputs) in dep_outputs {
713 let dep_upper = dep_name.to_uppercase().replace('-', "_");
714 for (prop, value) in outputs {
715 let prop_upper = prop.to_uppercase().replace('-', "_");
716 let key = format!("LSH_{dep_upper}_{prop_upper}");
717 resolved_env.entry(key).or_insert_with(|| value.clone());
718 }
719 }
720 spec.env = resolved_env;
721
722 if let Some(args) = spec.command.as_mut() {
724 for arg in args.iter_mut() {
725 *arg = interpolator.resolve(arg).map_err(|e| e.to_string())?;
726 }
727 }
728
729 Ok(spec)
730}
731
732#[cfg(unix)]
733async fn wait_for_shutdown_signal() {
734 use tokio::signal::unix::{SignalKind, signal};
735 let mut sigterm = match signal(SignalKind::terminate()) {
736 Ok(s) => s,
737 Err(e) => {
738 warn!("failed to install SIGTERM handler: {e}");
739 let _ = tokio::signal::ctrl_c().await;
740 return;
741 }
742 };
743 tokio::select! {
744 _ = tokio::signal::ctrl_c() => info!("received SIGINT"),
745 _ = sigterm.recv() => info!("received SIGTERM"),
746 }
747}
748
749#[cfg(windows)]
750async fn wait_for_shutdown_signal() {
751 let _ = tokio::signal::ctrl_c().await;
752 info!("received Ctrl+C");
753}