1use crate::{
2 CommandSpec, DevPlan, ProcessPlan, ReadinessProbe, ServicePlan, is_sensitive_environment_name,
3};
4use serde::{Deserialize, Serialize};
5use std::{
6 collections::BTreeMap,
7 future::{Future, pending},
8 path::{Path, PathBuf},
9 pin::Pin,
10 process::{ExitStatus, Stdio},
11 time::Duration,
12};
13use thiserror::Error;
14use tokio::{
15 io::{AsyncBufReadExt, AsyncRead, BufReader},
16 process::Child,
17 sync::mpsc,
18 task::JoinHandle,
19 time,
20};
21
22struct ManagedChild {
23 id: String,
24 child: Child,
25 log_tasks: Vec<JoinHandle<()>>,
26 #[cfg(unix)]
27 process_group: u32,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31enum CompletedCommand {
32 Finished,
33 Shutdown,
34}
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
37#[serde(rename_all = "snake_case")]
38pub enum DevStream {
39 Stdout,
40 Stderr,
41}
42
43#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
44#[serde(tag = "kind", rename_all = "snake_case")]
45pub enum DevEvent {
46 Starting {
47 id: String,
48 },
49 Ready {
50 id: String,
51 },
52 Log {
53 id: String,
54 stream: DevStream,
55 line: String,
56 },
57 Stopping {
58 id: String,
59 },
60 Stopped {
61 id: String,
62 },
63 Failed {
64 id: String,
65 },
66}
67
68#[derive(Debug, Clone)]
69pub struct Supervisor {
70 root: PathBuf,
71 poll_interval: Duration,
72 readiness_timeout: Duration,
73 shutdown_grace: Duration,
74}
75
76impl Supervisor {
77 #[must_use]
78 pub fn new(root: impl AsRef<Path>) -> Self {
79 Self {
80 root: root.as_ref().to_path_buf(),
81 poll_interval: Duration::from_millis(50),
82 readiness_timeout: Duration::from_secs(30),
83 shutdown_grace: Duration::from_secs(3),
84 }
85 }
86
87 #[must_use]
88 pub const fn with_poll_interval(mut self, poll_interval: Duration) -> Self {
89 self.poll_interval = poll_interval;
90 self
91 }
92
93 #[must_use]
94 pub const fn with_shutdown_grace(mut self, shutdown_grace: Duration) -> Self {
95 self.shutdown_grace = shutdown_grace;
96 self
97 }
98
99 #[must_use]
100 pub const fn with_readiness_timeout(mut self, readiness_timeout: Duration) -> Self {
101 self.readiness_timeout = readiness_timeout;
102 self
103 }
104
105 pub async fn run_until<S>(
106 &self,
107 plan: &DevPlan,
108 runtime_environment: &BTreeMap<String, String>,
109 shutdown: S,
110 events: mpsc::UnboundedSender<DevEvent>,
111 ) -> Result<(), SupervisorError>
112 where
113 S: Future<Output = ()>,
114 {
115 tokio::pin!(shutdown);
116 let mut started_services = Vec::new();
117 for service in &plan.services {
118 if let Some(start) = &service.start {
119 events.send(DevEvent::Starting {
120 id: service.id.clone(),
121 })?;
122 let result = self
123 .run_completed_until(
124 &service.id,
125 start,
126 runtime_environment,
127 &events,
128 shutdown.as_mut(),
129 )
130 .await;
131 match result {
132 Ok(CompletedCommand::Finished) => {
133 started_services.push(service);
134 }
135 Ok(CompletedCommand::Shutdown) => {
136 started_services.push(service);
137 return self
138 .stop_services(&started_services, runtime_environment, &events)
139 .await;
140 }
141 Err(error) => {
142 if !matches!(error, SupervisorError::Spawn { .. }) {
143 started_services.push(service);
144 }
145 return match self
146 .stop_services(&started_services, runtime_environment, &events)
147 .await
148 {
149 Ok(()) => Err(error),
150 Err(cleanup) => Err(cleanup),
151 };
152 }
153 }
154 events.send(DevEvent::Ready {
155 id: service.id.clone(),
156 })?;
157 }
158 }
159
160 for lifecycle in &plan.lifecycle {
161 events.send(DevEvent::Starting {
162 id: lifecycle.id.clone(),
163 })?;
164 match self
165 .run_completed_until(
166 &lifecycle.id,
167 &lifecycle.command,
168 runtime_environment,
169 &events,
170 shutdown.as_mut(),
171 )
172 .await
173 {
174 Ok(CompletedCommand::Finished) => {}
175 Ok(CompletedCommand::Shutdown) => {
176 return self
177 .stop_services(&started_services, runtime_environment, &events)
178 .await;
179 }
180 Err(error) => {
181 return match self
182 .stop_services(&started_services, runtime_environment, &events)
183 .await
184 {
185 Ok(()) => Err(error),
186 Err(cleanup) => Err(cleanup),
187 };
188 }
189 }
190 events.send(DevEvent::Ready {
191 id: lifecycle.id.clone(),
192 })?;
193 }
194
195 let mut children = Vec::new();
196 for process in &plan.processes {
197 events.send(DevEvent::Starting {
198 id: process.id.clone(),
199 })?;
200 match self.spawn_process(process, runtime_environment, &events) {
201 Ok(child) => {
202 children.push(child);
203 let readiness = tokio::select! {
204 result = self.wait_for_readiness(
205 process,
206 children.last_mut().expect("child was just inserted"),
207 &events,
208 ) => Some(result),
209 () = &mut shutdown => None,
210 };
211 let Some(readiness) = readiness else {
212 return self
213 .stop_topology(
214 &mut children,
215 &started_services,
216 runtime_environment,
217 &events,
218 )
219 .await;
220 };
221 if let Err(error) = readiness {
222 return match self
223 .stop_topology(
224 &mut children,
225 &started_services,
226 runtime_environment,
227 &events,
228 )
229 .await
230 {
231 Ok(()) => Err(error),
232 Err(cleanup) => Err(cleanup),
233 };
234 }
235 }
236 Err(error) => {
237 return match self
238 .stop_topology(
239 &mut children,
240 &started_services,
241 runtime_environment,
242 &events,
243 )
244 .await
245 {
246 Ok(()) => Err(error),
247 Err(cleanup) => Err(cleanup),
248 };
249 }
250 }
251 }
252
253 let mut ticker = time::interval(self.poll_interval);
254 let outcome = loop {
255 tokio::select! {
256 () = &mut shutdown => break Ok(()),
257 _ = ticker.tick() => {
258 let mut exited = None;
259 for managed in &mut children {
260 if let Some(status) = managed.child.try_wait().map_err(|source| {
261 SupervisorError::Inspect {
262 id: managed.id.clone(),
263 source,
264 }
265 })? {
266 exited = Some((managed.id.clone(), status));
267 break;
268 }
269 }
270 if let Some((id, status)) = exited {
271 break Err(SupervisorError::ProcessExited { id, status });
272 }
273 }
274 }
275 };
276
277 match self
278 .stop_topology(
279 &mut children,
280 &started_services,
281 runtime_environment,
282 &events,
283 )
284 .await
285 {
286 Ok(()) => outcome,
287 Err(cleanup) => Err(cleanup),
288 }
289 }
290
291 async fn wait_for_readiness(
292 &self,
293 process: &ProcessPlan,
294 managed: &mut ManagedChild,
295 events: &mpsc::UnboundedSender<DevEvent>,
296 ) -> Result<(), SupervisorError> {
297 let ReadinessProbe::Http { url } = &process.readiness else {
298 events.send(DevEvent::Ready {
299 id: process.id.clone(),
300 })?;
301 return Ok(());
302 };
303 let url = local_readiness_url(url, &process.id)?;
304 let client = reqwest::Client::builder()
305 .redirect(reqwest::redirect::Policy::none())
306 .build()
307 .map_err(|source| SupervisorError::HttpClient {
308 id: process.id.clone(),
309 source,
310 })?;
311 let deadline = time::Instant::now() + self.readiness_timeout;
312 loop {
313 if let Some(status) =
314 managed
315 .child
316 .try_wait()
317 .map_err(|source| SupervisorError::Inspect {
318 id: process.id.clone(),
319 source,
320 })?
321 {
322 return Err(SupervisorError::ProcessExited {
323 id: process.id.clone(),
324 status,
325 });
326 }
327 let remaining = deadline.saturating_duration_since(time::Instant::now());
328 if remaining.is_zero() {
329 return Err(SupervisorError::ReadinessTimeout {
330 id: process.id.clone(),
331 });
332 }
333 let request_timeout = remaining.min(Duration::from_secs(1));
334 let response = time::timeout(request_timeout, client.get(url.clone()).send()).await;
335 if matches!(response, Ok(Ok(response)) if response.status().is_success()) {
336 events.send(DevEvent::Ready {
337 id: process.id.clone(),
338 })?;
339 return Ok(());
340 }
341 time::sleep(self.poll_interval.min(remaining)).await;
342 }
343 }
344
345 async fn run_completed_until<S>(
346 &self,
347 id: &str,
348 command: &CommandSpec,
349 runtime_environment: &BTreeMap<String, String>,
350 events: &mpsc::UnboundedSender<DevEvent>,
351 mut shutdown: Pin<&mut S>,
352 ) -> Result<CompletedCommand, SupervisorError>
353 where
354 S: Future<Output = ()>,
355 {
356 let mut managed = self.spawn_command(id, command, runtime_environment, events)?;
357 let status = tokio::select! {
358 result = managed.child.wait() => Some(result),
359 () = shutdown.as_mut() => None,
360 };
361 let Some(status) = status else {
362 self.terminate_managed(&mut managed).await?;
363 return Ok(CompletedCommand::Shutdown);
364 };
365 self.finish_log_tasks(&mut managed).await;
366 let status = status.map_err(|source| SupervisorError::Inspect {
367 id: id.into(),
368 source,
369 })?;
370 if status.success() {
371 Ok(CompletedCommand::Finished)
372 } else {
373 Err(SupervisorError::CommandFailed {
374 id: id.into(),
375 status,
376 })
377 }
378 }
379
380 async fn run_completed(
381 &self,
382 id: &str,
383 command: &CommandSpec,
384 runtime_environment: &BTreeMap<String, String>,
385 events: &mpsc::UnboundedSender<DevEvent>,
386 ) -> Result<(), SupervisorError> {
387 let never = pending();
388 tokio::pin!(never);
389 match self
390 .run_completed_until(id, command, runtime_environment, events, never.as_mut())
391 .await?
392 {
393 CompletedCommand::Finished => Ok(()),
394 CompletedCommand::Shutdown => unreachable!("pending shutdown cannot resolve"),
395 }
396 }
397
398 fn spawn_command(
399 &self,
400 id: &str,
401 command: &CommandSpec,
402 runtime_environment: &BTreeMap<String, String>,
403 events: &mpsc::UnboundedSender<DevEvent>,
404 ) -> Result<ManagedChild, SupervisorError> {
405 let mut configured = configured_command(&self.root, command, runtime_environment);
406 configured.stdout(Stdio::piped()).stderr(Stdio::piped());
407 #[cfg(unix)]
408 configured.process_group(0);
409 let mut child = configured
410 .spawn()
411 .map_err(|source| SupervisorError::Spawn {
412 id: id.into(),
413 source,
414 })?;
415 #[cfg(unix)]
416 let process_group = child
417 .id()
418 .ok_or_else(|| SupervisorError::MissingProcessId { id: id.into() })?;
419 let redactions = redaction_values(runtime_environment, &command.environment);
420 let mut log_tasks = Vec::new();
421 if let Some(stdout) = child.stdout.take() {
422 log_tasks.push(spawn_log_reader(
423 stdout,
424 id.into(),
425 DevStream::Stdout,
426 redactions.clone(),
427 events.clone(),
428 ));
429 }
430 if let Some(stderr) = child.stderr.take() {
431 log_tasks.push(spawn_log_reader(
432 stderr,
433 id.into(),
434 DevStream::Stderr,
435 redactions,
436 events.clone(),
437 ));
438 }
439 Ok(ManagedChild {
440 id: id.into(),
441 child,
442 log_tasks,
443 #[cfg(unix)]
444 process_group,
445 })
446 }
447
448 fn spawn_process(
449 &self,
450 process: &ProcessPlan,
451 runtime_environment: &BTreeMap<String, String>,
452 events: &mpsc::UnboundedSender<DevEvent>,
453 ) -> Result<ManagedChild, SupervisorError> {
454 self.spawn_command(&process.id, &process.command, runtime_environment, events)
455 }
456
457 async fn terminate_managed(&self, managed: &mut ManagedChild) -> Result<(), SupervisorError> {
458 #[cfg(unix)]
459 signal_process_group(managed.process_group, rustix::process::Signal::TERM);
460 #[cfg(not(unix))]
461 let _ = managed.child.start_kill();
462 self.reap_managed_until(managed, time::Instant::now() + self.shutdown_grace)
463 .await
464 }
465
466 async fn reap_managed_until(
467 &self,
468 managed: &mut ManagedChild,
469 deadline: time::Instant,
470 ) -> Result<(), SupervisorError> {
471 let wait = time::timeout_at(deadline, managed.child.wait()).await;
472 let result = if let Ok(result) = wait {
473 result
474 } else {
475 #[cfg(unix)]
476 signal_process_group(managed.process_group, rustix::process::Signal::KILL);
477 let _ = managed.child.start_kill();
478 if let Ok(result) = time::timeout(self.shutdown_grace, managed.child.wait()).await {
479 result
480 } else {
481 for task in managed.log_tasks.drain(..) {
482 task.abort();
483 }
484 return Err(SupervisorError::ShutdownTimeout {
485 id: managed.id.clone(),
486 });
487 }
488 };
489 self.finish_log_tasks(managed).await;
490 result
491 .map(|_| ())
492 .map_err(|source| SupervisorError::Inspect {
493 id: managed.id.clone(),
494 source,
495 })
496 }
497
498 async fn finish_log_tasks(&self, managed: &mut ManagedChild) {
499 #[cfg(unix)]
500 signal_process_group(managed.process_group, rustix::process::Signal::KILL);
501 for task in managed.log_tasks.drain(..) {
502 let _ = task.await;
503 }
504 }
505
506 async fn stop_children(
507 &self,
508 children: &mut [ManagedChild],
509 events: &mpsc::UnboundedSender<DevEvent>,
510 ) -> Result<(), SupervisorError> {
511 for managed in children.iter_mut() {
512 let _ = events.send(DevEvent::Stopping {
513 id: managed.id.clone(),
514 });
515 #[cfg(unix)]
516 signal_process_group(managed.process_group, rustix::process::Signal::TERM);
517 #[cfg(not(unix))]
518 let _ = managed.child.start_kill();
519 }
520 let deadline = time::Instant::now() + self.shutdown_grace;
521 let mut first_error = None;
522 for managed in children {
523 match self.reap_managed_until(managed, deadline).await {
524 Ok(()) => {
525 let _ = events.send(DevEvent::Stopped {
526 id: managed.id.clone(),
527 });
528 }
529 Err(error) => {
530 let _ = events.send(DevEvent::Failed {
531 id: managed.id.clone(),
532 });
533 if first_error.is_none() {
534 first_error = Some(error);
535 }
536 }
537 }
538 }
539 match first_error {
540 Some(error) => Err(error),
541 None => Ok(()),
542 }
543 }
544
545 async fn stop_topology(
546 &self,
547 children: &mut [ManagedChild],
548 services: &[&ServicePlan],
549 runtime_environment: &BTreeMap<String, String>,
550 events: &mpsc::UnboundedSender<DevEvent>,
551 ) -> Result<(), SupervisorError> {
552 let child_result = self.stop_children(children, events).await;
553 let service_result = self
554 .stop_services(services, runtime_environment, events)
555 .await;
556 match (child_result, service_result) {
557 (Err(error), _) | (Ok(()), Err(error)) => Err(error),
558 (Ok(()), Ok(())) => Ok(()),
559 }
560 }
561
562 async fn stop_services(
563 &self,
564 services: &[&ServicePlan],
565 runtime_environment: &BTreeMap<String, String>,
566 events: &mpsc::UnboundedSender<DevEvent>,
567 ) -> Result<(), SupervisorError> {
568 let mut first_error = None;
569 for service in services.iter().rev() {
570 let Some(stop) = &service.stop else {
571 continue;
572 };
573 let _ = events.send(DevEvent::Stopping {
574 id: service.id.clone(),
575 });
576 let result = self
577 .run_completed(&service.id, stop, runtime_environment, events)
578 .await;
579 match result {
580 Ok(()) => {
581 let _ = events.send(DevEvent::Stopped {
582 id: service.id.clone(),
583 });
584 }
585 Err(error) => {
586 let _ = events.send(DevEvent::Failed {
587 id: service.id.clone(),
588 });
589 if first_error.is_none() {
590 first_error = Some(error);
591 }
592 }
593 }
594 }
595 match first_error {
596 Some(error) => Err(error),
597 None => Ok(()),
598 }
599 }
600}
601
602fn spawn_log_reader<R>(
603 reader: R,
604 id: String,
605 stream: DevStream,
606 redactions: Vec<String>,
607 events: mpsc::UnboundedSender<DevEvent>,
608) -> JoinHandle<()>
609where
610 R: AsyncRead + Unpin + Send + 'static,
611{
612 tokio::spawn(async move {
613 let mut lines = BufReader::new(reader).lines();
614 while let Ok(Some(line)) = lines.next_line().await {
615 let line = redact_line(line, &redactions);
616 if events
617 .send(DevEvent::Log {
618 id: id.clone(),
619 stream,
620 line,
621 })
622 .is_err()
623 {
624 break;
625 }
626 }
627 })
628}
629
630fn redaction_values(
631 runtime_environment: &BTreeMap<String, String>,
632 command_environment: &BTreeMap<String, String>,
633) -> Vec<String> {
634 let mut values = runtime_environment
635 .iter()
636 .chain(command_environment.iter())
637 .filter(|(name, value)| is_sensitive_environment_name(name) && !value.is_empty())
638 .map(|(_, value)| value.clone())
639 .collect::<Vec<_>>();
640 values.sort_by_key(|value| std::cmp::Reverse(value.len()));
641 values.dedup();
642 values
643}
644
645fn redact_line(mut line: String, redactions: &[String]) -> String {
646 for value in redactions {
647 line = line.replace(value, "<redacted>");
648 }
649 line
650}
651
652fn local_readiness_url(value: &str, id: &str) -> Result<reqwest::Url, SupervisorError> {
653 let url = reqwest::Url::parse(value)
654 .map_err(|_| SupervisorError::InvalidReadinessUrl { id: id.into() })?;
655 let local_host = match url.host_str() {
656 Some("localhost") => true,
657 Some(host) => host
658 .parse::<std::net::IpAddr>()
659 .is_ok_and(|address| address.is_loopback()),
660 None => false,
661 };
662 if url.scheme() != "http"
663 || !local_host
664 || !url.username().is_empty()
665 || url.password().is_some()
666 || url.query().is_some()
667 || url.fragment().is_some()
668 {
669 return Err(SupervisorError::InvalidReadinessUrl { id: id.into() });
670 }
671 Ok(url)
672}
673
674fn configured_command(
675 root: &Path,
676 command: &CommandSpec,
677 runtime_environment: &BTreeMap<String, String>,
678) -> tokio::process::Command {
679 let mut configured = tokio::process::Command::new(&command.program);
680 configured
681 .args(&command.arguments)
682 .current_dir(root)
683 .envs(runtime_environment)
684 .envs(&command.environment)
685 .kill_on_drop(true);
686 configured
687}
688
689#[cfg(unix)]
690fn signal_process_group(process_group: u32, signal: rustix::process::Signal) {
691 let Ok(raw) = i32::try_from(process_group) else {
692 return;
693 };
694 let Some(process_group) = rustix::process::Pid::from_raw(raw) else {
695 return;
696 };
697 let _ = rustix::process::kill_process_group(process_group, signal);
698}
699
700#[derive(Debug, Error)]
701pub enum SupervisorError {
702 #[error("development event receiver closed")]
703 EventReceiverClosed,
704 #[error("failed to start `{id}`: {source}")]
705 Spawn {
706 id: String,
707 #[source]
708 source: std::io::Error,
709 },
710 #[error("development command `{id}` exited with {status}")]
711 CommandFailed { id: String, status: ExitStatus },
712 #[error("development process `{id}` exited unexpectedly with {status}")]
713 ProcessExited { id: String, status: ExitStatus },
714 #[error("failed to inspect development process `{id}`: {source}")]
715 Inspect {
716 id: String,
717 #[source]
718 source: std::io::Error,
719 },
720 #[error("development process `{id}` did not expose a process identifier")]
721 MissingProcessId { id: String },
722 #[error("development process `{id}` could not be reaped before the shutdown timeout")]
723 ShutdownTimeout { id: String },
724 #[error("development process `{id}` has a non-local or invalid readiness URL")]
725 InvalidReadinessUrl { id: String },
726 #[error("failed to create the readiness client for `{id}`: {source}")]
727 HttpClient {
728 id: String,
729 #[source]
730 source: reqwest::Error,
731 },
732 #[error("development process `{id}` did not become ready before the timeout")]
733 ReadinessTimeout { id: String },
734}
735
736impl From<mpsc::error::SendError<DevEvent>> for SupervisorError {
737 fn from(_: mpsc::error::SendError<DevEvent>) -> Self {
738 Self::EventReceiverClosed
739 }
740}