mz_orchestrator_process/
lib.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use 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/// Configures a [`ProcessOrchestrator`].
60#[derive(Debug, Clone)]
61pub struct ProcessOrchestratorConfig {
62    /// The directory in which the orchestrator should look for executable
63    /// images.
64    pub image_dir: PathBuf,
65    /// Whether to supress output from spawned subprocesses.
66    pub suppress_output: bool,
67    /// The ID of the environment under orchestration.
68    pub environment_id: String,
69    /// The directory in which to store secrets.
70    pub secrets_dir: PathBuf,
71    /// A command to wrap the child command invocation
72    pub command_wrapper: Vec<String>,
73    /// Whether to crash this process if a child process crashes.
74    pub propagate_crashes: bool,
75    /// TCP proxy configuration.
76    ///
77    /// When enabled, for each named port of each created service, the process
78    /// orchestrator will bind a TCP listener that proxies incoming connections
79    /// to the underlying Unix domain socket. Each bound TCP address will be
80    /// emitted as a tracing event.
81    ///
82    /// The primary use is live debugging the running child services via tools
83    /// that do not support Unix domain sockets (e.g., Prometheus, web
84    /// browsers).
85    pub tcp_proxy: Option<ProcessOrchestratorTcpProxyConfig>,
86    /// A scratch directory that orchestrated processes can use for ephemeral storage.
87    pub scratch_directory: PathBuf,
88}
89
90/// Configures the TCP proxy for a [`ProcessOrchestrator`].
91///
92/// See [`ProcessOrchestratorConfig::tcp_proxy`].
93#[derive(Debug, Clone)]
94pub struct ProcessOrchestratorTcpProxyConfig {
95    /// The IP address on which to bind TCP listeners.
96    pub listen_addr: IpAddr,
97    /// A directory in which to write Prometheus scrape targets, for use with
98    /// Prometheus's file-based service discovery.
99    ///
100    /// Each [`NamespacedOrchestrator`] will maintain a single JSON file into
101    /// the directory named `NAMESPACE.json` containing the scrape targets for
102    /// all extant services. The scrape targets will use the TCP proxy address,
103    /// as Prometheus does not support scraping over Unix domain sockets.
104    ///
105    /// See also: <https://prometheus.io/docs/guides/file-sd/>
106    pub prometheus_service_discovery_dir: Option<PathBuf>,
107}
108
109/// An orchestrator backed by processes on the local machine.
110///
111/// **This orchestrator is for development only.** Due to limitations in the
112/// Unix process API, it does not exactly conform to the documented semantics
113/// of `Orchestrator`.
114#[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    /// Directly execute the provided binary
131    Direct,
132    /// Use Systemd to start the binary
133    Systemd,
134}
135
136impl LaunchSpec {
137    fn determine_implementation() -> Result<Self, anyhow::Error> {
138        // According to https://www.freedesktop.org/software/systemd/man/latest/sd_booted.html
139        // checking for `/run/systemd/system/` is the canonical way to determine if the system
140        // was booted up with systemd.
141        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                    // TODO: We could set `-p MemorySwapMax=0` here to disable regular swap.
173                }
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    /// Creates a new process orchestrator from the provided configuration.
191    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/// Configuration for a [`NamespacedProcessOrchestrator`].
282#[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        // Enable disk if the size does not disable it.
338        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
418/// Commands sent from a [`NamespacedProcessOrchestrator`] to its
419/// [`OrchestratorWorker`].
420///
421/// Commands for which the caller expects a result include a `result_tx` on which the
422/// [`OrchestratorWorker`] will deliver the result.
423enum 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
440/// Describes the desired state of a process.
441struct EnsureServiceConfig {
442    /// An opaque identifier for the executable or container image to run.
443    ///
444    /// Often names a container on Docker Hub or a path on the local machine.
445    pub image: String,
446    /// A function that generates the arguments for each process of the service
447    /// given the assigned listen addresses for each named port.
448    pub args: Box<dyn Fn(ServiceAssignments) -> Vec<String> + Send + Sync>,
449    /// Ports to expose.
450    pub ports: Vec<ServicePort>,
451    /// An optional limit on the memory that the service can use.
452    pub memory_limit: Option<MemoryLimit>,
453    /// An optional limit on the CPU that the service can use.
454    pub cpu_limit: Option<CpuLimit>,
455    /// The number of copies of this service to run.
456    pub scale: u16,
457    /// Arbitrary key–value pairs to attach to the service in the orchestrator
458    /// backend.
459    ///
460    /// The orchestrator backend may apply a prefix to the key if appropriate.
461    pub labels: BTreeMap<String, String>,
462    /// Whether scratch disk space should be allocated for the service.
463    pub disk: bool,
464}
465
466/// A task executing blocking work for a [`NamespacedProcessOrchestrator`] in the background.
467///
468/// This type exists to enable making [`NamespacedProcessOrchestrator::ensure_service`] and
469/// [`NamespacedProcessOrchestrator::drop_service`] non-blocking, allowing invocation of these
470/// methods in latency-sensitive contexts.
471///
472/// Note that, apart from `ensure_service` and `drop_service`, this worker also handles blocking
473/// orchestrator calls that query service state (such as `list_services`). These need to be
474/// sequenced through the worker loop to ensure they linearize as expected. For example, we want to
475/// ensure that a `list_services` result contains exactly those services that were previously
476/// created with `ensure_service` and not yet dropped with `drop_service`.
477struct 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                            // Justification for `unwrap`:
536                            //
537                            // `u64::try_cast_from(f: f64)`
538                            // will always succeed if 0 <= f < 2^64.
539                            // Since the max value of `process.cpu_usage()` is
540                            // 100.0 * num_of_cores, this will be true whenever there
541                            // are less than 2^64 / 10^9 logical cores, or about
542                            // 18 billion.
543                            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                // Process orchestrator does not support this right now.
557                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        // The service might already exist. If it has the same config as requested (currently we
594        // check only the scale), we have nothing to do. Otherwise we need to drop and recreate it.
595        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        // Create sockets for all processes in the service.
606        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            // Create the state for new processes.
622            let mut process_states = vec![];
623            for i in 0..scale.into() {
624                let listen_addrs = &peer_addrs[i];
625
626                // Fill out placeholders in the command wrapper for this process.
627                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                // Allocate listeners for each TCP proxy, if requested.
635                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                // Launch supervisor process.
676                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            // Update the in-memory process state. We do this after we've created
702            // all process states to avoid partially updating our in-memory state.
703            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        // Drop the supervisor for the service, if it exists. If this service
717        // was under supervision, this will kill all processes associated with
718        // it.
719        {
720            let mut supervisors = self.services.lock().expect("lock poisoned");
721            supervisors.remove(id);
722        }
723
724        // If the service was orphaned by a prior incarnation of the
725        // orchestrator, it won't have been under supervision and therefore will
726        // still be running. So kill any process that we have state for in the
727        // run directory.
728        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        // Clean up the on-disk state of the service.
744        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
944/// Supervises an existing process, if it exists.
945async 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    // Kill the process if the future is dropped.
961    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    // Periodically check if the process has terminated.
971    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    // The process has crashed. Exit the function without attempting to
977    // kill it.
978    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                // Give the process a bit of time to react to the signal
1010                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    // Immediately write out a file containing the PID of the child process and
1019    // its start time. We'll use this state to rediscover our children if we
1020    // crash and restart. There's a very small window where we can crash after
1021    // having spawned the child but before writing this file, in which case we
1022    // might orphan the process. We accept this risk, though. It's hard to do
1023    // anything more robust given the Unix APIs available to us, and the
1024    // solution here is good enough given that the process orchestrator is only
1025    // used in development/testing.
1026    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    // Likely not exhaustive. Feel free to add additional tests for other
1034    // indications of a crashed child process, as those conditions are
1035    // discovered.
1036    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    // Checking the start time protects against killing an unrelated process due
1070    // to PID reuse.
1071    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        // Unix socket addresses have a very low maximum length of around 100
1182        // bytes on most platforms.
1183        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}