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::Debug;
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::builder_async::PressOnDropButton;
27use mz_timely_util::columnation::ColumnationStack;
28use pin_project::pin_project;
29use serde::{Deserialize, Serialize};
30use timely::dataflow::{Scope, StreamVec};
31use timely::progress::Antichain;
32use tokio::sync::Semaphore;
33use tokio_util::sync::PollSemaphore;
34
35use crate::healthcheck::{HealthStatusMessage, StatusNamespace};
36use crate::source::RawSourceCreationConfig;
37
38/// An update produced by implementors of `SourceRender` that presents an _aggregated_
39/// description of the number of _offset_committed_ and _offset_known_ for the given
40/// source.
41///
42/// The aggregate is required to be a 64 bit unsigned integer, whose units are
43/// implementation-defined.
44#[derive(Clone, Debug)]
45pub enum ProgressStatisticsUpdate {
46    SteadyState {
47        offset_known: u64,
48        offset_committed: u64,
49    },
50    Snapshot {
51        records_known: u64,
52        records_staged: u64,
53    },
54}
55
56pub type StackedCollection<'scope, T, D> = Collection<'scope, T, ColumnationStack<(D, T, Diff)>>;
57
58/// Describes a source that can render itself in a timely scope.
59pub trait SourceRender {
60    type Time: SourceTimestamp;
61    const STATUS_NAMESPACE: StatusNamespace;
62
63    /// Renders the source in the provided timely scope.
64    ///
65    /// The `resume_uppers` stream can be used by the source to observe the overall progress of the
66    /// ingestion. When a frontier appears in this stream the source implementation can be certain
67    /// that future ingestion instances will request to read the external data only at times beyond
68    /// that frontier. Therefore, the source implementation can react to this stream by e.g
69    /// committing offsets upstream or advancing the LSN of a replication slot. It is safe to
70    /// ignore this argument.
71    ///
72    /// Rendering a source is expected to return four things.
73    ///
74    /// First, a source must produce a collection that is produced by the rendered dataflow and
75    /// must contain *definite*[^1] data for all times beyond the resumption frontier.
76    ///
77    /// Second, a source must produce a stream of health status updates.
78    ///
79    /// Third, a source must produce a probe stream that periodically reports the upstream
80    /// frontier. This is used to drive reclocking and mint new bindings.
81    ///
82    /// Finally, the source is expected to return an opaque token that when dropped will cause the
83    /// source to immediately drop all capabilities and advance its frontier to the empty antichain.
84    ///
85    /// [^1]: <https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/20210831_correctness.md#describing-definite-data>
86    fn render<'scope>(
87        self,
88        scope: Scope<'scope, Self::Time>,
89        config: &RawSourceCreationConfig,
90        resume_uppers: impl futures::Stream<Item = Antichain<Self::Time>> + 'static,
91        start_signal: impl std::future::Future<Output = ()> + 'static,
92    ) -> (
93        BTreeMap<
94            GlobalId,
95            StackedCollection<'scope, Self::Time, Result<SourceMessage, DataflowError>>,
96        >,
97        StreamVec<'scope, Self::Time, HealthStatusMessage>,
98        StreamVec<'scope, Self::Time, Probe<Self::Time>>,
99        Vec<PressOnDropButton>,
100    );
101}
102
103/// Source-agnostic wrapper for messages. Each source must implement a
104/// conversion to Message.
105#[derive(Debug, Clone)]
106pub struct SourceMessage {
107    /// The message key
108    pub key: Row,
109    /// The message value
110    pub value: Row,
111    /// Additional metadata columns requested by the user
112    pub metadata: Row,
113}
114
115/// The result of probing an upstream system for its write frontier.
116#[derive(Debug, Clone, Serialize, Deserialize)]
117pub struct Probe<T> {
118    /// The timestamp at which this probe was initiated.
119    pub probe_ts: mz_repr::Timestamp,
120    /// The frontier obtain from the upstream system.
121    pub upstream_frontier: Antichain<T>,
122}
123
124mod columnation {
125    use columnation::{Columnation, Region};
126    use mz_repr::Row;
127
128    use super::SourceMessage;
129
130    impl Columnation for SourceMessage {
131        type InnerRegion = SourceMessageRegion;
132    }
133
134    #[derive(Default)]
135    pub struct SourceMessageRegion {
136        inner: <Row as Columnation>::InnerRegion,
137    }
138
139    impl Region for SourceMessageRegion {
140        type Item = SourceMessage;
141
142        unsafe fn copy(&mut self, item: &Self::Item) -> Self::Item {
143            SourceMessage {
144                key: unsafe { self.inner.copy(&item.key) },
145                value: unsafe { self.inner.copy(&item.value) },
146                metadata: unsafe { self.inner.copy(&item.metadata) },
147            }
148        }
149
150        fn clear(&mut self) {
151            self.inner.clear()
152        }
153
154        fn reserve_items<'a, I>(&mut self, items: I)
155        where
156            Self: 'a,
157            I: Iterator<Item = &'a Self::Item> + Clone,
158        {
159            self.inner.reserve_items(
160                items
161                    .map(|item| [&item.key, &item.value, &item.metadata])
162                    .flatten(),
163            )
164        }
165
166        fn reserve_regions<'a, I>(&mut self, regions: I)
167        where
168            Self: 'a,
169            I: Iterator<Item = &'a Self> + Clone,
170        {
171            self.inner.reserve_regions(regions.map(|r| &r.inner))
172        }
173
174        fn heap_size(&self, callback: impl FnMut(usize, usize)) {
175            self.inner.heap_size(callback)
176        }
177    }
178}
179
180/// A record produced by a source
181#[derive(
182    Debug,
183    PartialEq,
184    Eq,
185    Hash,
186    PartialOrd,
187    Ord,
188    Clone,
189    Serialize,
190    Deserialize
191)]
192pub struct SourceOutput<FromTime> {
193    /// The record's key (or some empty/default value for sources without the concept of key)
194    pub key: Row,
195    /// The record's value
196    pub value: Row,
197    /// Additional metadata columns requested by the user
198    pub metadata: Row,
199    /// The original timestamp of this message
200    pub from_time: FromTime,
201}
202
203/// The output of the decoding operator
204#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
205pub struct DecodeResult<FromTime> {
206    /// The decoded key
207    pub key: Option<Result<Row, DecodeError>>,
208    /// The decoded value, as well as the the
209    /// differential `diff` value for this value, if the value
210    /// is present and not and error.
211    pub value: Option<Result<Row, DecodeError>>,
212    /// Additional metadata requested by the user
213    pub metadata: Row,
214    /// The original timestamp of this message
215    pub from_time: FromTime,
216}
217
218#[pin_project]
219pub struct SignaledFuture<F> {
220    #[pin]
221    fut: F,
222    semaphore: PollSemaphore,
223}
224
225impl<F: Future> SignaledFuture<F> {
226    pub fn new(semaphore: Arc<Semaphore>, fut: F) -> Self {
227        Self {
228            fut,
229            semaphore: PollSemaphore::new(semaphore),
230        }
231    }
232}
233
234impl<F: Future> Future for SignaledFuture<F> {
235    type Output = F::Output;
236
237    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
238        let this = self.project();
239        let permit = ready!(this.semaphore.poll_acquire(cx));
240        let ret = this.fut.poll(cx);
241        drop(permit);
242        ret
243    }
244}
mz_storage/source/types.rs

mz_storage/source/
types.rs
