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}