1use std::collections::BTreeMap;
11use std::env;
12use std::ffi::OsStr;
13use std::fmt::Debug;
14use std::fs::Permissions;
15use std::future::Future;
16use std::net::{IpAddr, SocketAddr, TcpListener as StdTcpListener};
17use std::os::unix::fs::PermissionsExt;
18use std::os::unix::process::ExitStatusExt;
19use std::path::{Path, PathBuf};
20use std::process::{ExitStatus, Stdio};
21use std::str::FromStr;
22use std::sync::atomic::{AtomicBool, Ordering};
23use std::sync::{Arc, Mutex};
24
25use anyhow::{Context, anyhow, bail};
26use async_stream::stream;
27use async_trait::async_trait;
28use chrono::{DateTime, Utc};
29use futures::StreamExt;
30use futures::stream::{BoxStream, FuturesUnordered};
31use itertools::Itertools;
32use libc::{SIGABRT, SIGBUS, SIGILL, SIGSEGV, SIGTRAP};
33use maplit::btreemap;
34use mz_orchestrator::scheduling_config::ServiceSchedulingConfig;
35use mz_orchestrator::{
36 CpuLimit, DiskLimit, MemoryLimit, NamespacedOrchestrator, Orchestrator, Service,
37 ServiceAssignments, ServiceConfig, ServiceEvent, ServicePort, ServiceProcessMetrics,
38 ServiceStatus,
39};
40use mz_ore::cast::{CastFrom, TryCastFrom};
41use mz_ore::error::ErrorExt;
42use mz_ore::netio::UnixSocketAddr;
43use mz_ore::result::ResultExt;
44use mz_ore::task::AbortOnDropHandle;
45use scopeguard::defer;
46use serde::Serialize;
47use sha1::{Digest, Sha1};
48use sysinfo::{Pid, PidExt, Process, ProcessExt, ProcessRefreshKind, System, SystemExt};
49use tokio::fs::remove_dir_all;
50use tokio::net::{TcpListener, UnixStream};
51use tokio::process::{Child, Command};
52use tokio::sync::{broadcast, mpsc, oneshot};
53use tokio::time::{self, Duration};
54use tokio::{fs, io, select};
55use tracing::{debug, error, info, warn};
56
57pub mod secrets;
58
59#[derive(Debug, Clone)]
61pub struct ProcessOrchestratorConfig {
62 pub image_dir: PathBuf,
65 pub suppress_output: bool,
67 pub environment_id: String,
69 pub secrets_dir: PathBuf,
71 pub command_wrapper: Vec<String>,
73 pub propagate_crashes: bool,
75 pub tcp_proxy: Option<ProcessOrchestratorTcpProxyConfig>,
86 pub scratch_directory: PathBuf,
88}
89
90#[derive(Debug, Clone)]
94pub struct ProcessOrchestratorTcpProxyConfig {
95 pub listen_addr: IpAddr,
97 pub prometheus_service_discovery_dir: Option<PathBuf>,
107}
108
109#[derive(Debug)]
115pub struct ProcessOrchestrator {
116 image_dir: PathBuf,
117 suppress_output: bool,
118 namespaces: Mutex<BTreeMap<String, Arc<dyn NamespacedOrchestrator>>>,
119 metadata_dir: PathBuf,
120 secrets_dir: PathBuf,
121 command_wrapper: Vec<String>,
122 propagate_crashes: bool,
123 tcp_proxy: Option<ProcessOrchestratorTcpProxyConfig>,
124 scratch_directory: PathBuf,
125 launch_spec: LaunchSpec,
126}
127
128#[derive(Debug, Clone, Copy)]
129enum LaunchSpec {
130 Direct,
132 Systemd,
134}
135
136impl LaunchSpec {
137 fn determine_implementation() -> Result<Self, anyhow::Error> {
138 match Path::new("/run/systemd/system/").try_exists()? {
142 true => Ok(Self::Systemd),
143 false => Ok(Self::Direct),
144 }
145 }
146
147 fn refine_command(
148 &self,
149 image: impl AsRef<OsStr>,
150 args: &[impl AsRef<OsStr>],
151 wrapper: &[String],
152 memory_limit: Option<&MemoryLimit>,
153 cpu_limit: Option<&CpuLimit>,
154 ) -> Command {
155 let mut cmd = match self {
156 Self::Direct => {
157 if let Some((program, wrapper_args)) = wrapper.split_first() {
158 let mut cmd = Command::new(program);
159 cmd.args(wrapper_args);
160 cmd.arg(image);
161 cmd
162 } else {
163 Command::new(image)
164 }
165 }
166 Self::Systemd => {
167 let mut cmd = Command::new("systemd-run");
168 cmd.args(["--user", "--scope", "--quiet"]);
169 if let Some(memory_limit) = memory_limit {
170 let memory_limit = memory_limit.0.as_u64();
171 cmd.args(["-p", &format!("MemoryMax={memory_limit}")]);
172 }
174 if let Some(cpu_limit) = cpu_limit {
175 let cpu_limit = (cpu_limit.as_millicpus() + 9) / 10;
176 cmd.args(["-p", &format!("CPUQuota={cpu_limit}%")]);
177 }
178
179 cmd.args(wrapper);
180 cmd.arg(image);
181 cmd
182 }
183 };
184 cmd.args(args);
185 cmd
186 }
187}
188
189impl ProcessOrchestrator {
190 pub async fn new(
192 ProcessOrchestratorConfig {
193 image_dir,
194 suppress_output,
195 environment_id,
196 secrets_dir,
197 command_wrapper,
198 propagate_crashes,
199 tcp_proxy,
200 scratch_directory,
201 }: ProcessOrchestratorConfig,
202 ) -> Result<ProcessOrchestrator, anyhow::Error> {
203 let metadata_dir = env::temp_dir().join(format!("environmentd-{environment_id}"));
204 fs::create_dir_all(&metadata_dir)
205 .await
206 .context("creating metadata directory")?;
207 fs::create_dir_all(&secrets_dir)
208 .await
209 .context("creating secrets directory")?;
210 fs::set_permissions(&secrets_dir, Permissions::from_mode(0o700))
211 .await
212 .context("setting secrets directory permissions")?;
213 if let Some(prometheus_dir) = tcp_proxy
214 .as_ref()
215 .and_then(|p| p.prometheus_service_discovery_dir.as_ref())
216 {
217 fs::create_dir_all(&prometheus_dir)
218 .await
219 .context("creating prometheus directory")?;
220 }
221
222 let launch_spec = LaunchSpec::determine_implementation()?;
223 info!(driver = ?launch_spec, "Process orchestrator launch spec");
224
225 Ok(ProcessOrchestrator {
226 image_dir: fs::canonicalize(image_dir).await?,
227 suppress_output,
228 namespaces: Mutex::new(BTreeMap::new()),
229 metadata_dir: fs::canonicalize(metadata_dir).await?,
230 secrets_dir: fs::canonicalize(secrets_dir).await?,
231 command_wrapper,
232 propagate_crashes,
233 tcp_proxy,
234 scratch_directory,
235 launch_spec,
236 })
237 }
238}
239
240impl Orchestrator for ProcessOrchestrator {
241 fn namespace(&self, namespace: &str) -> Arc<dyn NamespacedOrchestrator> {
242 let mut namespaces = self.namespaces.lock().expect("lock poisoned");
243 Arc::clone(namespaces.entry(namespace.into()).or_insert_with(|| {
244 let config = Arc::new(NamespacedProcessOrchestratorConfig {
245 namespace: namespace.into(),
246 image_dir: self.image_dir.clone(),
247 suppress_output: self.suppress_output,
248 metadata_dir: self.metadata_dir.clone(),
249 command_wrapper: self.command_wrapper.clone(),
250 propagate_crashes: self.propagate_crashes,
251 tcp_proxy: self.tcp_proxy.clone(),
252 scratch_directory: self.scratch_directory.clone(),
253 launch_spec: self.launch_spec,
254 });
255
256 let services = Arc::new(Mutex::new(BTreeMap::new()));
257 let (service_event_tx, service_event_rx) = broadcast::channel(16384);
258 let (command_tx, command_rx) = mpsc::unbounded_channel();
259
260 let worker = OrchestratorWorker {
261 config: Arc::clone(&config),
262 services: Arc::clone(&services),
263 service_event_tx,
264 system: System::new(),
265 command_rx,
266 }
267 .spawn();
268
269 Arc::new(NamespacedProcessOrchestrator {
270 config,
271 services,
272 service_event_rx,
273 command_tx,
274 scheduling_config: Default::default(),
275 _worker: worker,
276 })
277 }))
278 }
279}
280
281#[derive(Debug)]
283struct NamespacedProcessOrchestratorConfig {
284 namespace: String,
285 image_dir: PathBuf,
286 suppress_output: bool,
287 metadata_dir: PathBuf,
288 command_wrapper: Vec<String>,
289 propagate_crashes: bool,
290 tcp_proxy: Option<ProcessOrchestratorTcpProxyConfig>,
291 scratch_directory: PathBuf,
292 launch_spec: LaunchSpec,
293}
294
295impl NamespacedProcessOrchestratorConfig {
296 fn full_id(&self, id: &str) -> String {
297 format!("{}-{}", self.namespace, id)
298 }
299
300 fn service_run_dir(&self, id: &str) -> PathBuf {
301 self.metadata_dir.join(&self.full_id(id))
302 }
303
304 fn service_scratch_dir(&self, id: &str) -> PathBuf {
305 self.scratch_directory.join(&self.full_id(id))
306 }
307}
308
309#[derive(Debug)]
310struct NamespacedProcessOrchestrator {
311 config: Arc<NamespacedProcessOrchestratorConfig>,
312 services: Arc<Mutex<BTreeMap<String, Vec<ProcessState>>>>,
313 service_event_rx: broadcast::Receiver<ServiceEvent>,
314 command_tx: mpsc::UnboundedSender<WorkerCommand>,
315 scheduling_config: std::sync::RwLock<ServiceSchedulingConfig>,
316 _worker: AbortOnDropHandle<()>,
317}
318
319impl NamespacedProcessOrchestrator {
320 fn send_command(&self, cmd: WorkerCommand) {
321 self.command_tx.send(cmd).expect("worker task not dropped");
322 }
323}
324
325#[async_trait]
326impl NamespacedOrchestrator for NamespacedProcessOrchestrator {
327 fn ensure_service(
328 &self,
329 id: &str,
330 config: ServiceConfig,
331 ) -> Result<Box<dyn Service>, anyhow::Error> {
332 let service = ProcessService {
333 run_dir: self.config.service_run_dir(id),
334 scale: config.scale,
335 };
336
337 let disk = config.disk_limit != Some(DiskLimit::ZERO);
339
340 let config = EnsureServiceConfig {
341 image: config.image,
342 args: config.args,
343 ports: config.ports,
344 memory_limit: config.memory_limit,
345 cpu_limit: config.cpu_limit,
346 scale: config.scale,
347 labels: config.labels,
348 disk,
349 };
350
351 self.send_command(WorkerCommand::EnsureService {
352 id: id.to_string(),
353 config,
354 });
355
356 Ok(Box::new(service))
357 }
358
359 fn drop_service(&self, id: &str) -> Result<(), anyhow::Error> {
360 self.send_command(WorkerCommand::DropService { id: id.to_string() });
361 Ok(())
362 }
363
364 async fn list_services(&self) -> Result<Vec<String>, anyhow::Error> {
365 let (result_tx, result_rx) = oneshot::channel();
366 self.send_command(WorkerCommand::ListServices { result_tx });
367
368 result_rx.await.expect("worker task not dropped")
369 }
370
371 fn watch_services(&self) -> BoxStream<'static, Result<ServiceEvent, anyhow::Error>> {
372 let mut initial_events = vec![];
373 let mut service_event_rx = {
374 let services = self.services.lock().expect("lock poisoned");
375 for (service_id, process_states) in &*services {
376 for (process_id, process_state) in process_states.iter().enumerate() {
377 initial_events.push(ServiceEvent {
378 service_id: service_id.clone(),
379 process_id: u64::cast_from(process_id),
380 status: process_state.status.into(),
381 time: process_state.status_time,
382 });
383 }
384 }
385 self.service_event_rx.resubscribe()
386 };
387 Box::pin(stream! {
388 for event in initial_events {
389 yield Ok(event);
390 }
391 loop {
392 yield service_event_rx.recv().await.err_into();
393 }
394 })
395 }
396
397 async fn fetch_service_metrics(
398 &self,
399 id: &str,
400 ) -> Result<Vec<ServiceProcessMetrics>, anyhow::Error> {
401 let (result_tx, result_rx) = oneshot::channel();
402 self.send_command(WorkerCommand::FetchServiceMetrics {
403 id: id.to_string(),
404 result_tx,
405 });
406
407 result_rx.await.expect("worker task not dropped")
408 }
409
410 fn update_scheduling_config(
411 &self,
412 config: mz_orchestrator::scheduling_config::ServiceSchedulingConfig,
413 ) {
414 *self.scheduling_config.write().expect("poisoned") = config;
415 }
416}
417
418enum WorkerCommand {
424 EnsureService {
425 id: String,
426 config: EnsureServiceConfig,
427 },
428 DropService {
429 id: String,
430 },
431 ListServices {
432 result_tx: oneshot::Sender<Result<Vec<String>, anyhow::Error>>,
433 },
434 FetchServiceMetrics {
435 id: String,
436 result_tx: oneshot::Sender<Result<Vec<ServiceProcessMetrics>, anyhow::Error>>,
437 },
438}
439
440struct EnsureServiceConfig {
442 pub image: String,
446 pub args: Box<dyn Fn(ServiceAssignments) -> Vec<String> + Send + Sync>,
449 pub ports: Vec<ServicePort>,
451 pub memory_limit: Option<MemoryLimit>,
453 pub cpu_limit: Option<CpuLimit>,
455 pub scale: u16,
457 pub labels: BTreeMap<String, String>,
462 pub disk: bool,
464}
465
466struct OrchestratorWorker {
478 config: Arc<NamespacedProcessOrchestratorConfig>,
479 services: Arc<Mutex<BTreeMap<String, Vec<ProcessState>>>>,
480 service_event_tx: broadcast::Sender<ServiceEvent>,
481 system: System,
482 command_rx: mpsc::UnboundedReceiver<WorkerCommand>,
483}
484
485impl OrchestratorWorker {
486 fn spawn(self) -> AbortOnDropHandle<()> {
487 let name = format!("process-orchestrator:{}", self.config.namespace);
488 mz_ore::task::spawn(|| name, self.run()).abort_on_drop()
489 }
490
491 async fn run(mut self) {
492 while let Some(cmd) = self.command_rx.recv().await {
493 use WorkerCommand::*;
494 let result = match cmd {
495 EnsureService { id, config } => self.ensure_service(id, config).await,
496 DropService { id } => self.drop_service(&id).await,
497 ListServices { result_tx } => {
498 let _ = result_tx.send(self.list_services().await);
499 Ok(())
500 }
501 FetchServiceMetrics { id, result_tx } => {
502 let _ = result_tx.send(self.fetch_service_metrics(&id));
503 Ok(())
504 }
505 };
506
507 if let Err(error) = result {
508 panic!("process orchestrator worker failed: {error}");
509 }
510 }
511 }
512
513 fn fetch_service_metrics(
514 &mut self,
515 id: &str,
516 ) -> Result<Vec<ServiceProcessMetrics>, anyhow::Error> {
517 let pids: Vec<_> = {
518 let services = self.services.lock().expect("lock poisoned");
519 let Some(service) = services.get(id) else {
520 bail!("unknown service {id}")
521 };
522 service.iter().map(|p| p.pid()).collect()
523 };
524
525 let mut metrics = vec![];
526 for pid in pids {
527 let (cpu_nano_cores, memory_bytes) = match pid {
528 None => (None, None),
529 Some(pid) => {
530 self.system
531 .refresh_process_specifics(pid, ProcessRefreshKind::new().with_cpu());
532 match self.system.process(pid) {
533 None => (None, None),
534 Some(process) => {
535 let cpu = u64::try_cast_from(
544 (f64::from(process.cpu_usage()) * 10_000_000.0).trunc(),
545 )
546 .expect("sane value of process.cpu_usage()");
547 let memory = process.memory();
548 (Some(cpu), Some(memory))
549 }
550 }
551 }
552 };
553 metrics.push(ServiceProcessMetrics {
554 cpu_nano_cores,
555 memory_bytes,
556 disk_usage_bytes: None,
558 });
559 }
560 Ok(metrics)
561 }
562
563 async fn ensure_service(
564 &self,
565 id: String,
566 EnsureServiceConfig {
567 image,
568 args,
569 ports: ports_in,
570 memory_limit,
571 cpu_limit,
572 scale,
573 labels,
574 disk,
575 }: EnsureServiceConfig,
576 ) -> Result<(), anyhow::Error> {
577 let full_id = self.config.full_id(&id);
578
579 let run_dir = self.config.service_run_dir(&id);
580 fs::create_dir_all(&run_dir)
581 .await
582 .context("creating run directory")?;
583 let scratch_dir = if disk {
584 let scratch_dir = self.config.service_scratch_dir(&id);
585 fs::create_dir_all(&scratch_dir)
586 .await
587 .context("creating scratch directory")?;
588 Some(fs::canonicalize(&scratch_dir).await?)
589 } else {
590 None
591 };
592
593 let old_scale = {
596 let services = self.services.lock().expect("poisoned");
597 services.get(&id).map(|states| states.len())
598 };
599 match old_scale {
600 Some(old) if old == usize::from(scale) => return Ok(()),
601 Some(_) => self.drop_service(&id).await?,
602 None => (),
603 }
604
605 let mut peer_addrs = Vec::new();
607 for i in 0..scale.into() {
608 let addresses = ports_in
609 .iter()
610 .map(|port| {
611 let addr = socket_path(&run_dir, &port.name, i);
612 (port.name.clone(), addr)
613 })
614 .collect();
615 peer_addrs.push(addresses);
616 }
617
618 {
619 let mut services = self.services.lock().expect("lock poisoned");
620
621 let mut process_states = vec![];
623 for i in 0..scale.into() {
624 let listen_addrs = &peer_addrs[i];
625
626 let mut command_wrapper = self.config.command_wrapper.clone();
628 if let Some(parts) = command_wrapper.get_mut(1..) {
629 for part in parts {
630 *part = interpolate_command(&part[..], &full_id, listen_addrs);
631 }
632 }
633
634 let mut ports = vec![];
636 let mut tcp_proxy_addrs = BTreeMap::new();
637 for port in &ports_in {
638 let tcp_proxy_listener = match &self.config.tcp_proxy {
639 None => None,
640 Some(tcp_proxy) => {
641 let listener = StdTcpListener::bind((tcp_proxy.listen_addr, 0))
642 .with_context(|| format!("binding to {}", tcp_proxy.listen_addr))?;
643 listener.set_nonblocking(true)?;
644 let listener = TcpListener::from_std(listener)?;
645 let local_addr = listener.local_addr()?;
646 tcp_proxy_addrs.insert(port.name.clone(), local_addr);
647 Some(AddressedTcpListener {
648 listener,
649 local_addr,
650 })
651 }
652 };
653 ports.push(ServiceProcessPort {
654 name: port.name.clone(),
655 listen_addr: listen_addrs[&port.name].clone(),
656 tcp_proxy_listener,
657 });
658 }
659
660 let mut args = args(ServiceAssignments {
661 listen_addrs,
662 peer_addrs: &peer_addrs,
663 });
664 args.push(format!("--process={i}"));
665 if disk {
666 if let Some(scratch) = &scratch_dir {
667 args.push(format!("--scratch-directory={}", scratch.display()));
668 } else {
669 panic!(
670 "internal error: service requested disk but no scratch directory was configured"
671 );
672 }
673 }
674
675 let handle = mz_ore::task::spawn(
677 || format!("process-orchestrator:{full_id}-{i}"),
678 self.supervise_service_process(ServiceProcessConfig {
679 id: id.to_string(),
680 run_dir: run_dir.clone(),
681 i,
682 image: image.clone(),
683 args,
684 command_wrapper,
685 ports,
686 memory_limit,
687 cpu_limit,
688 launch_spec: self.config.launch_spec,
689 }),
690 );
691
692 process_states.push(ProcessState {
693 _handle: handle.abort_on_drop(),
694 status: ProcessStatus::NotReady,
695 status_time: Utc::now(),
696 labels: labels.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
697 tcp_proxy_addrs,
698 });
699 }
700
701 services.insert(id, process_states);
704 }
705
706 self.maybe_write_prometheus_service_discovery_file().await;
707
708 Ok(())
709 }
710
711 async fn drop_service(&self, id: &str) -> Result<(), anyhow::Error> {
712 let full_id = self.config.full_id(id);
713 let run_dir = self.config.service_run_dir(id);
714 let scratch_dir = self.config.service_scratch_dir(id);
715
716 {
720 let mut supervisors = self.services.lock().expect("lock poisoned");
721 supervisors.remove(id);
722 }
723
724 if let Ok(mut entries) = fs::read_dir(&run_dir).await {
729 while let Some(entry) = entries.next_entry().await? {
730 let path = entry.path();
731 if path.extension() == Some(OsStr::new("pid")) {
732 let mut system = System::new();
733 let Some(process) = find_process_from_pid_file(&mut system, &path).await else {
734 continue;
735 };
736 let pid = process.pid();
737 info!("terminating orphaned process for {full_id} with PID {pid}");
738 process.kill();
739 }
740 }
741 }
742
743 if let Err(e) = remove_dir_all(run_dir).await {
745 if e.kind() != io::ErrorKind::NotFound {
746 warn!(
747 "error cleaning up run directory for {full_id}: {}",
748 e.display_with_causes()
749 );
750 }
751 }
752 if let Err(e) = remove_dir_all(scratch_dir).await {
753 if e.kind() != io::ErrorKind::NotFound {
754 warn!(
755 "error cleaning up scratch directory for {full_id}: {}",
756 e.display_with_causes()
757 );
758 }
759 }
760
761 self.maybe_write_prometheus_service_discovery_file().await;
762 Ok(())
763 }
764
765 async fn list_services(&self) -> Result<Vec<String>, anyhow::Error> {
766 let mut services = vec![];
767 let namespace_prefix = format!("{}-", self.config.namespace);
768 let mut entries = fs::read_dir(&self.config.metadata_dir).await?;
769 while let Some(entry) = entries.next_entry().await? {
770 let filename = entry
771 .file_name()
772 .into_string()
773 .map_err(|_| anyhow!("unable to convert filename to string"))?;
774 if let Some(id) = filename.strip_prefix(&namespace_prefix) {
775 services.push(id.to_string());
776 }
777 }
778 Ok(services)
779 }
780
781 fn supervise_service_process(
782 &self,
783 ServiceProcessConfig {
784 id,
785 run_dir,
786 i,
787 image,
788 args,
789 command_wrapper,
790 ports,
791 memory_limit,
792 cpu_limit,
793 launch_spec,
794 }: ServiceProcessConfig,
795 ) -> impl Future<Output = ()> + use<> {
796 let suppress_output = self.config.suppress_output;
797 let propagate_crashes = self.config.propagate_crashes;
798 let image = self.config.image_dir.join(image);
799 let pid_file = run_dir.join(format!("{i}.pid"));
800 let full_id = self.config.full_id(&id);
801
802 let state_updater = ProcessStateUpdater {
803 namespace: self.config.namespace.clone(),
804 id,
805 i,
806 services: Arc::clone(&self.services),
807 service_event_tx: self.service_event_tx.clone(),
808 };
809
810 async move {
811 let mut proxy_handles = vec![];
812 for port in ports {
813 if let Some(tcp_listener) = port.tcp_proxy_listener {
814 info!(
815 "{full_id}-{i}: {} tcp proxy listening on {}",
816 port.name, tcp_listener.local_addr,
817 );
818 let uds_path = port.listen_addr;
819 let handle = mz_ore::task::spawn(
820 || format!("{full_id}-{i}-proxy-{}", port.name),
821 tcp_proxy(TcpProxyConfig {
822 name: format!("{full_id}-{i}-{}", port.name),
823 tcp_listener,
824 uds_path: uds_path.clone(),
825 }),
826 );
827 proxy_handles.push(handle.abort_on_drop());
828 }
829 }
830
831 supervise_existing_process(&state_updater, &pid_file).await;
832
833 loop {
834 let mut cmd = launch_spec.refine_command(
835 &image,
836 &args,
837 &command_wrapper,
838 memory_limit.as_ref(),
839 cpu_limit.as_ref(),
840 );
841 info!(
842 "launching {full_id}-{i} via {} {}...",
843 cmd.as_std().get_program().to_string_lossy(),
844 cmd.as_std()
845 .get_args()
846 .map(|arg| arg.to_string_lossy())
847 .join(" ")
848 );
849 if suppress_output {
850 cmd.stdout(Stdio::null());
851 cmd.stderr(Stdio::null());
852 }
853 match spawn_process(&state_updater, cmd, &pid_file, !command_wrapper.is_empty())
854 .await
855 {
856 Ok(status) => {
857 if propagate_crashes && did_process_crash(status) {
858 panic!(
859 "{full_id}-{i} crashed; aborting because propagate_crashes is enabled"
860 );
861 }
862 error!("{full_id}-{i} exited: {:?}; relaunching in 5s", status);
863 }
864 Err(e) => {
865 error!("{full_id}-{i} failed to spawn: {}; relaunching in 5s", e);
866 }
867 };
868 state_updater.update_state(ProcessStatus::NotReady);
869 time::sleep(Duration::from_secs(5)).await;
870 }
871 }
872 }
873
874 async fn maybe_write_prometheus_service_discovery_file(&self) {
875 #[derive(Serialize)]
876 struct StaticConfig {
877 labels: BTreeMap<String, String>,
878 targets: Vec<String>,
879 }
880
881 let Some(tcp_proxy) = &self.config.tcp_proxy else {
882 return;
883 };
884 let Some(dir) = &tcp_proxy.prometheus_service_discovery_dir else {
885 return;
886 };
887
888 let mut static_configs = vec![];
889 {
890 let services = self.services.lock().expect("lock poisoned");
891 for (id, states) in &*services {
892 for (i, state) in states.iter().enumerate() {
893 for (name, addr) in &state.tcp_proxy_addrs {
894 let mut labels = btreemap! {
895 "mz_orchestrator_namespace".into() => self.config.namespace.clone(),
896 "mz_orchestrator_service_id".into() => id.clone(),
897 "mz_orchestrator_port".into() => name.clone(),
898 "mz_orchestrator_ordinal".into() => i.to_string(),
899 };
900 for (k, v) in &state.labels {
901 let k = format!("mz_orchestrator_{}", k.replace('-', "_"));
902 labels.insert(k, v.clone());
903 }
904 static_configs.push(StaticConfig {
905 labels,
906 targets: vec![addr.to_string()],
907 })
908 }
909 }
910 }
911 }
912
913 let path = dir.join(Path::new(&self.config.namespace).with_extension("json"));
914 let contents = serde_json::to_vec_pretty(&static_configs).expect("valid json");
915 if let Err(e) = fs::write(&path, &contents).await {
916 warn!(
917 "{}: failed to write prometheus service discovery file: {}",
918 self.config.namespace,
919 e.display_with_causes()
920 );
921 }
922 }
923}
924
925struct ServiceProcessConfig {
926 id: String,
927 run_dir: PathBuf,
928 i: usize,
929 image: String,
930 args: Vec<String>,
931 command_wrapper: Vec<String>,
932 ports: Vec<ServiceProcessPort>,
933 memory_limit: Option<MemoryLimit>,
934 cpu_limit: Option<CpuLimit>,
935 launch_spec: LaunchSpec,
936}
937
938struct ServiceProcessPort {
939 name: String,
940 listen_addr: String,
941 tcp_proxy_listener: Option<AddressedTcpListener>,
942}
943
944async fn supervise_existing_process(state_updater: &ProcessStateUpdater, pid_file: &Path) {
946 let name = format!(
947 "{}-{}-{}",
948 state_updater.namespace, state_updater.id, state_updater.i
949 );
950
951 let mut system = System::new();
952 let Some(process) = find_process_from_pid_file(&mut system, pid_file).await else {
953 return;
954 };
955 let pid = process.pid();
956
957 info!(%pid, "discovered existing process for {name}");
958 state_updater.update_state(ProcessStatus::Ready { pid });
959
960 let need_kill = AtomicBool::new(true);
962 defer! {
963 state_updater.update_state(ProcessStatus::NotReady);
964 if need_kill.load(Ordering::SeqCst) {
965 info!(%pid, "terminating existing process for {name}");
966 process.kill();
967 }
968 }
969
970 let mut system = System::new();
972 while system.refresh_process_specifics(pid, ProcessRefreshKind::new()) {
973 time::sleep(Duration::from_secs(5)).await;
974 }
975
976 warn!(%pid, "process for {name} has crashed; will reboot");
979 need_kill.store(false, Ordering::SeqCst)
980}
981
982fn interpolate_command(
983 command_part: &str,
984 full_id: &str,
985 ports: &BTreeMap<String, String>,
986) -> String {
987 let mut command_part = command_part.replace("%N", full_id);
988 for (endpoint, port) in ports {
989 command_part = command_part.replace(&format!("%P:{endpoint}"), port);
990 }
991 command_part
992}
993
994async fn spawn_process(
995 state_updater: &ProcessStateUpdater,
996 mut cmd: Command,
997 pid_file: &Path,
998 send_sigterm: bool,
999) -> Result<ExitStatus, anyhow::Error> {
1000 struct KillOnDropChild(Child, bool);
1001
1002 impl Drop for KillOnDropChild {
1003 fn drop(&mut self) {
1004 if let (Some(pid), true) = (self.0.id().and_then(|id| i32::try_from(id).ok()), self.1) {
1005 let _ = nix::sys::signal::kill(
1006 nix::unistd::Pid::from_raw(pid),
1007 nix::sys::signal::Signal::SIGTERM,
1008 );
1009 tokio::task::block_in_place(|| std::thread::sleep(Duration::from_millis(500)));
1011 }
1012 let _ = self.0.start_kill();
1013 }
1014 }
1015
1016 let mut child = KillOnDropChild(cmd.spawn()?, send_sigterm);
1017
1018 let pid = Pid::from_u32(child.0.id().unwrap());
1027 write_pid_file(pid_file, pid).await?;
1028 state_updater.update_state(ProcessStatus::Ready { pid });
1029 Ok(child.0.wait().await?)
1030}
1031
1032fn did_process_crash(status: ExitStatus) -> bool {
1033 matches!(
1037 status.signal(),
1038 Some(SIGABRT | SIGBUS | SIGSEGV | SIGTRAP | SIGILL)
1039 )
1040}
1041
1042async fn write_pid_file(pid_file: &Path, pid: Pid) -> Result<(), anyhow::Error> {
1043 let mut system = System::new();
1044 system.refresh_process_specifics(pid, ProcessRefreshKind::new());
1045 let start_time = system.process(pid).map_or(0, |p| p.start_time());
1046 fs::write(pid_file, format!("{pid}\n{start_time}\n")).await?;
1047 Ok(())
1048}
1049
1050async fn find_process_from_pid_file<'a>(
1051 system: &'a mut System,
1052 pid_file: &Path,
1053) -> Option<&'a Process> {
1054 let Ok(contents) = fs::read_to_string(pid_file).await else {
1055 return None;
1056 };
1057 let lines = contents.trim().split('\n').collect::<Vec<_>>();
1058 let [pid, start_time] = lines.as_slice() else {
1059 return None;
1060 };
1061 let Ok(pid) = Pid::from_str(pid) else {
1062 return None;
1063 };
1064 let Ok(start_time) = u64::from_str(start_time) else {
1065 return None;
1066 };
1067 system.refresh_process_specifics(pid, ProcessRefreshKind::new());
1068 let process = system.process(pid)?;
1069 if process.start_time() != start_time {
1072 return None;
1073 }
1074 Some(process)
1075}
1076
1077struct TcpProxyConfig {
1078 name: String,
1079 tcp_listener: AddressedTcpListener,
1080 uds_path: String,
1081}
1082
1083async fn tcp_proxy(
1084 TcpProxyConfig {
1085 name,
1086 tcp_listener,
1087 uds_path,
1088 }: TcpProxyConfig,
1089) {
1090 let mut conns = FuturesUnordered::new();
1091 loop {
1092 select! {
1093 res = tcp_listener.listener.accept() => {
1094 debug!("{name}: accepting tcp proxy connection");
1095 let uds_path = uds_path.clone();
1096 conns.push(Box::pin(async move {
1097 let (mut tcp_conn, _) = res.context("accepting tcp connection")?;
1098 let mut uds_conn = UnixStream::connect(uds_path)
1099 .await
1100 .context("making uds connection")?;
1101 io::copy_bidirectional(&mut tcp_conn, &mut uds_conn)
1102 .await
1103 .context("proxying")
1104 }));
1105 }
1106 Some(Err(e)) = conns.next() => {
1107 warn!("{name}: tcp proxy connection failed: {}", e.display_with_causes());
1108 }
1109 }
1110 }
1111}
1112
1113struct ProcessStateUpdater {
1114 namespace: String,
1115 id: String,
1116 i: usize,
1117 services: Arc<Mutex<BTreeMap<String, Vec<ProcessState>>>>,
1118 service_event_tx: broadcast::Sender<ServiceEvent>,
1119}
1120
1121impl ProcessStateUpdater {
1122 fn update_state(&self, status: ProcessStatus) {
1123 let mut services = self.services.lock().expect("lock poisoned");
1124 let Some(process_states) = services.get_mut(&self.id) else {
1125 return;
1126 };
1127 let Some(process_state) = process_states.get_mut(self.i) else {
1128 return;
1129 };
1130 let status_time = Utc::now();
1131 process_state.status = status;
1132 process_state.status_time = status_time;
1133 let _ = self.service_event_tx.send(ServiceEvent {
1134 service_id: self.id.to_string(),
1135 process_id: u64::cast_from(self.i),
1136 status: status.into(),
1137 time: status_time,
1138 });
1139 }
1140}
1141
1142#[derive(Debug)]
1143struct ProcessState {
1144 _handle: AbortOnDropHandle<()>,
1145 status: ProcessStatus,
1146 status_time: DateTime<Utc>,
1147 labels: BTreeMap<String, String>,
1148 tcp_proxy_addrs: BTreeMap<String, SocketAddr>,
1149}
1150
1151impl ProcessState {
1152 fn pid(&self) -> Option<Pid> {
1153 match &self.status {
1154 ProcessStatus::NotReady => None,
1155 ProcessStatus::Ready { pid } => Some(*pid),
1156 }
1157 }
1158}
1159
1160#[derive(Debug, Clone, Copy)]
1161enum ProcessStatus {
1162 NotReady,
1163 Ready { pid: Pid },
1164}
1165
1166impl From<ProcessStatus> for ServiceStatus {
1167 fn from(status: ProcessStatus) -> ServiceStatus {
1168 match status {
1169 ProcessStatus::NotReady => ServiceStatus::Offline(None),
1170 ProcessStatus::Ready { .. } => ServiceStatus::Online,
1171 }
1172 }
1173}
1174
1175fn socket_path(run_dir: &Path, port: &str, process: usize) -> String {
1176 let desired = run_dir
1177 .join(format!("{port}-{process}"))
1178 .to_string_lossy()
1179 .into_owned();
1180 if UnixSocketAddr::from_pathname(&desired).is_err() {
1181 env::temp_dir()
1184 .join(hex::encode(Sha1::digest(desired)))
1185 .display()
1186 .to_string()
1187 } else {
1188 desired
1189 }
1190}
1191
1192struct AddressedTcpListener {
1193 listener: TcpListener,
1194 local_addr: SocketAddr,
1195}
1196
1197#[derive(Debug, Clone)]
1198struct ProcessService {
1199 run_dir: PathBuf,
1200 scale: u16,
1201}
1202
1203impl Service for ProcessService {
1204 fn addresses(&self, port: &str) -> Vec<String> {
1205 (0..self.scale)
1206 .map(|i| socket_path(&self.run_dir, port, i.into()))
1207 .collect()
1208 }
1209}