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::TraceReader;
20use differential_dataflow::trace::wrappers::frontier::TraceFrontier;
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 batches_through(&mut self, upper: AntichainRef<Self::Time>) -> Option<Vec<Self::Batch>> {
169        self.trace.batches_through(upper)
170    }
171
172    fn set_logical_compaction(&mut self, frontier: AntichainRef<Self::Time>) {
173        let Some(padded_since) = &mut self.padded_since else {
174            self.trace.set_logical_compaction(frontier);
175            return;
176        };
177
178        // If a padded trace is compacted to some frontier less than the inner trace's compaction
179        // frontier, advance the `padded_since`. Otherwise discard the padding and apply the
180        // compaction to the inner trace instead.
181        let trace_since = self.trace.get_logical_compaction();
182        if PartialOrder::less_than(&frontier, &trace_since) {
183            if PartialOrder::less_than(&padded_since.borrow(), &frontier) {
184                *padded_since = frontier.to_owned();
185            }
186        } else {
187            self.padded_since = None;
188            self.trace.set_logical_compaction(frontier);
189        }
190    }
191
192    fn get_logical_compaction(&mut self) -> AntichainRef<'_, Self::Time> {
193        match &self.padded_since {
194            Some(since) => since.borrow(),
195            None => self.trace.get_logical_compaction(),
196        }
197    }
198
199    fn set_physical_compaction(&mut self, frontier: AntichainRef<Self::Time>) {
200        self.trace.set_physical_compaction(frontier);
201    }
202
203    fn get_physical_compaction(&mut self) -> AntichainRef<'_, Self::Time> {
204        self.trace.get_physical_compaction()
205    }
206
207    fn map_batches<F: FnMut(&Self::Batch)>(&self, f: F) {
208        self.trace.map_batches(f)
209    }
210}
211
212impl<Tr> PaddedTrace<TraceAgent<Tr>>
213where
214    Tr: TraceReader<Time = Timestamp> + 'static,
215{
216    /// Import a trace restricted to a specific time interval `[since, until)`.
217    pub fn import_frontier_core<'scope>(
218        &mut self,
219        scope: Scope<'scope, Tr::Time>,
220        name: &str,
221        since: Antichain<Tr::Time>,
222        until: Antichain<Tr::Time>,
223    ) -> (
224        Arranged<'scope, TraceFrontier<TraceAgent<Tr>>>,
225        ShutdownButton<CapabilitySet<Tr::Time>>,
226    ) {
227        self.trace.import_frontier_core(scope, name, since, until)
228    }
229}
230
231/// Bundles together traces for the successful computations (`oks`), the
232/// failed computations (`errs`), additional tokens that should share
233/// the lifetime of the bundled traces (`to_drop`).
234#[derive(Clone)]
235pub struct TraceBundle {
236    oks: PaddedTrace<RowRowAgent<Timestamp, Diff>>,
237    errs: PaddedTrace<ErrAgent<Timestamp, Diff>>,
238    to_drop: Option<Rc<dyn Any>>,
239}
240
241impl TraceBundle {
242    /// Constructs a new trace bundle out of an `oks` trace and `errs` trace.
243    pub fn new<O, E>(oks: O, errs: E) -> TraceBundle
244    where
245        O: Into<PaddedTrace<RowRowAgent<Timestamp, Diff>>>,
246        E: Into<PaddedTrace<ErrAgent<Timestamp, Diff>>>,
247    {
248        TraceBundle {
249            oks: oks.into(),
250            errs: errs.into(),
251            to_drop: None,
252        }
253    }
254
255    /// Adds tokens to be dropped when the trace bundle is dropped.
256    pub fn with_drop<T>(self, to_drop: T) -> TraceBundle
257    where
258        T: 'static,
259    {
260        TraceBundle {
261            to_drop: Some(Rc::new(Box::new(to_drop))),
262            ..self
263        }
264    }
265
266    /// Returns a mutable reference to the `oks` trace.
267    pub fn oks_mut(&mut self) -> &mut PaddedTrace<RowRowAgent<Timestamp, Diff>> {
268        &mut self.oks
269    }
270
271    /// Returns a mutable reference to the `errs` trace.
272    pub fn errs_mut(&mut self) -> &mut PaddedTrace<ErrAgent<Timestamp, Diff>> {
273        &mut self.errs
274    }
275
276    /// Returns a reference to the `to_drop` tokens.
277    pub fn to_drop(&self) -> &Option<Rc<dyn Any>> {
278        &self.to_drop
279    }
280
281    /// Returns the frontier up to which the traces have been allowed to compact.
282    pub fn compaction_frontier(&mut self) -> Antichain<Timestamp> {
283        antichain_join(
284            &self.oks.get_logical_compaction(),
285            &self.errs.get_logical_compaction(),
286        )
287    }
288
289    /// Turns this trace bundle into a padded version that reports empty data for all times less
290    /// than the traces' current logical compaction frontier.
291    ///
292    /// Note that the padded bundle represents a different TVC than the original one, it is unsound
293    /// to use it to "uncompact" an existing TVC. The only valid use of the padded bundle is to
294    /// initializa a new TVC.
295    pub fn into_padded(self) -> Self {
296        Self {
297            oks: self.oks.into_padded(),
298            errs: self.errs.into_padded(),
299            to_drop: self.to_drop,
300        }
301    }
302}
mz_compute/arrangement/manager.rs

mz_compute/arrangement/
manager.rs
