mz_persist_client/
async_runtime.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
10//! Async runtime extensions.
11
12use std::future::Future;
13use std::sync::atomic::{AtomicUsize, Ordering};
14
15use mz_ore::metrics::{MetricsRegistry, register_runtime_metrics};
16use mz_ore::task::{JoinHandle, RuntimeExt};
17use tokio::runtime::{Builder, Runtime};
18
19/// A reasonable number of threads to use in tests: enough to reproduce nontrivial
20/// orderings if necessary and avoid blocking, but not too many.
21// This was done as a workaround for https://sourceware.org/bugzilla/show_bug.cgi?id=19951
22// in tests, but seems useful in general.
23pub const TEST_THREADS: usize = 4;
24
25/// An isolated runtime for asynchronous tasks, particularly work
26/// that may be CPU intensive such as encoding/decoding and shard
27/// maintenance.
28///
29/// Using a separate runtime allows Persist to isolate its expensive
30/// workloads on its own OS threads as an insurance policy against
31/// tasks that erroneously fail to yield for a long time. By using
32/// separate OS threads, the scheduler is able to context switch
33/// out of any problematic tasks, preserving liveness for the rest
34/// of the process.
35///
36/// Note: Even though the work done by this runtime might be "blocking" or
37/// CPU bound we should not use the [`tokio::task::spawn_blocking`] API.
38/// There can be issues during shutdown if tasks are currently running on the
39/// blocking thread pool [1], and the blocking thread pool generally creates
40/// many more threads than are physically available. This can pin CPU usage
41/// to 100% starving other important threads like the Coordinator.
42///
43/// [1]: <https://github.com/MaterializeInc/materialize/pull/13955>
44#[derive(Debug)]
45pub struct IsolatedRuntime {
46    inner: Option<Runtime>,
47}
48
49impl IsolatedRuntime {
50    /// Creates a new isolated runtime.
51    pub fn new(metrics: &MetricsRegistry, worker_threads: Option<usize>) -> IsolatedRuntime {
52        let mut runtime = Builder::new_multi_thread();
53        if let Some(worker_threads) = worker_threads {
54            runtime.worker_threads(worker_threads);
55        }
56        let runtime = runtime
57            .thread_name_fn(|| {
58                static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
59                let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
60                // This will wrap around eventually, which is not ideal, but it's important that
61                // it stays small to fit within OS limits.
62                format!("persist:{:04x}", id % 0x10000)
63            })
64            .enable_all()
65            .build()
66            .expect("known to be valid");
67        register_runtime_metrics("persist", runtime.metrics(), metrics);
68        IsolatedRuntime {
69            inner: Some(runtime),
70        }
71    }
72
73    /// Create an isolated runtime with appropriate values for tests.
74    pub fn new_for_tests() -> Self {
75        IsolatedRuntime::new(&MetricsRegistry::new(), Some(TEST_THREADS))
76    }
77
78    /// Spawns a task onto this runtime.
79    ///
80    /// Note: We purposefully do not use the [`tokio::task::spawn_blocking`] API here, see the doc
81    /// comment on [`IsolatedRuntime`] for explanation.
82    pub fn spawn_named<N, S, F>(&self, name: N, fut: F) -> JoinHandle<F::Output>
83    where
84        S: AsRef<str>,
85        N: FnOnce() -> S,
86        F: Future + Send + 'static,
87        F::Output: Send + 'static,
88    {
89        self.inner
90            .as_ref()
91            .expect("exists until drop")
92            .spawn_named(name, fut)
93    }
94}
95
96impl Drop for IsolatedRuntime {
97    fn drop(&mut self) {
98        // We don't need to worry about `shutdown_background` leaking
99        // blocking tasks (i.e., tasks spawned with `spawn_blocking`) because
100        // the `IsolatedRuntime` wrapper prevents access to `spawn_blocking`.
101        self.inner
102            .take()
103            .expect("cannot drop twice")
104            .shutdown_background()
105    }
106}
mz_persist_client/async_runtime.rs

mz_persist_client/
async_runtime.rs
