Skip to main content

mz_persist_client/internal/
state.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
10use anyhow::ensure;
11use async_stream::{stream, try_stream};
12use differential_dataflow::difference::Monoid;
13use mz_persist::metrics::ColumnarMetrics;
14use proptest::prelude::{Arbitrary, Strategy};
15use std::borrow::Cow;
16use std::cmp::Ordering;
17use std::collections::BTreeMap;
18use std::fmt::{Debug, Formatter};
19use std::marker::PhantomData;
20use std::ops::ControlFlow::{self, Break, Continue};
21use std::ops::{Deref, DerefMut};
22use std::time::Duration;
23
24use arrow::array::{Array, ArrayData, make_array};
25use arrow::datatypes::DataType;
26use bytes::Bytes;
27use differential_dataflow::Hashable;
28use differential_dataflow::lattice::Lattice;
29use differential_dataflow::trace::Description;
30use differential_dataflow::trace::implementations::BatchContainer;
31use futures::Stream;
32use futures_util::StreamExt;
33use itertools::Itertools;
34use mz_dyncfg::{Config, ParameterScope};
35use mz_ore::cast::CastFrom;
36use mz_ore::now::EpochMillis;
37use mz_ore::soft_panic_or_log;
38use mz_ore::vec::PartialOrdVecExt;
39use mz_persist::indexed::encoding::{BatchColumnarFormat, BlobTraceUpdates};
40use mz_persist::location::{Blob, SeqNo};
41use mz_persist_types::arrow::{ArrayBound, ProtoArrayData};
42use mz_persist_types::columnar::{ColumnEncoder, Schema};
43use mz_persist_types::schema::{SchemaId, backward_compatible};
44use mz_persist_types::{Codec, Codec64};
45use mz_proto::ProtoType;
46use mz_proto::RustType;
47use proptest_derive::Arbitrary;
48use semver::Version;
49use serde::ser::SerializeStruct;
50use serde::{Serialize, Serializer};
51use timely::PartialOrder;
52use timely::order::TotalOrder;
53use timely::progress::{Antichain, Timestamp};
54use tracing::info;
55use uuid::Uuid;
56
57use crate::critical::{CriticalReaderId, Opaque};
58use crate::error::InvalidUsage;
59use crate::internal::encoding::{
60    LazyInlineBatchPart, LazyPartStats, LazyProto, MetadataKey, MetadataMap, parse_id,
61};
62use crate::internal::gc::GcReq;
63use crate::internal::machine::retry_external;
64use crate::internal::paths::{BlobKey, PartId, PartialBatchKey, PartialRollupKey, WriterKey};
65use crate::internal::trace::{
66    ActiveCompaction, ApplyMergeResult, FueledMergeReq, FueledMergeRes, Trace,
67};
68use crate::metrics::Metrics;
69use crate::read::LeasedReaderId;
70use crate::schema::CaESchema;
71use crate::write::WriterId;
72use crate::{PersistConfig, ShardId};
73
74include!(concat!(
75    env!("OUT_DIR"),
76    "/mz_persist_client.internal.state.rs"
77));
78
79include!(concat!(
80    env!("OUT_DIR"),
81    "/mz_persist_client.internal.diff.rs"
82));
83
84/// Determines how often to write rollups, assigning a maintenance task after
85/// `rollup_threshold` seqnos have passed since the last rollup.
86///
87/// Tuning note: in the absence of a long reader seqno hold, and with
88/// incremental GC, this threshold will determine about how many live diffs are
89/// held in Consensus. Lowering this value decreases the live diff count at the
90/// cost of more maintenance work + blob writes.
91pub(crate) const ROLLUP_THRESHOLD: Config<usize> = Config::new(
92    "persist_rollup_threshold",
93    128,
94    "The number of seqnos between rollups.",
95    ParameterScope::Environment,
96);
97
98/// Determines how long to wait before an active rollup is considered
99/// "stuck" and a new rollup is started.
100pub(crate) const ROLLUP_FALLBACK_THRESHOLD_MS: Config<usize> = Config::new(
101    "persist_rollup_fallback_threshold_ms",
102    5000,
103    "The number of milliseconds before a worker claims an already claimed rollup.",
104    ParameterScope::Environment,
105);
106
107/// Feature flag the new active rollup tracking mechanism.
108/// We musn't enable this until we are fully deployed on the new version.
109pub(crate) const ROLLUP_USE_ACTIVE_ROLLUP: Config<bool> = Config::new(
110    "persist_rollup_use_active_rollup",
111    true,
112    "Whether to use the new active rollup tracking mechanism.",
113    ParameterScope::Environment,
114);
115
116/// Determines how long to wait before an active GC is considered
117/// "stuck" and a new GC is started.
118pub(crate) const GC_FALLBACK_THRESHOLD_MS: Config<usize> = Config::new(
119    "persist_gc_fallback_threshold_ms",
120    900000,
121    "The number of milliseconds before a worker claims an already claimed GC.",
122    ParameterScope::Environment,
123);
124
125/// See the config description string.
126pub(crate) const GC_MIN_VERSIONS: Config<usize> = Config::new(
127    "persist_gc_min_versions",
128    32,
129    "The number of un-GCd versions that may exist in state before we'll trigger a GC.",
130    ParameterScope::Environment,
131);
132
133/// See the config description string.
134pub(crate) const GC_MAX_VERSIONS: Config<usize> = Config::new(
135    "persist_gc_max_versions",
136    128_000,
137    "The maximum number of versions to GC in a single GC run.",
138    ParameterScope::Environment,
139);
140
141/// Feature flag the new active GC tracking mechanism.
142/// We musn't enable this until we are fully deployed on the new version.
143pub(crate) const GC_USE_ACTIVE_GC: Config<bool> = Config::new(
144    "persist_gc_use_active_gc",
145    false,
146    "Whether to use the new active GC tracking mechanism.",
147    ParameterScope::Environment,
148);
149
150pub(crate) const ENABLE_INCREMENTAL_COMPACTION: Config<bool> = Config::new(
151    "persist_enable_incremental_compaction",
152    false,
153    "Whether to enable incremental compaction.",
154    ParameterScope::Environment,
155);
156
157/// A token to disambiguate state commands that could not otherwise be
158/// idempotent.
159#[derive(Arbitrary, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
160#[serde(into = "String")]
161pub struct IdempotencyToken(pub(crate) [u8; 16]);
162
163impl std::fmt::Display for IdempotencyToken {
164    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
165        write!(f, "i{}", Uuid::from_bytes(self.0))
166    }
167}
168
169impl std::fmt::Debug for IdempotencyToken {
170    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
171        write!(f, "IdempotencyToken({})", Uuid::from_bytes(self.0))
172    }
173}
174
175impl std::str::FromStr for IdempotencyToken {
176    type Err = String;
177
178    fn from_str(s: &str) -> Result<Self, Self::Err> {
179        parse_id("i", "IdempotencyToken", s).map(IdempotencyToken)
180    }
181}
182
183impl From<IdempotencyToken> for String {
184    fn from(x: IdempotencyToken) -> Self {
185        x.to_string()
186    }
187}
188
189impl IdempotencyToken {
190    pub(crate) fn new() -> Self {
191        IdempotencyToken(*Uuid::new_v4().as_bytes())
192    }
193    pub(crate) const SENTINEL: IdempotencyToken = IdempotencyToken([17u8; 16]);
194}
195
196#[derive(Clone, Debug, PartialEq, Serialize)]
197pub struct LeasedReaderState<T> {
198    /// The seqno capability of this reader.
199    pub seqno: SeqNo,
200    /// The since capability of this reader.
201    pub since: Antichain<T>,
202    /// UNIX_EPOCH timestamp (in millis) of this reader's most recent heartbeat
203    pub last_heartbeat_timestamp_ms: u64,
204    /// Duration (in millis) allowed after [Self::last_heartbeat_timestamp_ms]
205    /// after which this reader may be expired
206    pub lease_duration_ms: u64,
207    /// For debugging.
208    pub debug: HandleDebugState,
209}
210
211#[derive(Clone, Debug, PartialEq, Serialize)]
212pub struct CriticalReaderState<T> {
213    /// The since capability of this reader.
214    pub since: Antichain<T>,
215    /// An opaque token matched on by compare_and_downgrade_since.
216    pub opaque: Opaque,
217    /// For debugging.
218    pub debug: HandleDebugState,
219}
220
221#[derive(Clone, Debug, PartialEq, Serialize)]
222pub struct WriterState<T> {
223    /// UNIX_EPOCH timestamp (in millis) of this writer's most recent heartbeat
224    pub last_heartbeat_timestamp_ms: u64,
225    /// Duration (in millis) allowed after [Self::last_heartbeat_timestamp_ms]
226    /// after which this writer may be expired
227    pub lease_duration_ms: u64,
228    /// The idempotency token of the most recent successful compare_and_append
229    /// by this writer.
230    pub most_recent_write_token: IdempotencyToken,
231    /// The upper of the most recent successful compare_and_append by this
232    /// writer.
233    pub most_recent_write_upper: Antichain<T>,
234    /// For debugging.
235    pub debug: HandleDebugState,
236}
237
238/// Debugging info for a reader or writer.
239#[derive(Arbitrary, Clone, Debug, Default, PartialEq, Serialize)]
240pub struct HandleDebugState {
241    /// Hostname of the persist user that registered this writer or reader. For
242    /// critical readers, this is the _most recent_ registration.
243    pub hostname: String,
244    /// Plaintext description of this writer or reader's intent.
245    pub purpose: String,
246}
247
248/// Part of the updates in a Batch.
249///
250/// Either a pointer to ones stored in Blob or the updates themselves inlined.
251#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
252#[serde(tag = "type")]
253pub enum BatchPart<T> {
254    Hollow(HollowBatchPart<T>),
255    Inline {
256        updates: LazyInlineBatchPart,
257        ts_rewrite: Option<Antichain<T>>,
258        schema_id: Option<SchemaId>,
259
260        /// ID of a schema that has since been deprecated and exists only to cleanly roundtrip.
261        deprecated_schema_id: Option<SchemaId>,
262    },
263}
264
265fn decode_structured_lower(lower: &LazyProto<ProtoArrayData>) -> Option<ArrayBound> {
266    let try_decode = |lower: &LazyProto<ProtoArrayData>| {
267        let proto = lower.decode()?;
268        let data = ArrayData::from_proto(proto)?;
269        ensure!(data.len() == 1);
270        Ok(ArrayBound::new(make_array(data), 0))
271    };
272
273    let decoded: anyhow::Result<ArrayBound> = try_decode(lower);
274
275    match decoded {
276        Ok(bound) => Some(bound),
277        Err(e) => {
278            soft_panic_or_log!("failed to decode bound: {e:#?}");
279            None
280        }
281    }
282}
283
284impl<T> BatchPart<T> {
285    pub fn hollow_bytes(&self) -> usize {
286        match self {
287            BatchPart::Hollow(x) => x.encoded_size_bytes,
288            BatchPart::Inline { .. } => 0,
289        }
290    }
291
292    pub fn is_inline(&self) -> bool {
293        matches!(self, BatchPart::Inline { .. })
294    }
295
296    pub fn inline_bytes(&self) -> usize {
297        match self {
298            BatchPart::Hollow(_) => 0,
299            BatchPart::Inline { updates, .. } => updates.encoded_size_bytes(),
300        }
301    }
302
303    pub fn writer_key(&self) -> Option<WriterKey> {
304        match self {
305            BatchPart::Hollow(x) => x.key.split().map(|(writer, _part)| writer),
306            BatchPart::Inline { .. } => None,
307        }
308    }
309
310    pub fn encoded_size_bytes(&self) -> usize {
311        match self {
312            BatchPart::Hollow(x) => x.encoded_size_bytes,
313            BatchPart::Inline { updates, .. } => updates.encoded_size_bytes(),
314        }
315    }
316
317    // A user-interpretable identifier or description of the part (for logs and
318    // such).
319    pub fn printable_name(&self) -> &str {
320        match self {
321            BatchPart::Hollow(x) => x.key.0.as_str(),
322            BatchPart::Inline { .. } => "<inline>",
323        }
324    }
325
326    pub fn stats(&self) -> Option<&LazyPartStats> {
327        match self {
328            BatchPart::Hollow(x) => x.stats.as_ref(),
329            BatchPart::Inline { .. } => None,
330        }
331    }
332
333    pub fn key_lower(&self) -> &[u8] {
334        match self {
335            BatchPart::Hollow(x) => x.key_lower.as_slice(),
336            // We don't duplicate the lowest key because this can be
337            // considerable overhead for small parts.
338            //
339            // The empty key might not be a tight lower bound, but it is a valid
340            // lower bound. If a caller is interested in a tighter lower bound,
341            // the data is inline.
342            BatchPart::Inline { .. } => &[],
343        }
344    }
345
346    pub fn structured_key_lower(&self) -> Option<ArrayBound> {
347        let part = match self {
348            BatchPart::Hollow(part) => part,
349            BatchPart::Inline { .. } => return None,
350        };
351
352        decode_structured_lower(part.structured_key_lower.as_ref()?)
353    }
354
355    pub fn ts_rewrite(&self) -> Option<&Antichain<T>> {
356        match self {
357            BatchPart::Hollow(x) => x.ts_rewrite.as_ref(),
358            BatchPart::Inline { ts_rewrite, .. } => ts_rewrite.as_ref(),
359        }
360    }
361
362    pub fn schema_id(&self) -> Option<SchemaId> {
363        match self {
364            BatchPart::Hollow(x) => x.schema_id,
365            BatchPart::Inline { schema_id, .. } => *schema_id,
366        }
367    }
368
369    pub fn deprecated_schema_id(&self) -> Option<SchemaId> {
370        match self {
371            BatchPart::Hollow(x) => x.deprecated_schema_id,
372            BatchPart::Inline {
373                deprecated_schema_id,
374                ..
375            } => *deprecated_schema_id,
376        }
377    }
378}
379
380impl<T: Timestamp + Codec64> BatchPart<T> {
381    pub fn is_structured_only(&self, metrics: &ColumnarMetrics) -> bool {
382        match self {
383            BatchPart::Hollow(x) => matches!(x.format, Some(BatchColumnarFormat::Structured)),
384            BatchPart::Inline { updates, .. } => {
385                let inline_part = updates.decode::<T>(metrics).expect("valid inline part");
386                matches!(inline_part.updates, BlobTraceUpdates::Structured { .. })
387            }
388        }
389    }
390
391    pub fn diffs_sum<D: Codec64 + Monoid>(&self, metrics: &ColumnarMetrics) -> Option<D> {
392        match self {
393            BatchPart::Hollow(x) => x.diffs_sum.map(D::decode),
394            BatchPart::Inline { updates, .. } => Some(
395                updates
396                    .decode::<T>(metrics)
397                    .expect("valid inline part")
398                    .updates
399                    .diffs_sum(),
400            ),
401        }
402    }
403}
404
405/// An ordered list of parts, generally stored as part of a larger run.
406#[derive(Debug, Clone)]
407pub struct HollowRun<T> {
408    /// Pointers usable to retrieve the updates.
409    pub(crate) parts: Vec<RunPart<T>>,
410}
411
412/// A reference to a [HollowRun], including the key in the blob store and some denormalized
413/// metadata.
414#[derive(Debug, Eq, PartialEq, Clone, Serialize)]
415pub struct HollowRunRef<T> {
416    pub key: PartialBatchKey,
417
418    /// The size of the referenced run object, plus all of the hollow objects it contains.
419    pub hollow_bytes: usize,
420
421    /// The size of the largest individual part in the run; useful for sizing compaction.
422    pub max_part_bytes: usize,
423
424    /// The lower bound of the data in this part, ordered by the codec ordering.
425    pub key_lower: Vec<u8>,
426
427    /// The lower bound of the data in this part, ordered by the structured ordering.
428    pub structured_key_lower: Option<LazyProto<ProtoArrayData>>,
429
430    pub diffs_sum: Option<[u8; 8]>,
431
432    pub(crate) _phantom_data: PhantomData<T>,
433}
434impl<T: Eq> PartialOrd<Self> for HollowRunRef<T> {
435    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
436        Some(self.cmp(other))
437    }
438}
439
440impl<T: Eq> Ord for HollowRunRef<T> {
441    fn cmp(&self, other: &Self) -> Ordering {
442        self.key.cmp(&other.key)
443    }
444}
445
446impl<T> HollowRunRef<T> {
447    pub fn writer_key(&self) -> Option<WriterKey> {
448        Some(self.key.split()?.0)
449    }
450}
451
452impl<T: Timestamp + Codec64> HollowRunRef<T> {
453    /// Stores the given runs and returns a [HollowRunRef] that points to them.
454    pub async fn set<D: Codec64 + Monoid>(
455        shard_id: ShardId,
456        blob: &dyn Blob,
457        writer: &WriterKey,
458        data: HollowRun<T>,
459        metrics: &Metrics,
460    ) -> Self {
461        let hollow_bytes = data.parts.iter().map(|p| p.hollow_bytes()).sum();
462        let max_part_bytes = data
463            .parts
464            .iter()
465            .map(|p| p.max_part_bytes())
466            .max()
467            .unwrap_or(0);
468        let key_lower = data
469            .parts
470            .first()
471            .map_or(vec![], |p| p.key_lower().to_vec());
472        let structured_key_lower = match data.parts.first() {
473            Some(RunPart::Many(r)) => r.structured_key_lower.clone(),
474            Some(RunPart::Single(BatchPart::Hollow(p))) => p.structured_key_lower.clone(),
475            Some(RunPart::Single(BatchPart::Inline { .. })) | None => None,
476        };
477        let diffs_sum = data
478            .parts
479            .iter()
480            .map(|p| {
481                p.diffs_sum::<D>(&metrics.columnar)
482                    .expect("valid diffs sum")
483            })
484            .reduce(|mut a, b| {
485                a.plus_equals(&b);
486                a
487            })
488            .expect("valid diffs sum")
489            .encode();
490
491        let key = PartialBatchKey::new(writer, &PartId::new());
492        let blob_key = key.complete(&shard_id);
493        let bytes = Bytes::from(prost::Message::encode_to_vec(&data.into_proto()));
494        let () = retry_external(&metrics.retries.external.hollow_run_set, || {
495            blob.set(&blob_key, bytes.clone())
496        })
497        .await;
498        Self {
499            key,
500            hollow_bytes,
501            max_part_bytes,
502            key_lower,
503            structured_key_lower,
504            diffs_sum: Some(diffs_sum),
505            _phantom_data: Default::default(),
506        }
507    }
508
509    /// Retrieve the [HollowRun] that this reference points to.
510    /// The caller is expected to ensure that this ref is the result of calling [HollowRunRef::set]
511    /// with the same shard id and backing store.
512    pub async fn get(
513        &self,
514        shard_id: ShardId,
515        blob: &dyn Blob,
516        metrics: &Metrics,
517    ) -> Option<HollowRun<T>> {
518        let blob_key = self.key.complete(&shard_id);
519        let mut bytes = retry_external(&metrics.retries.external.hollow_run_get, || {
520            blob.get(&blob_key)
521        })
522        .await?;
523        let proto_runs: ProtoHollowRun =
524            prost::Message::decode(&mut bytes).expect("illegal state: invalid proto bytes");
525        let runs = proto_runs
526            .into_rust()
527            .expect("illegal state: invalid encoded runs proto");
528        Some(runs)
529    }
530}
531
532/// Part of the updates in a run.
533///
534/// Either a pointer to ones stored in Blob or a single part stored inline.
535#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
536#[serde(untagged)]
537pub enum RunPart<T> {
538    Single(BatchPart<T>),
539    Many(HollowRunRef<T>),
540}
541
542impl<T: Ord> PartialOrd<Self> for RunPart<T> {
543    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
544        Some(self.cmp(other))
545    }
546}
547
548impl<T: Ord> Ord for RunPart<T> {
549    fn cmp(&self, other: &Self) -> Ordering {
550        match (self, other) {
551            (RunPart::Single(a), RunPart::Single(b)) => a.cmp(b),
552            (RunPart::Single(_), RunPart::Many(_)) => Ordering::Less,
553            (RunPart::Many(_), RunPart::Single(_)) => Ordering::Greater,
554            (RunPart::Many(a), RunPart::Many(b)) => a.cmp(b),
555        }
556    }
557}
558
559impl<T> RunPart<T> {
560    #[cfg(test)]
561    pub fn expect_hollow_part(&self) -> &HollowBatchPart<T> {
562        match self {
563            RunPart::Single(BatchPart::Hollow(hollow)) => hollow,
564            _ => panic!("expected hollow part!"),
565        }
566    }
567
568    pub fn hollow_bytes(&self) -> usize {
569        match self {
570            Self::Single(p) => p.hollow_bytes(),
571            Self::Many(r) => r.hollow_bytes,
572        }
573    }
574
575    pub fn is_inline(&self) -> bool {
576        match self {
577            Self::Single(p) => p.is_inline(),
578            Self::Many(_) => false,
579        }
580    }
581
582    pub fn inline_bytes(&self) -> usize {
583        match self {
584            Self::Single(p) => p.inline_bytes(),
585            Self::Many(_) => 0,
586        }
587    }
588
589    pub fn max_part_bytes(&self) -> usize {
590        match self {
591            Self::Single(p) => p.encoded_size_bytes(),
592            Self::Many(r) => r.max_part_bytes,
593        }
594    }
595
596    pub fn writer_key(&self) -> Option<WriterKey> {
597        match self {
598            Self::Single(p) => p.writer_key(),
599            Self::Many(r) => r.writer_key(),
600        }
601    }
602
603    pub fn encoded_size_bytes(&self) -> usize {
604        match self {
605            Self::Single(p) => p.encoded_size_bytes(),
606            Self::Many(r) => r.hollow_bytes,
607        }
608    }
609
610    pub fn schema_id(&self) -> Option<SchemaId> {
611        match self {
612            Self::Single(p) => p.schema_id(),
613            Self::Many(_) => None,
614        }
615    }
616
617    // A user-interpretable identifier or description of the part (for logs and
618    // such).
619    pub fn printable_name(&self) -> &str {
620        match self {
621            Self::Single(p) => p.printable_name(),
622            Self::Many(r) => r.key.0.as_str(),
623        }
624    }
625
626    pub fn stats(&self) -> Option<&LazyPartStats> {
627        match self {
628            Self::Single(p) => p.stats(),
629            // TODO: if we kept stats we could avoid fetching the metadata here.
630            Self::Many(_) => None,
631        }
632    }
633
634    pub fn key_lower(&self) -> &[u8] {
635        match self {
636            Self::Single(p) => p.key_lower(),
637            Self::Many(r) => r.key_lower.as_slice(),
638        }
639    }
640
641    pub fn structured_key_lower(&self) -> Option<ArrayBound> {
642        match self {
643            Self::Single(p) => p.structured_key_lower(),
644            Self::Many(_) => None,
645        }
646    }
647
648    pub fn ts_rewrite(&self) -> Option<&Antichain<T>> {
649        match self {
650            Self::Single(p) => p.ts_rewrite(),
651            Self::Many(_) => None,
652        }
653    }
654}
655
656impl<T> RunPart<T>
657where
658    T: Timestamp + Codec64,
659{
660    pub fn diffs_sum<D: Codec64 + Monoid>(&self, metrics: &ColumnarMetrics) -> Option<D> {
661        match self {
662            Self::Single(p) => p.diffs_sum(metrics),
663            Self::Many(hollow_run) => hollow_run.diffs_sum.map(D::decode),
664        }
665    }
666}
667
668/// A blob was missing!
669#[derive(Clone, Debug)]
670pub struct MissingBlob(BlobKey);
671
672impl std::fmt::Display for MissingBlob {
673    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
674        write!(f, "unexpectedly missing key: {}", self.0)
675    }
676}
677
678impl std::error::Error for MissingBlob {}
679
680impl<T: Timestamp + Codec64 + Sync> RunPart<T> {
681    pub fn part_stream<'a>(
682        &'a self,
683        shard_id: ShardId,
684        blob: &'a dyn Blob,
685        metrics: &'a Metrics,
686    ) -> impl Stream<Item = Result<Cow<'a, BatchPart<T>>, MissingBlob>> + Send + 'a {
687        try_stream! {
688            match self {
689                RunPart::Single(p) => {
690                    yield Cow::Borrowed(p);
691                }
692                RunPart::Many(r) => {
693                    let fetched = r.get(shard_id, blob, metrics).await
694                        .ok_or_else(|| MissingBlob(r.key.complete(&shard_id)))?;
695                    for run_part in fetched.parts {
696                        for await batch_part in
697                            run_part.part_stream(shard_id, blob, metrics).boxed()
698                        {
699                            yield Cow::Owned(batch_part?.into_owned());
700                        }
701                    }
702                }
703            }
704        }
705    }
706}
707
708impl<T: Ord> PartialOrd for BatchPart<T> {
709    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
710        Some(self.cmp(other))
711    }
712}
713
714impl<T: Ord> Ord for BatchPart<T> {
715    fn cmp(&self, other: &Self) -> Ordering {
716        match (self, other) {
717            (BatchPart::Hollow(s), BatchPart::Hollow(o)) => s.cmp(o),
718            (
719                BatchPart::Inline {
720                    updates: s_updates,
721                    ts_rewrite: s_ts_rewrite,
722                    schema_id: s_schema_id,
723                    deprecated_schema_id: s_deprecated_schema_id,
724                },
725                BatchPart::Inline {
726                    updates: o_updates,
727                    ts_rewrite: o_ts_rewrite,
728                    schema_id: o_schema_id,
729                    deprecated_schema_id: o_deprecated_schema_id,
730                },
731            ) => (
732                s_updates,
733                s_ts_rewrite.as_ref().map(|x| x.elements()),
734                s_schema_id,
735                s_deprecated_schema_id,
736            )
737                .cmp(&(
738                    o_updates,
739                    o_ts_rewrite.as_ref().map(|x| x.elements()),
740                    o_schema_id,
741                    o_deprecated_schema_id,
742                )),
743            (BatchPart::Hollow(_), BatchPart::Inline { .. }) => Ordering::Less,
744            (BatchPart::Inline { .. }, BatchPart::Hollow(_)) => Ordering::Greater,
745        }
746    }
747}
748
749/// What order are the parts in this run in?
750#[derive(Clone, Copy, Debug, PartialEq, Eq, Ord, PartialOrd, Serialize)]
751pub(crate) enum RunOrder {
752    /// They're in no particular order.
753    Unordered,
754    /// They're ordered based on the codec-encoded K/V bytes.
755    Codec,
756    /// They're ordered by the natural ordering of the structured data.
757    Structured,
758}
759
760#[derive(Clone, PartialEq, Eq, Ord, PartialOrd, Serialize, Copy, Hash)]
761pub struct RunId(pub(crate) [u8; 16]);
762
763impl std::fmt::Display for RunId {
764    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
765        write!(f, "ri{}", Uuid::from_bytes(self.0))
766    }
767}
768
769impl std::fmt::Debug for RunId {
770    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
771        write!(f, "RunId({})", Uuid::from_bytes(self.0))
772    }
773}
774
775impl std::str::FromStr for RunId {
776    type Err = String;
777
778    fn from_str(s: &str) -> Result<Self, Self::Err> {
779        parse_id("ri", "RunId", s).map(RunId)
780    }
781}
782
783impl From<RunId> for String {
784    fn from(x: RunId) -> Self {
785        x.to_string()
786    }
787}
788
789impl RunId {
790    pub(crate) fn new() -> Self {
791        RunId(*Uuid::new_v4().as_bytes())
792    }
793}
794
795impl Arbitrary for RunId {
796    type Parameters = ();
797    type Strategy = proptest::strategy::BoxedStrategy<Self>;
798    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
799        Strategy::prop_map(proptest::prelude::any::<u128>(), |n| {
800            RunId(*Uuid::from_u128(n).as_bytes())
801        })
802        .boxed()
803    }
804}
805
806/// Metadata shared across a run.
807#[derive(Clone, Debug, Default, PartialEq, Eq, Ord, PartialOrd, Serialize)]
808pub struct RunMeta {
809    /// If none, Persist should infer the order based on the proto metadata.
810    pub(crate) order: Option<RunOrder>,
811    /// All parts in a run should have the same schema.
812    pub(crate) schema: Option<SchemaId>,
813
814    /// ID of a schema that has since been deprecated and exists only to cleanly roundtrip.
815    pub(crate) deprecated_schema: Option<SchemaId>,
816
817    /// If set, a UUID that uniquely identifies this run.
818    pub(crate) id: Option<RunId>,
819
820    /// The number of updates in this run, or `None` if the number is unknown.
821    pub(crate) len: Option<usize>,
822
823    /// Additional unstructured metadata.
824    #[serde(skip_serializing_if = "MetadataMap::is_empty")]
825    pub(crate) meta: MetadataMap,
826}
827
828/// Metadata key for [RunMeta::bounds_truncated].
829const RUN_META_BOUNDS_TRUNCATED: MetadataKey<bool> = MetadataKey::new("truncated");
830
831impl RunMeta {
832    /// Whether this run's parts may hold updates outside the registered desc
833    /// of the batch that contains them.
834    ///
835    /// Set when a batch is appended under a desc narrower than the one it was
836    /// written with (truncation). Readers filter such updates out against the
837    /// registered desc, but per-part statistics like `diffs_sum` are computed
838    /// at write time over everything physically in the part, so accounting
839    /// that compares those statistics against data seen through a read must
840    /// skip runs with this bit set.
841    pub(crate) fn bounds_truncated(&self) -> bool {
842        self.meta.get(RUN_META_BOUNDS_TRUNCATED).unwrap_or(false)
843    }
844
845    /// Marks this run as possibly holding updates outside its batch's
846    /// registered desc. See [Self::bounds_truncated].
847    pub(crate) fn set_bounds_truncated(&mut self) {
848        self.meta.set(RUN_META_BOUNDS_TRUNCATED, true);
849    }
850}
851
852/// A subset of a [HollowBatch] corresponding 1:1 to a blob.
853#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
854pub struct HollowBatchPart<T> {
855    /// Pointer usable to retrieve the updates.
856    pub key: PartialBatchKey,
857    /// Miscellaneous metadata.
858    #[serde(skip_serializing_if = "MetadataMap::is_empty")]
859    pub meta: MetadataMap,
860    /// The encoded size of this part.
861    pub encoded_size_bytes: usize,
862    /// A lower bound on the keys in the part. (By default, this the minimum
863    /// possible key: `vec![]`.)
864    #[serde(serialize_with = "serialize_part_bytes")]
865    pub key_lower: Vec<u8>,
866    /// A lower bound on the keys in the part, stored as structured data.
867    #[serde(serialize_with = "serialize_lazy_proto")]
868    pub structured_key_lower: Option<LazyProto<ProtoArrayData>>,
869    /// Aggregate statistics about data contained in this part.
870    #[serde(serialize_with = "serialize_part_stats")]
871    pub stats: Option<LazyPartStats>,
872    /// A frontier to which timestamps in this part are advanced on read, if
873    /// set.
874    ///
875    /// A value of `Some([T::minimum()])` is functionally the same as `None`,
876    /// but we maintain the distinction between the two for some internal sanity
877    /// checking of invariants as well as metrics. If this ever becomes an
878    /// issue, everything still works with this as just `Antichain<T>`.
879    pub ts_rewrite: Option<Antichain<T>>,
880    /// A Codec64 encoded sum of all diffs in this part, if known.
881    ///
882    /// This is `None` if this part was written before we started storing this
883    /// information, or if it was written when the dyncfg was off.
884    ///
885    /// It could also make sense to model this as part of the pushdown stats, if
886    /// we later decide that's of some benefit.
887    #[serde(serialize_with = "serialize_diffs_sum")]
888    pub diffs_sum: Option<[u8; 8]>,
889    /// Columnar format that this batch was written in.
890    ///
891    /// This is `None` if this part was written before we started writing structured
892    /// columnar data.
893    pub format: Option<BatchColumnarFormat>,
894    /// The schemas used to encode the data in this batch part.
895    ///
896    /// Or None for historical data written before the schema registry was
897    /// added.
898    pub schema_id: Option<SchemaId>,
899
900    /// ID of a schema that has since been deprecated and exists only to cleanly roundtrip.
901    pub deprecated_schema_id: Option<SchemaId>,
902}
903
904/// A [Batch] but with the updates themselves stored externally.
905///
906/// [Batch]: differential_dataflow::trace::BatchReader
907#[derive(Clone, PartialEq, Eq)]
908pub struct HollowBatch<T> {
909    /// Describes the times of the updates in the batch.
910    pub desc: Description<T>,
911    /// The number of updates in the batch.
912    pub len: usize,
913    /// Pointers usable to retrieve the updates.
914    pub(crate) parts: Vec<RunPart<T>>,
915    /// Runs of sequential sorted batch parts, stored as indices into `parts`.
916    /// ex.
917    /// ```text
918    ///     parts=[p1, p2, p3], runs=[]     --> run  is  [p1, p2, p2]
919    ///     parts=[p1, p2, p3], runs=[1]    --> runs are [p1] and [p2, p3]
920    ///     parts=[p1, p2, p3], runs=[1, 2] --> runs are [p1], [p2], [p3]
921    /// ```
922    pub(crate) run_splits: Vec<usize>,
923    /// Run-level metadata: the first entry has metadata for the first run, and so on.
924    /// If there's no corresponding entry for a particular run, it's assumed to be [RunMeta::default()].
925    pub(crate) run_meta: Vec<RunMeta>,
926}
927
928impl<T: Debug> Debug for HollowBatch<T> {
929    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
930        let HollowBatch {
931            desc,
932            parts,
933            len,
934            run_splits: runs,
935            run_meta,
936        } = self;
937        f.debug_struct("HollowBatch")
938            .field(
939                "desc",
940                &(
941                    desc.lower().elements(),
942                    desc.upper().elements(),
943                    desc.since().elements(),
944                ),
945            )
946            .field("parts", &parts)
947            .field("len", &len)
948            .field("runs", &runs)
949            .field("run_meta", &run_meta)
950            .finish()
951    }
952}
953
954impl<T: Serialize> serde::Serialize for HollowBatch<T> {
955    fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
956        let HollowBatch {
957            desc,
958            len,
959            // Both parts and runs are covered by the self.runs call.
960            parts: _,
961            run_splits: _,
962            run_meta: _,
963        } = self;
964        let mut s = s.serialize_struct("HollowBatch", 5)?;
965        let () = s.serialize_field("lower", &desc.lower().elements())?;
966        let () = s.serialize_field("upper", &desc.upper().elements())?;
967        let () = s.serialize_field("since", &desc.since().elements())?;
968        let () = s.serialize_field("len", len)?;
969        let () = s.serialize_field("part_runs", &self.runs().collect::<Vec<_>>())?;
970        s.end()
971    }
972}
973
974impl<T: Ord> PartialOrd for HollowBatch<T> {
975    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
976        Some(self.cmp(other))
977    }
978}
979
980impl<T: Ord> Ord for HollowBatch<T> {
981    fn cmp(&self, other: &Self) -> Ordering {
982        // Deconstruct self and other so we get a compile failure if new fields
983        // are added.
984        let HollowBatch {
985            desc: self_desc,
986            parts: self_parts,
987            len: self_len,
988            run_splits: self_runs,
989            run_meta: self_run_meta,
990        } = self;
991        let HollowBatch {
992            desc: other_desc,
993            parts: other_parts,
994            len: other_len,
995            run_splits: other_runs,
996            run_meta: other_run_meta,
997        } = other;
998        (
999            self_desc.lower().elements(),
1000            self_desc.upper().elements(),
1001            self_desc.since().elements(),
1002            self_parts,
1003            self_len,
1004            self_runs,
1005            self_run_meta,
1006        )
1007            .cmp(&(
1008                other_desc.lower().elements(),
1009                other_desc.upper().elements(),
1010                other_desc.since().elements(),
1011                other_parts,
1012                other_len,
1013                other_runs,
1014                other_run_meta,
1015            ))
1016    }
1017}
1018
1019impl<T: Timestamp + Codec64 + Sync> HollowBatch<T> {
1020    pub(crate) fn part_stream<'a>(
1021        &'a self,
1022        shard_id: ShardId,
1023        blob: &'a dyn Blob,
1024        metrics: &'a Metrics,
1025    ) -> impl Stream<Item = Result<Cow<'a, BatchPart<T>>, MissingBlob>> + 'a {
1026        stream! {
1027            for part in &self.parts {
1028                for await part in part.part_stream(shard_id, blob, metrics) {
1029                    yield part;
1030                }
1031            }
1032        }
1033    }
1034}
1035impl<T> HollowBatch<T> {
1036    /// Construct an in-memory hollow batch from the given metadata.
1037    ///
1038    /// This method checks that `runs` is a sequence of valid indices into `parts`. The caller
1039    /// is responsible for ensuring that the defined runs are valid.
1040    ///
1041    /// `len` should represent the number of valid updates in the referenced parts.
1042    pub(crate) fn new(
1043        desc: Description<T>,
1044        parts: Vec<RunPart<T>>,
1045        len: usize,
1046        run_meta: Vec<RunMeta>,
1047        run_splits: Vec<usize>,
1048    ) -> Self {
1049        debug_assert!(
1050            run_splits.is_strictly_sorted(),
1051            "run indices should be strictly increasing"
1052        );
1053        mz_ore::soft_assert_no_log!(
1054            run_splits.first().map_or(true, |i| *i > 0),
1055            "run indices should be positive"
1056        );
1057        mz_ore::soft_assert_no_log!(
1058            run_splits.last().map_or(true, |i| *i < parts.len()),
1059            "run indices should be valid indices into parts"
1060        );
1061        mz_ore::soft_assert_no_log!(
1062            parts.is_empty() || run_meta.len() == run_splits.len() + 1,
1063            "all metadata should correspond to a run"
1064        );
1065
1066        Self {
1067            desc,
1068            len,
1069            parts,
1070            run_splits,
1071            run_meta,
1072        }
1073    }
1074
1075    /// Construct a batch of a single run with default metadata. Mostly interesting for tests.
1076    pub(crate) fn new_run(desc: Description<T>, parts: Vec<RunPart<T>>, len: usize) -> Self {
1077        let run_meta = if parts.is_empty() {
1078            vec![]
1079        } else {
1080            vec![RunMeta::default()]
1081        };
1082        Self {
1083            desc,
1084            len,
1085            parts,
1086            run_splits: vec![],
1087            run_meta,
1088        }
1089    }
1090
1091    #[cfg(test)]
1092    pub(crate) fn new_run_for_test(
1093        desc: Description<T>,
1094        parts: Vec<RunPart<T>>,
1095        len: usize,
1096        run_id: RunId,
1097    ) -> Self {
1098        let run_meta = if parts.is_empty() {
1099            vec![]
1100        } else {
1101            let mut meta = RunMeta::default();
1102            meta.id = Some(run_id);
1103            vec![meta]
1104        };
1105        Self {
1106            desc,
1107            len,
1108            parts,
1109            run_splits: vec![],
1110            run_meta,
1111        }
1112    }
1113
1114    /// An empty hollow batch, representing no updates over the given desc.
1115    pub(crate) fn empty(desc: Description<T>) -> Self {
1116        Self {
1117            desc,
1118            len: 0,
1119            parts: vec![],
1120            run_splits: vec![],
1121            run_meta: vec![],
1122        }
1123    }
1124
1125    pub(crate) fn runs(&self) -> impl Iterator<Item = (&RunMeta, &[RunPart<T>])> {
1126        let run_ends = self
1127            .run_splits
1128            .iter()
1129            .copied()
1130            .chain(std::iter::once(self.parts.len()));
1131        let run_metas = self.run_meta.iter();
1132        let run_parts = run_ends
1133            .scan(0, |start, end| {
1134                let range = *start..end;
1135                *start = end;
1136                Some(range)
1137            })
1138            .filter(|range| !range.is_empty())
1139            .map(|range| &self.parts[range]);
1140        run_metas.zip_eq(run_parts)
1141    }
1142
1143    pub(crate) fn inline_bytes(&self) -> usize {
1144        self.parts.iter().map(|x| x.inline_bytes()).sum()
1145    }
1146
1147    pub(crate) fn is_empty(&self) -> bool {
1148        self.parts.is_empty()
1149    }
1150
1151    pub(crate) fn part_count(&self) -> usize {
1152        self.parts.len()
1153    }
1154
1155    /// The sum of the encoded sizes of all parts in the batch.
1156    pub fn encoded_size_bytes(&self) -> usize {
1157        self.parts.iter().map(|p| p.encoded_size_bytes()).sum()
1158    }
1159}
1160
1161// See the comment on [Batch::rewrite_ts] for why this is TotalOrder.
1162impl<T: Timestamp + TotalOrder> HollowBatch<T> {
1163    pub(crate) fn rewrite_ts(
1164        &mut self,
1165        frontier: &Antichain<T>,
1166        new_upper: Antichain<T>,
1167    ) -> Result<(), String> {
1168        if !PartialOrder::less_than(frontier, &new_upper) {
1169            return Err(format!(
1170                "rewrite frontier {:?} !< rewrite upper {:?}",
1171                frontier.elements(),
1172                new_upper.elements(),
1173            ));
1174        }
1175        if PartialOrder::less_than(&new_upper, self.desc.upper()) {
1176            return Err(format!(
1177                "rewrite upper {:?} < batch upper {:?}",
1178                new_upper.elements(),
1179                self.desc.upper().elements(),
1180            ));
1181        }
1182
1183        // The following are things that it seems like we could support, but
1184        // initially we don't because we don't have a use case for them.
1185        if PartialOrder::less_than(frontier, self.desc.lower()) {
1186            return Err(format!(
1187                "rewrite frontier {:?} < batch lower {:?}",
1188                frontier.elements(),
1189                self.desc.lower().elements(),
1190            ));
1191        }
1192        if self.desc.since() != &Antichain::from_elem(T::minimum()) {
1193            return Err(format!(
1194                "batch since {:?} != minimum antichain {:?}",
1195                self.desc.since().elements(),
1196                [T::minimum()],
1197            ));
1198        }
1199        for part in self.parts.iter() {
1200            let Some(ts_rewrite) = part.ts_rewrite() else {
1201                continue;
1202            };
1203            if PartialOrder::less_than(frontier, ts_rewrite) {
1204                return Err(format!(
1205                    "rewrite frontier {:?} < batch rewrite {:?}",
1206                    frontier.elements(),
1207                    ts_rewrite.elements(),
1208                ));
1209            }
1210        }
1211
1212        self.desc = Description::new(
1213            self.desc.lower().clone(),
1214            new_upper,
1215            self.desc.since().clone(),
1216        );
1217        for part in &mut self.parts {
1218            match part {
1219                RunPart::Single(BatchPart::Hollow(part)) => {
1220                    part.ts_rewrite = Some(frontier.clone())
1221                }
1222                RunPart::Single(BatchPart::Inline { ts_rewrite, .. }) => {
1223                    *ts_rewrite = Some(frontier.clone())
1224                }
1225                RunPart::Many(runs) => {
1226                    // Currently unreachable: we only apply rewrites to user batches, and we don't
1227                    // ever generate runs of >1 part for those.
1228                    panic!("unexpected rewrite of a hollow runs ref: {runs:?}");
1229                }
1230            }
1231        }
1232        Ok(())
1233    }
1234}
1235
1236impl<T: Ord> PartialOrd for HollowBatchPart<T> {
1237    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1238        Some(self.cmp(other))
1239    }
1240}
1241
1242impl<T: Ord> Ord for HollowBatchPart<T> {
1243    fn cmp(&self, other: &Self) -> Ordering {
1244        // Deconstruct self and other so we get a compile failure if new fields
1245        // are added.
1246        let HollowBatchPart {
1247            key: self_key,
1248            meta: self_meta,
1249            encoded_size_bytes: self_encoded_size_bytes,
1250            key_lower: self_key_lower,
1251            structured_key_lower: self_structured_key_lower,
1252            stats: self_stats,
1253            ts_rewrite: self_ts_rewrite,
1254            diffs_sum: self_diffs_sum,
1255            format: self_format,
1256            schema_id: self_schema_id,
1257            deprecated_schema_id: self_deprecated_schema_id,
1258        } = self;
1259        let HollowBatchPart {
1260            key: other_key,
1261            meta: other_meta,
1262            encoded_size_bytes: other_encoded_size_bytes,
1263            key_lower: other_key_lower,
1264            structured_key_lower: other_structured_key_lower,
1265            stats: other_stats,
1266            ts_rewrite: other_ts_rewrite,
1267            diffs_sum: other_diffs_sum,
1268            format: other_format,
1269            schema_id: other_schema_id,
1270            deprecated_schema_id: other_deprecated_schema_id,
1271        } = other;
1272        (
1273            self_key,
1274            self_meta,
1275            self_encoded_size_bytes,
1276            self_key_lower,
1277            self_structured_key_lower,
1278            self_stats,
1279            self_ts_rewrite.as_ref().map(|x| x.elements()),
1280            self_diffs_sum,
1281            self_format,
1282            self_schema_id,
1283            self_deprecated_schema_id,
1284        )
1285            .cmp(&(
1286                other_key,
1287                other_meta,
1288                other_encoded_size_bytes,
1289                other_key_lower,
1290                other_structured_key_lower,
1291                other_stats,
1292                other_ts_rewrite.as_ref().map(|x| x.elements()),
1293                other_diffs_sum,
1294                other_format,
1295                other_schema_id,
1296                other_deprecated_schema_id,
1297            ))
1298    }
1299}
1300
1301/// A pointer to a rollup stored externally.
1302#[derive(Arbitrary, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize)]
1303pub struct HollowRollup {
1304    /// Pointer usable to retrieve the rollup.
1305    pub key: PartialRollupKey,
1306    /// The encoded size of this rollup, if known.
1307    pub encoded_size_bytes: Option<usize>,
1308}
1309
1310/// A pointer to a blob stored externally.
1311#[derive(Debug)]
1312pub enum HollowBlobRef<'a, T> {
1313    Batch(&'a HollowBatch<T>),
1314    Rollup(&'a HollowRollup),
1315}
1316
1317/// A rollup that is currently being computed.
1318#[derive(
1319    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Arbitrary, Serialize
1320)]
1321pub struct ActiveRollup {
1322    pub seqno: SeqNo,
1323    pub start_ms: u64,
1324}
1325
1326/// A garbage collection request that is currently being computed.
1327#[derive(
1328    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Arbitrary, Serialize
1329)]
1330pub struct ActiveGc {
1331    pub seqno: SeqNo,
1332    pub start_ms: u64,
1333}
1334
1335/// A sentinel for a state transition that was a no-op.
1336///
1337/// Critically, this also indicates that the no-op state transition was not
1338/// committed through compare_and_append and thus is _not linearized_.
1339#[derive(Debug)]
1340#[cfg_attr(any(test, debug_assertions), derive(PartialEq))]
1341pub struct NoOpStateTransition<T>(pub T);
1342
1343// TODO: Document invariants.
1344#[derive(Debug, Clone)]
1345#[cfg_attr(any(test, debug_assertions), derive(PartialEq))]
1346pub struct StateCollections<T> {
1347    /// The version of this state. This is typically identical to the version of the code
1348    /// that wrote it, but may diverge during 0dt upgrades and similar operations when a
1349    /// new version of code is intentionally interoperating with an older state format.
1350    pub(crate) version: Version,
1351
1352    // - Invariant: `<= all reader.since`
1353    // - Invariant: Doesn't regress across state versions.
1354    pub(crate) last_gc_req: SeqNo,
1355
1356    // - Invariant: There is a rollup with `seqno <= self.seqno_since`.
1357    pub(crate) rollups: BTreeMap<SeqNo, HollowRollup>,
1358
1359    /// The rollup that is currently being computed.
1360    pub(crate) active_rollup: Option<ActiveRollup>,
1361    /// The gc request that is currently being computed.
1362    pub(crate) active_gc: Option<ActiveGc>,
1363
1364    pub(crate) leased_readers: BTreeMap<LeasedReaderId, LeasedReaderState<T>>,
1365    pub(crate) critical_readers: BTreeMap<CriticalReaderId, CriticalReaderState<T>>,
1366    pub(crate) writers: BTreeMap<WriterId, WriterState<T>>,
1367    pub(crate) schemas: BTreeMap<SchemaId, EncodedSchemas>,
1368
1369    // - Invariant: `trace.since == meet(all reader.since)`
1370    // - Invariant: `trace.since` doesn't regress across state versions.
1371    // - Invariant: `trace.upper` doesn't regress across state versions.
1372    // - Invariant: `trace` upholds its own invariants.
1373    pub(crate) trace: Trace<T>,
1374}
1375
1376/// A key and val [Codec::Schema] encoded via [Codec::encode_schema].
1377///
1378/// This strategy of directly serializing the schema objects requires that
1379/// persist users do the right thing. Specifically, that an encoded schema
1380/// doesn't in some later version of mz decode to an in-mem object that acts
1381/// differently. In a sense, the current system (before the introduction of the
1382/// schema registry) where schemas are passed in unchecked to reader and writer
1383/// registration calls also has the same defect, so seems fine.
1384///
1385/// An alternative is to write down here some persist-specific representation of
1386/// the schema (e.g. the arrow DataType). This is a lot more work and also has
1387/// the potential to lead down a similar failure mode to the mz_persist_types
1388/// `Data` trait, where the boilerplate isn't worth the safety. Given that we
1389/// can always migrate later by rehydrating these, seems fine to start with the
1390/// easy thing.
1391#[derive(Debug, Clone, Serialize, PartialEq)]
1392pub struct EncodedSchemas {
1393    /// A full in-mem `K::Schema` impl encoded via [Codec::encode_schema].
1394    pub key: Bytes,
1395    /// The arrow `DataType` produced by this `K::Schema` at the time it was
1396    /// registered, encoded as a `ProtoDataType`.
1397    pub key_data_type: Bytes,
1398    /// A full in-mem `V::Schema` impl encoded via [Codec::encode_schema].
1399    pub val: Bytes,
1400    /// The arrow `DataType` produced by this `V::Schema` at the time it was
1401    /// registered, encoded as a `ProtoDataType`.
1402    pub val_data_type: Bytes,
1403}
1404
1405impl EncodedSchemas {
1406    pub(crate) fn decode_data_type(buf: &[u8]) -> DataType {
1407        let proto = prost::Message::decode(buf).expect("valid ProtoDataType");
1408        DataType::from_proto(proto).expect("valid DataType")
1409    }
1410}
1411
1412#[derive(Debug)]
1413#[cfg_attr(test, derive(PartialEq))]
1414pub enum CompareAndAppendBreak<T> {
1415    AlreadyCommitted,
1416    Upper {
1417        shard_upper: Antichain<T>,
1418        writer_upper: Antichain<T>,
1419    },
1420    InvalidUsage(InvalidUsage<T>),
1421    InlineBackpressure,
1422}
1423
1424#[derive(Debug)]
1425#[cfg_attr(test, derive(PartialEq))]
1426pub enum SnapshotErr<T> {
1427    AsOfNotYetAvailable(SeqNo, Upper<T>),
1428    AsOfHistoricalDistinctionsLost(Since<T>),
1429}
1430
1431impl<T> StateCollections<T>
1432where
1433    T: Timestamp + Lattice + Codec64,
1434{
1435    pub fn add_rollup(
1436        &mut self,
1437        add_rollup: (SeqNo, &HollowRollup),
1438    ) -> ControlFlow<NoOpStateTransition<bool>, bool> {
1439        let (rollup_seqno, rollup) = add_rollup;
1440        let applied = match self.rollups.get(&rollup_seqno) {
1441            Some(x) => x.key == rollup.key,
1442            None => {
1443                // PER-16: refuse to insert a rollup at a seqno that GC has
1444                // already physically removed. `apply_unbatched_idempotent_cmd`
1445                // replays the captured `(rollup_seqno, rollup)` tuple on
1446                // every retry. If the original CAS committed indeterminately
1447                // and a concurrent GC then removed the rollup, the retry
1448                // would otherwise observe an empty map entry and re-insert
1449                // the same `PartialRollupKey` at the same seqno, producing
1450                // a second Insert (and, later, a second Delete) for that
1451                // key in the diff stream and tripping the `assert!` in
1452                // `find_removable_blobs`.
1453                //
1454                // We use the smallest seqno currently in `self.rollups` as
1455                // the GC watermark. GC's `remove_rollups` always keeps the
1456                // latest rollup `<= seqno_since` and removes the older
1457                // ones, so the surviving minimum is strictly above every
1458                // seqno that has been removed. (Unlike `last_gc_req`, this
1459                // is only advanced by an actual physical removal, not by
1460                // `maybe_gc` deciding to *request* one — `last_gc_req` runs
1461                // ahead of removals and would falsely reject legitimate
1462                // late `add_rollup` calls for older seqnos.)
1463                if let Some(min_kept) = self.rollups.keys().next() {
1464                    if rollup_seqno < *min_kept {
1465                        return Continue(false);
1466                    }
1467                }
1468                self.active_rollup = None;
1469                self.rollups.insert(rollup_seqno, rollup.to_owned());
1470                true
1471            }
1472        };
1473        // This state transition is a no-op if applied is false but we
1474        // still commit the state change so that this gets linearized
1475        // (maybe we're looking at old state).
1476        Continue(applied)
1477    }
1478
1479    pub fn remove_rollups(
1480        &mut self,
1481        remove_rollups: &[(SeqNo, PartialRollupKey)],
1482    ) -> ControlFlow<NoOpStateTransition<Vec<SeqNo>>, Vec<SeqNo>> {
1483        if self.is_tombstone() {
1484            return Break(NoOpStateTransition(vec![]));
1485        }
1486
1487        // This state transition is called at the end of the GC process, so we
1488        // need to unset the `active_gc` field.
1489        let active_gc_was_set = self.active_gc.take().is_some();
1490
1491        if remove_rollups.is_empty() {
1492            return if active_gc_was_set {
1493                Continue(vec![])
1494            } else {
1495                Break(NoOpStateTransition(vec![]))
1496            };
1497        }
1498
1499        let mut removed = vec![];
1500        for (seqno, key) in remove_rollups {
1501            let removed_key = self.rollups.remove(seqno);
1502            mz_ore::soft_assert_no_log!(
1503                removed_key.as_ref().map_or(true, |x| &x.key == key),
1504                "rollup at {} to be removed has key {:?} in state, but GC asked to remove {}",
1505                seqno,
1506                removed_key,
1507                key
1508            );
1509
1510            if removed_key.is_some() {
1511                removed.push(*seqno);
1512            }
1513        }
1514
1515        Continue(removed)
1516    }
1517
1518    pub fn register_leased_reader(
1519        &mut self,
1520        hostname: &str,
1521        reader_id: &LeasedReaderId,
1522        purpose: &str,
1523        seqno: SeqNo,
1524        lease_duration: Duration,
1525        heartbeat_timestamp_ms: u64,
1526        use_critical_since: bool,
1527    ) -> ControlFlow<
1528        NoOpStateTransition<(LeasedReaderState<T>, SeqNo)>,
1529        (LeasedReaderState<T>, SeqNo),
1530    > {
1531        let since = if use_critical_since {
1532            self.critical_since()
1533                .unwrap_or_else(|| self.trace.since().clone())
1534        } else {
1535            self.trace.since().clone()
1536        };
1537        let reader_state = LeasedReaderState {
1538            debug: HandleDebugState {
1539                hostname: hostname.to_owned(),
1540                purpose: purpose.to_owned(),
1541            },
1542            seqno,
1543            since,
1544            last_heartbeat_timestamp_ms: heartbeat_timestamp_ms,
1545            lease_duration_ms: u64::try_from(lease_duration.as_millis())
1546                .expect("lease duration as millis should fit within u64"),
1547        };
1548
1549        // If the shard-global upper and since are both the empty antichain,
1550        // then no further writes can ever commit and no further reads can be
1551        // served. Optimize this by no-op-ing reader registration so that we can
1552        // settle the shard into a final unchanging tombstone state.
1553        if self.is_tombstone() {
1554            return Break(NoOpStateTransition((reader_state, self.seqno_since(seqno))));
1555        }
1556
1557        // TODO: Handle if the reader or writer already exists.
1558        self.leased_readers
1559            .insert(reader_id.clone(), reader_state.clone());
1560        Continue((reader_state, self.seqno_since(seqno)))
1561    }
1562
1563    pub fn register_critical_reader(
1564        &mut self,
1565        hostname: &str,
1566        reader_id: &CriticalReaderId,
1567        opaque: Opaque,
1568        purpose: &str,
1569    ) -> ControlFlow<NoOpStateTransition<CriticalReaderState<T>>, CriticalReaderState<T>> {
1570        let state = CriticalReaderState {
1571            debug: HandleDebugState {
1572                hostname: hostname.to_owned(),
1573                purpose: purpose.to_owned(),
1574            },
1575            since: self.trace.since().clone(),
1576            opaque,
1577        };
1578
1579        // We expire all readers if the upper and since both advance to the
1580        // empty antichain. Gracefully handle this. At the same time,
1581        // short-circuit the cmd application so we don't needlessly create new
1582        // SeqNos.
1583        if self.is_tombstone() {
1584            return Break(NoOpStateTransition(state));
1585        }
1586
1587        let state = match self.critical_readers.get_mut(reader_id) {
1588            Some(existing_state) => {
1589                existing_state.debug = state.debug;
1590                existing_state.clone()
1591            }
1592            None => {
1593                self.critical_readers
1594                    .insert(reader_id.clone(), state.clone());
1595                state
1596            }
1597        };
1598        Continue(state)
1599    }
1600
1601    pub fn register_schema<K: Codec, V: Codec>(
1602        &mut self,
1603        key_schema: &K::Schema,
1604        val_schema: &V::Schema,
1605    ) -> ControlFlow<NoOpStateTransition<Option<SchemaId>>, Option<SchemaId>> {
1606        fn encode_data_type(data_type: &DataType) -> Bytes {
1607            let proto = data_type.into_proto();
1608            prost::Message::encode_to_vec(&proto).into()
1609        }
1610
1611        // Look for an existing registered SchemaId for these schemas.
1612        //
1613        // The common case is that this should be a recent one, so as a minor
1614        // optimization, do this search in reverse order.
1615        //
1616        // TODO: Note that this impl is `O(schemas)`. Combined with the
1617        // possibility of cmd retries, it's possible but unlikely for this to
1618        // get expensive. We could maintain a reverse map to speed this up in
1619        // necessary. This would either need to work on the encoded
1620        // representation (which, we'd have to fall back to the linear scan) or
1621        // we'd need to add a Hash/Ord bound to Schema.
1622        let existing_id = self.schemas.iter().rev().find(|(_, x)| {
1623            K::decode_schema(&x.key) == *key_schema && V::decode_schema(&x.val) == *val_schema
1624        });
1625        match existing_id {
1626            Some((schema_id, _)) => {
1627                // TODO: Validate that the decoded schemas still produce records
1628                // of the recorded DataType, to detect shenanigans. Probably
1629                // best to wait until we've turned on Schema2 in prod and thus
1630                // committed to the current mappings.
1631                Break(NoOpStateTransition(Some(*schema_id)))
1632            }
1633            None if self.is_tombstone() => {
1634                // TODO: Is this right?
1635                Break(NoOpStateTransition(None))
1636            }
1637            None if self.schemas.is_empty() => {
1638                // We'll have to do something more sophisticated here to
1639                // generate the next id if/when we start supporting the removal
1640                // of schemas.
1641                let id = SchemaId(self.schemas.len());
1642                let key_data_type = mz_persist_types::columnar::data_type::<K>(key_schema)
1643                    .expect("valid key schema");
1644                let val_data_type = mz_persist_types::columnar::data_type::<V>(val_schema)
1645                    .expect("valid val schema");
1646                let prev = self.schemas.insert(
1647                    id,
1648                    EncodedSchemas {
1649                        key: K::encode_schema(key_schema),
1650                        key_data_type: encode_data_type(&key_data_type),
1651                        val: V::encode_schema(val_schema),
1652                        val_data_type: encode_data_type(&val_data_type),
1653                    },
1654                );
1655                assert_eq!(prev, None);
1656                Continue(Some(id))
1657            }
1658            None => {
1659                info!(
1660                    "register_schemas got {:?} expected {:?}",
1661                    key_schema,
1662                    self.schemas
1663                        .iter()
1664                        .map(|(id, x)| (id, K::decode_schema(&x.key)))
1665                        .collect::<Vec<_>>()
1666                );
1667                // Until we implement persist schema changes, only allow at most
1668                // one registered schema.
1669                Break(NoOpStateTransition(None))
1670            }
1671        }
1672    }
1673
1674    pub fn compare_and_evolve_schema<K: Codec, V: Codec>(
1675        &mut self,
1676        expected: SchemaId,
1677        key_schema: &K::Schema,
1678        val_schema: &V::Schema,
1679    ) -> ControlFlow<NoOpStateTransition<CaESchema<K, V>>, CaESchema<K, V>> {
1680        fn data_type<T>(schema: &impl Schema<T>) -> DataType {
1681            // To be defensive, create an empty batch and inspect the resulting
1682            // data type (as opposed to something like allowing the `Schema` to
1683            // declare the DataType).
1684            let array = Schema::encoder(schema).expect("valid schema").finish();
1685            Array::data_type(&array).clone()
1686        }
1687
1688        let (current_id, current) = self
1689            .schemas
1690            .last_key_value()
1691            .expect("all shards have a schema");
1692
1693        let current_key = K::decode_schema(&current.key);
1694        let current_key_dt = EncodedSchemas::decode_data_type(&current.key_data_type);
1695        let current_val = V::decode_schema(&current.val);
1696        let current_val_dt = EncodedSchemas::decode_data_type(&current.val_data_type);
1697
1698        let key_dt = data_type(key_schema);
1699        let val_dt = data_type(val_schema);
1700
1701        // If the schema is exactly the same as the current one, no-op. NOTE:
1702        // this check has to come before the `expected` one, otherwise the
1703        // command is not idempotent: `apply_unbatched_idempotent_cmd` retries
1704        // indeterminate errors, so a CaS that committed but lost its response
1705        // gets re-run against state that already carries its own evolution.
1706        //
1707        // Returning Ok when `*current_id != expected` is safe: register_schema
1708        // never mints two ids for the same content, so a content match is the
1709        // current schema and a stale `expected` just lags the committed retry.
1710        if current_key == *key_schema
1711            && current_key_dt == key_dt
1712            && current_val == *val_schema
1713            && current_val_dt == val_dt
1714        {
1715            return Break(NoOpStateTransition(CaESchema::Ok(*current_id)));
1716        }
1717
1718        if *current_id != expected {
1719            return Break(NoOpStateTransition(CaESchema::ExpectedMismatch {
1720                schema_id: *current_id,
1721                key: current_key,
1722                val: current_val,
1723            }));
1724        }
1725
1726        let key_fn = backward_compatible(&current_key_dt, &key_dt);
1727        let val_fn = backward_compatible(&current_val_dt, &val_dt);
1728        let (Some(key_fn), Some(val_fn)) = (key_fn, val_fn) else {
1729            return Break(NoOpStateTransition(CaESchema::Incompatible));
1730        };
1731        // Persist initially disallows dropping columns. This would require a
1732        // bunch more work (e.g. not safe to use the latest schema in
1733        // compaction) and isn't initially necessary in mz.
1734        if key_fn.contains_drop() || val_fn.contains_drop() {
1735            return Break(NoOpStateTransition(CaESchema::Incompatible));
1736        }
1737
1738        // We'll have to do something more sophisticated here to
1739        // generate the next id if/when we start supporting the removal
1740        // of schemas.
1741        let id = SchemaId(self.schemas.len());
1742        self.schemas.insert(
1743            id,
1744            EncodedSchemas {
1745                key: K::encode_schema(key_schema),
1746                key_data_type: prost::Message::encode_to_vec(&key_dt.into_proto()).into(),
1747                val: V::encode_schema(val_schema),
1748                val_data_type: prost::Message::encode_to_vec(&val_dt.into_proto()).into(),
1749            },
1750        );
1751        Continue(CaESchema::Ok(id))
1752    }
1753
1754    pub fn compare_and_append(
1755        &mut self,
1756        batch: &HollowBatch<T>,
1757        writer_id: &WriterId,
1758        heartbeat_timestamp_ms: u64,
1759        lease_duration_ms: u64,
1760        idempotency_token: &IdempotencyToken,
1761        debug_info: &HandleDebugState,
1762        inline_writes_total_max_bytes: usize,
1763        claim_compaction_percent: usize,
1764        claim_compaction_min_version: Option<&Version>,
1765    ) -> ControlFlow<CompareAndAppendBreak<T>, Vec<FueledMergeReq<T>>> {
1766        // We expire all writers if the upper and since both advance to the
1767        // empty antichain. Gracefully handle this. At the same time,
1768        // short-circuit the cmd application so we don't needlessly create new
1769        // SeqNos.
1770        if self.is_tombstone() {
1771            assert_eq!(self.trace.upper(), &Antichain::new());
1772            return Break(CompareAndAppendBreak::Upper {
1773                shard_upper: Antichain::new(),
1774                // This writer might have been registered before the shard upper
1775                // was advanced, which would make this pessimistic in the
1776                // Indeterminate handling of compare_and_append at the machine
1777                // level, but that's fine.
1778                writer_upper: Antichain::new(),
1779            });
1780        }
1781
1782        let writer_state = self
1783            .writers
1784            .entry(writer_id.clone())
1785            .or_insert_with(|| WriterState {
1786                last_heartbeat_timestamp_ms: heartbeat_timestamp_ms,
1787                lease_duration_ms,
1788                most_recent_write_token: IdempotencyToken::SENTINEL,
1789                most_recent_write_upper: Antichain::from_elem(T::minimum()),
1790                debug: debug_info.clone(),
1791            });
1792
1793        if PartialOrder::less_than(batch.desc.upper(), batch.desc.lower()) {
1794            return Break(CompareAndAppendBreak::InvalidUsage(
1795                InvalidUsage::InvalidBounds {
1796                    lower: batch.desc.lower().clone(),
1797                    upper: batch.desc.upper().clone(),
1798                },
1799            ));
1800        }
1801
1802        // If the time interval is empty, the list of updates must also be
1803        // empty.
1804        if batch.desc.upper() == batch.desc.lower() && !batch.is_empty() {
1805            return Break(CompareAndAppendBreak::InvalidUsage(
1806                InvalidUsage::InvalidEmptyTimeInterval {
1807                    lower: batch.desc.lower().clone(),
1808                    upper: batch.desc.upper().clone(),
1809                    keys: batch
1810                        .parts
1811                        .iter()
1812                        .map(|x| x.printable_name().to_owned())
1813                        .collect(),
1814                },
1815            ));
1816        }
1817
1818        if idempotency_token == &writer_state.most_recent_write_token {
1819            // If the last write had the same idempotency_token, then this must
1820            // have already committed. Sanity check that the most recent write
1821            // upper matches and that the shard upper is at least the write
1822            // upper, if it's not something very suspect is going on.
1823            assert_eq!(batch.desc.upper(), &writer_state.most_recent_write_upper);
1824            assert!(
1825                PartialOrder::less_equal(batch.desc.upper(), self.trace.upper()),
1826                "{:?} vs {:?}",
1827                batch.desc.upper(),
1828                self.trace.upper()
1829            );
1830            return Break(CompareAndAppendBreak::AlreadyCommitted);
1831        }
1832
1833        let shard_upper = self.trace.upper();
1834        if shard_upper != batch.desc.lower() {
1835            return Break(CompareAndAppendBreak::Upper {
1836                shard_upper: shard_upper.clone(),
1837                writer_upper: writer_state.most_recent_write_upper.clone(),
1838            });
1839        }
1840
1841        let new_inline_bytes = batch.inline_bytes();
1842        if new_inline_bytes > 0 {
1843            let mut existing_inline_bytes = 0;
1844            self.trace
1845                .map_batches(|x| existing_inline_bytes += x.inline_bytes());
1846            // TODO: For very small batches, it may actually _increase_ the size
1847            // of state to flush them out. Consider another threshold under
1848            // which an inline part can be appended no matter what.
1849            if existing_inline_bytes + new_inline_bytes >= inline_writes_total_max_bytes {
1850                return Break(CompareAndAppendBreak::InlineBackpressure);
1851            }
1852        }
1853
1854        let mut merge_reqs = if batch.desc.upper() != batch.desc.lower() {
1855            self.trace.push_batch(batch.clone())
1856        } else {
1857            Vec::new()
1858        };
1859
1860        // NB: we don't claim unclaimed compactions when the recording flag is off, even if we'd
1861        // otherwise be allowed to, to avoid triggering the same compactions in every writer.
1862        let all_empty_reqs = merge_reqs
1863            .iter()
1864            .all(|req| req.inputs.iter().all(|b| b.batch.is_empty()));
1865        if all_empty_reqs && !batch.is_empty() {
1866            let mut reqs_to_take = claim_compaction_percent / 100;
1867            if (usize::cast_from(idempotency_token.hashed()) % 100)
1868                < (claim_compaction_percent % 100)
1869            {
1870                reqs_to_take += 1;
1871            }
1872            let threshold_ms = heartbeat_timestamp_ms.saturating_sub(lease_duration_ms);
1873            let min_writer = claim_compaction_min_version.map(WriterKey::for_version);
1874            merge_reqs.extend(
1875                // We keep the oldest `reqs_to_take` batches, under the theory that they're least
1876                // likely to be compacted soon for other reasons.
1877                self.trace
1878                    .fueled_merge_reqs_before_ms(threshold_ms, min_writer)
1879                    .take(reqs_to_take),
1880            )
1881        }
1882
1883        for req in &merge_reqs {
1884            self.trace.claim_compaction(
1885                req.id,
1886                ActiveCompaction {
1887                    start_ms: heartbeat_timestamp_ms,
1888                },
1889            )
1890        }
1891
1892        mz_ore::soft_assert_eq_no_log!(self.trace.upper(), batch.desc.upper());
1893        writer_state.most_recent_write_token = idempotency_token.clone();
1894        // The writer's most recent upper should only go forward.
1895        assert!(
1896            PartialOrder::less_equal(&writer_state.most_recent_write_upper, batch.desc.upper()),
1897            "{:?} vs {:?}",
1898            writer_state.most_recent_write_upper,
1899            batch.desc.upper()
1900        );
1901        writer_state
1902            .most_recent_write_upper
1903            .clone_from(batch.desc.upper());
1904
1905        // Heartbeat the writer state to keep our idempotency token alive.
1906        writer_state.last_heartbeat_timestamp_ms = std::cmp::max(
1907            heartbeat_timestamp_ms,
1908            writer_state.last_heartbeat_timestamp_ms,
1909        );
1910
1911        Continue(merge_reqs)
1912    }
1913
1914    pub fn apply_merge_res<D: Codec64 + Monoid + PartialEq>(
1915        &mut self,
1916        res: &FueledMergeRes<T>,
1917        metrics: &ColumnarMetrics,
1918    ) -> ControlFlow<NoOpStateTransition<ApplyMergeResult>, ApplyMergeResult> {
1919        // We expire all writers if the upper and since both advance to the
1920        // empty antichain. Gracefully handle this. At the same time,
1921        // short-circuit the cmd application so we don't needlessly create new
1922        // SeqNos.
1923        if self.is_tombstone() {
1924            return Break(NoOpStateTransition(ApplyMergeResult::NotAppliedNoMatch));
1925        }
1926
1927        let apply_merge_result = self.trace.apply_merge_res_checked::<D>(res, metrics);
1928        Continue(apply_merge_result)
1929    }
1930
1931    pub fn spine_exert(
1932        &mut self,
1933        fuel: usize,
1934    ) -> ControlFlow<NoOpStateTransition<Vec<FueledMergeReq<T>>>, Vec<FueledMergeReq<T>>> {
1935        let (merge_reqs, did_work) = self.trace.exert(fuel);
1936        if did_work {
1937            Continue(merge_reqs)
1938        } else {
1939            assert!(merge_reqs.is_empty());
1940            // Break if we have nothing useful to do to save the seqno (and
1941            // resulting crdb traffic)
1942            Break(NoOpStateTransition(Vec::new()))
1943        }
1944    }
1945
1946    pub fn downgrade_since(
1947        &mut self,
1948        reader_id: &LeasedReaderId,
1949        seqno: SeqNo,
1950        outstanding_seqno: SeqNo,
1951        new_since: &Antichain<T>,
1952        heartbeat_timestamp_ms: u64,
1953    ) -> ControlFlow<NoOpStateTransition<Since<T>>, Since<T>> {
1954        // We expire all readers if the upper and since both advance to the
1955        // empty antichain. Gracefully handle this. At the same time,
1956        // short-circuit the cmd application so we don't needlessly create new
1957        // SeqNos.
1958        if self.is_tombstone() {
1959            return Break(NoOpStateTransition(Since(Antichain::new())));
1960        }
1961
1962        // The only way to have a missing reader in state is if it's been expired... and in that
1963        // case, we behave the same as though that reader had been downgraded to the empty antichain.
1964        let Some(reader_state) = self.leased_reader(reader_id) else {
1965            tracing::warn!(
1966                "Leased reader {reader_id} was expired due to inactivity. Did the machine go to sleep?",
1967            );
1968            return Break(NoOpStateTransition(Since(Antichain::new())));
1969        };
1970
1971        // Also use this as an opportunity to heartbeat the reader and downgrade
1972        // the seqno capability.
1973        reader_state.last_heartbeat_timestamp_ms = std::cmp::max(
1974            heartbeat_timestamp_ms,
1975            reader_state.last_heartbeat_timestamp_ms,
1976        );
1977
1978        let seqno = {
1979            assert!(
1980                outstanding_seqno >= reader_state.seqno,
1981                "SeqNos cannot go backward; however, oldest leased SeqNo ({:?}) \
1982                    is behind current reader_state ({:?})",
1983                outstanding_seqno,
1984                reader_state.seqno,
1985            );
1986            std::cmp::min(outstanding_seqno, seqno)
1987        };
1988
1989        reader_state.seqno = seqno;
1990
1991        let reader_current_since = if PartialOrder::less_than(&reader_state.since, new_since) {
1992            reader_state.since.clone_from(new_since);
1993            self.update_since();
1994            new_since.clone()
1995        } else {
1996            // No-op, but still commit the state change so that this gets
1997            // linearized.
1998            reader_state.since.clone()
1999        };
2000
2001        Continue(Since(reader_current_since))
2002    }
2003
2004    pub fn compare_and_downgrade_since(
2005        &mut self,
2006        reader_id: &CriticalReaderId,
2007        expected_opaque: &Opaque,
2008        (new_opaque, new_since): (&Opaque, &Antichain<T>),
2009    ) -> ControlFlow<
2010        NoOpStateTransition<Result<Since<T>, (Opaque, Since<T>)>>,
2011        Result<Since<T>, (Opaque, Since<T>)>,
2012    > {
2013        // We expire all readers if the upper and since both advance to the
2014        // empty antichain. Gracefully handle this. At the same time,
2015        // short-circuit the cmd application so we don't needlessly create new
2016        // SeqNos.
2017        if self.is_tombstone() {
2018            // Match the idempotence behavior below of ignoring the token if
2019            // since is already advanced enough (in this case, because it's a
2020            // tombstone, we know it's the empty antichain).
2021            return Break(NoOpStateTransition(Ok(Since(Antichain::new()))));
2022        }
2023
2024        let reader_state = self.critical_reader(reader_id);
2025
2026        if reader_state.opaque != *expected_opaque {
2027            // No-op, but still commit the state change so that this gets
2028            // linearized.
2029            return Continue(Err((
2030                reader_state.opaque.clone(),
2031                Since(reader_state.since.clone()),
2032            )));
2033        }
2034
2035        reader_state.opaque = new_opaque.clone();
2036        if PartialOrder::less_equal(&reader_state.since, new_since) {
2037            reader_state.since.clone_from(new_since);
2038            self.update_since();
2039            Continue(Ok(Since(new_since.clone())))
2040        } else {
2041            // no work to be done -- the reader state's `since` is already sufficiently
2042            // advanced. we may someday need to revisit this branch when it's possible
2043            // for two `since` frontiers to be incomparable.
2044            Continue(Ok(Since(reader_state.since.clone())))
2045        }
2046    }
2047
2048    pub fn expire_leased_reader(
2049        &mut self,
2050        reader_id: &LeasedReaderId,
2051    ) -> ControlFlow<NoOpStateTransition<bool>, bool> {
2052        // We expire all readers if the upper and since both advance to the
2053        // empty antichain. Gracefully handle this. At the same time,
2054        // short-circuit the cmd application so we don't needlessly create new
2055        // SeqNos.
2056        if self.is_tombstone() {
2057            return Break(NoOpStateTransition(false));
2058        }
2059
2060        let existed = self.leased_readers.remove(reader_id).is_some();
2061        if existed {
2062            // TODO(database-issues#6885): Re-enable this
2063            //
2064            // Temporarily disabling this because we think it might be the cause
2065            // of the remap since bug. Specifically, a clusterd process has a
2066            // ReadHandle for maintaining the once and one inside a Listen. If
2067            // we crash and stay down for longer than the read lease duration,
2068            // it's possible that an expiry of them both in quick succession
2069            // jumps the since forward to the Listen one.
2070            //
2071            // Don't forget to update the downgrade_since when this gets
2072            // switched back on.
2073            //
2074            // self.update_since();
2075        }
2076        // No-op if existed is false, but still commit the state change so that
2077        // this gets linearized.
2078        Continue(existed)
2079    }
2080
2081    pub fn expire_critical_reader(
2082        &mut self,
2083        reader_id: &CriticalReaderId,
2084    ) -> ControlFlow<NoOpStateTransition<bool>, bool> {
2085        // We expire all readers if the upper and since both advance to the
2086        // empty antichain. Gracefully handle this. At the same time,
2087        // short-circuit the cmd application so we don't needlessly create new
2088        // SeqNos.
2089        if self.is_tombstone() {
2090            return Break(NoOpStateTransition(false));
2091        }
2092
2093        let existed = self.critical_readers.remove(reader_id).is_some();
2094        if existed {
2095            // TODO(database-issues#6885): Re-enable this
2096            //
2097            // Temporarily disabling this because we think it might be the cause
2098            // of the remap since bug. Specifically, a clusterd process has a
2099            // ReadHandle for maintaining the once and one inside a Listen. If
2100            // we crash and stay down for longer than the read lease duration,
2101            // it's possible that an expiry of them both in quick succession
2102            // jumps the since forward to the Listen one.
2103            //
2104            // Don't forget to update the downgrade_since when this gets
2105            // switched back on.
2106            //
2107            // self.update_since();
2108        }
2109        // This state transition is a no-op if existed is false, but we still
2110        // commit the state change so that this gets linearized (maybe we're
2111        // looking at old state).
2112        Continue(existed)
2113    }
2114
2115    pub fn expire_writer(
2116        &mut self,
2117        writer_id: &WriterId,
2118    ) -> ControlFlow<NoOpStateTransition<bool>, bool> {
2119        // We expire all writers if the upper and since both advance to the
2120        // empty antichain. Gracefully handle this. At the same time,
2121        // short-circuit the cmd application so we don't needlessly create new
2122        // SeqNos.
2123        if self.is_tombstone() {
2124            return Break(NoOpStateTransition(false));
2125        }
2126
2127        let existed = self.writers.remove(writer_id).is_some();
2128        // This state transition is a no-op if existed is false, but we still
2129        // commit the state change so that this gets linearized (maybe we're
2130        // looking at old state).
2131        Continue(existed)
2132    }
2133
2134    fn leased_reader(&mut self, id: &LeasedReaderId) -> Option<&mut LeasedReaderState<T>> {
2135        self.leased_readers.get_mut(id)
2136    }
2137
2138    fn critical_reader(&mut self, id: &CriticalReaderId) -> &mut CriticalReaderState<T> {
2139        self.critical_readers
2140            .get_mut(id)
2141            .unwrap_or_else(|| {
2142                panic!(
2143                    "Unknown CriticalReaderId({}). It was either never registered, or has been manually expired.",
2144                    id
2145                )
2146            })
2147    }
2148
2149    fn critical_since(&self) -> Option<Antichain<T>> {
2150        let mut critical_sinces = self.critical_readers.values().map(|r| &r.since);
2151        let mut since = critical_sinces.next().cloned()?;
2152        for s in critical_sinces {
2153            since.meet_assign(s);
2154        }
2155        Some(since)
2156    }
2157
2158    fn update_since(&mut self) {
2159        let mut sinces_iter = self
2160            .leased_readers
2161            .values()
2162            .map(|x| &x.since)
2163            .chain(self.critical_readers.values().map(|x| &x.since));
2164        let mut since = match sinces_iter.next() {
2165            Some(since) => since.clone(),
2166            None => {
2167                // If there are no current readers, leave `since` unchanged so
2168                // it doesn't regress.
2169                return;
2170            }
2171        };
2172        while let Some(s) = sinces_iter.next() {
2173            since.meet_assign(s);
2174        }
2175        self.trace.downgrade_since(&since);
2176    }
2177
2178    fn seqno_since(&self, seqno: SeqNo) -> SeqNo {
2179        let mut seqno_since = seqno;
2180        for cap in self.leased_readers.values() {
2181            seqno_since = std::cmp::min(seqno_since, cap.seqno);
2182        }
2183        // critical_readers don't hold a seqno capability.
2184        seqno_since
2185    }
2186
2187    fn tombstone_batch() -> HollowBatch<T> {
2188        HollowBatch::empty(Description::new(
2189            Antichain::from_elem(T::minimum()),
2190            Antichain::new(),
2191            Antichain::new(),
2192        ))
2193    }
2194
2195    pub(crate) fn is_tombstone(&self) -> bool {
2196        self.trace.upper().is_empty()
2197            && self.trace.since().is_empty()
2198            && self.writers.is_empty()
2199            && self.leased_readers.is_empty()
2200            && self.critical_readers.is_empty()
2201    }
2202
2203    pub(crate) fn is_single_empty_batch(&self) -> bool {
2204        let mut batch_count = 0;
2205        let mut is_empty = true;
2206        self.trace.map_batches(|b| {
2207            batch_count += 1;
2208            is_empty &= b.is_empty()
2209        });
2210        batch_count <= 1 && is_empty
2211    }
2212
2213    pub fn become_tombstone_and_shrink(&mut self) -> ControlFlow<NoOpStateTransition<()>, ()> {
2214        assert_eq!(self.trace.upper(), &Antichain::new());
2215        assert_eq!(self.trace.since(), &Antichain::new());
2216
2217        // Remember our current state, so we can decide whether we have to
2218        // record a transition in durable state.
2219        let was_tombstone = self.is_tombstone();
2220
2221        // Enter the "tombstone" state, if we're not in it already.
2222        self.writers.clear();
2223        self.leased_readers.clear();
2224        self.critical_readers.clear();
2225
2226        mz_ore::soft_assert_no_log!(self.is_tombstone());
2227
2228        // Now that we're in a "tombstone" state -- ie. nobody can read the data from a shard or write to
2229        // it -- the actual contents of our batches no longer matter.
2230        // This method progressively replaces batches in our state with simpler versions, to allow
2231        // freeing up resources and to reduce the state size. (Since the state is unreadable, this
2232        // is not visible to clients.) We do this a little bit at a time to avoid really large state
2233        // transitions... most operations happen incrementally, and large single writes can overwhelm
2234        // a backing store. See comments for why we believe the relevant diffs are reasonably small.
2235
2236        let mut to_replace = None;
2237        let mut batch_count = 0;
2238        self.trace.map_batches(|b| {
2239            batch_count += 1;
2240            if !b.is_empty() && to_replace.is_none() {
2241                to_replace = Some(b.desc.clone());
2242            }
2243        });
2244        if let Some(desc) = to_replace {
2245            // We have a nonempty batch: replace it with an empty batch and return.
2246            // This should not produce an excessively large diff: if it did, we wouldn't have been
2247            // able to append that batch in the first place.
2248            let result = self.trace.apply_tombstone_merge(&desc);
2249            assert!(
2250                result.matched(),
2251                "merge with a matching desc should always match"
2252            );
2253            Continue(())
2254        } else if batch_count > 1 {
2255            // All our batches are empty, but we have more than one of them. Replace the whole set
2256            // with a new single-batch trace.
2257            // This produces a diff with a size proportional to the number of batches, but since
2258            // Spine keeps a logarithmic number of batches this should never be excessively large.
2259            let mut new_trace = Trace::default();
2260            new_trace.downgrade_since(&Antichain::new());
2261            let merge_reqs = new_trace.push_batch(Self::tombstone_batch());
2262            assert_eq!(merge_reqs, Vec::new());
2263            self.trace = new_trace;
2264            Continue(())
2265        } else if !was_tombstone {
2266            // We were not tombstoned before, so have to make sure this state
2267            // transition is recorded.
2268            Continue(())
2269        } else {
2270            // All our batches are empty, and there's only one... there's no shrinking this
2271            // tombstone further.
2272            Break(NoOpStateTransition(()))
2273        }
2274    }
2275}
2276
2277// TODO: Document invariants.
2278#[derive(Debug)]
2279#[cfg_attr(any(test, debug_assertions), derive(Clone, PartialEq))]
2280pub struct State<T> {
2281    pub(crate) shard_id: ShardId,
2282
2283    pub(crate) seqno: SeqNo,
2284    /// A strictly increasing wall time of when this state was written, in
2285    /// milliseconds since the unix epoch.
2286    pub(crate) walltime_ms: u64,
2287    /// Hostname of the persist user that created this version of state. For
2288    /// debugging.
2289    pub(crate) hostname: String,
2290    pub(crate) collections: StateCollections<T>,
2291}
2292
2293/// A newtype wrapper of State that guarantees the K, V, and D codecs match the
2294/// ones in durable storage.
2295pub struct TypedState<K, V, T, D> {
2296    pub(crate) state: State<T>,
2297
2298    // According to the docs, PhantomData is to "mark things that act like they
2299    // own a T". State doesn't actually own K, V, or D, just the ability to
2300    // produce them. Using the `fn() -> T` pattern gets us the same variance as
2301    // T [1], but also allows State to correctly derive Send+Sync.
2302    //
2303    // [1]:
2304    //     https://doc.rust-lang.org/nomicon/phantom-data.html#table-of-phantomdata-patterns
2305    pub(crate) _phantom: PhantomData<fn() -> (K, V, D)>,
2306}
2307
2308impl<K, V, T: Clone, D> TypedState<K, V, T, D> {
2309    #[cfg(any(test, debug_assertions))]
2310    pub(crate) fn clone(&self, hostname: String) -> Self {
2311        TypedState {
2312            state: State {
2313                shard_id: self.shard_id.clone(),
2314                seqno: self.seqno.clone(),
2315                walltime_ms: self.walltime_ms,
2316                hostname,
2317                collections: self.collections.clone(),
2318            },
2319            _phantom: PhantomData,
2320        }
2321    }
2322
2323    pub(crate) fn clone_for_rollup(&self) -> Self {
2324        TypedState {
2325            state: State {
2326                shard_id: self.shard_id.clone(),
2327                seqno: self.seqno.clone(),
2328                walltime_ms: self.walltime_ms,
2329                hostname: self.hostname.clone(),
2330                collections: self.collections.clone(),
2331            },
2332            _phantom: PhantomData,
2333        }
2334    }
2335}
2336
2337impl<K, V, T: Debug, D> Debug for TypedState<K, V, T, D> {
2338    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
2339        // Deconstruct self so we get a compile failure if new fields
2340        // are added.
2341        let TypedState { state, _phantom } = self;
2342        f.debug_struct("TypedState").field("state", state).finish()
2343    }
2344}
2345
2346// Impl PartialEq regardless of the type params.
2347#[cfg(any(test, debug_assertions))]
2348impl<K, V, T: PartialEq, D> PartialEq for TypedState<K, V, T, D> {
2349    fn eq(&self, other: &Self) -> bool {
2350        // Deconstruct self and other so we get a compile failure if new fields
2351        // are added.
2352        let TypedState {
2353            state: self_state,
2354            _phantom,
2355        } = self;
2356        let TypedState {
2357            state: other_state,
2358            _phantom,
2359        } = other;
2360        self_state == other_state
2361    }
2362}
2363
2364impl<K, V, T, D> Deref for TypedState<K, V, T, D> {
2365    type Target = State<T>;
2366
2367    fn deref(&self) -> &Self::Target {
2368        &self.state
2369    }
2370}
2371
2372impl<K, V, T, D> DerefMut for TypedState<K, V, T, D> {
2373    fn deref_mut(&mut self) -> &mut Self::Target {
2374        &mut self.state
2375    }
2376}
2377
2378impl<K, V, T, D> TypedState<K, V, T, D>
2379where
2380    K: Codec,
2381    V: Codec,
2382    T: Timestamp + Lattice + Codec64,
2383    D: Codec64,
2384{
2385    pub fn new(
2386        applier_version: Version,
2387        shard_id: ShardId,
2388        hostname: String,
2389        walltime_ms: u64,
2390    ) -> Self {
2391        let state = State {
2392            shard_id,
2393            seqno: SeqNo::minimum(),
2394            walltime_ms,
2395            hostname,
2396            collections: StateCollections {
2397                version: applier_version,
2398                last_gc_req: SeqNo::minimum(),
2399                rollups: BTreeMap::new(),
2400                active_rollup: None,
2401                active_gc: None,
2402                leased_readers: BTreeMap::new(),
2403                critical_readers: BTreeMap::new(),
2404                writers: BTreeMap::new(),
2405                schemas: BTreeMap::new(),
2406                trace: Trace::default(),
2407            },
2408        };
2409        TypedState {
2410            state,
2411            _phantom: PhantomData,
2412        }
2413    }
2414
2415    pub fn clone_apply<R, E, WorkFn>(
2416        &self,
2417        cfg: &PersistConfig,
2418        work_fn: &mut WorkFn,
2419    ) -> ControlFlow<E, (R, Self)>
2420    where
2421        WorkFn: FnMut(SeqNo, &PersistConfig, &mut StateCollections<T>) -> ControlFlow<E, R>,
2422    {
2423        // We do not increment the version by default, though work_fn can if it chooses to.
2424        let mut new_state = State {
2425            shard_id: self.shard_id,
2426            seqno: self.seqno.next(),
2427            walltime_ms: (cfg.now)(),
2428            hostname: cfg.hostname.clone(),
2429            collections: self.collections.clone(),
2430        };
2431
2432        // Make sure walltime_ms is strictly increasing, in case clocks are
2433        // offset.
2434        if new_state.walltime_ms <= self.walltime_ms {
2435            new_state.walltime_ms = self.walltime_ms + 1;
2436        }
2437
2438        let work_ret = work_fn(new_state.seqno, cfg, &mut new_state.collections)?;
2439        let new_state = TypedState {
2440            state: new_state,
2441            _phantom: PhantomData,
2442        };
2443        Continue((work_ret, new_state))
2444    }
2445}
2446
2447#[derive(Copy, Clone, Debug)]
2448pub struct GcConfig {
2449    pub use_active_gc: bool,
2450    pub fallback_threshold_ms: u64,
2451    pub min_versions: usize,
2452    pub max_versions: usize,
2453}
2454
2455impl<T> State<T>
2456where
2457    T: Timestamp + Lattice + Codec64,
2458{
2459    pub fn shard_id(&self) -> ShardId {
2460        self.shard_id
2461    }
2462
2463    pub fn seqno(&self) -> SeqNo {
2464        self.seqno
2465    }
2466
2467    pub fn since(&self) -> &Antichain<T> {
2468        self.collections.trace.since()
2469    }
2470
2471    pub fn upper(&self) -> &Antichain<T> {
2472        self.collections.trace.upper()
2473    }
2474
2475    pub fn spine_batch_count(&self) -> usize {
2476        self.collections.trace.num_spine_batches()
2477    }
2478
2479    pub fn size_metrics(&self) -> StateSizeMetrics {
2480        let mut ret = StateSizeMetrics::default();
2481        self.blobs().for_each(|x| match x {
2482            HollowBlobRef::Batch(x) => {
2483                ret.hollow_batch_count += 1;
2484                ret.batch_part_count += x.part_count();
2485                ret.num_updates += x.len;
2486
2487                let batch_size = x.encoded_size_bytes();
2488                for x in x.parts.iter() {
2489                    if x.ts_rewrite().is_some() {
2490                        ret.rewrite_part_count += 1;
2491                    }
2492                    if x.is_inline() {
2493                        ret.inline_part_count += 1;
2494                        ret.inline_part_bytes += x.inline_bytes();
2495                    }
2496                }
2497                ret.largest_batch_bytes = std::cmp::max(ret.largest_batch_bytes, batch_size);
2498                ret.state_batches_bytes += batch_size;
2499            }
2500            HollowBlobRef::Rollup(x) => {
2501                ret.state_rollup_count += 1;
2502                ret.state_rollups_bytes += x.encoded_size_bytes.unwrap_or_default()
2503            }
2504        });
2505        ret
2506    }
2507
2508    pub fn latest_rollup(&self) -> (&SeqNo, &HollowRollup) {
2509        // We maintain the invariant that every version of state has at least
2510        // one rollup.
2511        self.collections
2512            .rollups
2513            .iter()
2514            .rev()
2515            .next()
2516            .expect("State should have at least one rollup if seqno > minimum")
2517    }
2518
2519    pub(crate) fn seqno_since(&self) -> SeqNo {
2520        self.collections.seqno_since(self.seqno)
2521    }
2522
2523    // Returns whether the cmd proposing this state has been selected to perform
2524    // background garbage collection work.
2525    //
2526    // If it was selected, this information is recorded in the state itself for
2527    // commit along with the cmd's state transition. This helps us to avoid
2528    // redundant work.
2529    //
2530    // Correctness does not depend on a gc assignment being executed, nor on
2531    // them being executed in the order they are given. But it is expected that
2532    // gc assignments are best-effort respected. In practice, cmds like
2533    // register_foo or expire_foo, where it would be awkward, ignore gc.
2534    pub fn maybe_gc(&mut self, is_write: bool, now: u64, cfg: GcConfig) -> Option<GcReq> {
2535        let GcConfig {
2536            use_active_gc,
2537            fallback_threshold_ms,
2538            min_versions,
2539            max_versions,
2540        } = cfg;
2541        // This is an arbitrary-ish threshold that scales with seqno, but never
2542        // gets particularly big. It probably could be much bigger and certainly
2543        // could use a tuning pass at some point.
2544        let gc_threshold = if use_active_gc {
2545            u64::cast_from(min_versions)
2546        } else {
2547            std::cmp::max(
2548                1,
2549                u64::cast_from(self.seqno.0.next_power_of_two().trailing_zeros()),
2550            )
2551        };
2552        let new_seqno_since = self.seqno_since();
2553        // Collect until the new seqno since... or the old since plus the max number of versions,
2554        // whatever is less.
2555        let gc_until_seqno = new_seqno_since.min(SeqNo(
2556            self.collections
2557                .last_gc_req
2558                .0
2559                .saturating_add(u64::cast_from(max_versions)),
2560        ));
2561        let should_gc = new_seqno_since
2562            .0
2563            .saturating_sub(self.collections.last_gc_req.0)
2564            >= gc_threshold;
2565
2566        // If we wouldn't otherwise gc, check if we have an active gc. If we do, and
2567        // it's been a while since it started, we should gc.
2568        let should_gc = if use_active_gc && !should_gc {
2569            match self.collections.active_gc {
2570                Some(active_gc) => now.saturating_sub(active_gc.start_ms) > fallback_threshold_ms,
2571                None => false,
2572            }
2573        } else {
2574            should_gc
2575        };
2576        // Assign GC traffic preferentially to writers, falling back to anyone
2577        // generating new state versions if there are no writers.
2578        let should_gc = should_gc && (is_write || self.collections.writers.is_empty());
2579        // Always assign GC work to a tombstoned shard to have the chance to
2580        // clean up any residual blobs. This is safe (won't cause excess gc)
2581        // as the only allowed command after becoming a tombstone is to write
2582        // the final rollup.
2583        let tombstone_needs_gc = self.collections.is_tombstone();
2584        let should_gc = should_gc || tombstone_needs_gc;
2585        let should_gc = if use_active_gc {
2586            // If we have an active gc, we should only gc if the active gc is
2587            // sufficiently old. This is to avoid doing more gc work than
2588            // necessary.
2589            should_gc
2590                && match self.collections.active_gc {
2591                    Some(active) => now.saturating_sub(active.start_ms) > fallback_threshold_ms,
2592                    None => true,
2593                }
2594        } else {
2595            should_gc
2596        };
2597        if should_gc {
2598            self.collections.last_gc_req = gc_until_seqno;
2599            Some(GcReq {
2600                shard_id: self.shard_id,
2601                new_seqno_since: gc_until_seqno,
2602            })
2603        } else {
2604            None
2605        }
2606    }
2607
2608    /// Return the number of gc-ineligible state versions.
2609    pub fn seqnos_held(&self) -> usize {
2610        usize::cast_from(self.seqno.0.saturating_sub(self.seqno_since().0))
2611    }
2612
2613    /// Expire all readers and writers up to the given walltime_ms.
2614    pub fn expire_at(&mut self, walltime_ms: EpochMillis) -> ExpiryMetrics {
2615        let mut metrics = ExpiryMetrics::default();
2616        let shard_id = self.shard_id();
2617        self.collections.leased_readers.retain(|id, state| {
2618            let retain = state.last_heartbeat_timestamp_ms + state.lease_duration_ms >= walltime_ms;
2619            if !retain {
2620                info!(
2621                    "Force expiring reader {id} ({}) of shard {shard_id} due to inactivity",
2622                    state.debug.purpose
2623                );
2624                metrics.readers_expired += 1;
2625            }
2626            retain
2627        });
2628        // critical_readers don't need forced expiration. (In fact, that's the point!)
2629        self.collections.writers.retain(|id, state| {
2630            let retain =
2631                (state.last_heartbeat_timestamp_ms + state.lease_duration_ms) >= walltime_ms;
2632            if !retain {
2633                info!(
2634                    "Force expiring writer {id} ({}) of shard {shard_id} due to inactivity",
2635                    state.debug.purpose
2636                );
2637                metrics.writers_expired += 1;
2638            }
2639            retain
2640        });
2641        metrics
2642    }
2643
2644    /// Returns the batches that contain updates up to (and including) the given `as_of`. The
2645    /// result `Vec` contains blob keys, along with a [`Description`] of what updates in the
2646    /// referenced parts are valid to read.
2647    pub fn snapshot(&self, as_of: &Antichain<T>) -> Result<Vec<HollowBatch<T>>, SnapshotErr<T>> {
2648        if PartialOrder::less_than(as_of, self.collections.trace.since()) {
2649            return Err(SnapshotErr::AsOfHistoricalDistinctionsLost(Since(
2650                self.collections.trace.since().clone(),
2651            )));
2652        }
2653        let upper = self.collections.trace.upper();
2654        if PartialOrder::less_equal(upper, as_of) {
2655            return Err(SnapshotErr::AsOfNotYetAvailable(
2656                self.seqno,
2657                Upper(upper.clone()),
2658            ));
2659        }
2660
2661        let batches = self
2662            .collections
2663            .trace
2664            .batches()
2665            .filter(|b| !PartialOrder::less_than(as_of, b.desc.lower()))
2666            .cloned()
2667            .collect();
2668        Ok(batches)
2669    }
2670
2671    // NB: Unlike the other methods here, this one is read-only.
2672    pub fn verify_listen(&self, as_of: &Antichain<T>) -> Result<(), Since<T>> {
2673        if PartialOrder::less_than(as_of, self.collections.trace.since()) {
2674            return Err(Since(self.collections.trace.since().clone()));
2675        }
2676        Ok(())
2677    }
2678
2679    pub fn next_listen_batch(&self, frontier: &Antichain<T>) -> Result<HollowBatch<T>, SeqNo> {
2680        // TODO: Avoid the O(n^2) here: `next_listen_batch` is called once per
2681        // batch and this iterates through all batches to find the next one.
2682        self.collections
2683            .trace
2684            .batches()
2685            .find(|b| {
2686                PartialOrder::less_equal(b.desc.lower(), frontier)
2687                    && PartialOrder::less_than(frontier, b.desc.upper())
2688            })
2689            .cloned()
2690            .ok_or(self.seqno)
2691    }
2692
2693    pub fn active_rollup(&self) -> Option<ActiveRollup> {
2694        self.collections.active_rollup
2695    }
2696
2697    pub fn need_rollup(
2698        &self,
2699        threshold: usize,
2700        use_active_rollup: bool,
2701        fallback_threshold_ms: u64,
2702        now: u64,
2703    ) -> Option<SeqNo> {
2704        let (latest_rollup_seqno, _) = self.latest_rollup();
2705
2706        // Tombstoned shards require one final rollup. However, because we
2707        // write a rollup as of SeqNo X and then link it in using a state
2708        // transition (in this case from X to X+1), the minimum number of
2709        // live diffs is actually two. Detect when we're in this minimal
2710        // two diff state and stop the (otherwise) infinite iteration.
2711        if self.collections.is_tombstone() && latest_rollup_seqno.next() < self.seqno {
2712            return Some(self.seqno);
2713        }
2714
2715        let seqnos_since_last_rollup = self.seqno.0.saturating_sub(latest_rollup_seqno.0);
2716
2717        if use_active_rollup {
2718            // If sequnos_since_last_rollup>threshold, and there is no existing rollup in progress,
2719            // we should start a new rollup.
2720            // If there is an active rollup, we should check if it has been running too long.
2721            // If it has, we should start a new rollup.
2722            // This is to guard against a worker dying/taking too long/etc.
2723            if seqnos_since_last_rollup > u64::cast_from(threshold) {
2724                match self.active_rollup() {
2725                    Some(active_rollup) => {
2726                        if now.saturating_sub(active_rollup.start_ms) > fallback_threshold_ms {
2727                            return Some(self.seqno);
2728                        }
2729                    }
2730                    None => {
2731                        return Some(self.seqno);
2732                    }
2733                }
2734            }
2735        } else {
2736            // every `threshold` seqnos since the latest rollup, assign rollup maintenance.
2737            // we avoid assigning rollups to every seqno past the threshold to avoid handles
2738            // racing / performing redundant work.
2739            if seqnos_since_last_rollup > 0
2740                && seqnos_since_last_rollup % u64::cast_from(threshold) == 0
2741            {
2742                return Some(self.seqno);
2743            }
2744
2745            // however, since maintenance is best-effort and could fail, do assign rollup
2746            // work to every seqno after a fallback threshold to ensure one is written.
2747            if seqnos_since_last_rollup
2748                > u64::cast_from(
2749                    threshold * PersistConfig::DEFAULT_FALLBACK_ROLLUP_THRESHOLD_MULTIPLIER,
2750                )
2751            {
2752                return Some(self.seqno);
2753            }
2754        }
2755
2756        None
2757    }
2758
2759    pub(crate) fn blobs(&self) -> impl Iterator<Item = HollowBlobRef<'_, T>> {
2760        let batches = self.collections.trace.batches().map(HollowBlobRef::Batch);
2761        let rollups = self.collections.rollups.values().map(HollowBlobRef::Rollup);
2762        batches.chain(rollups)
2763    }
2764}
2765
2766fn serialize_part_bytes<S: Serializer>(val: &[u8], s: S) -> Result<S::Ok, S::Error> {
2767    let val = hex::encode(val);
2768    val.serialize(s)
2769}
2770
2771fn serialize_lazy_proto<S: Serializer, T: prost::Message + Default>(
2772    val: &Option<LazyProto<T>>,
2773    s: S,
2774) -> Result<S::Ok, S::Error> {
2775    val.as_ref()
2776        .map(|lazy| hex::encode(&lazy.into_proto()))
2777        .serialize(s)
2778}
2779
2780fn serialize_part_stats<S: Serializer>(
2781    val: &Option<LazyPartStats>,
2782    s: S,
2783) -> Result<S::Ok, S::Error> {
2784    // These bytes come from blob and are never validated on the way in, so a
2785    // malformed or newer-version encoding reaches here intact. Report it as
2786    // absent rather than panicking, and keep the field's shape stable for
2787    // consumers of the inspect-state output by logging the failure instead of
2788    // serializing a differently typed value in its place.
2789    let stats = val.as_ref().and_then(|x| match x.try_decode() {
2790        Ok(stats) => Some(stats.key),
2791        Err(err) => {
2792            tracing::warn!("undecodable part stats, reporting as absent: {err}");
2793            None
2794        }
2795    });
2796    stats.serialize(s)
2797}
2798
2799fn serialize_diffs_sum<S: Serializer>(val: &Option<[u8; 8]>, s: S) -> Result<S::Ok, S::Error> {
2800    // This is only used for debugging, so hack to assume that D is i64.
2801    let val = val.map(i64::decode);
2802    val.serialize(s)
2803}
2804
2805// This Serialize impl is used for debugging/testing and exposed via SQL. It's
2806// intentionally gated from users, so not strictly subject to our backward
2807// compatibility guarantees, but still probably best to be thoughtful about
2808// making unnecessary changes. Additionally, it's nice to make the output as
2809// nice to use as possible without tying our hands for the actual code usages.
2810impl<T: Serialize + Timestamp + Lattice> Serialize for State<T> {
2811    fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
2812        let State {
2813            shard_id,
2814            seqno,
2815            walltime_ms,
2816            hostname,
2817            collections:
2818                StateCollections {
2819                    version: applier_version,
2820                    last_gc_req,
2821                    rollups,
2822                    active_rollup,
2823                    active_gc,
2824                    leased_readers,
2825                    critical_readers,
2826                    writers,
2827                    schemas,
2828                    trace,
2829                },
2830        } = self;
2831        let mut s = s.serialize_struct("State", 13)?;
2832        let () = s.serialize_field("applier_version", &applier_version.to_string())?;
2833        let () = s.serialize_field("shard_id", shard_id)?;
2834        let () = s.serialize_field("seqno", seqno)?;
2835        let () = s.serialize_field("walltime_ms", walltime_ms)?;
2836        let () = s.serialize_field("hostname", hostname)?;
2837        let () = s.serialize_field("last_gc_req", last_gc_req)?;
2838        let () = s.serialize_field("rollups", rollups)?;
2839        let () = s.serialize_field("active_rollup", active_rollup)?;
2840        let () = s.serialize_field("active_gc", active_gc)?;
2841        let () = s.serialize_field("leased_readers", leased_readers)?;
2842        let () = s.serialize_field("critical_readers", critical_readers)?;
2843        let () = s.serialize_field("writers", writers)?;
2844        let () = s.serialize_field("schemas", schemas)?;
2845        let () = s.serialize_field("since", &trace.since().elements())?;
2846        let () = s.serialize_field("upper", &trace.upper().elements())?;
2847        let trace = trace.flatten();
2848        let () = s.serialize_field("batches", &trace.legacy_batches.keys().collect::<Vec<_>>())?;
2849        let () = s.serialize_field("hollow_batches", &trace.hollow_batches)?;
2850        let () = s.serialize_field("spine_batches", &trace.spine_batches)?;
2851        let () = s.serialize_field("merges", &trace.merges)?;
2852        s.end()
2853    }
2854}
2855
2856#[derive(Debug, Default)]
2857pub struct StateSizeMetrics {
2858    pub hollow_batch_count: usize,
2859    pub batch_part_count: usize,
2860    pub rewrite_part_count: usize,
2861    pub num_updates: usize,
2862    pub largest_batch_bytes: usize,
2863    pub state_batches_bytes: usize,
2864    pub state_rollups_bytes: usize,
2865    pub state_rollup_count: usize,
2866    pub inline_part_count: usize,
2867    pub inline_part_bytes: usize,
2868}
2869
2870#[derive(Default)]
2871pub struct ExpiryMetrics {
2872    pub(crate) readers_expired: usize,
2873    pub(crate) writers_expired: usize,
2874}
2875
2876/// Wrapper for Antichain that represents a Since
2877#[derive(Debug, Clone, PartialEq)]
2878pub struct Since<T>(pub Antichain<T>);
2879
2880/// Wrapper for Antichain that represents an Upper
2881#[derive(Debug, PartialEq)]
2882pub struct Upper<T>(pub Antichain<T>);
2883
2884#[cfg(test)]
2885pub(crate) mod tests {
2886    use std::ops::Range;
2887    use std::str::FromStr;
2888
2889    use bytes::Bytes;
2890    use mz_build_info::DUMMY_BUILD_INFO;
2891    use mz_dyncfg::ConfigUpdates;
2892    use mz_ore::now::SYSTEM_TIME;
2893    use mz_ore::{assert_none, assert_ok};
2894    use mz_proto::RustType;
2895    use proptest::prelude::*;
2896    use proptest::strategy::ValueTree;
2897
2898    use crate::InvalidUsage::{InvalidBounds, InvalidEmptyTimeInterval};
2899    use crate::cache::PersistClientCache;
2900    use crate::internal::encoding::any_some_lazy_part_stats;
2901    use crate::internal::paths::RollupId;
2902    use crate::internal::trace::tests::any_trace;
2903    use crate::tests::new_test_client_cache;
2904    use crate::{Diagnostics, PersistLocation};
2905
2906    use super::*;
2907
2908    const LEASE_DURATION_MS: u64 = 900 * 1000;
2909    fn debug_state() -> HandleDebugState {
2910        HandleDebugState {
2911            hostname: "debug".to_owned(),
2912            purpose: "finding the bugs".to_owned(),
2913        }
2914    }
2915
2916    pub fn any_hollow_batch_with_exact_runs<T: Arbitrary + Timestamp>(
2917        num_runs: usize,
2918    ) -> impl Strategy<Value = HollowBatch<T>> {
2919        (
2920            any::<T>(),
2921            any::<T>(),
2922            any::<T>(),
2923            proptest::collection::vec(any_run_part::<T>(), num_runs + 1..20),
2924            any::<usize>(),
2925        )
2926            .prop_map(move |(t0, t1, since, parts, len)| {
2927                let (lower, upper) = if t0 <= t1 {
2928                    (Antichain::from_elem(t0), Antichain::from_elem(t1))
2929                } else {
2930                    (Antichain::from_elem(t1), Antichain::from_elem(t0))
2931                };
2932                let since = Antichain::from_elem(since);
2933
2934                let run_splits = (1..num_runs)
2935                    .map(|i| i * parts.len() / num_runs)
2936                    .collect::<Vec<_>>();
2937
2938                let run_meta = (0..num_runs)
2939                    .map(|_| {
2940                        let mut meta = RunMeta::default();
2941                        meta.id = Some(RunId::new());
2942                        meta
2943                    })
2944                    .collect::<Vec<_>>();
2945
2946                HollowBatch::new(
2947                    Description::new(lower, upper, since),
2948                    parts,
2949                    len % 10,
2950                    run_meta,
2951                    run_splits,
2952                )
2953            })
2954    }
2955
2956    pub fn any_hollow_batch<T: Arbitrary + Timestamp>() -> impl Strategy<Value = HollowBatch<T>> {
2957        Strategy::prop_map(
2958            (
2959                any::<T>(),
2960                any::<T>(),
2961                any::<T>(),
2962                proptest::collection::vec(any_run_part::<T>(), 0..20),
2963                any::<usize>(),
2964                0..=10usize,
2965                proptest::collection::vec(any::<RunId>(), 10),
2966            ),
2967            |(t0, t1, since, parts, len, num_runs, run_ids)| {
2968                let (lower, upper) = if t0 <= t1 {
2969                    (Antichain::from_elem(t0), Antichain::from_elem(t1))
2970                } else {
2971                    (Antichain::from_elem(t1), Antichain::from_elem(t0))
2972                };
2973                let since = Antichain::from_elem(since);
2974                if num_runs > 0 && parts.len() > 2 && num_runs < parts.len() {
2975                    let run_splits = (1..num_runs)
2976                        .map(|i| i * parts.len() / num_runs)
2977                        .collect::<Vec<_>>();
2978
2979                    let run_meta = (0..num_runs)
2980                        .enumerate()
2981                        .map(|(i, _)| {
2982                            let mut meta = RunMeta::default();
2983                            meta.id = Some(run_ids[i]);
2984                            meta
2985                        })
2986                        .collect::<Vec<_>>();
2987
2988                    HollowBatch::new(
2989                        Description::new(lower, upper, since),
2990                        parts,
2991                        len % 10,
2992                        run_meta,
2993                        run_splits,
2994                    )
2995                } else {
2996                    HollowBatch::new_run_for_test(
2997                        Description::new(lower, upper, since),
2998                        parts,
2999                        len % 10,
3000                        run_ids[0],
3001                    )
3002                }
3003            },
3004        )
3005    }
3006
3007    pub fn any_batch_part<T: Arbitrary + Timestamp>() -> impl Strategy<Value = BatchPart<T>> {
3008        Strategy::prop_map(
3009            (
3010                any::<bool>(),
3011                any_hollow_batch_part(),
3012                any::<Option<T>>(),
3013                any::<Option<SchemaId>>(),
3014                any::<Option<SchemaId>>(),
3015            ),
3016            |(is_hollow, hollow, ts_rewrite, schema_id, deprecated_schema_id)| {
3017                if is_hollow {
3018                    BatchPart::Hollow(hollow)
3019                } else {
3020                    let updates = LazyInlineBatchPart::from_proto(Bytes::new()).unwrap();
3021                    let ts_rewrite = ts_rewrite.map(Antichain::from_elem);
3022                    BatchPart::Inline {
3023                        updates,
3024                        ts_rewrite,
3025                        schema_id,
3026                        deprecated_schema_id,
3027                    }
3028                }
3029            },
3030        )
3031    }
3032
3033    pub fn any_run_part<T: Arbitrary + Timestamp>() -> impl Strategy<Value = RunPart<T>> {
3034        Strategy::prop_map(any_batch_part(), |part| RunPart::Single(part))
3035    }
3036
3037    pub fn any_hollow_batch_part<T: Arbitrary + Timestamp>()
3038    -> impl Strategy<Value = HollowBatchPart<T>> {
3039        Strategy::prop_map(
3040            (
3041                any::<PartialBatchKey>(),
3042                any::<usize>(),
3043                any::<Vec<u8>>(),
3044                any_some_lazy_part_stats(),
3045                any::<Option<T>>(),
3046                any::<[u8; 8]>(),
3047                any::<Option<BatchColumnarFormat>>(),
3048                any::<Option<SchemaId>>(),
3049                any::<Option<SchemaId>>(),
3050            ),
3051            |(
3052                key,
3053                encoded_size_bytes,
3054                key_lower,
3055                stats,
3056                ts_rewrite,
3057                diffs_sum,
3058                format,
3059                schema_id,
3060                deprecated_schema_id,
3061            )| {
3062                HollowBatchPart {
3063                    key,
3064                    meta: Default::default(),
3065                    encoded_size_bytes,
3066                    key_lower,
3067                    structured_key_lower: None,
3068                    stats,
3069                    ts_rewrite: ts_rewrite.map(Antichain::from_elem),
3070                    diffs_sum: Some(diffs_sum),
3071                    format,
3072                    schema_id,
3073                    deprecated_schema_id,
3074                }
3075            },
3076        )
3077    }
3078
3079    pub fn any_leased_reader_state<T: Arbitrary>() -> impl Strategy<Value = LeasedReaderState<T>> {
3080        Strategy::prop_map(
3081            (
3082                any::<SeqNo>(),
3083                any::<Option<T>>(),
3084                any::<u64>(),
3085                any::<u64>(),
3086                any::<HandleDebugState>(),
3087            ),
3088            |(seqno, since, last_heartbeat_timestamp_ms, mut lease_duration_ms, debug)| {
3089                // lease_duration_ms of 0 means this state was written by an old
3090                // version of code, which means we'll migrate it in the decode
3091                // path. Avoid.
3092                if lease_duration_ms == 0 {
3093                    lease_duration_ms += 1;
3094                }
3095                LeasedReaderState {
3096                    seqno,
3097                    since: since.map_or_else(Antichain::new, Antichain::from_elem),
3098                    last_heartbeat_timestamp_ms,
3099                    lease_duration_ms,
3100                    debug,
3101                }
3102            },
3103        )
3104    }
3105
3106    pub fn any_critical_reader_state<T>() -> impl Strategy<Value = CriticalReaderState<T>>
3107    where
3108        T: Arbitrary,
3109    {
3110        Strategy::prop_map(
3111            (
3112                any::<Option<T>>(),
3113                any::<Opaque>(),
3114                any::<HandleDebugState>(),
3115            ),
3116            |(since, opaque, debug)| CriticalReaderState {
3117                since: since.map_or_else(Antichain::new, Antichain::from_elem),
3118                opaque,
3119                debug,
3120            },
3121        )
3122    }
3123
3124    pub fn any_writer_state<T: Arbitrary>() -> impl Strategy<Value = WriterState<T>> {
3125        Strategy::prop_map(
3126            (
3127                any::<u64>(),
3128                any::<u64>(),
3129                any::<IdempotencyToken>(),
3130                any::<Option<T>>(),
3131                any::<HandleDebugState>(),
3132            ),
3133            |(
3134                last_heartbeat_timestamp_ms,
3135                lease_duration_ms,
3136                most_recent_write_token,
3137                most_recent_write_upper,
3138                debug,
3139            )| WriterState {
3140                last_heartbeat_timestamp_ms,
3141                lease_duration_ms,
3142                most_recent_write_token,
3143                most_recent_write_upper: most_recent_write_upper
3144                    .map_or_else(Antichain::new, Antichain::from_elem),
3145                debug,
3146            },
3147        )
3148    }
3149
3150    pub fn any_encoded_schemas() -> impl Strategy<Value = EncodedSchemas> {
3151        Strategy::prop_map(
3152            (
3153                any::<Vec<u8>>(),
3154                any::<Vec<u8>>(),
3155                any::<Vec<u8>>(),
3156                any::<Vec<u8>>(),
3157            ),
3158            |(key, key_data_type, val, val_data_type)| EncodedSchemas {
3159                key: Bytes::from(key),
3160                key_data_type: Bytes::from(key_data_type),
3161                val: Bytes::from(val),
3162                val_data_type: Bytes::from(val_data_type),
3163            },
3164        )
3165    }
3166
3167    pub fn any_state<T: Arbitrary + Timestamp + Lattice>(
3168        num_trace_batches: Range<usize>,
3169    ) -> impl Strategy<Value = State<T>> {
3170        let part1 = (
3171            any::<ShardId>(),
3172            any::<SeqNo>(),
3173            any::<u64>(),
3174            any::<String>(),
3175            any::<SeqNo>(),
3176            proptest::collection::btree_map(any::<SeqNo>(), any::<HollowRollup>(), 1..3),
3177            proptest::option::of(any::<ActiveRollup>()),
3178        );
3179
3180        let part2 = (
3181            proptest::option::of(any::<ActiveGc>()),
3182            proptest::collection::btree_map(
3183                any::<LeasedReaderId>(),
3184                any_leased_reader_state::<T>(),
3185                1..3,
3186            ),
3187            proptest::collection::btree_map(
3188                any::<CriticalReaderId>(),
3189                any_critical_reader_state::<T>(),
3190                1..3,
3191            ),
3192            proptest::collection::btree_map(any::<WriterId>(), any_writer_state::<T>(), 0..3),
3193            proptest::collection::btree_map(any::<SchemaId>(), any_encoded_schemas(), 0..3),
3194            any_trace::<T>(num_trace_batches),
3195        );
3196
3197        (part1, part2).prop_map(
3198            |(
3199                (shard_id, seqno, walltime_ms, hostname, last_gc_req, rollups, active_rollup),
3200                (active_gc, leased_readers, critical_readers, writers, schemas, trace),
3201            )| State {
3202                shard_id,
3203                seqno,
3204                walltime_ms,
3205                hostname,
3206                collections: StateCollections {
3207                    version: Version::new(1, 2, 3),
3208                    last_gc_req,
3209                    rollups,
3210                    active_rollup,
3211                    active_gc,
3212                    leased_readers,
3213                    critical_readers,
3214                    writers,
3215                    schemas,
3216                    trace,
3217                },
3218            },
3219        )
3220    }
3221
3222    pub(crate) fn hollow<T: Timestamp>(
3223        lower: T,
3224        upper: T,
3225        keys: &[&str],
3226        len: usize,
3227    ) -> HollowBatch<T> {
3228        HollowBatch::new_run(
3229            Description::new(
3230                Antichain::from_elem(lower),
3231                Antichain::from_elem(upper),
3232                Antichain::from_elem(T::minimum()),
3233            ),
3234            keys.iter()
3235                .map(|x| {
3236                    RunPart::Single(BatchPart::Hollow(HollowBatchPart {
3237                        key: PartialBatchKey((*x).to_owned()),
3238                        meta: Default::default(),
3239                        encoded_size_bytes: 0,
3240                        key_lower: vec![],
3241                        structured_key_lower: None,
3242                        stats: None,
3243                        ts_rewrite: None,
3244                        diffs_sum: None,
3245                        format: None,
3246                        schema_id: None,
3247                        deprecated_schema_id: None,
3248                    }))
3249                })
3250                .collect(),
3251            len,
3252        )
3253    }
3254
3255    #[mz_ore::test]
3256    fn downgrade_since() {
3257        let mut state = TypedState::<(), (), u64, i64>::new(
3258            DUMMY_BUILD_INFO.semver_version(),
3259            ShardId::new(),
3260            "".to_owned(),
3261            0,
3262        );
3263        let reader = LeasedReaderId::new();
3264        let seqno = SeqNo::minimum();
3265        let now = SYSTEM_TIME.clone();
3266        let _ = state.collections.register_leased_reader(
3267            "",
3268            &reader,
3269            "",
3270            seqno,
3271            Duration::from_secs(10),
3272            now(),
3273            false,
3274        );
3275
3276        // The shard global since == 0 initially.
3277        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(0));
3278
3279        // Greater
3280        assert_eq!(
3281            state.collections.downgrade_since(
3282                &reader,
3283                seqno,
3284                seqno,
3285                &Antichain::from_elem(2),
3286                now()
3287            ),
3288            Continue(Since(Antichain::from_elem(2)))
3289        );
3290        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(2));
3291        // Equal (no-op)
3292        assert_eq!(
3293            state.collections.downgrade_since(
3294                &reader,
3295                seqno,
3296                seqno,
3297                &Antichain::from_elem(2),
3298                now()
3299            ),
3300            Continue(Since(Antichain::from_elem(2)))
3301        );
3302        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(2));
3303        // Less (no-op)
3304        assert_eq!(
3305            state.collections.downgrade_since(
3306                &reader,
3307                seqno,
3308                seqno,
3309                &Antichain::from_elem(1),
3310                now()
3311            ),
3312            Continue(Since(Antichain::from_elem(2)))
3313        );
3314        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(2));
3315
3316        // Create a second reader.
3317        let reader2 = LeasedReaderId::new();
3318        let _ = state.collections.register_leased_reader(
3319            "",
3320            &reader2,
3321            "",
3322            seqno,
3323            Duration::from_secs(10),
3324            now(),
3325            false,
3326        );
3327
3328        // Shard since doesn't change until the meet (min) of all reader sinces changes.
3329        assert_eq!(
3330            state.collections.downgrade_since(
3331                &reader2,
3332                seqno,
3333                seqno,
3334                &Antichain::from_elem(3),
3335                now()
3336            ),
3337            Continue(Since(Antichain::from_elem(3)))
3338        );
3339        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(2));
3340        // Shard since == 3 when all readers have since >= 3.
3341        assert_eq!(
3342            state.collections.downgrade_since(
3343                &reader,
3344                seqno,
3345                seqno,
3346                &Antichain::from_elem(5),
3347                now()
3348            ),
3349            Continue(Since(Antichain::from_elem(5)))
3350        );
3351        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(3));
3352
3353        // Shard since unaffected readers with since > shard since expiring.
3354        assert_eq!(
3355            state.collections.expire_leased_reader(&reader),
3356            Continue(true)
3357        );
3358        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(3));
3359
3360        // Create a third reader.
3361        let reader3 = LeasedReaderId::new();
3362        let _ = state.collections.register_leased_reader(
3363            "",
3364            &reader3,
3365            "",
3366            seqno,
3367            Duration::from_secs(10),
3368            now(),
3369            false,
3370        );
3371
3372        // Shard since doesn't change until the meet (min) of all reader sinces changes.
3373        assert_eq!(
3374            state.collections.downgrade_since(
3375                &reader3,
3376                seqno,
3377                seqno,
3378                &Antichain::from_elem(10),
3379                now()
3380            ),
3381            Continue(Since(Antichain::from_elem(10)))
3382        );
3383        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(3));
3384
3385        // Shard since advances when reader with the minimal since expires.
3386        assert_eq!(
3387            state.collections.expire_leased_reader(&reader2),
3388            Continue(true)
3389        );
3390        // TODO(database-issues#6885): expiry temporarily doesn't advance since
3391        // Switch this assertion back when we re-enable this.
3392        //
3393        // assert_eq!(state.collections.trace.since(), &Antichain::from_elem(10));
3394        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(3));
3395
3396        // Shard since unaffected when all readers are expired.
3397        assert_eq!(
3398            state.collections.expire_leased_reader(&reader3),
3399            Continue(true)
3400        );
3401        // TODO(database-issues#6885): expiry temporarily doesn't advance since
3402        // Switch this assertion back when we re-enable this.
3403        //
3404        // assert_eq!(state.collections.trace.since(), &Antichain::from_elem(10));
3405        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(3));
3406    }
3407
3408    #[mz_ore::test]
3409    fn compare_and_downgrade_since() {
3410        let mut state = TypedState::<(), (), u64, i64>::new(
3411            DUMMY_BUILD_INFO.semver_version(),
3412            ShardId::new(),
3413            "".to_owned(),
3414            0,
3415        );
3416        let reader = CriticalReaderId::new();
3417        let _ = state
3418            .collections
3419            .register_critical_reader("", &reader, Opaque::encode(&0u64), "");
3420
3421        // The shard global since == 0 initially.
3422        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(0));
3423        // The initial opaque value should be set.
3424        assert_eq!(
3425            state
3426                .collections
3427                .critical_reader(&reader)
3428                .opaque
3429                .decode::<u64>(),
3430            u64::MIN
3431        );
3432
3433        // Greater
3434        assert_eq!(
3435            state.collections.compare_and_downgrade_since(
3436                &reader,
3437                &Opaque::encode(&0u64),
3438                (&Opaque::encode(&1u64), &Antichain::from_elem(2)),
3439            ),
3440            Continue(Ok(Since(Antichain::from_elem(2))))
3441        );
3442        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(2));
3443        assert_eq!(
3444            state
3445                .collections
3446                .critical_reader(&reader)
3447                .opaque
3448                .decode::<u64>(),
3449            1
3450        );
3451        // Equal (no-op)
3452        assert_eq!(
3453            state.collections.compare_and_downgrade_since(
3454                &reader,
3455                &Opaque::encode(&1u64),
3456                (&Opaque::encode(&2u64), &Antichain::from_elem(2)),
3457            ),
3458            Continue(Ok(Since(Antichain::from_elem(2))))
3459        );
3460        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(2));
3461        assert_eq!(
3462            state
3463                .collections
3464                .critical_reader(&reader)
3465                .opaque
3466                .decode::<u64>(),
3467            2
3468        );
3469        // Less (no-op)
3470        assert_eq!(
3471            state.collections.compare_and_downgrade_since(
3472                &reader,
3473                &Opaque::encode(&2u64),
3474                (&Opaque::encode(&3u64), &Antichain::from_elem(1)),
3475            ),
3476            Continue(Ok(Since(Antichain::from_elem(2))))
3477        );
3478        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(2));
3479        assert_eq!(
3480            state
3481                .collections
3482                .critical_reader(&reader)
3483                .opaque
3484                .decode::<u64>(),
3485            3
3486        );
3487    }
3488
3489    #[mz_ore::test]
3490    fn compare_and_append() {
3491        let state = &mut TypedState::<String, String, u64, i64>::new(
3492            DUMMY_BUILD_INFO.semver_version(),
3493            ShardId::new(),
3494            "".to_owned(),
3495            0,
3496        )
3497        .collections;
3498
3499        let writer_id = WriterId::new();
3500        let now = SYSTEM_TIME.clone();
3501
3502        // State is initially empty.
3503        assert_eq!(state.trace.num_spine_batches(), 0);
3504        assert_eq!(state.trace.num_hollow_batches(), 0);
3505        assert_eq!(state.trace.num_updates(), 0);
3506
3507        // Cannot insert a batch with a lower != current shard upper.
3508        assert_eq!(
3509            state.compare_and_append(
3510                &hollow(1, 2, &["key1"], 1),
3511                &writer_id,
3512                now(),
3513                LEASE_DURATION_MS,
3514                &IdempotencyToken::new(),
3515                &debug_state(),
3516                0,
3517                100,
3518                None
3519            ),
3520            Break(CompareAndAppendBreak::Upper {
3521                shard_upper: Antichain::from_elem(0),
3522                writer_upper: Antichain::from_elem(0)
3523            })
3524        );
3525
3526        // Insert an empty batch with an upper > lower..
3527        assert!(
3528            state
3529                .compare_and_append(
3530                    &hollow(0, 5, &[], 0),
3531                    &writer_id,
3532                    now(),
3533                    LEASE_DURATION_MS,
3534                    &IdempotencyToken::new(),
3535                    &debug_state(),
3536                    0,
3537                    100,
3538                    None
3539                )
3540                .is_continue()
3541        );
3542
3543        // Cannot insert a batch with a upper less than the lower.
3544        assert_eq!(
3545            state.compare_and_append(
3546                &hollow(5, 4, &["key1"], 1),
3547                &writer_id,
3548                now(),
3549                LEASE_DURATION_MS,
3550                &IdempotencyToken::new(),
3551                &debug_state(),
3552                0,
3553                100,
3554                None
3555            ),
3556            Break(CompareAndAppendBreak::InvalidUsage(InvalidBounds {
3557                lower: Antichain::from_elem(5),
3558                upper: Antichain::from_elem(4)
3559            }))
3560        );
3561
3562        // Cannot insert a nonempty batch with an upper equal to lower.
3563        assert_eq!(
3564            state.compare_and_append(
3565                &hollow(5, 5, &["key1"], 1),
3566                &writer_id,
3567                now(),
3568                LEASE_DURATION_MS,
3569                &IdempotencyToken::new(),
3570                &debug_state(),
3571                0,
3572                100,
3573                None
3574            ),
3575            Break(CompareAndAppendBreak::InvalidUsage(
3576                InvalidEmptyTimeInterval {
3577                    lower: Antichain::from_elem(5),
3578                    upper: Antichain::from_elem(5),
3579                    keys: vec!["key1".to_owned()],
3580                }
3581            ))
3582        );
3583
3584        // Can insert an empty batch with an upper equal to lower.
3585        assert!(
3586            state
3587                .compare_and_append(
3588                    &hollow(5, 5, &[], 0),
3589                    &writer_id,
3590                    now(),
3591                    LEASE_DURATION_MS,
3592                    &IdempotencyToken::new(),
3593                    &debug_state(),
3594                    0,
3595                    100,
3596                    None
3597                )
3598                .is_continue()
3599        );
3600    }
3601
3602    #[mz_ore::test]
3603    fn snapshot() {
3604        let now = SYSTEM_TIME.clone();
3605
3606        let mut state = TypedState::<String, String, u64, i64>::new(
3607            DUMMY_BUILD_INFO.semver_version(),
3608            ShardId::new(),
3609            "".to_owned(),
3610            0,
3611        );
3612        // Cannot take a snapshot with as_of == shard upper.
3613        assert_eq!(
3614            state.snapshot(&Antichain::from_elem(0)),
3615            Err(SnapshotErr::AsOfNotYetAvailable(
3616                SeqNo(0),
3617                Upper(Antichain::from_elem(0))
3618            ))
3619        );
3620
3621        // Cannot take a snapshot with as_of > shard upper.
3622        assert_eq!(
3623            state.snapshot(&Antichain::from_elem(5)),
3624            Err(SnapshotErr::AsOfNotYetAvailable(
3625                SeqNo(0),
3626                Upper(Antichain::from_elem(0))
3627            ))
3628        );
3629
3630        let writer_id = WriterId::new();
3631
3632        // Advance upper to 5.
3633        assert!(
3634            state
3635                .collections
3636                .compare_and_append(
3637                    &hollow(0, 5, &["key1"], 1),
3638                    &writer_id,
3639                    now(),
3640                    LEASE_DURATION_MS,
3641                    &IdempotencyToken::new(),
3642                    &debug_state(),
3643                    0,
3644                    100,
3645                    None
3646                )
3647                .is_continue()
3648        );
3649
3650        // Can take a snapshot with as_of < upper.
3651        assert_eq!(
3652            state.snapshot(&Antichain::from_elem(0)),
3653            Ok(vec![hollow(0, 5, &["key1"], 1)])
3654        );
3655
3656        // Can take a snapshot with as_of >= shard since, as long as as_of < shard_upper.
3657        assert_eq!(
3658            state.snapshot(&Antichain::from_elem(4)),
3659            Ok(vec![hollow(0, 5, &["key1"], 1)])
3660        );
3661
3662        // Cannot take a snapshot with as_of >= upper.
3663        assert_eq!(
3664            state.snapshot(&Antichain::from_elem(5)),
3665            Err(SnapshotErr::AsOfNotYetAvailable(
3666                SeqNo(0),
3667                Upper(Antichain::from_elem(5))
3668            ))
3669        );
3670        assert_eq!(
3671            state.snapshot(&Antichain::from_elem(6)),
3672            Err(SnapshotErr::AsOfNotYetAvailable(
3673                SeqNo(0),
3674                Upper(Antichain::from_elem(5))
3675            ))
3676        );
3677
3678        let reader = LeasedReaderId::new();
3679        // Advance the since to 2.
3680        let _ = state.collections.register_leased_reader(
3681            "",
3682            &reader,
3683            "",
3684            SeqNo::minimum(),
3685            Duration::from_secs(10),
3686            now(),
3687            false,
3688        );
3689        assert_eq!(
3690            state.collections.downgrade_since(
3691                &reader,
3692                SeqNo::minimum(),
3693                SeqNo::minimum(),
3694                &Antichain::from_elem(2),
3695                now()
3696            ),
3697            Continue(Since(Antichain::from_elem(2)))
3698        );
3699        assert_eq!(state.collections.trace.since(), &Antichain::from_elem(2));
3700        // Cannot take a snapshot with as_of < shard_since.
3701        assert_eq!(
3702            state.snapshot(&Antichain::from_elem(1)),
3703            Err(SnapshotErr::AsOfHistoricalDistinctionsLost(Since(
3704                Antichain::from_elem(2)
3705            )))
3706        );
3707
3708        // Advance the upper to 10 via an empty batch.
3709        assert!(
3710            state
3711                .collections
3712                .compare_and_append(
3713                    &hollow(5, 10, &[], 0),
3714                    &writer_id,
3715                    now(),
3716                    LEASE_DURATION_MS,
3717                    &IdempotencyToken::new(),
3718                    &debug_state(),
3719                    0,
3720                    100,
3721                    None
3722                )
3723                .is_continue()
3724        );
3725
3726        // Can still take snapshots at times < upper.
3727        assert_eq!(
3728            state.snapshot(&Antichain::from_elem(7)),
3729            Ok(vec![hollow(0, 5, &["key1"], 1), hollow(5, 10, &[], 0)])
3730        );
3731
3732        // Cannot take snapshots with as_of >= upper.
3733        assert_eq!(
3734            state.snapshot(&Antichain::from_elem(10)),
3735            Err(SnapshotErr::AsOfNotYetAvailable(
3736                SeqNo(0),
3737                Upper(Antichain::from_elem(10))
3738            ))
3739        );
3740
3741        // Advance upper to 15.
3742        assert!(
3743            state
3744                .collections
3745                .compare_and_append(
3746                    &hollow(10, 15, &["key2"], 1),
3747                    &writer_id,
3748                    now(),
3749                    LEASE_DURATION_MS,
3750                    &IdempotencyToken::new(),
3751                    &debug_state(),
3752                    0,
3753                    100,
3754                    None
3755                )
3756                .is_continue()
3757        );
3758
3759        // Filter out batches whose lowers are less than the requested as of (the
3760        // batches that are too far in the future for the requested as_of).
3761        assert_eq!(
3762            state.snapshot(&Antichain::from_elem(9)),
3763            Ok(vec![hollow(0, 5, &["key1"], 1), hollow(5, 10, &[], 0)])
3764        );
3765
3766        // Don't filter out batches whose lowers are <= the requested as_of.
3767        assert_eq!(
3768            state.snapshot(&Antichain::from_elem(10)),
3769            Ok(vec![
3770                hollow(0, 5, &["key1"], 1),
3771                hollow(5, 10, &[], 0),
3772                hollow(10, 15, &["key2"], 1)
3773            ])
3774        );
3775
3776        assert_eq!(
3777            state.snapshot(&Antichain::from_elem(11)),
3778            Ok(vec![
3779                hollow(0, 5, &["key1"], 1),
3780                hollow(5, 10, &[], 0),
3781                hollow(10, 15, &["key2"], 1)
3782            ])
3783        );
3784    }
3785
3786    #[mz_ore::test]
3787    fn next_listen_batch() {
3788        let mut state = TypedState::<String, String, u64, i64>::new(
3789            DUMMY_BUILD_INFO.semver_version(),
3790            ShardId::new(),
3791            "".to_owned(),
3792            0,
3793        );
3794
3795        // Empty collection never has any batches to listen for, regardless of the
3796        // current frontier.
3797        assert_eq!(
3798            state.next_listen_batch(&Antichain::from_elem(0)),
3799            Err(SeqNo(0))
3800        );
3801        assert_eq!(state.next_listen_batch(&Antichain::new()), Err(SeqNo(0)));
3802
3803        let writer_id = WriterId::new();
3804        let now = SYSTEM_TIME.clone();
3805
3806        // Add two batches of data, one from [0, 5) and then another from [5, 10).
3807        assert!(
3808            state
3809                .collections
3810                .compare_and_append(
3811                    &hollow(0, 5, &["key1"], 1),
3812                    &writer_id,
3813                    now(),
3814                    LEASE_DURATION_MS,
3815                    &IdempotencyToken::new(),
3816                    &debug_state(),
3817                    0,
3818                    100,
3819                    None
3820                )
3821                .is_continue()
3822        );
3823        assert!(
3824            state
3825                .collections
3826                .compare_and_append(
3827                    &hollow(5, 10, &["key2"], 1),
3828                    &writer_id,
3829                    now(),
3830                    LEASE_DURATION_MS,
3831                    &IdempotencyToken::new(),
3832                    &debug_state(),
3833                    0,
3834                    100,
3835                    None
3836                )
3837                .is_continue()
3838        );
3839
3840        // All frontiers in [0, 5) return the first batch.
3841        for t in 0..=4 {
3842            assert_eq!(
3843                state.next_listen_batch(&Antichain::from_elem(t)),
3844                Ok(hollow(0, 5, &["key1"], 1))
3845            );
3846        }
3847
3848        // All frontiers in [5, 10) return the second batch.
3849        for t in 5..=9 {
3850            assert_eq!(
3851                state.next_listen_batch(&Antichain::from_elem(t)),
3852                Ok(hollow(5, 10, &["key2"], 1))
3853            );
3854        }
3855
3856        // There is no batch currently available for t = 10.
3857        assert_eq!(
3858            state.next_listen_batch(&Antichain::from_elem(10)),
3859            Err(SeqNo(0))
3860        );
3861
3862        // By definition, there is no frontier ever at the empty antichain which
3863        // is the time after all possible times.
3864        assert_eq!(state.next_listen_batch(&Antichain::new()), Err(SeqNo(0)));
3865    }
3866
3867    #[mz_ore::test]
3868    fn expire_writer() {
3869        let mut state = TypedState::<String, String, u64, i64>::new(
3870            DUMMY_BUILD_INFO.semver_version(),
3871            ShardId::new(),
3872            "".to_owned(),
3873            0,
3874        );
3875        let now = SYSTEM_TIME.clone();
3876
3877        let writer_id_one = WriterId::new();
3878
3879        let writer_id_two = WriterId::new();
3880
3881        // Writer is eligible to write
3882        assert!(
3883            state
3884                .collections
3885                .compare_and_append(
3886                    &hollow(0, 2, &["key1"], 1),
3887                    &writer_id_one,
3888                    now(),
3889                    LEASE_DURATION_MS,
3890                    &IdempotencyToken::new(),
3891                    &debug_state(),
3892                    0,
3893                    100,
3894                    None
3895                )
3896                .is_continue()
3897        );
3898
3899        assert!(
3900            state
3901                .collections
3902                .expire_writer(&writer_id_one)
3903                .is_continue()
3904        );
3905
3906        // Other writers should still be able to write
3907        assert!(
3908            state
3909                .collections
3910                .compare_and_append(
3911                    &hollow(2, 5, &["key2"], 1),
3912                    &writer_id_two,
3913                    now(),
3914                    LEASE_DURATION_MS,
3915                    &IdempotencyToken::new(),
3916                    &debug_state(),
3917                    0,
3918                    100,
3919                    None
3920                )
3921                .is_continue()
3922        );
3923    }
3924
3925    #[mz_ore::test]
3926    fn maybe_gc_active_gc() {
3927        const GC_CONFIG: GcConfig = GcConfig {
3928            use_active_gc: true,
3929            fallback_threshold_ms: 5000,
3930            min_versions: 99,
3931            max_versions: 500,
3932        };
3933        let now_fn = SYSTEM_TIME.clone();
3934
3935        let mut state = TypedState::<String, String, u64, i64>::new(
3936            DUMMY_BUILD_INFO.semver_version(),
3937            ShardId::new(),
3938            "".to_owned(),
3939            0,
3940        );
3941
3942        let now = now_fn();
3943        // Empty state doesn't need gc, regardless of is_write.
3944        assert_eq!(state.maybe_gc(true, now, GC_CONFIG), None);
3945        assert_eq!(state.maybe_gc(false, now, GC_CONFIG), None);
3946
3947        // Artificially advance the seqno so the seqno_since advances past our
3948        // internal gc_threshold.
3949        state.seqno = SeqNo(100);
3950        assert_eq!(state.seqno_since(), SeqNo(100));
3951
3952        // When a writer is present, non-writes don't gc.
3953        let writer_id = WriterId::new();
3954        let _ = state.collections.compare_and_append(
3955            &hollow(1, 2, &["key1"], 1),
3956            &writer_id,
3957            now,
3958            LEASE_DURATION_MS,
3959            &IdempotencyToken::new(),
3960            &debug_state(),
3961            0,
3962            100,
3963            None,
3964        );
3965        assert_eq!(state.maybe_gc(false, now, GC_CONFIG), None);
3966
3967        // A write will gc though.
3968        assert_eq!(
3969            state.maybe_gc(true, now, GC_CONFIG),
3970            Some(GcReq {
3971                shard_id: state.shard_id,
3972                new_seqno_since: SeqNo(100)
3973            })
3974        );
3975
3976        // But if we write down an active gc, we won't gc.
3977        state.collections.active_gc = Some(ActiveGc {
3978            seqno: state.seqno,
3979            start_ms: now,
3980        });
3981
3982        state.seqno = SeqNo(200);
3983        assert_eq!(state.seqno_since(), SeqNo(200));
3984
3985        assert_eq!(state.maybe_gc(true, now, GC_CONFIG), None);
3986
3987        state.seqno = SeqNo(300);
3988        assert_eq!(state.seqno_since(), SeqNo(300));
3989        // But if we advance the time past the threshold, we will gc.
3990        let new_now = now + GC_CONFIG.fallback_threshold_ms + 1;
3991        assert_eq!(
3992            state.maybe_gc(true, new_now, GC_CONFIG),
3993            Some(GcReq {
3994                shard_id: state.shard_id,
3995                new_seqno_since: SeqNo(300)
3996            })
3997        );
3998
3999        // Even if the sequence number doesn't pass the threshold, if the
4000        // active gc is expired, we will gc.
4001
4002        state.seqno = SeqNo(301);
4003        assert_eq!(state.seqno_since(), SeqNo(301));
4004        assert_eq!(
4005            state.maybe_gc(true, new_now, GC_CONFIG),
4006            Some(GcReq {
4007                shard_id: state.shard_id,
4008                new_seqno_since: SeqNo(301)
4009            })
4010        );
4011
4012        state.collections.active_gc = None;
4013
4014        // Artificially advance the seqno (again) so the seqno_since advances
4015        // past our internal gc_threshold (again).
4016        state.seqno = SeqNo(400);
4017        assert_eq!(state.seqno_since(), SeqNo(400));
4018
4019        let now = now_fn();
4020
4021        // If there are no writers, even a non-write will gc.
4022        let _ = state.collections.expire_writer(&writer_id);
4023        assert_eq!(
4024            state.maybe_gc(false, now, GC_CONFIG),
4025            Some(GcReq {
4026                shard_id: state.shard_id,
4027                new_seqno_since: SeqNo(400)
4028            })
4029        );
4030
4031        // Upper-bound the number of seqnos we'll attempt to collect in one go.
4032        let previous_seqno = state.seqno;
4033        state.seqno = SeqNo(10_000);
4034        assert_eq!(state.seqno_since(), SeqNo(10_000));
4035
4036        let now = now_fn();
4037        assert_eq!(
4038            state.maybe_gc(true, now, GC_CONFIG),
4039            Some(GcReq {
4040                shard_id: state.shard_id,
4041                new_seqno_since: SeqNo(previous_seqno.0 + u64::cast_from(GC_CONFIG.max_versions))
4042            })
4043        );
4044    }
4045
4046    #[mz_ore::test]
4047    fn maybe_gc_classic() {
4048        const GC_CONFIG: GcConfig = GcConfig {
4049            use_active_gc: false,
4050            fallback_threshold_ms: 5000,
4051            min_versions: 16,
4052            max_versions: 128,
4053        };
4054        const NOW_MS: u64 = 0;
4055
4056        let mut state = TypedState::<String, String, u64, i64>::new(
4057            DUMMY_BUILD_INFO.semver_version(),
4058            ShardId::new(),
4059            "".to_owned(),
4060            0,
4061        );
4062
4063        // Empty state doesn't need gc, regardless of is_write.
4064        assert_eq!(state.maybe_gc(true, NOW_MS, GC_CONFIG), None);
4065        assert_eq!(state.maybe_gc(false, NOW_MS, GC_CONFIG), None);
4066
4067        // Artificially advance the seqno so the seqno_since advances past our
4068        // internal gc_threshold.
4069        state.seqno = SeqNo(100);
4070        assert_eq!(state.seqno_since(), SeqNo(100));
4071
4072        // When a writer is present, non-writes don't gc.
4073        let writer_id = WriterId::new();
4074        let now = SYSTEM_TIME.clone();
4075        let _ = state.collections.compare_and_append(
4076            &hollow(1, 2, &["key1"], 1),
4077            &writer_id,
4078            now(),
4079            LEASE_DURATION_MS,
4080            &IdempotencyToken::new(),
4081            &debug_state(),
4082            0,
4083            100,
4084            None,
4085        );
4086        assert_eq!(state.maybe_gc(false, NOW_MS, GC_CONFIG), None);
4087
4088        // A write will gc though.
4089        assert_eq!(
4090            state.maybe_gc(true, NOW_MS, GC_CONFIG),
4091            Some(GcReq {
4092                shard_id: state.shard_id,
4093                new_seqno_since: SeqNo(100)
4094            })
4095        );
4096
4097        // Artificially advance the seqno (again) so the seqno_since advances
4098        // past our internal gc_threshold (again).
4099        state.seqno = SeqNo(200);
4100        assert_eq!(state.seqno_since(), SeqNo(200));
4101
4102        // If there are no writers, even a non-write will gc.
4103        let _ = state.collections.expire_writer(&writer_id);
4104        assert_eq!(
4105            state.maybe_gc(false, NOW_MS, GC_CONFIG),
4106            Some(GcReq {
4107                shard_id: state.shard_id,
4108                new_seqno_since: SeqNo(200)
4109            })
4110        );
4111    }
4112
4113    #[mz_ore::test]
4114    fn need_rollup_active_rollup() {
4115        const ROLLUP_THRESHOLD: usize = 3;
4116        const ROLLUP_USE_ACTIVE_ROLLUP: bool = true;
4117        const ROLLUP_FALLBACK_THRESHOLD_MS: u64 = 5000;
4118        let now = SYSTEM_TIME.clone();
4119
4120        mz_ore::test::init_logging();
4121        let mut state = TypedState::<String, String, u64, i64>::new(
4122            DUMMY_BUILD_INFO.semver_version(),
4123            ShardId::new(),
4124            "".to_owned(),
4125            0,
4126        );
4127
4128        let rollup_seqno = SeqNo(5);
4129        let rollup = HollowRollup {
4130            key: PartialRollupKey::new(rollup_seqno, &RollupId::new()),
4131            encoded_size_bytes: None,
4132        };
4133
4134        assert!(
4135            state
4136                .collections
4137                .add_rollup((rollup_seqno, &rollup))
4138                .is_continue()
4139        );
4140
4141        // shouldn't need a rollup at the seqno of the rollup
4142        state.seqno = SeqNo(5);
4143        assert_none!(state.need_rollup(
4144            ROLLUP_THRESHOLD,
4145            ROLLUP_USE_ACTIVE_ROLLUP,
4146            ROLLUP_FALLBACK_THRESHOLD_MS,
4147            now()
4148        ));
4149
4150        // shouldn't need a rollup at seqnos less than our threshold
4151        state.seqno = SeqNo(6);
4152        assert_none!(state.need_rollup(
4153            ROLLUP_THRESHOLD,
4154            ROLLUP_USE_ACTIVE_ROLLUP,
4155            ROLLUP_FALLBACK_THRESHOLD_MS,
4156            now()
4157        ));
4158        state.seqno = SeqNo(7);
4159        assert_none!(state.need_rollup(
4160            ROLLUP_THRESHOLD,
4161            ROLLUP_USE_ACTIVE_ROLLUP,
4162            ROLLUP_FALLBACK_THRESHOLD_MS,
4163            now()
4164        ));
4165        state.seqno = SeqNo(8);
4166        assert_none!(state.need_rollup(
4167            ROLLUP_THRESHOLD,
4168            ROLLUP_USE_ACTIVE_ROLLUP,
4169            ROLLUP_FALLBACK_THRESHOLD_MS,
4170            now()
4171        ));
4172
4173        let mut current_time = now();
4174        // hit our threshold! we should need a rollup
4175        state.seqno = SeqNo(9);
4176        assert_eq!(
4177            state
4178                .need_rollup(
4179                    ROLLUP_THRESHOLD,
4180                    ROLLUP_USE_ACTIVE_ROLLUP,
4181                    ROLLUP_FALLBACK_THRESHOLD_MS,
4182                    current_time
4183                )
4184                .expect("rollup"),
4185            SeqNo(9)
4186        );
4187
4188        state.collections.active_rollup = Some(ActiveRollup {
4189            seqno: SeqNo(9),
4190            start_ms: current_time,
4191        });
4192
4193        // There is now an active rollup, so we shouldn't need a rollup.
4194        assert_none!(state.need_rollup(
4195            ROLLUP_THRESHOLD,
4196            ROLLUP_USE_ACTIVE_ROLLUP,
4197            ROLLUP_FALLBACK_THRESHOLD_MS,
4198            current_time
4199        ));
4200
4201        state.seqno = SeqNo(10);
4202        // We still don't need a rollup, even though the seqno is greater than
4203        // the rollup threshold.
4204        assert_none!(state.need_rollup(
4205            ROLLUP_THRESHOLD,
4206            ROLLUP_USE_ACTIVE_ROLLUP,
4207            ROLLUP_FALLBACK_THRESHOLD_MS,
4208            current_time
4209        ));
4210
4211        // But if we wait long enough, we should need a rollup again.
4212        current_time += u64::cast_from(ROLLUP_FALLBACK_THRESHOLD_MS) + 1;
4213        assert_eq!(
4214            state
4215                .need_rollup(
4216                    ROLLUP_THRESHOLD,
4217                    ROLLUP_USE_ACTIVE_ROLLUP,
4218                    ROLLUP_FALLBACK_THRESHOLD_MS,
4219                    current_time
4220                )
4221                .expect("rollup"),
4222            SeqNo(10)
4223        );
4224
4225        state.seqno = SeqNo(9);
4226        // Clear the active rollup and ensure we need a rollup again.
4227        state.collections.active_rollup = None;
4228        let rollup_seqno = SeqNo(9);
4229        let rollup = HollowRollup {
4230            key: PartialRollupKey::new(rollup_seqno, &RollupId::new()),
4231            encoded_size_bytes: None,
4232        };
4233        assert!(
4234            state
4235                .collections
4236                .add_rollup((rollup_seqno, &rollup))
4237                .is_continue()
4238        );
4239
4240        state.seqno = SeqNo(11);
4241        // We shouldn't need a rollup at seqnos less than our threshold
4242        assert_none!(state.need_rollup(
4243            ROLLUP_THRESHOLD,
4244            ROLLUP_USE_ACTIVE_ROLLUP,
4245            ROLLUP_FALLBACK_THRESHOLD_MS,
4246            current_time
4247        ));
4248        // hit our threshold! we should need a rollup
4249        state.seqno = SeqNo(13);
4250        assert_eq!(
4251            state
4252                .need_rollup(
4253                    ROLLUP_THRESHOLD,
4254                    ROLLUP_USE_ACTIVE_ROLLUP,
4255                    ROLLUP_FALLBACK_THRESHOLD_MS,
4256                    current_time
4257                )
4258                .expect("rollup"),
4259            SeqNo(13)
4260        );
4261    }
4262
4263    #[mz_ore::test]
4264    fn need_rollup_classic() {
4265        const ROLLUP_THRESHOLD: usize = 3;
4266        const ROLLUP_USE_ACTIVE_ROLLUP: bool = false;
4267        const ROLLUP_FALLBACK_THRESHOLD_MS: u64 = 0;
4268        const NOW: u64 = 0;
4269
4270        mz_ore::test::init_logging();
4271        let mut state = TypedState::<String, String, u64, i64>::new(
4272            DUMMY_BUILD_INFO.semver_version(),
4273            ShardId::new(),
4274            "".to_owned(),
4275            0,
4276        );
4277
4278        let rollup_seqno = SeqNo(5);
4279        let rollup = HollowRollup {
4280            key: PartialRollupKey::new(rollup_seqno, &RollupId::new()),
4281            encoded_size_bytes: None,
4282        };
4283
4284        assert!(
4285            state
4286                .collections
4287                .add_rollup((rollup_seqno, &rollup))
4288                .is_continue()
4289        );
4290
4291        // shouldn't need a rollup at the seqno of the rollup
4292        state.seqno = SeqNo(5);
4293        assert_none!(state.need_rollup(
4294            ROLLUP_THRESHOLD,
4295            ROLLUP_USE_ACTIVE_ROLLUP,
4296            ROLLUP_FALLBACK_THRESHOLD_MS,
4297            NOW
4298        ));
4299
4300        // shouldn't need a rollup at seqnos less than our threshold
4301        state.seqno = SeqNo(6);
4302        assert_none!(state.need_rollup(
4303            ROLLUP_THRESHOLD,
4304            ROLLUP_USE_ACTIVE_ROLLUP,
4305            ROLLUP_FALLBACK_THRESHOLD_MS,
4306            NOW
4307        ));
4308        state.seqno = SeqNo(7);
4309        assert_none!(state.need_rollup(
4310            ROLLUP_THRESHOLD,
4311            ROLLUP_USE_ACTIVE_ROLLUP,
4312            ROLLUP_FALLBACK_THRESHOLD_MS,
4313            NOW
4314        ));
4315
4316        // hit our threshold! we should need a rollup
4317        state.seqno = SeqNo(8);
4318        assert_eq!(
4319            state
4320                .need_rollup(
4321                    ROLLUP_THRESHOLD,
4322                    ROLLUP_USE_ACTIVE_ROLLUP,
4323                    ROLLUP_FALLBACK_THRESHOLD_MS,
4324                    NOW
4325                )
4326                .expect("rollup"),
4327            SeqNo(8)
4328        );
4329
4330        // but we don't need rollups for every seqno > the threshold
4331        state.seqno = SeqNo(9);
4332        assert_none!(state.need_rollup(
4333            ROLLUP_THRESHOLD,
4334            ROLLUP_USE_ACTIVE_ROLLUP,
4335            ROLLUP_FALLBACK_THRESHOLD_MS,
4336            NOW
4337        ));
4338
4339        // we only need a rollup each `ROLLUP_THRESHOLD` beyond our current seqno
4340        state.seqno = SeqNo(11);
4341        assert_eq!(
4342            state
4343                .need_rollup(
4344                    ROLLUP_THRESHOLD,
4345                    ROLLUP_USE_ACTIVE_ROLLUP,
4346                    ROLLUP_FALLBACK_THRESHOLD_MS,
4347                    NOW
4348                )
4349                .expect("rollup"),
4350            SeqNo(11)
4351        );
4352
4353        // add another rollup and ensure we're always picking the latest
4354        let rollup_seqno = SeqNo(6);
4355        let rollup = HollowRollup {
4356            key: PartialRollupKey::new(rollup_seqno, &RollupId::new()),
4357            encoded_size_bytes: None,
4358        };
4359        assert!(
4360            state
4361                .collections
4362                .add_rollup((rollup_seqno, &rollup))
4363                .is_continue()
4364        );
4365
4366        state.seqno = SeqNo(8);
4367        assert_none!(state.need_rollup(
4368            ROLLUP_THRESHOLD,
4369            ROLLUP_USE_ACTIVE_ROLLUP,
4370            ROLLUP_FALLBACK_THRESHOLD_MS,
4371            NOW
4372        ));
4373        state.seqno = SeqNo(9);
4374        assert_eq!(
4375            state
4376                .need_rollup(
4377                    ROLLUP_THRESHOLD,
4378                    ROLLUP_USE_ACTIVE_ROLLUP,
4379                    ROLLUP_FALLBACK_THRESHOLD_MS,
4380                    NOW
4381                )
4382                .expect("rollup"),
4383            SeqNo(9)
4384        );
4385
4386        // and ensure that after a fallback point, we assign every seqno work
4387        let fallback_seqno = SeqNo(
4388            rollup_seqno.0
4389                * u64::cast_from(PersistConfig::DEFAULT_FALLBACK_ROLLUP_THRESHOLD_MULTIPLIER),
4390        );
4391        state.seqno = fallback_seqno;
4392        assert_eq!(
4393            state
4394                .need_rollup(
4395                    ROLLUP_THRESHOLD,
4396                    ROLLUP_USE_ACTIVE_ROLLUP,
4397                    ROLLUP_FALLBACK_THRESHOLD_MS,
4398                    NOW
4399                )
4400                .expect("rollup"),
4401            fallback_seqno
4402        );
4403        state.seqno = fallback_seqno.next();
4404        assert_eq!(
4405            state
4406                .need_rollup(
4407                    ROLLUP_THRESHOLD,
4408                    ROLLUP_USE_ACTIVE_ROLLUP,
4409                    ROLLUP_FALLBACK_THRESHOLD_MS,
4410                    NOW
4411                )
4412                .expect("rollup"),
4413            fallback_seqno.next()
4414        );
4415    }
4416
4417    #[mz_ore::test]
4418    fn idempotency_token_sentinel() {
4419        assert_eq!(
4420            IdempotencyToken::SENTINEL.to_string(),
4421            "i11111111-1111-1111-1111-111111111111"
4422        );
4423    }
4424
4425    /// This test generates an "arbitrary" State, but uses a fixed seed for the
4426    /// randomness, so that it's deterministic. This lets us assert the
4427    /// serialization of that State against a golden file that's committed,
4428    /// making it easy to see what the serialization (used in an upcoming
4429    /// INSPECT feature) looks like.
4430    ///
4431    /// This golden will have to be updated each time we change State, but
4432    /// that's a feature, not a bug.
4433    #[mz_ore::test]
4434    #[cfg_attr(miri, ignore)] // too slow
4435    fn state_inspect_serde_json() {
4436        const STATE_SERDE_JSON: &str = include_str!("state_serde.json");
4437        let mut runner = proptest::test_runner::TestRunner::deterministic();
4438        let tree = any_state::<u64>(6..8).new_tree(&mut runner).unwrap();
4439        let json = serde_json::to_string_pretty(&tree.current()).unwrap();
4440        assert_eq!(
4441            json.trim(),
4442            STATE_SERDE_JSON.trim(),
4443            "\n\nNEW GOLDEN\n{}\n",
4444            json
4445        );
4446    }
4447
4448    #[mz_persist_proc::test(tokio::test)]
4449    #[cfg_attr(miri, ignore)] // too slow
4450    async fn sneaky_downgrades(dyncfgs: ConfigUpdates) {
4451        let mut clients = new_test_client_cache(&dyncfgs);
4452        let shard_id = ShardId::new();
4453
4454        async fn open_and_write(
4455            clients: &mut PersistClientCache,
4456            version: semver::Version,
4457            shard_id: ShardId,
4458        ) -> Result<(), tokio::task::JoinError> {
4459            clients.cfg.build_version = version.clone();
4460            clients.clear_state_cache();
4461            let client = clients.open(PersistLocation::new_in_mem()).await.unwrap();
4462            // Run in a task so we can catch the panic.
4463            mz_ore::task::spawn(|| version.to_string(), async move {
4464                let () = client
4465                    .upgrade_version::<String, (), u64, i64>(shard_id, Diagnostics::for_tests())
4466                    .await
4467                    .expect("valid usage");
4468                let (mut write, _) = client.expect_open::<String, (), u64, i64>(shard_id).await;
4469                let current = *write.upper().as_option().unwrap();
4470                // Do a write so that we tag the state with the version.
4471                write
4472                    .expect_compare_and_append_batch(&mut [], current, current + 1)
4473                    .await;
4474            })
4475            .into_tokio_handle()
4476            .await
4477        }
4478
4479        // Start at v0.10.0.
4480        let res = open_and_write(&mut clients, Version::new(0, 10, 0), shard_id).await;
4481        assert_ok!(res);
4482
4483        // Upgrade to v0.11.0 is allowed.
4484        let res = open_and_write(&mut clients, Version::new(0, 11, 0), shard_id).await;
4485        assert_ok!(res);
4486
4487        // Downgrade to v0.10.0 is no longer allowed.
4488        let res = open_and_write(&mut clients, Version::new(0, 10, 0), shard_id).await;
4489        assert!(res.unwrap_err().is_panic());
4490
4491        // Downgrade to v0.9.0 is _NOT_ allowed.
4492        let res = open_and_write(&mut clients, Version::new(0, 9, 0), shard_id).await;
4493        assert!(res.unwrap_err().is_panic());
4494    }
4495
4496    #[mz_ore::test]
4497    fn runid_roundtrip() {
4498        proptest!(|(runid: RunId)| {
4499            let runid_str = runid.to_string();
4500            let parsed = RunId::from_str(&runid_str);
4501            prop_assert_eq!(parsed, Ok(runid));
4502        });
4503    }
4504
4505    /// Regression for PER-16. `add_rollup` must refuse to (re-)insert a
4506    /// rollup at a seqno that GC has already physically removed.
4507    /// Without this guard, an `apply_unbatched_idempotent_cmd` retry of
4508    /// `add_rollup` that runs after a concurrent GC removed the rollup's
4509    /// original map entry will re-insert it under the same key, producing
4510    /// duplicate Insert events in the diff stream and (later) duplicate
4511    /// Delete events that trip the `assert!` in `gc.rs`'s
4512    /// `find_removable_blobs`.
4513    ///
4514    /// This test reaches the bug state via the public
4515    /// `add_rollup` → `add_rollup` (newer) → `remove_rollups` (the older
4516    /// one) → `add_rollup` (retry of the older one) sequence, mirroring
4517    /// how a delayed retry of an indeterminate add_rollup commit can race
4518    /// a concurrent GC pass that has since added a newer rollup and
4519    /// removed the older one.
4520    #[mz_ore::test]
4521    fn add_rollup_idempotent_across_gc_removal() {
4522        let mut state = TypedState::<String, String, u64, i64>::new(
4523            DUMMY_BUILD_INFO.semver_version(),
4524            ShardId::new(),
4525            "".to_owned(),
4526            0,
4527        );
4528
4529        let older_seqno = SeqNo(10);
4530        let older = HollowRollup {
4531            key: PartialRollupKey::new(older_seqno, &RollupId::new()),
4532            encoded_size_bytes: None,
4533        };
4534        let newer_seqno = SeqNo(20);
4535        let newer = HollowRollup {
4536            key: PartialRollupKey::new(newer_seqno, &RollupId::new()),
4537            encoded_size_bytes: None,
4538        };
4539        let add_older = |state: &mut StateCollections<u64>| state.add_rollup((older_seqno, &older));
4540
4541        // First attempt commits the older rollup.
4542        assert_eq!(add_older(&mut state.collections), Continue(true));
4543        // A repeat at the same seqno with the same key is the pre-existing
4544        // idempotent no-op and should report applied=true.
4545        assert_eq!(add_older(&mut state.collections), Continue(true));
4546        assert_eq!(state.collections.rollups.len(), 1);
4547
4548        // The shard keeps making progress and a later command commits a
4549        // newer rollup. This is what gives GC something to keep when it
4550        // later removes the older entry.
4551        assert_eq!(
4552            state.collections.add_rollup((newer_seqno, &newer)),
4553            Continue(true),
4554        );
4555
4556        // The GC worker, having found a kept rollup at `newer_seqno`,
4557        // commits the removal of the older one. State.rollups now only
4558        // contains the kept rollup; the older entry is gone.
4559        let _ = state
4560            .collections
4561            .remove_rollups(&[(older_seqno, older.key.clone())]);
4562        assert!(!state.collections.rollups.contains_key(&older_seqno));
4563        assert!(state.collections.rollups.contains_key(&newer_seqno));
4564
4565        // The retry replays the same captured `(older_seqno, older)`
4566        // tuple. Even though no entry references the key anymore,
4567        // re-inserting it must be refused: `older_seqno` falls below the
4568        // smallest live rollup seqno, which is the GC watermark.
4569        assert_eq!(add_older(&mut state.collections), Continue(false));
4570        assert!(!state.collections.rollups.contains_key(&older_seqno));
4571        assert_eq!(state.collections.rollups.len(), 1);
4572    }
4573}