Skip to main content

mz_persist_client/
read.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//! Read capabilities and handles
11
12use async_stream::stream;
13use std::collections::BTreeMap;
14use std::fmt::Debug;
15use std::future::Future;
16use std::sync::Arc;
17use std::time::{Duration, Instant};
18
19use differential_dataflow::Hashable;
20use differential_dataflow::consolidation::consolidate_updates;
21use differential_dataflow::difference::Monoid;
22use differential_dataflow::lattice::Lattice;
23use futures::Stream;
24use futures_util::{StreamExt, stream};
25use mz_dyncfg::{Config, ParameterScope};
26use mz_ore::cast::CastLossy;
27use mz_ore::halt;
28use mz_ore::instrument;
29use mz_ore::task::JoinHandle;
30use mz_persist::location::{Blob, SeqNo};
31use mz_persist_types::columnar::{ColumnDecoder, Schema};
32use mz_persist_types::{Codec, Codec64};
33use proptest_derive::Arbitrary;
34use serde::{Deserialize, Serialize};
35use timely::PartialOrder;
36use timely::order::TotalOrder;
37use timely::progress::{Antichain, Timestamp};
38use tracing::warn;
39use uuid::Uuid;
40
41use crate::batch::BLOB_TARGET_SIZE;
42use crate::cfg::{COMPACTION_MEMORY_BOUND_BYTES, RetryParameters};
43use crate::fetch::FetchConfig;
44use crate::fetch::{FetchBatchFilter, FetchedPart, Lease, LeasedBatchPart, fetch_leased_part};
45use crate::internal::encoding::Schemas;
46use crate::internal::machine::{Machine, next_listen_batch_retry_params};
47use crate::internal::metrics::{Metrics, ReadMetrics, ShardMetrics};
48use crate::internal::state::{HollowBatch, LeasedReaderState, SnapshotErr};
49use crate::internal::watch::{AwaitableState, StateWatch};
50use crate::iter::{Consolidator, StructuredSort};
51use crate::schema::SchemaCache;
52use crate::stats::{SnapshotPartStats, SnapshotPartsStats, SnapshotStats};
53use crate::{GarbageCollector, PersistConfig, ShardId, parse_id};
54
55pub use crate::internal::encoding::LazyPartStats;
56pub use crate::internal::state::Since;
57
58/// An opaque identifier for a reader of a persist durable TVC (aka shard).
59#[derive(
60    Arbitrary,
61    Clone,
62    PartialEq,
63    Eq,
64    PartialOrd,
65    Ord,
66    Hash,
67    Serialize,
68    Deserialize
69)]
70#[serde(try_from = "String", into = "String")]
71pub struct LeasedReaderId(pub(crate) [u8; 16]);
72
73impl std::fmt::Display for LeasedReaderId {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        write!(f, "r{}", Uuid::from_bytes(self.0))
76    }
77}
78
79impl std::fmt::Debug for LeasedReaderId {
80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        write!(f, "LeasedReaderId({})", Uuid::from_bytes(self.0))
82    }
83}
84
85impl std::str::FromStr for LeasedReaderId {
86    type Err = String;
87
88    fn from_str(s: &str) -> Result<Self, Self::Err> {
89        parse_id("r", "LeasedReaderId", s).map(LeasedReaderId)
90    }
91}
92
93impl From<LeasedReaderId> for String {
94    fn from(reader_id: LeasedReaderId) -> Self {
95        reader_id.to_string()
96    }
97}
98
99impl TryFrom<String> for LeasedReaderId {
100    type Error = String;
101
102    fn try_from(s: String) -> Result<Self, Self::Error> {
103        s.parse()
104    }
105}
106
107impl LeasedReaderId {
108    pub(crate) fn new() -> Self {
109        LeasedReaderId(*Uuid::new_v4().as_bytes())
110    }
111}
112
113/// Capable of generating a snapshot of all data at `as_of`, followed by a
114/// listen of all updates.
115///
116/// For more details, see [`ReadHandle::snapshot`] and [`Listen`].
117#[derive(Debug)]
118pub struct Subscribe<K: Codec, V: Codec, T, D> {
119    snapshot: Option<Vec<LeasedBatchPart<T>>>,
120    listen: Listen<K, V, T, D>,
121}
122
123impl<K, V, T, D> Subscribe<K, V, T, D>
124where
125    K: Debug + Codec,
126    V: Debug + Codec,
127    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
128    D: Monoid + Codec64 + Send + Sync,
129{
130    fn new(snapshot_parts: Vec<LeasedBatchPart<T>>, listen: Listen<K, V, T, D>) -> Self {
131        Subscribe {
132            snapshot: Some(snapshot_parts),
133            listen,
134        }
135    }
136
137    /// Returns a `LeasedBatchPart` enriched with the proper metadata.
138    ///
139    /// First returns snapshot parts, until they're exhausted, at which point
140    /// begins returning listen parts.
141    ///
142    /// The returned `Antichain` represents the subscription progress as it will
143    /// be _after_ the returned parts are fetched.
144    #[instrument(level = "debug", fields(shard = %self.listen.handle.machine.shard_id()))]
145    pub async fn next(
146        &mut self,
147        // If Some, an override for the default listen sleep retry parameters.
148        listen_retry: Option<RetryParameters>,
149    ) -> Vec<ListenEvent<T, LeasedBatchPart<T>>> {
150        match self.snapshot.take() {
151            Some(parts) => vec![ListenEvent::Updates(parts)],
152            None => {
153                let (parts, upper) = self.listen.next(listen_retry).await;
154                vec![ListenEvent::Updates(parts), ListenEvent::Progress(upper)]
155            }
156        }
157    }
158}
159
160impl<K, V, T, D> Subscribe<K, V, T, D>
161where
162    K: Debug + Codec,
163    V: Debug + Codec,
164    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
165    D: Monoid + Codec64 + Send + Sync,
166{
167    /// Equivalent to `next`, but rather than returning a [`LeasedBatchPart`],
168    /// fetches and returns the data from within it.
169    #[instrument(level = "debug", fields(shard = %self.listen.handle.machine.shard_id()))]
170    pub async fn fetch_next(&mut self) -> Vec<ListenEvent<T, ((K, V), T, D)>> {
171        let events = self.next(None).await;
172        let new_len = events
173            .iter()
174            .map(|event| match event {
175                ListenEvent::Updates(parts) => parts.len(),
176                ListenEvent::Progress(_) => 1,
177            })
178            .sum();
179        let mut ret = Vec::with_capacity(new_len);
180        for event in events {
181            match event {
182                ListenEvent::Updates(parts) => {
183                    for part in parts {
184                        let fetched_part = self.listen.fetch_batch_part(part).await;
185                        let updates = fetched_part.collect::<Vec<_>>();
186                        if !updates.is_empty() {
187                            ret.push(ListenEvent::Updates(updates));
188                        }
189                    }
190                }
191                ListenEvent::Progress(progress) => ret.push(ListenEvent::Progress(progress)),
192            }
193        }
194        ret
195    }
196
197    /// Fetches the contents of `part` and returns its lease.
198    pub async fn fetch_batch_part(&mut self, part: LeasedBatchPart<T>) -> FetchedPart<K, V, T, D> {
199        self.listen.fetch_batch_part(part).await
200    }
201}
202
203impl<K, V, T, D> Subscribe<K, V, T, D>
204where
205    K: Debug + Codec,
206    V: Debug + Codec,
207    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
208    D: Monoid + Codec64 + Send + Sync,
209{
210    /// Politely expires this subscribe, releasing its lease.
211    ///
212    /// There is a best-effort impl in Drop to expire the
213    /// [`ReadHandle`] held by the subscribe that wasn't explicitly expired
214    /// with this method. When possible, explicit expiry is still preferred
215    /// because it also ensures that the background task is complete.
216    pub async fn expire(mut self) {
217        let _ = self.snapshot.take(); // Drop all leased parts.
218        self.listen.expire().await;
219    }
220}
221
222/// Data and progress events of a shard subscription.
223///
224/// TODO: Unify this with [timely::dataflow::operators::capture::event::Event].
225#[derive(Debug, PartialEq)]
226pub enum ListenEvent<T, D> {
227    /// Progress of the shard.
228    Progress(Antichain<T>),
229    /// Data of the shard.
230    Updates(Vec<D>),
231}
232
233/// An ongoing subscription of updates to a shard.
234#[derive(Debug)]
235pub struct Listen<K: Codec, V: Codec, T, D> {
236    handle: ReadHandle<K, V, T, D>,
237    as_of: Antichain<T>,
238    since: Antichain<T>,
239    frontier: Antichain<T>,
240}
241
242impl<K, V, T, D> Listen<K, V, T, D>
243where
244    K: Debug + Codec,
245    V: Debug + Codec,
246    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
247    D: Monoid + Codec64 + Send + Sync,
248{
249    async fn new(
250        mut handle: ReadHandle<K, V, T, D>,
251        as_of: Antichain<T>,
252    ) -> Result<Self, Since<T>> {
253        let () = handle.machine.verify_listen(&as_of)?;
254
255        let since = as_of.clone();
256        if !PartialOrder::less_equal(handle.since(), &since) {
257            // We can't guarantee that the as-of will be available by the time we start reading,
258            // since our handle is already downgraded too far.
259            return Err(Since(handle.since().clone()));
260        }
261        // This listen only needs to distinguish things after its frontier
262        // (initially as_of although the frontier is inclusive and the as_of
263        // isn't). Be a good citizen and downgrade early.
264        handle.downgrade_since(&since).await;
265        Ok(Listen {
266            handle,
267            since,
268            frontier: as_of.clone(),
269            as_of,
270        })
271    }
272
273    /// An exclusive upper bound on the progress of this Listen.
274    pub fn frontier(&self) -> &Antichain<T> {
275        &self.frontier
276    }
277
278    /// Attempt to pull out the next values of this subscription.
279    ///
280    /// The returned [`LeasedBatchPart`] is appropriate to use with
281    /// `crate::fetch::fetch_leased_part`.
282    ///
283    /// The returned `Antichain` represents the subscription progress as it will
284    /// be _after_ the returned parts are fetched.
285    pub async fn next(
286        &mut self,
287        // If Some, an override for the default listen sleep retry parameters.
288        retry: Option<RetryParameters>,
289    ) -> (Vec<LeasedBatchPart<T>>, Antichain<T>) {
290        // Wait until the upper is past our frontier - ie. there is another batch for us to process.
291        let retry = retry
292            .unwrap_or_else(|| next_listen_batch_retry_params(&self.handle.machine.applier.cfg));
293        self.handle
294            .machine
295            .wait_for_upper_past(
296                &self.frontier,
297                &mut self.handle.watch,
298                Some(&self.handle.reader_id),
299                &self.handle.metrics.retries.next_listen_batch,
300                retry,
301            )
302            .await;
303
304        // Obtain a lease before grabbing the upcoming batch from state.
305        let lease = self.handle.lease_seqno().await;
306        let batch = match self
307            .handle
308            .machine
309            .applier
310            .next_listen_batch(&self.frontier)
311        {
312            Ok(batch) => batch,
313            Err(seqno) => {
314                panic!(
315                    "waited for upper past {frontier:?}, but no listen batch was available at {seqno:?}!",
316                    frontier = self.frontier.elements()
317                );
318            }
319        };
320
321        // A lot of things across mz have to line up to hold the following
322        // invariant and violations only show up as subtle correctness errors,
323        // so explicitly validate it here. Better to panic and roll back a
324        // release than be incorrect (also potentially corrupting a sink).
325        //
326        // Note that the since check is intentionally less_than, not less_equal.
327        // If a batch's since is X, that means we can no longer distinguish X
328        // (beyond self.frontier) from X-1 (not beyond self.frontier) to keep
329        // former and filter out the latter.
330        let acceptable_desc = PartialOrder::less_than(batch.desc.since(), &self.frontier)
331            // Special case when the frontier == the as_of (i.e. the first
332            // time this is called on a new Listen). Because as_of is
333            // _exclusive_, we don't need to be able to distinguish X from
334            // X-1.
335            || (self.frontier == self.as_of
336            && PartialOrder::less_equal(batch.desc.since(), &self.frontier));
337        if !acceptable_desc {
338            let lease_state = self
339                .handle
340                .machine
341                .applier
342                .reader_lease(self.handle.reader_id.clone());
343            if let Some(lease) = lease_state {
344                panic!(
345                    "Listen on {} received a batch {:?} advanced past the listen frontier {:?}, but the lease has not expired: {:?}",
346                    self.handle.machine.shard_id(),
347                    batch.desc,
348                    self.frontier,
349                    lease
350                )
351            } else {
352                // Ideally we'd percolate this error up, so callers could eg. restart a dataflow
353                // instead of restarting a process...
354                halt!(
355                    "Listen on {} received a batch {:?} advanced past the listen frontier {:?} after the reader has expired. \
356                     This can happen in exceptional cases: a machine goes to sleep or is running out of memory or CPU, for example.",
357                    self.handle.machine.shard_id(),
358                    batch.desc,
359                    self.frontier
360                )
361            }
362        }
363
364        let new_frontier = batch.desc.upper().clone();
365
366        // We will have a new frontier, so this is an opportunity to downgrade our
367        // since capability. Go through `maybe_heartbeat` so we can rate limit
368        // this along with our heartbeats.
369        //
370        // HACK! Everything would be simpler if we could downgrade since to the
371        // new frontier, but we can't. The next call needs to be able to
372        // distinguish between the times T at the frontier (to emit updates with
373        // these times) and T-1 (to filter them). Advancing the since to
374        // frontier would erase the ability to distinguish between them. Ideally
375        // we'd use what is conceptually "batch.upper - 1" (the greatest
376        // elements that are still strictly less than batch.upper, which will be
377        // the new value of self.frontier after this call returns), but the
378        // trait bounds on T don't give us a way to compute that directly.
379        // Instead, we sniff out any elements in self.frontier (the upper of the
380        // batch the last time we called this) that are strictly less_than the
381        // batch upper to compute a new since. For totally ordered times
382        // (currently always the case in mz) self.frontier will always have a
383        // single element and it will be less_than upper. We
384        // could also abuse the fact that every time we actually emit is
385        // guaranteed by definition to be less_than upper to be a bit more
386        // prompt, but this would involve a lot more temporary antichains and
387        // it's unclear if that's worth it.
388        for x in self.frontier.elements().iter() {
389            let less_than_upper = batch.desc.upper().elements().iter().any(|u| x.less_than(u));
390            if less_than_upper {
391                self.since.join_assign(&Antichain::from_elem(x.clone()));
392            }
393        }
394
395        // IMPORTANT! Make sure this `lease_batch_parts` stays before the
396        // `maybe_downgrade_since` call. Otherwise, we might give up our
397        // capability on the batch's SeqNo before we lease it, which could lead
398        // to blobs that it references being GC'd.
399        let filter = FetchBatchFilter::Listen {
400            as_of: self.as_of.clone(),
401            lower: self.frontier.clone(),
402        };
403        let parts = self
404            .handle
405            .lease_batch_parts(lease, batch, filter)
406            .collect()
407            .await;
408
409        self.handle.maybe_downgrade_since(&self.since).await;
410
411        // NB: Keep this after we use self.frontier to join_assign self.since
412        // and also after we construct metadata.
413        self.frontier = new_frontier;
414
415        (parts, self.frontier.clone())
416    }
417}
418
419impl<K, V, T, D> Listen<K, V, T, D>
420where
421    K: Debug + Codec,
422    V: Debug + Codec,
423    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
424    D: Monoid + Codec64 + Send + Sync,
425{
426    /// Attempt to pull out the next values of this subscription.
427    ///
428    /// The updates received in [ListenEvent::Updates] should be assumed to be in arbitrary order
429    /// and not necessarily consolidated. However, the timestamp of each individual update will be
430    /// greater than or equal to the last received [ListenEvent::Progress] frontier (or this
431    /// [Listen]'s initial `as_of` frontier if no progress event has been emitted yet) and less
432    /// than the next [ListenEvent::Progress] frontier.
433    ///
434    /// If you have a use for consolidated listen output, given that snapshots can't be
435    /// consolidated, come talk to us!
436    #[instrument(level = "debug", name = "listen::next", fields(shard = %self.handle.machine.shard_id()))]
437    pub async fn fetch_next(&mut self) -> Vec<ListenEvent<T, ((K, V), T, D)>> {
438        let (parts, progress) = self.next(None).await;
439        let mut ret = Vec::with_capacity(parts.len() + 1);
440        for part in parts {
441            let fetched_part = self.fetch_batch_part(part).await;
442            let updates = fetched_part.collect::<Vec<_>>();
443            if !updates.is_empty() {
444                ret.push(ListenEvent::Updates(updates));
445            }
446        }
447        ret.push(ListenEvent::Progress(progress));
448        ret
449    }
450
451    /// Convert listener into futures::Stream
452    pub fn into_stream(mut self) -> impl Stream<Item = ListenEvent<T, ((K, V), T, D)>> {
453        async_stream::stream!({
454            loop {
455                for msg in self.fetch_next().await {
456                    yield msg;
457                }
458            }
459        })
460    }
461
462    /// Test helper to read from the listener until the given frontier is
463    /// reached. Because compaction can arbitrarily combine batches, we only
464    /// return the final progress info.
465    #[cfg(test)]
466    #[track_caller]
467    pub async fn read_until(&mut self, ts: &T) -> (Vec<((K, V), T, D)>, Antichain<T>) {
468        let mut updates = Vec::new();
469        let mut frontier = Antichain::from_elem(T::minimum());
470        while self.frontier.less_than(ts) {
471            for event in self.fetch_next().await {
472                match event {
473                    ListenEvent::Updates(mut x) => updates.append(&mut x),
474                    ListenEvent::Progress(x) => frontier = x,
475                }
476            }
477        }
478        // Unlike most tests, intentionally don't consolidate updates here
479        // because Listen replays them at the original fidelity.
480        (updates, frontier)
481    }
482}
483
484impl<K, V, T, D> Listen<K, V, T, D>
485where
486    K: Debug + Codec,
487    V: Debug + Codec,
488    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
489    D: Monoid + Codec64 + Send + Sync,
490{
491    /// Fetches the contents of `part` and returns its lease.
492    ///
493    /// This is broken out into its own function to provide a trivial means for
494    /// [`Subscribe`], which contains a [`Listen`], to fetch batches.
495    async fn fetch_batch_part(&mut self, part: LeasedBatchPart<T>) -> FetchedPart<K, V, T, D> {
496        let fetched_part = fetch_leased_part(
497            &self.handle.cfg,
498            &part,
499            self.handle.blob.as_ref(),
500            Arc::clone(&self.handle.metrics),
501            &self.handle.metrics.read.listen,
502            &self.handle.machine.applier.shard_metrics,
503            &self.handle.reader_id,
504            self.handle.read_schemas.clone(),
505            &mut self.handle.schema_cache,
506        )
507        .await;
508        fetched_part
509    }
510
511    /// Politely expires this listen, releasing its lease.
512    ///
513    /// There is a best-effort impl in to expire the
514    /// [`ReadHandle`] held by the listen that wasn't explicitly expired with
515    /// this method. When possible, explicit expiry is still preferred because
516    /// it also ensures that the background task is complete.
517    pub async fn expire(self) {
518        self.handle.expire().await
519    }
520}
521
522/// The state for read holds shared between the heartbeat task and the main client,
523/// including the since frontier and the seqno hold.
524#[derive(Debug)]
525pub(crate) struct ReadHolds<T> {
526    /// The frontier we should hold the time-based lease back to.
527    held_since: Antichain<T>,
528    /// The since hold we've actually committed to state. Should always be <=
529    /// the since hold.
530    applied_since: Antichain<T>,
531    /// The largest seqno we've observed in the state.
532    recent_seqno: SeqNo,
533    /// The set of active leases. We hold back the seqno to the minimum lease or
534    /// the recent_seqno, whichever is earlier.
535    leases: BTreeMap<SeqNo, Lease>,
536    /// True iff this state is expired.
537    expired: bool,
538    /// Used to trigger the background task to heartbeat state, instead of waiting for
539    /// the next tick.
540    request_sync: bool,
541}
542
543impl<T> ReadHolds<T>
544where
545    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
546{
547    pub fn downgrade_since(&mut self, since: &Antichain<T>) {
548        self.held_since.join_assign(since);
549    }
550
551    pub fn observe_seqno(&mut self, seqno: SeqNo) {
552        self.recent_seqno = seqno.max(self.recent_seqno);
553    }
554
555    pub fn lease_seqno(&mut self) -> Lease {
556        let seqno = self.recent_seqno;
557        let lease = self
558            .leases
559            .entry(seqno)
560            .or_insert_with(|| Lease::new(seqno));
561        lease.clone()
562    }
563
564    pub fn outstanding_seqno(&mut self) -> SeqNo {
565        while let Some(first) = self.leases.first_entry() {
566            if first.get().count() <= 1 {
567                first.remove();
568            } else {
569                return *first.key();
570            }
571        }
572        self.recent_seqno
573    }
574}
575
576/// A "capability" granting the ability to read the state of some shard at times
577/// greater or equal to `self.since()`.
578///
579/// Production users should call [Self::expire] before dropping a ReadHandle so
580/// that it can expire its leases. If/when rust gets AsyncDrop, this will be
581/// done automatically.
582///
583/// All async methods on ReadHandle retry for as long as they are able, but the
584/// returned [std::future::Future]s implement "cancel on drop" semantics. This
585/// means that callers can add a timeout using [tokio::time::timeout] or
586/// [tokio::time::timeout_at].
587///
588/// ```rust,no_run
589/// # let mut read: mz_persist_client::read::ReadHandle<String, String, u64, i64> = unimplemented!();
590/// # let timeout: std::time::Duration = unimplemented!();
591/// # let new_since: timely::progress::Antichain<u64> = unimplemented!();
592/// # async {
593/// tokio::time::timeout(timeout, read.downgrade_since(&new_since)).await
594/// # };
595/// ```
596#[derive(Debug)]
597pub struct ReadHandle<K: Codec, V: Codec, T, D> {
598    pub(crate) cfg: PersistConfig,
599    pub(crate) metrics: Arc<Metrics>,
600    pub(crate) machine: Machine<K, V, T, D>,
601    pub(crate) gc: GarbageCollector<K, V, T, D>,
602    pub(crate) blob: Arc<dyn Blob>,
603    watch: StateWatch<K, V, T, D>,
604
605    pub(crate) reader_id: LeasedReaderId,
606    pub(crate) read_schemas: Schemas<K, V>,
607    pub(crate) schema_cache: SchemaCache<K, V, T, D>,
608
609    since: Antichain<T>,
610    pub(crate) hold_state: AwaitableState<ReadHolds<T>>,
611    pub(crate) unexpired_state: Option<UnexpiredReadHandleState>,
612}
613
614/// Length of time after a reader's last operation after which the reader may be
615/// expired.
616pub(crate) const READER_LEASE_DURATION: Config<Duration> = Config::new(
617    "persist_reader_lease_duration",
618    Duration::from_secs(60 * 15),
619    "The time after which we'll clean up stale read leases",
620    ParameterScope::Environment,
621);
622
623impl<K, V, T, D> ReadHandle<K, V, T, D>
624where
625    K: Debug + Codec,
626    V: Debug + Codec,
627    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
628    D: Monoid + Codec64 + Send + Sync,
629{
630    #[allow(clippy::unused_async)]
631    pub(crate) async fn new(
632        cfg: PersistConfig,
633        metrics: Arc<Metrics>,
634        machine: Machine<K, V, T, D>,
635        gc: GarbageCollector<K, V, T, D>,
636        blob: Arc<dyn Blob>,
637        reader_id: LeasedReaderId,
638        read_schemas: Schemas<K, V>,
639        state: LeasedReaderState<T>,
640    ) -> Self {
641        let schema_cache = machine.applier.schema_cache();
642        let hold_state = AwaitableState::new(ReadHolds {
643            held_since: state.since.clone(),
644            applied_since: state.since.clone(),
645            recent_seqno: state.seqno,
646            leases: Default::default(),
647            expired: false,
648            request_sync: false,
649        });
650        ReadHandle {
651            cfg,
652            metrics: Arc::clone(&metrics),
653            machine: machine.clone(),
654            gc: gc.clone(),
655            blob,
656            watch: machine.applier.watch(),
657            reader_id: reader_id.clone(),
658            read_schemas,
659            schema_cache,
660            since: state.since,
661            hold_state: hold_state.clone(),
662            unexpired_state: Some(UnexpiredReadHandleState {
663                heartbeat_task: Self::start_reader_heartbeat_task(
664                    machine, reader_id, gc, hold_state,
665                ),
666            }),
667        }
668    }
669
670    fn start_reader_heartbeat_task(
671        machine: Machine<K, V, T, D>,
672        reader_id: LeasedReaderId,
673        gc: GarbageCollector<K, V, T, D>,
674        leased_seqnos: AwaitableState<ReadHolds<T>>,
675    ) -> JoinHandle<()> {
676        let metrics = Arc::clone(&machine.applier.metrics);
677        let name = format!(
678            "persist::heartbeat_read({},{})",
679            machine.shard_id(),
680            reader_id
681        );
682        mz_ore::task::spawn(|| name, {
683            metrics.tasks.heartbeat_read.instrument_task(async move {
684                Self::reader_heartbeat_task(machine, reader_id, gc, leased_seqnos).await
685            })
686        })
687    }
688
689    async fn reader_heartbeat_task(
690        machine: Machine<K, V, T, D>,
691        reader_id: LeasedReaderId,
692        gc: GarbageCollector<K, V, T, D>,
693        leased_seqnos: AwaitableState<ReadHolds<T>>,
694    ) {
695        let sleep_duration = READER_LEASE_DURATION.get(&machine.applier.cfg) / 4;
696        // Jitter the first tick to avoid a thundering herd when many readers are started around
697        // the same instant, like during deploys.
698        let jitter: f64 = f64::cast_lossy(reader_id.hashed()) / f64::cast_lossy(u64::MAX);
699        let mut interval = tokio::time::interval_at(
700            tokio::time::Instant::now() + sleep_duration.mul_f64(jitter),
701            sleep_duration,
702        );
703        let mut held_since = leased_seqnos.read(|s| s.held_since.clone());
704        loop {
705            let before_sleep = Instant::now();
706            let _woke_by_tick = tokio::select! {
707                _tick = interval.tick() => {
708                    true
709                }
710                _whatever = leased_seqnos.wait_while(|s| !s.request_sync) => {
711                    false
712                }
713            };
714
715            let elapsed_since_before_sleeping = before_sleep.elapsed();
716            if elapsed_since_before_sleeping > sleep_duration + Duration::from_secs(60) {
717                warn!(
718                    "reader ({}) of shard ({}) went {}s between heartbeats",
719                    reader_id,
720                    machine.shard_id(),
721                    elapsed_since_before_sleeping.as_secs_f64()
722                );
723            }
724
725            let before_heartbeat = Instant::now();
726            let current_seqno = machine.seqno();
727            let result = leased_seqnos.modify(|s| {
728                if s.expired {
729                    Err(())
730                } else {
731                    s.observe_seqno(current_seqno);
732                    s.request_sync = false;
733                    held_since.join_assign(&s.held_since);
734                    Ok(s.outstanding_seqno())
735                }
736            });
737            let actual_since = match result {
738                Ok(held_seqno) => {
739                    let (seqno, actual_since, maintenance) = machine
740                        .downgrade_since(&reader_id, held_seqno, &held_since)
741                        .await;
742                    leased_seqnos.modify(|s| {
743                        s.applied_since.clone_from(&actual_since.0);
744                        s.observe_seqno(seqno)
745                    });
746                    maintenance.start_performing(&machine, &gc);
747                    actual_since
748                }
749                Err(()) => {
750                    let (seqno, maintenance) = machine.expire_leased_reader(&reader_id).await;
751                    leased_seqnos.modify(|s| s.observe_seqno(seqno));
752                    maintenance.start_performing(&machine, &gc);
753                    break;
754                }
755            };
756
757            let elapsed_since_heartbeat = before_heartbeat.elapsed();
758            if elapsed_since_heartbeat > Duration::from_secs(60) {
759                warn!(
760                    "reader ({}) of shard ({}) heartbeat call took {}s",
761                    reader_id,
762                    machine.shard_id(),
763                    elapsed_since_heartbeat.as_secs_f64(),
764                );
765            }
766
767            if PartialOrder::less_than(&held_since, &actual_since.0) {
768                // If the read handle was intentionally expired, this task
769                // *should* be aborted before it observes the expiration. So if
770                // we get here, this task somehow failed to keep the read lease
771                // alive. Warn loudly, because there's now a live read handle to
772                // an expired shard that will panic if used, but don't panic,
773                // just in case there is some edge case that results in this
774                // task observing the intentional expiration of a read handle.
775                warn!(
776                    "heartbeat task for reader ({}) of shard ({}) exiting due to expired lease \
777                     while read handle is live",
778                    reader_id,
779                    machine.shard_id(),
780                );
781                return;
782            }
783        }
784    }
785
786    /// This handle's shard id.
787    pub fn shard_id(&self) -> ShardId {
788        self.machine.shard_id()
789    }
790
791    /// This handle's `since` frontier.
792    ///
793    /// This will always be greater or equal to the shard-global `since`.
794    pub fn since(&self) -> &Antichain<T> {
795        &self.since
796    }
797
798    /// A less-stale cached version of the shard-global `upper` frontier.
799    ///
800    /// This is the most recently known upper for this shard process-wide. It
801    /// requires a mutex and a clone. This will always be less or equal to the
802    /// shard-global `upper`.
803    pub fn shared_upper(&self) -> Antichain<T> {
804        self.machine.applier.clone_upper()
805    }
806
807    #[cfg(test)]
808    fn outstanding_seqno(&self) -> SeqNo {
809        let current_seqno = self.machine.seqno();
810        self.hold_state.modify(|s| {
811            s.observe_seqno(current_seqno);
812            s.outstanding_seqno()
813        })
814    }
815
816    /// Forwards the since frontier of this handle, giving up the ability to
817    /// read at times not greater or equal to `new_since`.
818    ///
819    /// This may trigger (asynchronous) compaction and consolidation in the
820    /// system. A `new_since` of the empty antichain "finishes" this shard,
821    /// promising that no more data will ever be read by this handle.
822    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
823    pub async fn downgrade_since(&mut self, new_since: &Antichain<T>) {
824        self.since = new_since.clone();
825        self.hold_state.modify(|s| {
826            s.downgrade_since(new_since);
827            s.request_sync = true;
828        });
829        self.hold_state
830            .wait_while(|s| PartialOrder::less_than(&s.applied_since, new_since))
831            .await;
832    }
833
834    /// Returns an ongoing subscription of updates to a shard.
835    ///
836    /// The stream includes all data at times greater than `as_of`. Combined
837    /// with [Self::snapshot] it will produce exactly correct results: the
838    /// snapshot is the TVCs contents at `as_of` and all subsequent updates
839    /// occur at exactly their indicated time. The recipient should only
840    /// downgrade their read capability when they are certain they have all data
841    /// through the frontier they would downgrade to.
842    ///
843    /// This takes ownership of the ReadHandle so the Listen can use it to
844    /// [Self::downgrade_since] as it progresses. If you need to keep this
845    /// handle, then [Self::clone] it before calling listen.
846    ///
847    /// The `Since` error indicates that the requested `as_of` cannot be served
848    /// (the caller has out of date information) and includes the smallest
849    /// `as_of` that would have been accepted.
850    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
851    pub async fn listen(self, as_of: Antichain<T>) -> Result<Listen<K, V, T, D>, Since<T>> {
852        Listen::new(self, as_of).await
853    }
854
855    async fn snapshot_batches(
856        &mut self,
857        as_of: Antichain<T>,
858    ) -> Result<(Lease, Vec<HollowBatch<T>>), Since<T>> {
859        self.machine
860            .wait_for_upper_past(
861                &as_of,
862                &mut self.watch,
863                Some(&self.reader_id),
864                &self.metrics.retries.snapshot,
865                RetryParameters::persist_defaults(),
866            )
867            .await;
868        let lease = self.lease_seqno().await;
869        let batches = match self.machine.applier.snapshot(&as_of) {
870            Ok(data) => data,
871            Err(SnapshotErr::AsOfHistoricalDistinctionsLost(since)) => return Err(since),
872            Err(SnapshotErr::AsOfNotYetAvailable(seqno, upper)) => {
873                panic!(
874                    "waited for upper past {as_of:?}, but at latest seqno {seqno:?} the frontier was only {upper:?}",
875                    as_of = as_of.elements(),
876                    upper = upper.0.elements(),
877                )
878            }
879        };
880        Ok((lease, batches))
881    }
882
883    /// Returns all of the contents of the shard TVC at `as_of` broken up into
884    /// [`LeasedBatchPart`]es. These parts can be "turned in" via
885    /// `crate::fetch::fetch_batch_part` to receive the data they contain.
886    ///
887    /// This command returns the contents of this shard as of `as_of` once they
888    /// are known. This may "block" (in an async-friendly way) if `as_of` is
889    /// greater or equal to the current `upper` of the shard. The recipient
890    /// should only downgrade their read capability when they are certain they
891    /// have all data through the frontier they would downgrade to.
892    ///
893    /// The `Since` error indicates that the requested `as_of` cannot be served
894    /// (the caller has out of date information) and includes the smallest
895    /// `as_of` that would have been accepted.
896    #[instrument(level = "trace", fields(shard = %self.machine.shard_id()))]
897    pub async fn snapshot(
898        &mut self,
899        as_of: Antichain<T>,
900    ) -> Result<Vec<LeasedBatchPart<T>>, Since<T>> {
901        let (lease, batches) = self.snapshot_batches(as_of.clone()).await?;
902
903        if !PartialOrder::less_equal(self.since(), &as_of) {
904            return Err(Since(self.since().clone()));
905        }
906
907        let filter = FetchBatchFilter::Snapshot { as_of };
908        let mut leased_parts = Vec::new();
909        for batch in batches {
910            // Flatten the HollowBatch into one LeasedBatchPart per key. Each key
911            // corresponds to a "part" or s3 object. This allows persist_source
912            // to distribute work by parts (smallish, more even size) instead of
913            // batches (arbitrarily large).
914            leased_parts.extend(
915                self.lease_batch_parts(lease.clone(), batch, filter.clone())
916                    .collect::<Vec<_>>()
917                    .await,
918            );
919        }
920        Ok(leased_parts)
921    }
922
923    /// Returns a snapshot of all of a shard's data using `as_of`, followed by
924    /// listening to any future updates.
925    ///
926    /// For more details on this operation's semantics, see [Self::snapshot] and
927    /// [Self::listen].
928    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
929    pub async fn subscribe(
930        mut self,
931        as_of: Antichain<T>,
932    ) -> Result<Subscribe<K, V, T, D>, Since<T>> {
933        let snapshot_parts = self.snapshot(as_of.clone()).await?;
934        let listen = self.listen(as_of.clone()).await?;
935        Ok(Subscribe::new(snapshot_parts, listen))
936    }
937
938    fn lease_batch_parts(
939        &mut self,
940        lease: Lease,
941        batch: HollowBatch<T>,
942        filter: FetchBatchFilter<T>,
943    ) -> impl Stream<Item = LeasedBatchPart<T>> + '_ {
944        stream! {
945            let blob = Arc::clone(&self.blob);
946            let metrics = Arc::clone(&self.metrics);
947            let desc = batch.desc.clone();
948            for await part in batch.part_stream(self.shard_id(), &*blob, &*metrics) {
949                yield LeasedBatchPart {
950                    metrics: Arc::clone(&self.metrics),
951                    shard_id: self.machine.shard_id(),
952                    filter: filter.clone(),
953                    desc: desc.clone(),
954                    part: part.expect("leased part").into_owned(),
955                    lease: lease.clone(),
956                    reader_id: self.reader_id.clone(),
957                    filter_pushdown_audit: false,
958                }
959            }
960        }
961    }
962
963    /// Tracks that the `ReadHandle`'s machine's current `SeqNo` is being
964    /// "leased out" to a `LeasedBatchPart`. This ensures that until the reader is expired or the
965    /// [Lease] is dropped, we will hold onto all versions of the state with this sequence number
966    /// or larger. In particular, this means that any batches that are present in the state at
967    /// the leased seqno _or in any future state_ will be preserved as long as the lease is active.
968    ///
969    /// Note that this doesn't present batches from earlier versions of state from being GCed,
970    /// even if they were observed just before the lease call. Callers should generally take a lease
971    /// first, and then inspect state to find any batches to return or process.
972    async fn lease_seqno(&mut self) -> Lease {
973        let current_seqno = self.machine.seqno();
974        let lease = self.hold_state.modify(|s| {
975            s.observe_seqno(current_seqno);
976            s.lease_seqno()
977        });
978        // The seqno we've leased may be the seqno observed by our heartbeat task, which could be
979        // ahead of the last state we saw. Ensure we only observe states in the future of our hold.
980        // (Since these are backed by the same state in the same process, this should all be pretty
981        // fast.)
982        self.watch.wait_for_seqno_ge(lease.seqno()).await;
983        lease
984    }
985
986    /// Returns an independent [ReadHandle] with a new [LeasedReaderId] but the
987    /// same `since`.
988    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
989    pub async fn clone(&self, purpose: &str) -> Self {
990        let new_reader_id = LeasedReaderId::new();
991        let machine = self.machine.clone();
992        let gc = self.gc.clone();
993        let (reader_state, maintenance) = machine
994            .register_leased_reader(
995                &new_reader_id,
996                purpose,
997                READER_LEASE_DURATION.get(&self.cfg),
998                false,
999            )
1000            .await;
1001        maintenance.start_performing(&machine, &gc);
1002        // The point of clone is that you're guaranteed to have the same (or
1003        // greater) since capability, verify that.
1004        // TODO: better if it's the same since capability exactly.
1005        assert!(PartialOrder::less_equal(&reader_state.since, &self.since));
1006        let new_reader = ReadHandle::new(
1007            self.cfg.clone(),
1008            Arc::clone(&self.metrics),
1009            machine,
1010            gc,
1011            Arc::clone(&self.blob),
1012            new_reader_id,
1013            self.read_schemas.clone(),
1014            reader_state,
1015        )
1016        .await;
1017        new_reader
1018    }
1019
1020    /// A rate-limited version of [Self::downgrade_since].
1021    ///
1022    /// Users can call this as frequently as they wish; changes will be periodically
1023    /// flushed to state by the heartbeat task.
1024    #[allow(clippy::unused_async)]
1025    pub async fn maybe_downgrade_since(&mut self, new_since: &Antichain<T>) {
1026        self.since = new_since.clone();
1027        self.hold_state.modify(|s| {
1028            s.downgrade_since(new_since);
1029        });
1030    }
1031
1032    /// Politely expires this reader, releasing its lease.
1033    ///
1034    /// There is a best-effort impl in Drop to expire a reader that wasn't
1035    /// explictly expired with this method. When possible, explicit expiry is
1036    /// still preferred because it also ensures that the background task is complete.
1037    #[instrument(level = "debug", fields(shard = %self.machine.shard_id()))]
1038    pub async fn expire(mut self) {
1039        self.hold_state.modify(|s| {
1040            s.expired = true;
1041            s.request_sync = true;
1042        });
1043        let Some(unexpired_state) = self.unexpired_state.take() else {
1044            return;
1045        };
1046        unexpired_state.heartbeat_task.await;
1047    }
1048
1049    /// Test helper for a [Self::listen] call that is expected to succeed.
1050    #[cfg(test)]
1051    #[track_caller]
1052    pub async fn expect_listen(self, as_of: T) -> Listen<K, V, T, D> {
1053        self.listen(Antichain::from_elem(as_of))
1054            .await
1055            .expect("cannot serve requested as_of")
1056    }
1057}
1058
1059/// State for a read handle that has not been explicitly expired.
1060#[derive(Debug)]
1061pub(crate) struct UnexpiredReadHandleState {
1062    pub(crate) heartbeat_task: JoinHandle<()>,
1063}
1064
1065/// An incremental cursor through a particular shard, returned from [ReadHandle::snapshot_cursor].
1066///
1067/// To read an entire dataset, the
1068/// client should call `next` until it returns `None`, which signals all data has been returned...
1069/// but it's also free to abandon the instance at any time if it eg. only needs a few entries.
1070#[derive(Debug)]
1071pub struct Cursor<K: Codec, V: Codec, T: Timestamp + Codec64, D: Codec64, L = Lease> {
1072    consolidator: Consolidator<T, D, StructuredSort<K, V, T, D>>,
1073    max_len: usize,
1074    max_bytes: usize,
1075    _lease: L,
1076    read_schemas: Schemas<K, V>,
1077}
1078
1079impl<K: Codec, V: Codec, T: Timestamp + Codec64, D: Codec64, L> Cursor<K, V, T, D, L> {
1080    /// Extracts and returns the lease from the cursor. Allowing the caller to
1081    /// do any necessary cleanup associated with the lease.
1082    pub fn into_lease(self: Self) -> L {
1083        self._lease
1084    }
1085}
1086
1087impl<K, V, T, D, L> Cursor<K, V, T, D, L>
1088where
1089    K: Debug + Codec + Ord,
1090    V: Debug + Codec + Ord,
1091    T: Timestamp + Lattice + Codec64 + Sync,
1092    D: Monoid + Ord + Codec64 + Send + Sync,
1093{
1094    /// Grab the next batch of consolidated data.
1095    pub async fn next(&mut self) -> Option<impl Iterator<Item = ((K, V), T, D)> + '_> {
1096        let Self {
1097            consolidator,
1098            max_len,
1099            max_bytes,
1100            _lease,
1101            read_schemas: _,
1102        } = self;
1103
1104        let part = consolidator
1105            .next_chunk(*max_len, *max_bytes)
1106            .await
1107            .expect("fetching a leased part")?;
1108        let key_decoder = self
1109            .read_schemas
1110            .key
1111            .decoder_any(part.key.as_ref())
1112            .expect("ok");
1113        let val_decoder = self
1114            .read_schemas
1115            .val
1116            .decoder_any(part.val.as_ref())
1117            .expect("ok");
1118        let iter = (0..part.len()).map(move |i| {
1119            let mut k = K::default();
1120            let mut v = V::default();
1121            key_decoder.decode(i, &mut k);
1122            val_decoder.decode(i, &mut v);
1123            let t = T::decode(part.time.value(i).to_le_bytes());
1124            let d = D::decode(part.diff.value(i).to_le_bytes());
1125            ((k, v), t, d)
1126        });
1127
1128        Some(iter)
1129    }
1130}
1131
1132impl<K, V, T, D> ReadHandle<K, V, T, D>
1133where
1134    K: Debug + Codec + Ord,
1135    V: Debug + Codec + Ord,
1136    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
1137    D: Monoid + Ord + Codec64 + Send + Sync,
1138{
1139    /// Generates a [Self::snapshot], and fetches all of the batches it
1140    /// contains.
1141    ///
1142    /// The output is consolidated. Furthermore, to keep memory usage down when
1143    /// reading a snapshot that consolidates well, this consolidates as it goes.
1144    ///
1145    /// Potential future improvements (if necessary):
1146    /// - Accept something like a `F: Fn(K,V) -> (K,V)` argument, which looks
1147    ///   like an MFP you might be pushing down. Reason being that if you are
1148    ///   projecting or transforming in a way that allows further consolidation,
1149    ///   amazing.
1150    /// - Reuse any code we write to streaming-merge consolidate in
1151    ///   persist_source here.
1152    pub async fn snapshot_and_fetch(
1153        &mut self,
1154        as_of: Antichain<T>,
1155    ) -> Result<Vec<((K, V), T, D)>, Since<T>> {
1156        let mut cursor = self.snapshot_cursor(as_of, |_| true).await?;
1157        let mut contents = Vec::new();
1158        while let Some(iter) = cursor.next().await {
1159            contents.extend(iter);
1160        }
1161
1162        // We don't currently guarantee that encoding is one-to-one, so we still need to
1163        // consolidate the decoded outputs. However, let's report if this isn't a noop.
1164        let old_len = contents.len();
1165        consolidate_updates(&mut contents);
1166        if old_len != contents.len() {
1167            // TODO(bkirwi): do we need more / finer-grained metrics for this?
1168            self.machine
1169                .applier
1170                .shard_metrics
1171                .unconsolidated_snapshot
1172                .inc();
1173        }
1174
1175        Ok(contents)
1176    }
1177
1178    /// Generates a [Self::snapshot], and fetches all of the batches it
1179    /// contains.
1180    ///
1181    /// To keep memory usage down when reading a snapshot that consolidates well, this consolidates
1182    /// as it goes. However, note that only the serialized data is consolidated: the deserialized
1183    /// data will only be consolidated if your K/V codecs are one-to-one.
1184    pub async fn snapshot_cursor(
1185        &mut self,
1186        as_of: Antichain<T>,
1187        should_fetch_part: impl for<'a> Fn(Option<&'a LazyPartStats>) -> bool,
1188    ) -> Result<Cursor<K, V, T, D>, Since<T>> {
1189        let (lease, batches) = self.snapshot_batches(as_of.clone()).await?;
1190
1191        Self::read_batches_consolidated(
1192            &self.cfg,
1193            Arc::clone(&self.metrics),
1194            Arc::clone(&self.machine.applier.shard_metrics),
1195            self.metrics.read.snapshot.clone(),
1196            Arc::clone(&self.blob),
1197            self.shard_id(),
1198            as_of,
1199            self.read_schemas.clone(),
1200            &batches,
1201            lease,
1202            should_fetch_part,
1203            COMPACTION_MEMORY_BOUND_BYTES.get(&self.cfg),
1204        )
1205    }
1206
1207    pub(crate) fn read_batches_consolidated<L>(
1208        persist_cfg: &PersistConfig,
1209        metrics: Arc<Metrics>,
1210        shard_metrics: Arc<ShardMetrics>,
1211        read_metrics: ReadMetrics,
1212        blob: Arc<dyn Blob>,
1213        shard_id: ShardId,
1214        as_of: Antichain<T>,
1215        schemas: Schemas<K, V>,
1216        batches: &[HollowBatch<T>],
1217        lease: L,
1218        should_fetch_part: impl for<'a> Fn(Option<&'a LazyPartStats>) -> bool,
1219        memory_budget_bytes: usize,
1220    ) -> Result<Cursor<K, V, T, D, L>, Since<T>> {
1221        let context = format!("{}[as_of={:?}]", shard_id, as_of.elements());
1222        let filter = FetchBatchFilter::Snapshot {
1223            as_of: as_of.clone(),
1224        };
1225
1226        let mut consolidator = Consolidator::new(
1227            context,
1228            FetchConfig::from_persist_config(persist_cfg),
1229            shard_id,
1230            StructuredSort::new(schemas.clone()),
1231            blob,
1232            metrics,
1233            shard_metrics,
1234            read_metrics,
1235            filter,
1236            None,
1237            memory_budget_bytes,
1238        );
1239        for batch in batches {
1240            for (meta, run) in batch.runs() {
1241                consolidator.enqueue_run(
1242                    &batch.desc,
1243                    meta,
1244                    run.into_iter()
1245                        .filter(|p| should_fetch_part(p.stats()))
1246                        .cloned(),
1247                );
1248            }
1249        }
1250        // This default may end up consolidating more records than previously
1251        // for cases like fast-path peeks, where only the first few entries are used.
1252        // If this is a noticeable performance impact, thread the max-len in from the caller.
1253        let max_len = persist_cfg.compaction_yield_after_n_updates;
1254        let max_bytes = BLOB_TARGET_SIZE.get(persist_cfg).max(1);
1255
1256        Ok(Cursor {
1257            consolidator,
1258            max_len,
1259            max_bytes,
1260            _lease: lease,
1261            read_schemas: schemas,
1262        })
1263    }
1264
1265    /// Returns aggregate statistics about the contents of the shard TVC at the
1266    /// given frontier.
1267    ///
1268    /// This command returns the contents of this shard as of `as_of` once they
1269    /// are known. This may "block" (in an async-friendly way) if `as_of` is
1270    /// greater or equal to the current `upper` of the shard. If `None` is given
1271    /// for `as_of`, then the latest stats known by this process are used.
1272    ///
1273    /// The `Since` error indicates that the requested `as_of` cannot be served
1274    /// (the caller has out of date information) and includes the smallest
1275    /// `as_of` that would have been accepted.
1276    pub fn snapshot_stats(
1277        &self,
1278        as_of: Option<Antichain<T>>,
1279    ) -> impl Future<Output = Result<SnapshotStats, Since<T>>> + Send + 'static {
1280        let machine = self.machine.clone();
1281        async move {
1282            let batches = match as_of {
1283                Some(as_of) => machine.unleased_snapshot(&as_of).await?,
1284                None => machine.applier.all_batches(),
1285            };
1286            let num_updates = batches.iter().map(|b| b.len).sum();
1287            Ok(SnapshotStats {
1288                shard_id: machine.shard_id(),
1289                num_updates,
1290            })
1291        }
1292    }
1293
1294    /// Returns aggregate statistics about the contents of the shard TVC at the
1295    /// given frontier.
1296    ///
1297    /// This command returns the contents of this shard as of `as_of` once they
1298    /// are known. This may "block" (in an async-friendly way) if `as_of` is
1299    /// greater or equal to the current `upper` of the shard.
1300    ///
1301    /// The `Since` error indicates that the requested `as_of` cannot be served
1302    /// (the caller has out of date information) and includes the smallest
1303    /// `as_of` that would have been accepted.
1304    pub async fn snapshot_parts_stats(
1305        &self,
1306        as_of: Antichain<T>,
1307    ) -> Result<SnapshotPartsStats, Since<T>> {
1308        let batches = self.machine.unleased_snapshot(&as_of).await?;
1309        let parts = stream::iter(&batches)
1310            .flat_map(|b| b.part_stream(self.shard_id(), &*self.blob, &*self.metrics))
1311            .map(|p| {
1312                let p = p.expect("live batch");
1313                SnapshotPartStats {
1314                    encoded_size_bytes: p.encoded_size_bytes(),
1315                    stats: p.stats().cloned(),
1316                }
1317            })
1318            .collect()
1319            .await;
1320        Ok(SnapshotPartsStats {
1321            metrics: Arc::clone(&self.machine.applier.metrics),
1322            shard_id: self.machine.shard_id(),
1323            parts,
1324        })
1325    }
1326}
1327
1328impl<K, V, T, D> ReadHandle<K, V, T, D>
1329where
1330    K: Debug + Codec + Ord,
1331    V: Debug + Codec + Ord,
1332    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
1333    D: Monoid + Codec64 + Send + Sync,
1334{
1335    /// Generates a [Self::snapshot], and streams out all of the updates
1336    /// it contains in bounded memory.
1337    ///
1338    /// The output is not consolidated.
1339    pub async fn snapshot_and_stream(
1340        &mut self,
1341        as_of: Antichain<T>,
1342    ) -> Result<impl Stream<Item = ((K, V), T, D)> + use<K, V, T, D>, Since<T>> {
1343        let snap = self.snapshot(as_of).await?;
1344
1345        let blob = Arc::clone(&self.blob);
1346        let metrics = Arc::clone(&self.metrics);
1347        let snapshot_metrics = self.metrics.read.snapshot.clone();
1348        let shard_metrics = Arc::clone(&self.machine.applier.shard_metrics);
1349        let reader_id = self.reader_id.clone();
1350        let schemas = self.read_schemas.clone();
1351        let mut schema_cache = self.schema_cache.clone();
1352        let persist_cfg = self.cfg.clone();
1353        let stream = async_stream::stream! {
1354            for part in snap {
1355                let mut fetched_part = fetch_leased_part(
1356                    &persist_cfg,
1357                    &part,
1358                    blob.as_ref(),
1359                    Arc::clone(&metrics),
1360                    &snapshot_metrics,
1361                    &shard_metrics,
1362                    &reader_id,
1363                    schemas.clone(),
1364                    &mut schema_cache,
1365                )
1366                .await;
1367
1368                while let Some(next) = fetched_part.next() {
1369                    yield next;
1370                }
1371            }
1372        };
1373
1374        Ok(stream)
1375    }
1376}
1377
1378impl<K, V, T, D> ReadHandle<K, V, T, D>
1379where
1380    K: Debug + Codec + Ord,
1381    V: Debug + Codec + Ord,
1382    T: Timestamp + TotalOrder + Lattice + Codec64 + Sync,
1383    D: Monoid + Ord + Codec64 + Send + Sync,
1384{
1385    /// Test helper to generate a [Self::snapshot] call that is expected to
1386    /// succeed, process its batches, and then return its data sorted.
1387    #[cfg(test)]
1388    #[track_caller]
1389    pub async fn expect_snapshot_and_fetch(&mut self, as_of: T) -> Vec<((K, V), T, D)> {
1390        let mut ret = self
1391            .snapshot_and_fetch(Antichain::from_elem(as_of))
1392            .await
1393            .expect("cannot serve requested as_of");
1394
1395        ret.sort();
1396        ret
1397    }
1398}
1399
1400impl<K: Codec, V: Codec, T, D> Drop for ReadHandle<K, V, T, D> {
1401    fn drop(&mut self) {
1402        self.hold_state.modify(|s| {
1403            s.expired = true;
1404            s.request_sync = true;
1405        });
1406    }
1407}
1408
1409#[cfg(test)]
1410mod tests {
1411    use std::pin;
1412    use std::str::FromStr;
1413
1414    use mz_dyncfg::ConfigUpdates;
1415    use mz_ore::cast::CastFrom;
1416    use mz_ore::metrics::MetricsRegistry;
1417    use mz_persist::mem::{MemBlob, MemBlobConfig, MemConsensus};
1418    use mz_persist::unreliable::{UnreliableConsensus, UnreliableHandle};
1419    use serde::{Deserialize, Serialize};
1420    use serde_json::json;
1421    use tokio_stream::StreamExt;
1422
1423    use crate::async_runtime::IsolatedRuntime;
1424    use crate::batch::BLOB_TARGET_SIZE;
1425    use crate::cache::StateCache;
1426    use crate::internal::metrics::Metrics;
1427    use crate::rpc::NoopPubSubSender;
1428    use crate::tests::{all_ok, new_test_client};
1429    use crate::{Diagnostics, PersistClient, PersistConfig, ShardId};
1430
1431    use super::*;
1432
1433    // Verifies `Subscribe` can be dropped while holding snapshot batches.
1434    #[mz_persist_proc::test(tokio::test)]
1435    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1436    async fn drop_unused_subscribe(dyncfgs: ConfigUpdates) {
1437        let data = [
1438            (("0".to_owned(), "zero".to_owned()), 0, 1),
1439            (("1".to_owned(), "one".to_owned()), 1, 1),
1440            (("2".to_owned(), "two".to_owned()), 2, 1),
1441        ];
1442
1443        let (mut write, read) = new_test_client(&dyncfgs)
1444            .await
1445            .expect_open::<String, String, u64, i64>(crate::ShardId::new())
1446            .await;
1447
1448        write.expect_compare_and_append(&data[0..1], 0, 1).await;
1449        write.expect_compare_and_append(&data[1..2], 1, 2).await;
1450        write.expect_compare_and_append(&data[2..3], 2, 3).await;
1451
1452        let subscribe = read
1453            .subscribe(timely::progress::Antichain::from_elem(2))
1454            .await
1455            .unwrap();
1456        assert!(
1457            !subscribe.snapshot.as_ref().unwrap().is_empty(),
1458            "snapshot must have batches for test to be meaningful"
1459        );
1460        drop(subscribe);
1461    }
1462
1463    // Verifies that we streaming-consolidate away identical key-values in the same batch.
1464    #[mz_persist_proc::test(tokio::test)]
1465    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1466    async fn streaming_consolidate(dyncfgs: ConfigUpdates) {
1467        let data = &[
1468            // Identical records should sum together...
1469            (("k".to_owned(), "v".to_owned()), 0, 1),
1470            (("k".to_owned(), "v".to_owned()), 1, 1),
1471            (("k".to_owned(), "v".to_owned()), 2, 1),
1472            // ...and when they cancel out entirely they should be omitted.
1473            (("k2".to_owned(), "v".to_owned()), 0, 1),
1474            (("k2".to_owned(), "v".to_owned()), 1, -1),
1475        ];
1476
1477        let (mut write, read) = {
1478            let client = new_test_client(&dyncfgs).await;
1479            client.cfg.set_config(&BLOB_TARGET_SIZE, 1000); // So our batch stays together!
1480            client
1481                .expect_open::<String, String, u64, i64>(crate::ShardId::new())
1482                .await
1483        };
1484
1485        write.expect_compare_and_append(data, 0, 5).await;
1486
1487        let mut snapshot = read
1488            .subscribe(timely::progress::Antichain::from_elem(4))
1489            .await
1490            .unwrap();
1491
1492        let mut updates = vec![];
1493        'outer: loop {
1494            for event in snapshot.fetch_next().await {
1495                match event {
1496                    ListenEvent::Progress(t) => {
1497                        if !t.less_than(&4) {
1498                            break 'outer;
1499                        }
1500                    }
1501                    ListenEvent::Updates(data) => {
1502                        updates.extend(data);
1503                    }
1504                }
1505            }
1506        }
1507        assert_eq!(updates, &[(("k".to_owned(), "v".to_owned()), 4u64, 3i64)],)
1508    }
1509
1510    #[mz_persist_proc::test(tokio::test)]
1511    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
1512    async fn snapshot_and_stream(dyncfgs: ConfigUpdates) {
1513        let data = &mut [
1514            (("k1".to_owned(), "v1".to_owned()), 0, 1),
1515            (("k2".to_owned(), "v2".to_owned()), 1, 1),
1516            (("k3".to_owned(), "v3".to_owned()), 2, 1),
1517            (("k4".to_owned(), "v4".to_owned()), 2, 1),
1518            (("k5".to_owned(), "v5".to_owned()), 3, 1),
1519        ];
1520
1521        let (mut write, mut read) = {
1522            let client = new_test_client(&dyncfgs).await;
1523            client.cfg.set_config(&BLOB_TARGET_SIZE, 0); // split batches across multiple parts
1524            client
1525                .expect_open::<String, String, u64, i64>(crate::ShardId::new())
1526                .await
1527        };
1528
1529        write.expect_compare_and_append(&data[0..2], 0, 2).await;
1530        write.expect_compare_and_append(&data[2..4], 2, 3).await;
1531        write.expect_compare_and_append(&data[4..], 3, 4).await;
1532
1533        let as_of = Antichain::from_elem(3);
1534        let mut snapshot = pin::pin!(read.snapshot_and_stream(as_of.clone()).await.unwrap());
1535
1536        let mut snapshot_rows = vec![];
1537        while let Some(((k, v), t, d)) = snapshot.next().await {
1538            snapshot_rows.push(((k, v), t, d));
1539        }
1540
1541        for ((_k, _v), t, _d) in data.as_mut_slice() {
1542            t.advance_by(as_of.borrow());
1543        }
1544
1545        assert_eq!(data.as_slice(), snapshot_rows.as_slice());
1546    }
1547
1548    // Verifies the semantics of `SeqNo` leases + checks dropping `LeasedBatchPart` semantics.
1549    #[mz_persist_proc::test(tokio::test)]
1550    #[cfg_attr(miri, ignore)] // https://github.com/MaterializeInc/database-issues/issues/5964
1551    async fn seqno_leases(dyncfgs: ConfigUpdates) {
1552        let mut data = vec![];
1553        for i in 0..20 {
1554            data.push(((i.to_string(), i.to_string()), i, 1))
1555        }
1556
1557        let shard_id = ShardId::new();
1558
1559        let client = new_test_client(&dyncfgs).await;
1560        let (mut write, read) = client
1561            .expect_open::<String, String, u64, i64>(shard_id)
1562            .await;
1563
1564        // Seed with some values
1565        let mut offset = 0;
1566        let mut width = 2;
1567
1568        for i in offset..offset + width {
1569            write
1570                .expect_compare_and_append(
1571                    &data[i..i + 1],
1572                    u64::cast_from(i),
1573                    u64::cast_from(i) + 1,
1574                )
1575                .await;
1576        }
1577        offset += width;
1578
1579        // Create machinery for subscribe + fetch
1580        let mut fetcher = client
1581            .create_batch_fetcher::<String, String, u64, i64>(
1582                shard_id,
1583                Default::default(),
1584                Default::default(),
1585                false,
1586                Diagnostics::for_tests(),
1587            )
1588            .await
1589            .unwrap();
1590
1591        let mut subscribe = read
1592            .subscribe(timely::progress::Antichain::from_elem(1))
1593            .await
1594            .expect("cannot serve requested as_of");
1595
1596        // Determine the sequence number held by our subscribe.
1597        let original_seqno_since = subscribe.listen.handle.outstanding_seqno();
1598        if let Some(snapshot) = &subscribe.snapshot {
1599            for part in snapshot {
1600                assert!(
1601                    part.lease.seqno() >= original_seqno_since,
1602                    "our seqno hold must cover all parts"
1603                );
1604            }
1605        }
1606
1607        let mut parts = vec![];
1608
1609        width = 4;
1610        // Collect parts while continuing to write values
1611        for i in offset..offset + width {
1612            for event in subscribe.next(None).await {
1613                if let ListenEvent::Updates(mut new_parts) = event {
1614                    parts.append(&mut new_parts);
1615                    // Here and elsewhere we "cheat" and immediately downgrade the since
1616                    // to demonstrate the effects of SeqNo leases immediately.
1617                    subscribe
1618                        .listen
1619                        .handle
1620                        .downgrade_since(&subscribe.listen.since)
1621                        .await;
1622                }
1623            }
1624
1625            write
1626                .expect_compare_and_append(
1627                    &data[i..i + 1],
1628                    u64::cast_from(i),
1629                    u64::cast_from(i) + 1,
1630                )
1631                .await;
1632
1633            // SeqNo is not downgraded
1634            assert_eq!(
1635                subscribe.listen.handle.machine.applier.seqno_since(),
1636                original_seqno_since
1637            );
1638        }
1639
1640        offset += width;
1641
1642        let mut seqno_since = subscribe.listen.handle.machine.applier.seqno_since();
1643
1644        // We're starting out with the original, non-downgraded SeqNo
1645        assert_eq!(seqno_since, original_seqno_since);
1646
1647        // We have to handle the parts we generate during the next loop to
1648        // ensure they don't panic.
1649        let mut subsequent_parts = vec![];
1650
1651        // Ensure monotonicity of seqnos we're processing, otherwise the
1652        // invariant we're testing (returning the last part of a seqno will
1653        // downgrade its since) will not hold.
1654        let mut this_seqno = SeqNo::minimum();
1655
1656        // Repeat the same process as above, more or less, while fetching + returning parts
1657        for (mut i, part) in parts.into_iter().enumerate() {
1658            let part_seqno = part.lease.seqno();
1659            let last_seqno = this_seqno;
1660            this_seqno = part_seqno;
1661            assert!(this_seqno >= last_seqno);
1662
1663            let (part, lease) = part.into_exchangeable_part();
1664            let _ = fetcher.fetch_leased_part(part).await;
1665            drop(lease);
1666
1667            // Simulates an exchange
1668            for event in subscribe.next(None).await {
1669                if let ListenEvent::Updates(parts) = event {
1670                    for part in parts {
1671                        let (_, lease) = part.into_exchangeable_part();
1672                        subsequent_parts.push(lease);
1673                    }
1674                }
1675            }
1676
1677            subscribe
1678                .listen
1679                .handle
1680                .downgrade_since(&subscribe.listen.since)
1681                .await;
1682
1683            // Write more new values
1684            i += offset;
1685            write
1686                .expect_compare_and_append(
1687                    &data[i..i + 1],
1688                    u64::cast_from(i),
1689                    u64::cast_from(i) + 1,
1690                )
1691                .await;
1692
1693            // We should expect the SeqNo to be downgraded if this part's SeqNo
1694            // is no longer leased to any other parts, either.
1695            let expect_downgrade = subscribe.listen.handle.outstanding_seqno() > part_seqno;
1696
1697            let new_seqno_since = subscribe.listen.handle.machine.applier.seqno_since();
1698            if expect_downgrade {
1699                assert!(new_seqno_since > seqno_since);
1700            } else {
1701                assert_eq!(new_seqno_since, seqno_since);
1702            }
1703            seqno_since = new_seqno_since;
1704        }
1705
1706        // SeqNo since was downgraded
1707        assert!(seqno_since > original_seqno_since);
1708
1709        // Return any outstanding parts, to prevent a panic!
1710        drop(subsequent_parts);
1711        drop(subscribe);
1712    }
1713
1714    #[mz_ore::test]
1715    fn reader_id_human_readable_serde() {
1716        #[derive(Debug, Serialize, Deserialize)]
1717        struct Container {
1718            reader_id: LeasedReaderId,
1719        }
1720
1721        // roundtrip through json
1722        let id =
1723            LeasedReaderId::from_str("r00000000-1234-5678-0000-000000000000").expect("valid id");
1724        assert_eq!(
1725            id,
1726            serde_json::from_value(serde_json::to_value(id.clone()).expect("serializable"))
1727                .expect("deserializable")
1728        );
1729
1730        // deserialize a serialized string directly
1731        assert_eq!(
1732            id,
1733            serde_json::from_str("\"r00000000-1234-5678-0000-000000000000\"")
1734                .expect("deserializable")
1735        );
1736
1737        // roundtrip id through a container type
1738        let json = json!({ "reader_id": id });
1739        assert_eq!(
1740            "{\"reader_id\":\"r00000000-1234-5678-0000-000000000000\"}",
1741            &json.to_string()
1742        );
1743        let container: Container = serde_json::from_value(json).expect("deserializable");
1744        assert_eq!(container.reader_id, id);
1745    }
1746
1747    // Verifies performance optimizations where a Listener doesn't fetch the
1748    // latest Consensus state if the one it currently has can serve the next
1749    // request.
1750    #[mz_ore::test(tokio::test)]
1751    #[cfg_attr(miri, ignore)] // too slow
1752    async fn skip_consensus_fetch_optimization() {
1753        let data = vec![
1754            (("0".to_owned(), "zero".to_owned()), 0, 1),
1755            (("1".to_owned(), "one".to_owned()), 1, 1),
1756            (("2".to_owned(), "two".to_owned()), 2, 1),
1757        ];
1758
1759        let cfg = PersistConfig::new_for_tests();
1760        let blob = Arc::new(MemBlob::open(MemBlobConfig::default()));
1761        let consensus = Arc::new(MemConsensus::default());
1762        let unreliable = UnreliableHandle::default();
1763        unreliable.totally_available();
1764        let consensus = Arc::new(UnreliableConsensus::new(consensus, unreliable.clone()));
1765        let metrics = Arc::new(Metrics::new(&cfg, &MetricsRegistry::new()));
1766        let pubsub_sender = Arc::new(NoopPubSubSender);
1767        let (mut write, mut read) = PersistClient::new(
1768            cfg,
1769            blob,
1770            consensus,
1771            metrics,
1772            Arc::new(IsolatedRuntime::new_for_tests()),
1773            Arc::new(StateCache::new_no_metrics()),
1774            pubsub_sender,
1775        )
1776        .expect("client construction failed")
1777        .expect_open::<String, String, u64, i64>(ShardId::new())
1778        .await;
1779
1780        write.expect_compare_and_append(&data[0..1], 0, 1).await;
1781        write.expect_compare_and_append(&data[1..2], 1, 2).await;
1782        write.expect_compare_and_append(&data[2..3], 2, 3).await;
1783
1784        let snapshot = read.expect_snapshot_and_fetch(2).await;
1785        let mut listen = read.expect_listen(0).await;
1786
1787        // Manually advance the listener's machine so that it has the latest
1788        // state by fetching the first events from next. This is awkward but
1789        // only necessary because we're about to do some weird things with
1790        // unreliable.
1791        let listen_actual = listen.fetch_next().await;
1792        let expected_events = vec![ListenEvent::Progress(Antichain::from_elem(1))];
1793        assert_eq!(listen_actual, expected_events);
1794
1795        // At this point, the snapshot and listen's state should have all the
1796        // writes. Test this by making consensus completely unavailable.
1797        unreliable.totally_unavailable();
1798        assert_eq!(snapshot, all_ok(&data, 2));
1799        assert_eq!(
1800            listen.read_until(&3).await,
1801            (all_ok(&data[1..], 1), Antichain::from_elem(3))
1802        );
1803    }
1804}