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}
258
259impl<T: Timestamp, CB: ContainerBuilder> SharedFuel<T, CB> {
260    /// Takes ownership of `handles`, which [`Self::give`] addresses by position.
261    pub fn new(handles: Vec<AsyncOutputHandle<T, CB>>) -> Self {
262        Self {
263            handles,
264            outstanding_bytes: 0,
265        }
266    }
267
268    /// Gives `data` at `cap` to the handle at `index`, charging its [`FuelSize`] against the
269    /// shared budget. Yields back to timely and resets the budget once it exceeds
270    /// [`MAX_OUTSTANDING_BYTES`].
271    ///
272    /// # Panics
273    ///
274    /// Panics if `index` is not the position of a handle passed to [`Self::new`].
275    pub async fn give<D>(&mut self, index: usize, cap: &Capability<T>, data: D)
276    where
277        CB: PushInto<D>,
278        D: FuelSize,
279    {
280        let size = data.fuel_size();
281        self.handles[index].give(cap, data);
282        if self.charge(size) {
283            tokio::task::yield_now().await;
284        }
285    }
286
287    /// Charges `size`, returning whether the budget was exceeded, in which case it resets.
288    fn charge(&mut self, size: usize) -> bool {
289        self.outstanding_bytes += size;
290        let exceeded = self.outstanding_bytes > MAX_OUTSTANDING_BYTES;
291        if exceeded {
292            self.outstanding_bytes = 0;
293        }
294        exceeded
295    }
296}
297
298/// The result of probing an upstream system for its write frontier.
299#[derive(Debug, Clone, Serialize, Deserialize)]
300pub struct Probe<T> {
301    /// The timestamp at which this probe was initiated.
302    pub probe_ts: mz_repr::Timestamp,
303    /// The frontier obtain from the upstream system.
304    pub upstream_frontier: Antichain<T>,
305}
306
307/// A record produced by a source
308#[derive(
309    Debug,
310    PartialEq,
311    Eq,
312    Hash,
313    PartialOrd,
314    Ord,
315    Clone,
316    Serialize,
317    Deserialize
318)]
319pub struct SourceOutput<FromTime> {
320    /// The record's key (or some empty/default value for sources without the concept of key)
321    pub key: Row,
322    /// The record's value
323    pub value: Row,
324    /// Additional metadata columns requested by the user
325    pub metadata: Row,
326    /// The original timestamp of this message
327    pub from_time: FromTime,
328}
329
330/// The output of the decoding operator
331#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
332pub struct DecodeResult<FromTime> {
333    /// The decoded key
334    pub key: Option<Result<Row, DecodeError>>,
335    /// The decoded value, as well as the the
336    /// differential `diff` value for this value, if the value
337    /// is present and not and error.
338    pub value: Option<Result<Row, DecodeError>>,
339    /// Additional metadata requested by the user
340    pub metadata: Row,
341    /// The original timestamp of this message
342    pub from_time: FromTime,
343}
344
345#[pin_project]
346pub struct SignaledFuture<F> {
347    #[pin]
348    fut: F,
349    semaphore: PollSemaphore,
350}
351
352impl<F: Future> SignaledFuture<F> {
353    pub fn new(semaphore: Arc<Semaphore>, fut: F) -> Self {
354        Self {
355            fut,
356            semaphore: PollSemaphore::new(semaphore),
357        }
358    }
359}
360
361impl<F: Future> Future for SignaledFuture<F> {
362    type Output = F::Output;
363
364    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
365        let this = self.project();
366        let permit = ready!(this.semaphore.poll_acquire(cx));
367        let ret = this.fut.poll(cx);
368        drop(permit);
369        ret
370    }
371}
372
373#[cfg(test)]
374mod tests {
375    use mz_storage_types::sources::MzOffset;
376
377    use super::*;
378
379    #[mz_ore::test]
380    fn resume_uppers_display() {
381        let offset = |o| Antichain::from_elem(MzOffset::from(o));
382        let uppers = ResumeUppers {
383            source: Some(offset(10)),
384            exports: BTreeMap::from([
385                (GlobalId::User(3), offset(30)),
386                (GlobalId::User(4), offset(10)),
387            ]),
388        };
389        assert_eq!(
390            uppers.to_string(),
391            "source={10} exports={u3: {30}, u4: {10}}"
392        );
393
394        let pending = ResumeUppers {
395            source: None,
396            exports: BTreeMap::from([(GlobalId::User(3), Antichain::<MzOffset>::new())]),
397        };
398        assert_eq!(pending.to_string(), "source=pending exports={u3: {}}");
399    }
400}
mz_storage/source/types.rs

mz_storage/source/
types.rs
