mz_compute/arrangement/
manager.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//! Management of arrangements across dataflows.
11
12use std::any::Any;
13use std::collections::BTreeMap;
14use std::rc::Rc;
15use std::time::Instant;
16
17use differential_dataflow::lattice::antichain_join;
18use differential_dataflow::operators::arrange::{Arranged, ShutdownButton, TraceAgent};
19use differential_dataflow::trace::wrappers::frontier::TraceFrontier;
20use differential_dataflow::trace::{Span, TraceReader};
21use mz_repr::{Diff, GlobalId, Timestamp};
22use timely::PartialOrder;
23use timely::dataflow::Scope;
24use timely::dataflow::operators::CapabilitySet;
25use timely::progress::Timestamp as _;
26use timely::progress::frontier::{Antichain, AntichainRef};
27
28use crate::metrics::WorkerMetrics;
29use crate::typedefs::{ErrAgent, RowRowAgent};
30
31/// A `TraceManager` stores maps from global identifiers to the primary arranged
32/// representation of that collection.
33pub struct TraceManager {
34    pub(crate) traces: BTreeMap<GlobalId, TraceBundle>,
35    metrics: WorkerMetrics,
36}
37
38impl TraceManager {
39    /// TODO(undocumented)
40    pub fn new(metrics: WorkerMetrics) -> Self {
41        TraceManager {
42            traces: BTreeMap::new(),
43            metrics,
44        }
45    }
46
47    /// performs maintenance work on the managed traces.
48    ///
49    /// In particular, this method enables the physical merging of batches, so that at most a logarithmic
50    /// number of batches need to be maintained. Any new batches introduced after this method is called
51    /// will not be physically merged until the method is called again. This is mostly due to limitations
52    /// of differential dataflow, which requires users to perform this explicitly; if that changes we may
53    /// be able to remove this code.
54    pub fn maintenance(&mut self) {
55        let start = Instant::now();
56        self.metrics.arrangement_maintenance_active_info.set(1);
57
58        let mut antichain = Antichain::new();
59        for bundle in self.traces.values_mut() {
60            bundle.oks.read_upper(&mut antichain);
61            bundle.oks.set_physical_compaction(antichain.borrow());
62            bundle.errs.read_upper(&mut antichain);
63            bundle.errs.set_physical_compaction(antichain.borrow());
64        }
65
66        let duration = start.elapsed().as_secs_f64();
67        self.metrics
68            .arrangement_maintenance_seconds_total
69            .inc_by(duration);
70        self.metrics.arrangement_maintenance_active_info.set(0);
71    }
72
73    /// Enables compaction of traces associated with the identifier.
74    ///
75    /// Compaction may not occur immediately, but once this method is called the
76    /// associated traces may not accumulate to the correct quantities for times
77    /// not in advance of `frontier`. Users should take care to only rely on
78    /// accumulations at times in advance of `frontier`.
79    pub fn allow_compaction(&mut self, id: GlobalId, frontier: AntichainRef<Timestamp>) {
80        if let Some(bundle) = self.traces.get_mut(&id) {
81            bundle.oks.set_logical_compaction(frontier);
82            bundle.errs.set_logical_compaction(frontier);
83        }
84    }
85
86    /// Returns a reference to the trace for `id`, should it exist.
87    pub fn get(&self, id: &GlobalId) -> Option<&TraceBundle> {
88        self.traces.get(id)
89    }
90
91    /// Returns a mutable reference to the trace for `id`, should it
92    /// exist.
93    pub fn get_mut(&mut self, id: &GlobalId) -> Option<&mut TraceBundle> {
94        self.traces.get_mut(id)
95    }
96
97    /// Binds the arrangement for `id` to `trace`.
98    pub fn set(&mut self, id: GlobalId, trace: TraceBundle) {
99        self.traces.insert(id, trace);
100    }
101
102    /// Removes the trace for `id`.
103    pub fn remove(&mut self, id: &GlobalId) -> Option<TraceBundle> {
104        self.traces.remove(id)
105    }
106}
107
108/// Handle to a trace that can be padded.
109///
110/// A padded trace contains empty data for all times greater than or equal to its `padded_since`
111/// and less than the logical compaction frontier of the inner `trace`.
112///
113/// This type is intentionally limited to only work with `mz_repr::Timestamp` times, because that
114/// is all that's required by `TraceManager`. It can be made to be more generic, at the cost of
115/// more complicated reasoning about the correct management of the involved frontiers.
116#[derive(Clone)]
117pub struct PaddedTrace<Tr>
118where
119    Tr: TraceReader,
120{
121    /// The wrapped trace.
122    trace: Tr,
123    /// The frontier from which the trace is padded, or `None` if it is not padded.
124    ///
125    /// Invariant: The contained frontier is less than the logical compaction frontier of `trace`.
126    ///
127    /// All methods of `PaddedTrace` are written to uphold this invariant. In particular,
128    /// `set_logical_compaction_frontier`  sets the `padded_since` to `None` if the new compaction
129    /// frontier is >= the previous compaction frontier of `trace`.
130    padded_since: Option<Antichain<Tr::Time>>,
131}
132
133impl<Tr> From<Tr> for PaddedTrace<Tr>
134where
135    Tr: TraceReader,
136{
137    fn from(trace: Tr) -> Self {
138        Self {
139            trace,
140            padded_since: None,
141        }
142    }
143}
144
145impl<Tr> PaddedTrace<Tr>
146where
147    Tr: TraceReader,
148{
149    /// Turns this trace into a padded version that reports empty data for all times less than the
150    /// trace's current logical compaction frontier.
151    fn into_padded(mut self) -> Self {
152        let trace_since = self.trace.get_logical_compaction();
153        let minimum_frontier = Antichain::from_elem(Tr::Time::minimum());
154        if PartialOrder::less_than(&minimum_frontier.borrow(), &trace_since) {
155            self.padded_since = Some(minimum_frontier);
156        }
157        self
158    }
159}
160
161impl<Tr> TraceReader for PaddedTrace<Tr>
162where
163    Tr: TraceReader,
164{
165    type Time = Tr::Time;
166    type Batch = Tr::Batch;
167
168    fn spans_through(
169        &mut self,
170        upper: AntichainRef<Self::Time>,
171    ) -> Option<Vec<Span<Self::Time, Self::Batch>>> {
172        self.trace.spans_through(upper)
173    }
174
175    fn batches_through(&mut self, upper: AntichainRef<Self::Time>) -> Option<Vec<Self::Batch>> {
176        self.trace.batches_through(upper)
177    }
178
179    fn set_logical_compaction(&mut self, frontier: AntichainRef<Self::Time>) {
180        let Some(padded_since) = &mut self.padded_since else {
181            self.trace.set_logical_compaction(frontier);
182            return;
183        };
184
185        // If a padded trace is compacted to some frontier less than the inner trace's compaction
186        // frontier, advance the `padded_since`. Otherwise discard the padding and apply the
187        // compaction to the inner trace instead.
188        let trace_since = self.trace.get_logical_compaction();
189        if PartialOrder::less_than(&frontier, &trace_since) {
190            if PartialOrder::less_than(&padded_since.borrow(), &frontier) {
191                *padded_since = frontier.to_owned();
192            }
193        } else {
194            self.padded_since = None;
195            self.trace.set_logical_compaction(frontier);
196        }
197    }
198
199    fn get_logical_compaction(&mut self) -> AntichainRef<'_, Self::Time> {
200        match &self.padded_since {
201            Some(since) => since.borrow(),
202            None => self.trace.get_logical_compaction(),
203        }
204    }
205
206    fn set_physical_compaction(&mut self, frontier: AntichainRef<Self::Time>) {
207        self.trace.set_physical_compaction(frontier);
208    }
209
210    fn get_physical_compaction(&mut self) -> AntichainRef<'_, Self::Time> {
211        self.trace.get_physical_compaction()
212    }
213
214    fn map_spans<F: FnMut(&Span<Self::Time, Self::Batch>)>(&self, f: F) {
215        self.trace.map_spans(f)
216    }
217}
218
219impl<Tr> PaddedTrace<TraceAgent<Tr>>
220where
221    Tr: TraceReader<Time = Timestamp> + 'static,
222{
223    /// Import a trace restricted to a specific time interval `[since, until)`.
224    pub fn import_frontier_core<'scope>(
225        &mut self,
226        scope: Scope<'scope, Tr::Time>,
227        name: &str,
228        since: Antichain<Tr::Time>,
229        until: Antichain<Tr::Time>,
230    ) -> (
231        Arranged<'scope, TraceFrontier<TraceAgent<Tr>>>,
232        ShutdownButton<CapabilitySet<Tr::Time>>,
233    ) {
234        self.trace.import_frontier_core(scope, name, since, until)
235    }
236}
237
238/// Bundles together traces for the successful computations (`oks`), the
239/// failed computations (`errs`), additional tokens that should share
240/// the lifetime of the bundled traces (`to_drop`).
241#[derive(Clone)]
242pub struct TraceBundle {
243    oks: PaddedTrace<RowRowAgent<Timestamp, Diff>>,
244    errs: PaddedTrace<ErrAgent<Timestamp, Diff>>,
245    to_drop: Option<Rc<dyn Any>>,
246}
247
248impl TraceBundle {
249    /// Constructs a new trace bundle out of an `oks` trace and `errs` trace.
250    pub fn new<O, E>(oks: O, errs: E) -> TraceBundle
251    where
252        O: Into<PaddedTrace<RowRowAgent<Timestamp, Diff>>>,
253        E: Into<PaddedTrace<ErrAgent<Timestamp, Diff>>>,
254    {
255        TraceBundle {
256            oks: oks.into(),
257            errs: errs.into(),
258            to_drop: None,
259        }
260    }
261
262    /// Adds tokens to be dropped when the trace bundle is dropped.
263    pub fn with_drop<T>(self, to_drop: T) -> TraceBundle
264    where
265        T: 'static,
266    {
267        TraceBundle {
268            to_drop: Some(Rc::new(Box::new(to_drop))),
269            ..self
270        }
271    }
272
273    /// Returns a mutable reference to the `oks` trace.
274    pub fn oks_mut(&mut self) -> &mut PaddedTrace<RowRowAgent<Timestamp, Diff>> {
275        &mut self.oks
276    }
277
278    /// Returns a mutable reference to the `errs` trace.
279    pub fn errs_mut(&mut self) -> &mut PaddedTrace<ErrAgent<Timestamp, Diff>> {
280        &mut self.errs
281    }
282
283    /// Returns mutable references to both traces.
284    ///
285    /// A reader of both, such as a peek that must rule out errors before it returns rows, cannot
286    /// take them one at a time, because each of [`TraceBundle::oks_mut`] and
287    /// [`TraceBundle::errs_mut`] borrows the whole bundle.
288    pub fn oks_errs_mut(
289        &mut self,
290    ) -> (
291        &mut PaddedTrace<RowRowAgent<Timestamp, Diff>>,
292        &mut PaddedTrace<ErrAgent<Timestamp, Diff>>,
293    ) {
294        (&mut self.oks, &mut self.errs)
295    }
296
297    /// Returns a reference to the `to_drop` tokens.
298    pub fn to_drop(&self) -> &Option<Rc<dyn Any>> {
299        &self.to_drop
300    }
301
302    /// Returns the frontier up to which the traces have been allowed to compact.
303    pub fn compaction_frontier(&mut self) -> Antichain<Timestamp> {
304        antichain_join(
305            &self.oks.get_logical_compaction(),
306            &self.errs.get_logical_compaction(),
307        )
308    }
309
310    /// Turns this trace bundle into a padded version that reports empty data for all times less
311    /// than the traces' current logical compaction frontier.
312    ///
313    /// Note that the padded bundle represents a different TVC than the original one, it is unsound
314    /// to use it to "uncompact" an existing TVC. The only valid use of the padded bundle is to
315    /// initializa a new TVC.
316    pub fn into_padded(self) -> Self {
317        Self {
318            oks: self.oks.into_padded(),
319            errs: self.errs.into_padded(),
320            to_drop: self.to_drop,
321        }
322    }
323}
mz_compute/arrangement/manager.rs

mz_compute/arrangement/
manager.rs
