mz_storage/source/
types.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//! Types related to the source ingestion pipeline/framework.
11
12// https://github.com/tokio-rs/prost/issues/237
13// #![allow(missing_docs)]
14
15use std::collections::BTreeMap;
16use std::fmt::{self, Debug, Display, Formatter};
17use std::future::Future;
18use std::pin::Pin;
19use std::sync::Arc;
20use std::task::{Context, Poll, ready};
21
22use differential_dataflow::Collection;
23use mz_repr::{Diff, GlobalId, Row};
24use mz_storage_types::errors::{DataflowError, DecodeError};
25use mz_storage_types::sources::SourceTimestamp;
26use mz_timely_util::antichain::AntichainExt;
27use mz_timely_util::builder_async::{AsyncOutputHandle, MAX_OUTSTANDING_BYTES, PressOnDropButton};
28use pin_project::pin_project;
29use serde::{Deserialize, Serialize};
30use timely::ContainerBuilder;
31use timely::container::PushInto;
32use timely::dataflow::operators::Capability;
33use timely::dataflow::{Scope, StreamVec};
34use timely::progress::{Antichain, Timestamp};
35use tokio::sync::Semaphore;
36use tokio_util::sync::PollSemaphore;
37
38use crate::healthcheck::{HealthStatusMessage, StatusNamespace};
39use crate::source::RawSourceCreationConfig;
40
41/// An update produced by implementors of `SourceRender` that presents an _aggregated_
42/// description of the number of _offset_committed_ and _offset_known_ for the given
43/// source.
44///
45/// The aggregate is required to be a 64 bit unsigned integer, whose units are
46/// implementation-defined.
47#[derive(Clone, Debug)]
48pub enum ProgressStatisticsUpdate {
49    SteadyState {
50        offset_known: u64,
51        offset_committed: u64,
52    },
53    Snapshot {
54        records_known: u64,
55        records_staged: u64,
56    },
57}
58
59pub type StackedCollection<'scope, T, D> = Collection<'scope, T, Vec<(D, T, Diff)>>;
60
61/// The durably committed progress of an ingestion, in the source's timestamp type.
62///
63/// Each value is a complete snapshot, so a consumer that skips values only misses intermediate
64/// progress.
65#[derive(Clone, Debug, PartialEq)]
66pub struct ResumeUppers<T> {
67    /// The frontier every export has committed through. `None` while some export has not yet
68    /// committed beyond the ingestion's as_of, for example while it snapshots.
69    pub source: Option<Antichain<T>>,
70    /// The frontier each export has committed through. An export is absent until it commits
71    /// beyond the ingestion's as_of.
72    // TODO(maz): Report `offset_committed` and `offset_known` from the source pipeline, derived from
73    // these frontiers and the probes, so sources only supply their seeds and we can maintain
74    // SourceStatistics outside source operators.
75    pub exports: BTreeMap<GlobalId, Antichain<T>>,
76}
77
78/// Formats as `source={..} exports={id: {..}, ..}`, with `source=pending` while
79/// [`ResumeUppers::source`] is `None`.
80impl<T: Display + 'static> Display for ResumeUppers<T> {
81    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
82        match &self.source {
83            Some(source) => write!(f, "source={}", source.pretty())?,
84            None => f.write_str("source=pending")?,
85        }
86        f.write_str(" exports={")?;
87        for (i, (id, upper)) in self.exports.iter().enumerate() {
88            if i > 0 {
89                f.write_str(", ")?;
90            }
91            write!(f, "{id}: {}", upper.pretty())?;
92        }
93        f.write_str("}")
94    }
95}
96
97/// Describes a source that can render itself in a timely scope.
98pub trait SourceRender {
99    type Time: SourceTimestamp;
100    const STATUS_NAMESPACE: StatusNamespace;
101
102    /// Renders the source in the provided timely scope.
103    ///
104    /// The `resume_uppers` stream can be used by the source to observe the progress of the
105    /// ingestion. When a frontier appears as [`ResumeUppers::source`] the source implementation
106    /// can be certain that future ingestion instances will request to read the external data only
107    /// at times beyond that frontier. Therefore, the source implementation can react to this
108    /// stream by e.g committing offsets upstream or advancing the LSN of a replication slot. The
109    /// [`ResumeUppers::exports`] are what the source reports as each export's `offset_committed`
110    /// statistic.
111    ///
112    /// Rendering a source is expected to return four things.
113    ///
114    /// First, a source must produce a collection that is produced by the rendered dataflow and
115    /// must contain *definite*[^1] data for all times beyond the resumption frontier.
116    ///
117    /// Second, a source must produce a stream of health status updates.
118    ///
119    /// Third, a source must produce a probe stream that periodically reports the upstream
120    /// frontier. This is used to drive reclocking and mint new bindings.
121    ///
122    /// Finally, the source is expected to return an opaque token that when dropped will cause the
123    /// source to immediately drop all capabilities and advance its frontier to the empty antichain.
124    ///
125    /// [^1]: <https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/20210831_correctness.md#describing-definite-data>
126    fn render<'scope>(
127        self,
128        scope: Scope<'scope, Self::Time>,
129        config: &RawSourceCreationConfig,
130        resume_uppers: impl futures::Stream<Item = ResumeUppers<Self::Time>> + 'static,
131        start_signal: impl std::future::Future<Output = ()> + 'static,
132    ) -> (
133        BTreeMap<
134            GlobalId,
135            StackedCollection<'scope, Self::Time, Result<SourceMessage, DataflowError>>,
136        >,
137        StreamVec<'scope, Self::Time, HealthStatusMessage>,
138        StreamVec<'scope, Self::Time, Probe<Self::Time>>,
139        Vec<PressOnDropButton>,
140    );
141}
142
143/// Source-agnostic wrapper for messages. Each source must implement a
144/// conversion to Message.
145#[derive(Debug, Clone)]
146pub struct SourceMessage {
147    /// The message key
148    pub key: Row,
149    /// The message value
150    pub value: Row,
151    /// Additional metadata columns requested by the user
152    pub metadata: Row,
153}
154
155impl SourceMessage {
156    /// Heap-size estimate used to drive fuel accounting in source operators.
157    pub fn byte_len(&self) -> usize {
158        self.key.byte_len() + self.value.byte_len() + self.metadata.byte_len()
159    }
160}
161
162/// Heap-size estimate used by source operators to drive `give_fueled` yielding.
163///
164/// Stack-only values may report `size_of_val(self)`; values with heap-allocated
165/// payloads (rows, byte buffers) should include that payload.
166pub trait FuelSize {
167    fn fuel_size(&self) -> usize;
168}
169
170impl FuelSize for SourceMessage {
171    fn fuel_size(&self) -> usize {
172        self.byte_len()
173    }
174}
175
176impl FuelSize for Row {
177    fn fuel_size(&self) -> usize {
178        self.byte_len()
179    }
180}
181
182impl FuelSize for Vec<u8> {
183    fn fuel_size(&self) -> usize {
184        self.len()
185    }
186}
187
188impl FuelSize for bytes::Bytes {
189    fn fuel_size(&self) -> usize {
190        self.len()
191    }
192}
193
194impl<T: FuelSize, E: FuelSize> FuelSize for Result<T, E> {
195    fn fuel_size(&self) -> usize {
196        match self {
197            Ok(t) => t.fuel_size(),
198            Err(e) => e.fuel_size(),
199        }
200    }
201}
202
203/// Tuples sum the fuel size of their elements. Trivial coordinate fields
204/// (output indices, timestamps, diffs) implement `FuelSize` to return their
205/// stack size, so an `update.fuel_size()` call charges only the heap-allocated
206/// payload like rows or byte buffers.
207impl<A: FuelSize, B: FuelSize> FuelSize for (A, B) {
208    fn fuel_size(&self) -> usize {
209        self.0.fuel_size() + self.1.fuel_size()
210    }
211}
212
213impl<A: FuelSize, B: FuelSize, C: FuelSize> FuelSize for (A, B, C) {
214    fn fuel_size(&self) -> usize {
215        self.0.fuel_size() + self.1.fuel_size() + self.2.fuel_size()
216    }
217}
218
219/// Convenience macro for declaring `FuelSize` on a stack-only type.
220macro_rules! impl_fuel_size_stack {
221    ($($t:ty),* $(,)?) => {
222        $(
223            impl FuelSize for $t {
224                fn fuel_size(&self) -> usize {
225                    std::mem::size_of_val(self)
226                }
227            }
228        )*
229    };
230}
231
232impl_fuel_size_stack!(
233    usize,
234    u32,
235    u64,
236    Diff,
237    mz_repr::Timestamp,
238    mz_storage_types::sources::MzOffset,
239    mz_sql_server_util::cdc::Lsn,
240    DataflowError,
241);
242
243impl<P, T> FuelSize for mz_timely_util::order::Partitioned<P, T> {
244    fn fuel_size(&self) -> usize {
245        std::mem::size_of_val(self)
246    }
247}
248
249/// Output handles of one operator that share a fuel budget, so that the operator yields back to
250/// timely once per [`MAX_OUTSTANDING_BYTES`] of its total emission across them.
251///
252/// [`AsyncOutputHandle::give_fueled`] counts bytes per handle, so an operator with many handles
253/// could emit up to the bound on each of them before yielding.
254pub struct SharedFuel<T: Timestamp, CB: ContainerBuilder> {
255    handles: Vec<AsyncOutputHandle<T, CB>>,
256    outstanding_bytes: usize,
257    yielded: bool,
258}
259
260impl<T: Timestamp, CB: ContainerBuilder> SharedFuel<T, CB> {
261    /// Takes ownership of `handles`, which [`Self::give`] addresses by position.
262    pub fn new(handles: Vec<AsyncOutputHandle<T, CB>>) -> Self {
263        Self {
264            handles,
265            outstanding_bytes: 0,
266            yielded: false,
267        }
268    }
269
270    /// Returns whether [`Self::give`] yielded since the last call.
271    pub fn take_yielded(&mut self) -> bool {
272        std::mem::take(&mut self.yielded)
273    }
274
275    /// Gives `data` at `cap` to the handle at `index`, charging its [`FuelSize`] against the
276    /// shared budget. Yields back to timely and resets the budget once it exceeds
277    /// [`MAX_OUTSTANDING_BYTES`].
278    ///
279    /// # Panics
280    ///
281    /// Panics if `index` is not the position of a handle passed to [`Self::new`].
282    pub async fn give<D>(&mut self, index: usize, cap: &Capability<T>, data: D)
283    where
284        CB: PushInto<D>,
285        D: FuelSize,
286    {
287        let size = data.fuel_size();
288        self.handles[index].give(cap, data);
289        if self.charge(size) {
290            tokio::task::yield_now().await;
291        }
292    }
293
294    /// Charges `size`, returning whether the budget was exceeded, in which case it resets.
295    fn charge(&mut self, size: usize) -> bool {
296        self.outstanding_bytes += size;
297        let exceeded = self.outstanding_bytes > MAX_OUTSTANDING_BYTES;
298        if exceeded {
299            self.outstanding_bytes = 0;
300            self.yielded = true;
301        }
302        exceeded
303    }
304}
305
306/// The result of probing an upstream system for its write frontier.
307#[derive(Debug, Clone, Serialize, Deserialize)]
308pub struct Probe<T> {
309    /// The timestamp at which this probe was initiated.
310    pub probe_ts: mz_repr::Timestamp,
311    /// The frontier obtain from the upstream system.
312    pub upstream_frontier: Antichain<T>,
313}
314
315/// A record produced by a source
316#[derive(
317    Debug,
318    PartialEq,
319    Eq,
320    Hash,
321    PartialOrd,
322    Ord,
323    Clone,
324    Serialize,
325    Deserialize
326)]
327pub struct SourceOutput<FromTime> {
328    /// The record's key (or some empty/default value for sources without the concept of key)
329    pub key: Row,
330    /// The record's value
331    pub value: Row,
332    /// Additional metadata columns requested by the user
333    pub metadata: Row,
334    /// The original timestamp of this message
335    pub from_time: FromTime,
336}
337
338/// The output of the decoding operator
339#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
340pub struct DecodeResult<FromTime> {
341    /// The decoded key
342    pub key: Option<Result<Row, DecodeError>>,
343    /// The decoded value, as well as the the
344    /// differential `diff` value for this value, if the value
345    /// is present and not and error.
346    pub value: Option<Result<Row, DecodeError>>,
347    /// Additional metadata requested by the user
348    pub metadata: Row,
349    /// The original timestamp of this message
350    pub from_time: FromTime,
351}
352
353#[pin_project]
354pub struct SignaledFuture<F> {
355    #[pin]
356    fut: F,
357    semaphore: PollSemaphore,
358}
359
360impl<F: Future> SignaledFuture<F> {
361    pub fn new(semaphore: Arc<Semaphore>, fut: F) -> Self {
362        Self {
363            fut,
364            semaphore: PollSemaphore::new(semaphore),
365        }
366    }
367}
368
369impl<F: Future> Future for SignaledFuture<F> {
370    type Output = F::Output;
371
372    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
373        let this = self.project();
374        let permit = ready!(this.semaphore.poll_acquire(cx));
375        let ret = this.fut.poll(cx);
376        drop(permit);
377        ret
378    }
379}
380
381#[cfg(test)]
382mod tests {
383    use mz_storage_types::sources::MzOffset;
384
385    use super::*;
386
387    #[mz_ore::test]
388    fn resume_uppers_display() {
389        let offset = |o| Antichain::from_elem(MzOffset::from(o));
390        let uppers = ResumeUppers {
391            source: Some(offset(10)),
392            exports: BTreeMap::from([
393                (GlobalId::User(3), offset(30)),
394                (GlobalId::User(4), offset(10)),
395            ]),
396        };
397        assert_eq!(
398            uppers.to_string(),
399            "source={10} exports={u3: {30}, u4: {10}}"
400        );
401
402        let pending = ResumeUppers {
403            source: None,
404            exports: BTreeMap::from([(GlobalId::User(3), Antichain::<MzOffset>::new())]),
405        };
406        assert_eq!(pending.to_string(), "source=pending exports={u3: {}}");
407    }
408}
mz_storage/source/types.rs

mz_storage/source/
types.rs
