mz_timestamp_oracle/
batching_oracle.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//! A timestamp oracle that wraps a `TimestampOracle` and batches calls
11//! to it.
12
13use std::sync::Arc;
14
15use async_trait::async_trait;
16use mz_ore::cast::CastFrom;
17use tokio::sync::mpsc::UnboundedSender;
18use tokio::sync::oneshot;
19
20use crate::metrics::Metrics;
21use crate::{TimestampOracle, WriteTimestamp};
22
23/// A batching [`TimestampOracle`] backed by a [`TimestampOracle`]
24///
25/// This will only batch calls to `read_ts` because the rest of the system
26/// already naturally does batching of write-related calls via the group commit
27/// mechanism. Write-related calls are passed straight through to the backing
28/// oracle.
29///
30/// For `read_ts` calls, we have to be careful to never cache results from the
31/// backing oracle: for the timestamp to be linearized we can never return a
32/// result as of an earlier moment, but batching them up is correct because this
33/// can only make it so that we return later timestamps. Those later timestamps
34/// still fall within the duration of the `read_ts` call and so are linearized.
35pub struct BatchingTimestampOracle<T> {
36    inner: Arc<dyn TimestampOracle<T> + Send + Sync>,
37    command_tx: UnboundedSender<Command<T>>,
38}
39
40/// A command on the internal batching command stream.
41enum Command<T> {
42    ReadTs(oneshot::Sender<T>),
43}
44
45impl<T> std::fmt::Debug for BatchingTimestampOracle<T> {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        f.debug_struct("BatchingTimestampOracle").finish()
48    }
49}
50
51impl<T> BatchingTimestampOracle<T>
52where
53    T: Clone + Send + Sync + 'static,
54{
55    /// Crates a [`BatchingTimestampOracle`] that uses the given inner oracle.
56    pub fn new(metrics: Arc<Metrics>, oracle: Arc<dyn TimestampOracle<T> + Send + Sync>) -> Self {
57        let (command_tx, mut command_rx) = tokio::sync::mpsc::unbounded_channel();
58
59        let task_oracle = Arc::clone(&oracle);
60
61        mz_ore::task::spawn(|| "BatchingTimestampOracle Worker Task", async move {
62            let read_ts_metrics = &metrics.batching.read_ts;
63
64            // See comment on `BatchingTimestampOracle` for why this batching is
65            // correct.
66            while let Some(cmd) = command_rx.recv().await {
67                let mut pending_cmds = vec![cmd];
68                while let Ok(cmd) = command_rx.try_recv() {
69                    pending_cmds.push(cmd);
70                }
71
72                read_ts_metrics
73                    .ops_count
74                    .inc_by(u64::cast_from(pending_cmds.len()));
75                read_ts_metrics.batches_count.inc();
76
77                let ts = task_oracle.read_ts().await;
78                for cmd in pending_cmds {
79                    match cmd {
80                        Command::ReadTs(response_tx) => {
81                            // It's okay if the receiver drops, just means
82                            // they're not interested anymore.
83                            let _ = response_tx.send(ts.clone());
84                        }
85                    }
86                }
87            }
88
89            tracing::debug!("shutting down BatchingTimestampOracle task");
90        });
91
92        Self {
93            inner: oracle,
94            command_tx,
95        }
96    }
97}
98
99/// Waits forever, for a `read_ts` whose worker task is no longer there to
100/// answer it.
101///
102/// The worker task owns both ends of the internal channels, and while a caller
103/// holds a sender the only thing that ends it is the Tokio runtime dropping it
104/// during shutdown: a panic in the task would abort the process instead
105/// (`mz_ore::panic::install_enhanced_handler`). There is no timestamp that
106/// could be invented here without breaking linearizability, so the caller waits
107/// and shutdown drops its task at this await point.
108async fn worker_task_gone<T>() -> T {
109    tracing::debug!("BatchingTimestampOracle worker task is gone, parking read_ts");
110    std::future::pending().await
111}
112
113#[async_trait]
114impl<T> TimestampOracle<T> for BatchingTimestampOracle<T>
115where
116    T: Send + Sync,
117{
118    async fn write_ts(&self) -> WriteTimestamp<T> {
119        self.inner.write_ts().await
120    }
121
122    async fn peek_write_ts(&self) -> T {
123        self.inner.peek_write_ts().await
124    }
125
126    async fn read_ts(&self) -> T {
127        let (tx, rx) = oneshot::channel();
128
129        if self.command_tx.send(Command::ReadTs(tx)).is_err() {
130            return worker_task_gone().await;
131        }
132
133        match rx.await {
134            Ok(ts) => ts,
135            Err(_) => worker_task_gone().await,
136        }
137    }
138
139    async fn apply_write(&self, write_ts: T) {
140        self.inner.apply_write(write_ts).await
141    }
142}
143
144#[cfg(test)]
145mod tests {
146
147    use mz_ore::metrics::MetricsRegistry;
148    use mz_repr::Timestamp;
149    use tracing::info;
150
151    use crate::postgres_oracle::{PostgresTimestampOracle, PostgresTimestampOracleConfig};
152
153    use super::*;
154
155    /// An oracle that answers nothing, for tests that only exercise the
156    /// batching wrapper's own plumbing.
157    #[derive(Debug)]
158    struct PendingOracle;
159
160    #[async_trait]
161    impl TimestampOracle<Timestamp> for PendingOracle {
162        async fn write_ts(&self) -> WriteTimestamp<Timestamp> {
163            std::future::pending().await
164        }
165
166        async fn peek_write_ts(&self) -> Timestamp {
167            std::future::pending().await
168        }
169
170        async fn read_ts(&self) -> Timestamp {
171            std::future::pending().await
172        }
173
174        async fn apply_write(&self, _write_ts: Timestamp) {
175            std::future::pending().await
176        }
177    }
178
179    /// Runtime shutdown drops the worker task while callers still hold the
180    /// oracle, so `read_ts` must wait rather than panic on the closed channel.
181    /// Dropping the runtime the worker was spawned on, while keeping the oracle
182    /// alive on another one, reproduces that state deterministically.
183    #[mz_ore::test]
184    fn test_read_ts_waits_when_worker_task_is_gone() {
185        let metrics = Arc::new(Metrics::new(&MetricsRegistry::new()));
186
187        let worker_runtime = tokio::runtime::Runtime::new().expect("can build runtime");
188        let oracle = worker_runtime
189            .block_on(async { BatchingTimestampOracle::new(metrics, Arc::new(PendingOracle)) });
190        drop(worker_runtime);
191
192        let caller_runtime = tokio::runtime::Runtime::new().expect("can build runtime");
193        caller_runtime.block_on(async {
194            let read_ts = std::pin::pin!(oracle.read_ts());
195            assert!(futures::poll!(read_ts).is_pending());
196        });
197    }
198
199    #[mz_ore::test(tokio::test)]
200    #[cfg_attr(miri, ignore)] // error: unsupported operation: can't call foreign function `TLS_client_method` on OS `linux`
201    async fn test_batching_timestamp_oracle() -> Result<(), anyhow::Error> {
202        let config = match PostgresTimestampOracleConfig::new_for_test() {
203            Some(config) => config,
204            None => {
205                info!(
206                    "{} env not set: skipping test that uses external service",
207                    PostgresTimestampOracleConfig::EXTERNAL_TESTS_POSTGRES_URL
208                );
209                return Ok(());
210            }
211        };
212        let metrics = Arc::new(Metrics::new(&MetricsRegistry::new()));
213
214        crate::tests::timestamp_oracle_impl_test(|timeline, now_fn, initial_ts| {
215            // We use the postgres oracle as the backing oracle.
216            let pg_oracle = PostgresTimestampOracle::open(
217                config.clone(),
218                timeline,
219                initial_ts,
220                now_fn,
221                false, /* read-only */
222            );
223
224            async {
225                let arced_pg_oracle: Arc<dyn TimestampOracle<Timestamp> + Send + Sync> =
226                    Arc::new(pg_oracle.await);
227
228                let batching_oracle =
229                    BatchingTimestampOracle::new(Arc::clone(&metrics), arced_pg_oracle);
230
231                let arced_oracle: Arc<dyn TimestampOracle<Timestamp> + Send + Sync> =
232                    Arc::new(batching_oracle);
233
234                arced_oracle
235            }
236        })
237        .await?;
238
239        Ok(())
240    }
241}
mz_timestamp_oracle/batching_oracle.rs

mz_timestamp_oracle/
batching_oracle.rs
