mz_row_spine/
arc_batch.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//! An `Arc`-backed batch newtype whose contents can be shared across timely runtimes.
11//!
12//! Differential's default spines reference-count their batches with `Rc`, which is worker-local.
13//! Sharing an arrangement with another runtime (a reader on a different worker thread) needs the
14//! batches behind an `Arc` so a batch whose contents are `Send + Sync` can be read from that other
15//! thread.
16//!
17//! The blanket `impl Trait for Arc<B>` that would express this lives outside this crate: both `Arc`
18//! and differential's `Batch`/`Builder`/`Merger`/`Cursor` traits are foreign, so the orphan rule
19//! forbids it here. [`ArcBatch`] is a local newtype around `Arc<B>` that carries those impls
20//! instead. The impls delegate straight through to the inner batch, so `ArcBatch<B>` behaves
21//! exactly like `B` except that its handle is atomically reference counted.
22//!
23//! This mirrors differential's own `rc_blanket_impls` (for `Rc<B>`), swapping in `Arc`. Keeping it
24//! here as a newtype lets cross-thread arrangement sharing build against a released
25//! differential-dataflow, with no differential-side `Arc` batch impls required.
26
27use std::sync::Arc;
28
29use differential_dataflow::trace::{
30    Batch, BatchReader, Builder, Cursor, Description, Merger, Navigable,
31};
32use timely::progress::{Antichain, frontier::AntichainRef};
33
34/// An `Arc`-backed batch, shareable across threads when `B`'s contents are `Send + Sync`.
35///
36/// A transparent newtype around `Arc<B>`. Cloning shares the underlying batch, exactly like the
37/// `Rc`-backed default, but with atomic reference counting.
38pub struct ArcBatch<B>(pub Arc<B>);
39
40// Hand-written rather than derived: `#[derive(Clone)]` would bound `B: Clone`, but `Arc<B>` is
41// `Clone` for any `B` (it clones the handle, not the batch). The derived bound would make
42// `ArcBatch<B>: Clone` fail for a non-`Clone` batch such as `OrdValBatch`, which in turn breaks
43// `Spine<ArcBatch<B>>: TraceReader`.
44impl<B> Clone for ArcBatch<B> {
45    fn clone(&self) -> Self {
46        ArcBatch(Arc::clone(&self.0))
47    }
48}
49
50impl<B> ArcBatch<B> {
51    /// Wraps a batch in an `Arc`.
52    pub fn new(batch: B) -> Self {
53        ArcBatch(Arc::new(batch))
54    }
55}
56
57impl<B> std::ops::Deref for ArcBatch<B> {
58    type Target = B;
59    fn deref(&self) -> &B {
60        &self.0
61    }
62}
63
64impl<B: BatchReader + Navigable> Navigable for ArcBatch<B> {
65    type Cursor = ArcBatchCursor<B::Cursor>;
66    fn cursor(&self) -> Self::Cursor {
67        // Disambiguate to the inner batch's cursor, reached through the `Deref`, so the wrapper's
68        // `Cursor` is `B`'s rather than any impl that might exist on `Arc<B>` itself.
69        ArcBatchCursor::new(<B as Navigable>::cursor(&self.0))
70    }
71}
72
73impl<B: BatchReader> BatchReader for ArcBatch<B> {
74    type Time = B::Time;
75    fn len(&self) -> usize {
76        self.0.len()
77    }
78    fn description(&self) -> &Description<Self::Time> {
79        self.0.description()
80    }
81}
82
83/// Cursor over an [`ArcBatch`], delegating to the inner batch's cursor.
84pub struct ArcBatchCursor<C> {
85    cursor: C,
86}
87
88impl<C> ArcBatchCursor<C> {
89    fn new(cursor: C) -> Self {
90        ArcBatchCursor { cursor }
91    }
92}
93
94impl<C: Cursor> Cursor for ArcBatchCursor<C> {
95    type Storage = ArcBatch<C::Storage>;
96
97    type Key<'a> = C::Key<'a>;
98    type ValOwn = C::ValOwn;
99    type Val<'a> = C::Val<'a>;
100    type Time = C::Time;
101    type TimeGat<'a> = C::TimeGat<'a>;
102    type Diff = C::Diff;
103    type DiffGat<'a> = C::DiffGat<'a>;
104    type KeyContainer = C::KeyContainer;
105    type ValContainer = C::ValContainer;
106    type TimeContainer = C::TimeContainer;
107    type DiffContainer = C::DiffContainer;
108
109    #[inline]
110    fn key_valid(&self, storage: &Self::Storage) -> bool {
111        self.cursor.key_valid(&storage.0)
112    }
113    #[inline]
114    fn val_valid(&self, storage: &Self::Storage) -> bool {
115        self.cursor.val_valid(&storage.0)
116    }
117
118    #[inline]
119    fn key<'a>(&self, storage: &'a Self::Storage) -> Self::Key<'a> {
120        self.cursor.key(&storage.0)
121    }
122    #[inline]
123    fn val<'a>(&self, storage: &'a Self::Storage) -> Self::Val<'a> {
124        self.cursor.val(&storage.0)
125    }
126
127    #[inline]
128    fn get_key<'a>(&self, storage: &'a Self::Storage) -> Option<Self::Key<'a>> {
129        self.cursor.get_key(&storage.0)
130    }
131    #[inline]
132    fn get_val<'a>(&self, storage: &'a Self::Storage) -> Option<Self::Val<'a>> {
133        self.cursor.get_val(&storage.0)
134    }
135
136    #[inline]
137    fn map_times<L: FnMut(Self::TimeGat<'_>, Self::DiffGat<'_>)>(
138        &mut self,
139        storage: &Self::Storage,
140        logic: L,
141    ) {
142        self.cursor.map_times(&storage.0, logic)
143    }
144
145    #[inline]
146    fn step_key(&mut self, storage: &Self::Storage) {
147        self.cursor.step_key(&storage.0)
148    }
149    #[inline]
150    fn seek_key(&mut self, storage: &Self::Storage, key: Self::Key<'_>) {
151        self.cursor.seek_key(&storage.0, key)
152    }
153
154    #[inline]
155    fn step_val(&mut self, storage: &Self::Storage) {
156        self.cursor.step_val(&storage.0)
157    }
158    #[inline]
159    fn seek_val(&mut self, storage: &Self::Storage, val: Self::Val<'_>) {
160        self.cursor.seek_val(&storage.0, val)
161    }
162
163    #[inline]
164    fn rewind_keys(&mut self, storage: &Self::Storage) {
165        self.cursor.rewind_keys(&storage.0)
166    }
167    #[inline]
168    fn rewind_vals(&mut self, storage: &Self::Storage) {
169        self.cursor.rewind_vals(&storage.0)
170    }
171}
172
173impl<B: Batch> Batch for ArcBatch<B> {
174    type Merger = ArcMerger<B>;
175    fn empty(lower: Antichain<Self::Time>, upper: Antichain<Self::Time>) -> Self {
176        ArcBatch::new(B::empty(lower, upper))
177    }
178}
179
180/// Builds [`ArcBatch`]es, delegating to the inner batch's builder.
181pub struct ArcBuilder<B: Builder> {
182    builder: B,
183}
184
185impl<B: Builder> Builder for ArcBuilder<B> {
186    type Input = B::Input;
187    type Time = B::Time;
188    type Output = ArcBatch<B::Output>;
189    fn with_capacity(keys: usize, vals: usize, upds: usize) -> Self {
190        ArcBuilder {
191            builder: B::with_capacity(keys, vals, upds),
192        }
193    }
194    fn push(&mut self, input: &mut Self::Input) {
195        self.builder.push(input)
196    }
197    fn done(self, description: Description<Self::Time>) -> ArcBatch<B::Output> {
198        ArcBatch::new(self.builder.done(description))
199    }
200    fn seal(chain: &mut Vec<Self::Input>, description: Description<Self::Time>) -> Self::Output {
201        ArcBatch::new(B::seal(chain, description))
202    }
203}
204
205/// Merges [`ArcBatch`]es, delegating to the inner batch's merger.
206pub struct ArcMerger<B: Batch> {
207    merger: B::Merger,
208}
209
210impl<B: Batch> Merger<ArcBatch<B>> for ArcMerger<B> {
211    fn new(
212        source1: &ArcBatch<B>,
213        source2: &ArcBatch<B>,
214        compaction_frontier: AntichainRef<B::Time>,
215    ) -> Self {
216        ArcMerger {
217            merger: B::begin_merge(&source1.0, &source2.0, compaction_frontier),
218        }
219    }
220    fn work(&mut self, source1: &ArcBatch<B>, source2: &ArcBatch<B>, fuel: &mut isize) {
221        self.merger.work(&source1.0, &source2.0, fuel)
222    }
223    fn done(self) -> ArcBatch<B> {
224        ArcBatch::new(self.merger.done())
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use differential_dataflow::trace::cursor::Cursor;
231    use differential_dataflow::trace::implementations::ord_neu::OrdValBatcher;
232    use differential_dataflow::trace::{Batcher, Builder, Navigable};
233    use timely::container::PushInto;
234    use timely::progress::Antichain;
235
236    use crate::ArcOrdValBuilder;
237
238    /// An `ArcBatch`'s cursor can be constructed and read from a thread other than the one that
239    /// built it, proving the newtype's batches are usable across a thread boundary. This is the
240    /// property that lets [`crate::ArcOrdValSpine`] (and the `RowRow`/`Err` spines built on
241    /// [`ArcBatch`]) back a cross-runtime shared trace; the default `Rc`-backed spines are
242    /// worker-local by design and do not have it.
243    ///
244    /// Mirrors differential-dataflow's own `tests/trace.rs` cross-thread batch read, over the local
245    /// [`ArcBatch`] newtype.
246    #[mz_ore::test]
247    fn arc_batch_reads_from_other_thread() {
248        fn assert_send_sync<T: Send + Sync>(_: &T) {}
249
250        let mut batcher = OrdValBatcher::<u64, u64, usize, i64>::new(None, 0);
251        batcher.push_into(vec![((1, 2), 0, 1), ((2, 3), 1, 1)]);
252        let (mut chain, description) = batcher.seal(Antichain::from_elem(2));
253        let batch = ArcOrdValBuilder::<u64, u64, usize, i64>::seal(&mut chain, description);
254
255        assert_send_sync(&batch);
256
257        let read = std::thread::spawn(move || {
258            let mut cursor = batch.cursor();
259            cursor.to_vec(&batch, |k| *k, |v| *v)
260        })
261        .join()
262        .expect("reader thread panicked");
263
264        assert_eq!(read, vec![((1, 2), vec![(0, 1)]), ((2, 3), vec![(1, 1)])]);
265    }
266}
mz_row_spine/arc_batch.rs

mz_row_spine/
arc_batch.rs
