mz_orchestrator/
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::fmt;
12use std::str::FromStr;
13use std::sync::Arc;
14
15use async_trait::async_trait;
16use bytesize::ByteSize;
17use chrono::{DateTime, Utc};
18use derivative::Derivative;
19use futures_core::stream::BoxStream;
20use mz_ore::cast::CastFrom;
21use serde::de::Unexpected;
22use serde::{Deserialize, Deserializer, Serialize};
23
24/// An orchestrator manages services.
25///
26/// A service is a set of one or more processes running the same image. See
27/// [`ServiceConfig`] for details.
28///
29/// All services live within a namespace. A namespace allows multiple users to
30/// share an orchestrator without conflicting: each user can only create,
31/// delete, and list the services within their namespace. Namespaces are not
32/// isolated at the network level, however: services in one namespace can
33/// communicate with services in another namespace with no restrictions.
34///
35/// Services **must** be tolerant of running as part of a distributed system. In
36/// particular, services **must** be prepared for the possibility that there are
37/// two live processes with the same identity. This can happen, for example,
38/// when the machine hosting a process *appears* to fail, from the perspective
39/// of the orchestrator, and so the orchestrator restarts the process on another
40/// machine, but in fact the original machine is still alive, just on the
41/// opposite side of a network partition. Be sure to design any communication
42/// with other services (e.g., an external database) to correctly handle
43/// competing communication from another incarnation of the service.
44///
45/// The intent is that you can implement `Orchestrator` with pods in Kubernetes,
46/// containers in Docker, or processes on your local machine.
47pub trait Orchestrator: fmt::Debug + Send + Sync {
48    /// Enter a namespace in the orchestrator.
49    fn namespace(&self, namespace: &str) -> Arc<dyn NamespacedOrchestrator>;
50}
51
52/// An orchestrator restricted to a single namespace.
53#[async_trait]
54pub trait NamespacedOrchestrator: fmt::Debug + Send + Sync {
55    /// Ensures that a service with the given configuration is running.
56    ///
57    /// If a service with the same ID already exists, its configuration is
58    /// updated to match `config`. This may or may not involve restarting the
59    /// service, depending on whether the existing service matches `config`.
60    fn ensure_service(
61        &self,
62        id: &str,
63        config: ServiceConfig,
64    ) -> Result<Box<dyn Service>, anyhow::Error>;
65
66    /// Drops the identified service, if it exists.
67    fn drop_service(&self, id: &str) -> Result<(), anyhow::Error>;
68
69    /// Lists the identifiers of all known services.
70    async fn list_services(&self) -> Result<Vec<String>, anyhow::Error>;
71
72    /// Watch for status changes of all known services.
73    fn watch_services(&self) -> BoxStream<'static, Result<ServiceEvent, anyhow::Error>>;
74
75    /// Gets resource usage metrics for all processes associated with a service.
76    ///
77    /// Returns `Err` if the entire process failed. Returns `Ok(v)` otherwise,
78    /// with one element in `v` for each process of the service,
79    /// even in not all metrics could be collected for all processes.
80    /// In such a case, the corresponding fields of `ServiceProcessMetrics` will be `None`.
81    async fn fetch_service_metrics(
82        &self,
83        id: &str,
84    ) -> Result<Vec<ServiceProcessMetrics>, anyhow::Error>;
85
86    fn update_scheduling_config(&self, config: scheduling_config::ServiceSchedulingConfig);
87}
88
89/// An event describing a status change of an orchestrated service.
90#[derive(Debug, Clone, Serialize)]
91pub struct ServiceEvent {
92    pub service_id: String,
93    pub process_id: u64,
94    pub status: ServiceStatus,
95    pub time: DateTime<Utc>,
96}
97
98/// Why the service is not ready, if known
99#[derive(Debug, Clone, Copy, Serialize, Eq, PartialEq)]
100pub enum OfflineReason {
101    OomKilled,
102    Initializing,
103}
104
105impl fmt::Display for OfflineReason {
106    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
107        match self {
108            OfflineReason::OomKilled => f.write_str("oom-killed"),
109            OfflineReason::Initializing => f.write_str("initializing"),
110        }
111    }
112}
113
114/// Describes the status of an orchestrated service.
115#[derive(Debug, Clone, Copy, Serialize, Eq, PartialEq)]
116pub enum ServiceStatus {
117    /// Service is ready to accept requests.
118    Online,
119    /// Service is not ready to accept requests.
120    /// The inner element is `None` if the reason
121    /// is unknown
122    Offline(Option<OfflineReason>),
123}
124
125impl ServiceStatus {
126    /// Returns the service status as a kebab-case string.
127    pub fn as_kebab_case_str(&self) -> &'static str {
128        match self {
129            ServiceStatus::Online => "online",
130            ServiceStatus::Offline(_) => "offline",
131        }
132    }
133}
134
135/// Describes a running service managed by an `Orchestrator`.
136pub trait Service: fmt::Debug + Send + Sync {
137    /// Given the name of a port, returns the addresses for each of the
138    /// service's processes, in order.
139    ///
140    /// Panics if `port` does not name a valid port.
141    fn addresses(&self, port: &str) -> Vec<String>;
142}
143
144#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize, Eq, PartialEq)]
145pub struct ServiceProcessMetrics {
146    pub cpu_nano_cores: Option<u64>,
147    pub memory_bytes: Option<u64>,
148    pub disk_usage_bytes: Option<u64>,
149}
150
151/// A simple language for describing assertions about a label's existence and value.
152///
153/// Used by [`LabelSelector`].
154#[derive(Clone, Debug)]
155pub enum LabelSelectionLogic {
156    /// The label exists and its value equals the given value.
157    /// Equivalent to `InSet { values: vec![value] }`
158    Eq { value: String },
159    /// Either the label does not exist, or it exists
160    /// but its value does not equal the given value.
161    /// Equivalent to `NotInSet { values: vec![value] }`
162    NotEq { value: String },
163    /// The label exists.
164    Exists,
165    /// The label does not exist.
166    NotExists,
167    /// The label exists and its value is one of the given values.
168    InSet { values: Vec<String> },
169    /// Either the label does not exist, or it exists
170    /// but its value is not one of the given values.
171    NotInSet { values: Vec<String> },
172}
173
174/// A simple language for describing whether a label
175/// exists and whether the value corresponding to it is in some set.
176/// Intended to correspond to the capabilities offered by Kubernetes label selectors,
177/// but without directly exposing Kubernetes API code to consumers of this module.
178#[derive(Clone, Debug)]
179pub struct LabelSelector {
180    /// The name of the label
181    pub label_name: String,
182    /// An assertion about the existence and value of a label
183    /// named `label_name`
184    pub logic: LabelSelectionLogic,
185}
186
187/// Describes the desired state of a service.
188#[derive(Derivative)]
189#[derivative(Debug)]
190pub struct ServiceConfig {
191    /// An opaque identifier for the executable or container image to run.
192    ///
193    /// Often names a container on Docker Hub or a path on the local machine.
194    pub image: String,
195    /// For the Kubernetes orchestrator, this is an init container to
196    /// configure for the pod running the service.
197    pub init_container_image: Option<String>,
198    /// A function that generates the arguments for each process of the service
199    /// given the assigned listen addresses for each named port.
200    #[derivative(Debug = "ignore")]
201    pub args: Box<dyn Fn(ServiceAssignments) -> Vec<String> + Send + Sync>,
202    /// Ports to expose.
203    pub ports: Vec<ServicePort>,
204    /// An optional limit on the memory that the service can use.
205    pub memory_limit: Option<MemoryLimit>,
206    /// An optional request on the memory that the service can use. If unspecified,
207    /// use the same value as `memory_limit`.
208    pub memory_request: Option<MemoryLimit>,
209    /// An optional limit on the CPU that the service can use.
210    pub cpu_limit: Option<CpuLimit>,
211    /// The number of copies of this service to run.
212    pub scale: u16,
213    /// Arbitrary key–value pairs to attach to the service in the orchestrator
214    /// backend.
215    ///
216    /// The orchestrator backend may apply a prefix to the key if appropriate.
217    pub labels: BTreeMap<String, String>,
218    /// Arbitrary key–value pairs to attach to the service as annotations in the
219    /// orchestrator backend.
220    ///
221    /// The orchestrator backend may apply a prefix to the key if appropriate.
222    pub annotations: BTreeMap<String, String>,
223    /// The availability zones the service can be run in. If no availability
224    /// zones are specified, the orchestrator is free to choose one.
225    pub availability_zones: Option<Vec<String>>,
226    /// A set of label selectors selecting all _other_ services that are replicas of this one.
227    ///
228    /// This may be used to implement anti-affinity. If _all_ such selectors
229    /// match for a given service, this service should not be co-scheduled on
230    /// a machine with that service.
231    ///
232    /// The orchestrator backend may or may not actually implement anti-affinity functionality.
233    pub other_replicas_selector: Vec<LabelSelector>,
234    /// A set of label selectors selecting all services that are replicas of this one,
235    /// including itself.
236    ///
237    /// This may be used to implement placement spread.
238    ///
239    /// The orchestrator backend may or may not actually implement placement spread functionality.
240    pub replicas_selector: Vec<LabelSelector>,
241
242    /// The maximum amount of scratch disk space that the service is allowed to consume.
243    pub disk_limit: Option<DiskLimit>,
244    /// Node selector for this service.
245    pub node_selector: BTreeMap<String, String>,
246}
247
248/// A named port associated with a service.
249#[derive(Debug, Clone, PartialEq, Eq)]
250pub struct ServicePort {
251    /// A descriptive name for the port.
252    ///
253    /// Note that not all orchestrator backends make use of port names.
254    pub name: String,
255    /// The desired port number.
256    ///
257    /// Not all orchestrator backends will make use of the hint.
258    pub port_hint: u16,
259}
260
261/// Assignments that the orchestrator has made for a process in a service.
262#[derive(Clone, Debug)]
263pub struct ServiceAssignments<'a> {
264    /// For each specified [`ServicePort`] name, a listen address.
265    pub listen_addrs: &'a BTreeMap<String, String>,
266    /// The listen addresses of each peer in the service.
267    ///
268    /// The order of peers is significant. Each peer is uniquely identified by its position in the
269    /// list.
270    pub peer_addrs: &'a [BTreeMap<String, String>],
271}
272
273impl ServiceAssignments<'_> {
274    /// Return the peer addresses for the specified [`ServicePort`] name.
275    pub fn peer_addresses(&self, name: &str) -> Vec<String> {
276        self.peer_addrs.iter().map(|a| a[name].clone()).collect()
277    }
278}
279
280/// Describes a limit on memory.
281#[derive(Copy, Clone, Debug, PartialOrd, Eq, Ord, PartialEq)]
282pub struct MemoryLimit(pub ByteSize);
283
284impl MemoryLimit {
285    pub const MAX: Self = Self(ByteSize(u64::MAX));
286}
287
288impl<'de> Deserialize<'de> for MemoryLimit {
289    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
290    where
291        D: Deserializer<'de>,
292    {
293        <String as Deserialize>::deserialize(deserializer)
294            .and_then(|s| {
295                ByteSize::from_str(&s).map_err(|_e| {
296                    use serde::de::Error;
297                    D::Error::invalid_value(serde::de::Unexpected::Str(&s), &"valid size in bytes")
298                })
299            })
300            .map(MemoryLimit)
301    }
302}
303
304impl Serialize for MemoryLimit {
305    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
306    where
307        S: serde::Serializer,
308    {
309        <String as Serialize>::serialize(&self.0.to_string(), serializer)
310    }
311}
312
313/// Describes a limit on CPU resources.
314#[derive(Debug, Copy, Clone, Eq, Ord, PartialEq, PartialOrd)]
315pub struct CpuLimit {
316    millicpus: usize,
317}
318
319impl CpuLimit {
320    pub const MAX: Self = Self::from_millicpus(usize::MAX / 1_000_000);
321
322    /// Constructs a new CPU limit from a number of millicpus.
323    pub const fn from_millicpus(millicpus: usize) -> CpuLimit {
324        CpuLimit { millicpus }
325    }
326
327    /// Returns the CPU limit in millicpus.
328    pub fn as_millicpus(&self) -> usize {
329        self.millicpus
330    }
331
332    /// Returns the CPU limit in nanocpus.
333    pub fn as_nanocpus(&self) -> u64 {
334        // The largest possible value of a u64 is
335        // 18_446_744_073_709_551_615,
336        // so we won't overflow this
337        // unless we have an instance with
338        // ~18.45 billion cores.
339        //
340        // Such an instance seems unrealistic,
341        // at least until we raise another few rounds
342        // of funding ...
343
344        u64::cast_from(self.millicpus)
345            .checked_mul(1_000_000)
346            .expect("Nano-CPUs must be representable")
347    }
348}
349
350impl<'de> Deserialize<'de> for CpuLimit {
351    // TODO(benesch): remove this once this function no longer makes use of
352    // potentially dangerous `as` conversions.
353    #[allow(clippy::as_conversions)]
354    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
355    where
356        D: serde::Deserializer<'de>,
357    {
358        // Note -- we just round off any precision beyond 0.001 here.
359        let float = f64::deserialize(deserializer)?;
360        let millicpus = (float * 1000.).round();
361        if millicpus < 0. || millicpus > (std::usize::MAX as f64) {
362            use serde::de::Error;
363            Err(D::Error::invalid_value(
364                Unexpected::Float(float),
365                &"a float representing a plausible number of CPUs",
366            ))
367        } else {
368            Ok(Self {
369                millicpus: millicpus as usize,
370            })
371        }
372    }
373}
374
375impl Serialize for CpuLimit {
376    // TODO(benesch): remove this once this function no longer makes use of
377    // potentially dangerous `as` conversions.
378    #[allow(clippy::as_conversions)]
379    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
380    where
381        S: serde::Serializer,
382    {
383        <f64 as Serialize>::serialize(&(self.millicpus as f64 / 1000.0), serializer)
384    }
385}
386
387/// Describes a limit on disk usage.
388#[derive(Copy, Clone, Debug, PartialOrd, Eq, Ord, PartialEq)]
389pub struct DiskLimit(pub ByteSize);
390
391impl DiskLimit {
392    pub const ZERO: Self = Self(ByteSize(0));
393    pub const MAX: Self = Self(ByteSize(u64::MAX));
394    pub const ARBITRARY: Self = Self(ByteSize::gib(1));
395}
396
397impl<'de> Deserialize<'de> for DiskLimit {
398    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
399    where
400        D: Deserializer<'de>,
401    {
402        <String as Deserialize>::deserialize(deserializer)
403            .and_then(|s| {
404                ByteSize::from_str(&s).map_err(|_e| {
405                    use serde::de::Error;
406                    D::Error::invalid_value(serde::de::Unexpected::Str(&s), &"valid size in bytes")
407                })
408            })
409            .map(DiskLimit)
410    }
411}
412
413impl Serialize for DiskLimit {
414    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
415    where
416        S: serde::Serializer,
417    {
418        <String as Serialize>::serialize(&self.0.to_string(), serializer)
419    }
420}
421
422/// Configuration for how services are scheduled. These may be ignored by orchestrator
423/// implementations.
424pub mod scheduling_config {
425    #[derive(Debug, Clone)]
426    pub struct ServiceTopologySpreadConfig {
427        /// If `true`, enable spread for replicated services.
428        ///
429        /// Defaults to `true`.
430        pub enabled: bool,
431        /// If `true`, ignore services with `scale` > 1 when expressing
432        /// spread constraints.
433        ///
434        /// Default to `true`.
435        pub ignore_non_singular_scale: bool,
436        /// The `maxSkew` for spread constraints.
437        /// See
438        /// <https://kubernetes.io/docs/concepts/scheduling-eviction/topology-spread-constraints/>
439        /// for more details.
440        ///
441        /// Defaults to `1`.
442        pub max_skew: i32,
443        /// The `minDomains` for spread constraints.
444        /// See
445        /// <https://kubernetes.io/docs/concepts/scheduling-eviction/topology-spread-constraints/>
446        /// for more details.
447        ///
448        /// Defaults to None.
449        pub min_domains: Option<i32>,
450        /// If `true`, make the spread constraints into a preference.
451        ///
452        /// Defaults to `false`.
453        pub soft: bool,
454    }
455
456    #[derive(Debug, Clone)]
457    pub struct ServiceSchedulingConfig {
458        /// If `Some`, add a affinity preference with the given
459        /// weight for services that horizontally scale.
460        ///
461        /// Defaults to `Some(100)`.
462        pub multi_pod_az_affinity_weight: Option<i32>,
463        /// If `true`, make the node-scope anti-affinity between
464        /// replicated services a preference over a constraint.
465        ///
466        /// Defaults to `false`.
467        pub soften_replication_anti_affinity: bool,
468        /// The weight for `soften_replication_anti_affinity.
469        ///
470        /// Defaults to `100`.
471        pub soften_replication_anti_affinity_weight: i32,
472        /// Configuration for `TopologySpreadConstraint`'s
473        pub topology_spread: ServiceTopologySpreadConfig,
474        /// If `true`, make the az-scope node affinity soft.
475        ///
476        /// Defaults to `false`.
477        pub soften_az_affinity: bool,
478        /// The weight for `soften_replication_anti_affinity.
479        ///
480        /// Defaults to `100`.
481        pub soften_az_affinity_weight: i32,
482        // Whether to enable security context for the service.
483        pub security_context_enabled: bool,
484    }
485
486    pub const DEFAULT_POD_AZ_AFFINITY_WEIGHT: Option<i32> = Some(100);
487    pub const DEFAULT_SOFTEN_REPLICATION_ANTI_AFFINITY: bool = false;
488    pub const DEFAULT_SOFTEN_REPLICATION_ANTI_AFFINITY_WEIGHT: i32 = 100;
489
490    pub const DEFAULT_TOPOLOGY_SPREAD_ENABLED: bool = true;
491    pub const DEFAULT_TOPOLOGY_SPREAD_IGNORE_NON_SINGULAR_SCALE: bool = true;
492    pub const DEFAULT_TOPOLOGY_SPREAD_MAX_SKEW: i32 = 1;
493    pub const DEFAULT_TOPOLOGY_SPREAD_MIN_DOMAIN: Option<i32> = None;
494    pub const DEFAULT_TOPOLOGY_SPREAD_SOFT: bool = false;
495
496    pub const DEFAULT_SOFTEN_AZ_AFFINITY: bool = false;
497    pub const DEFAULT_SOFTEN_AZ_AFFINITY_WEIGHT: i32 = 100;
498    pub const DEFAULT_SECURITY_CONTEXT_ENABLED: bool = true;
499
500    impl Default for ServiceSchedulingConfig {
501        fn default() -> Self {
502            ServiceSchedulingConfig {
503                multi_pod_az_affinity_weight: DEFAULT_POD_AZ_AFFINITY_WEIGHT,
504                soften_replication_anti_affinity: DEFAULT_SOFTEN_REPLICATION_ANTI_AFFINITY,
505                soften_replication_anti_affinity_weight:
506                    DEFAULT_SOFTEN_REPLICATION_ANTI_AFFINITY_WEIGHT,
507                topology_spread: ServiceTopologySpreadConfig {
508                    enabled: DEFAULT_TOPOLOGY_SPREAD_ENABLED,
509                    ignore_non_singular_scale: DEFAULT_TOPOLOGY_SPREAD_IGNORE_NON_SINGULAR_SCALE,
510                    max_skew: DEFAULT_TOPOLOGY_SPREAD_MAX_SKEW,
511                    min_domains: DEFAULT_TOPOLOGY_SPREAD_MIN_DOMAIN,
512                    soft: DEFAULT_TOPOLOGY_SPREAD_SOFT,
513                },
514                soften_az_affinity: DEFAULT_SOFTEN_AZ_AFFINITY,
515                soften_az_affinity_weight: DEFAULT_SOFTEN_AZ_AFFINITY_WEIGHT,
516                security_context_enabled: DEFAULT_SECURITY_CONTEXT_ENABLED,
517            }
518        }
519    }
520}