mz_storage/

upsert.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

use std::cell::RefCell;
use std::cmp::Reverse;
use std::convert::AsRef;
use std::fmt::Debug;
use std::hash::{Hash, Hasher};
use std::path::PathBuf;
use std::sync::Arc;

use differential_dataflow::hashable::Hashable;
use differential_dataflow::{AsCollection, VecCollection};
use futures::StreamExt;
use futures::future::FutureExt;
use indexmap::map::Entry;
use itertools::Itertools;
use mz_ore::error::ErrorExt;
use mz_repr::{Datum, DatumVec, Diff, GlobalId, Row};
use mz_rocksdb::ValueIterator;
use mz_sql_server_util::cdc::Lsn;
use mz_storage_types::configuration::StorageConfiguration;
use mz_storage_types::dyncfgs;
use mz_storage_types::errors::{DataflowError, EnvelopeError, UpsertError};
use mz_storage_types::sources::MzOffset;
use mz_storage_types::sources::envelope::UpsertEnvelope;
use mz_storage_types::sources::kafka::{KafkaTimestamp, RangeBound};
use mz_storage_types::sources::mysql::GtidPartition;
use mz_timely_util::builder_async::{
    AsyncOutputHandle, Event as AsyncEvent, OperatorBuilder as AsyncOperatorBuilder,
    PressOnDropButton,
};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use timely::dataflow::channels::pact::Exchange;
use timely::dataflow::operators::{Capability, InputCapability, Operator};
use timely::dataflow::{Scope, StreamVec};
use timely::order::{PartialOrder, TotalOrder};
use timely::progress::timestamp::Refines;
use timely::progress::{Antichain, Timestamp};

use crate::healthcheck::HealthStatusUpdate;
use crate::metrics::upsert::{UpsertBackpressureMetrics, UpsertMetrics};
use crate::storage_state::StorageInstanceContext;
use crate::{upsert_continual_feedback, upsert_continual_feedback_v2};
use types::{
    BincodeOpts, StateValue, UpsertState, UpsertStateBackend, consolidating_merge_function,
    upsert_bincode_opts,
};

#[cfg(any(test, feature = "fuzzing"))]
pub mod memory;
pub(crate) mod rocksdb;
// TODO(aljoscha): Move next to upsert module, rename to upsert_types.
pub(crate) mod types;

pub type UpsertValue = Result<Row, Box<UpsertError>>;

#[derive(
    Copy,
    Clone,
    Hash,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Serialize,
    Deserialize,
    bytemuck::AnyBitPattern,
    bytemuck::NoUninit
)]
#[repr(transparent)]
pub struct UpsertKey([u8; 32]);

impl columnation::Columnation for UpsertKey {
    type InnerRegion = columnation::CopyRegion<UpsertKey>;
}

/// Columnar (the `columnar` crate, distinct from `columnation`) support for
/// `UpsertKey`, so the upsert-v2 source stash can use a paged columnar merge
/// batcher keyed natively by `UpsertKey` (no `Row` packing). `UpsertKey` is a
/// POD `[u8; 32]` newtype, so the container is a fixed-stride byte column.
///
/// This is hand-rolled rather than `#[derive(Columnar)]`d for one load-bearing
/// reason: the reference type must be `&UpsertKey`. `&UpsertKey` is `Copy + Ord`
/// (the lexicographic `[u8; 32]` order the persist-feedback trace is keyed on),
/// which both satisfies the merge batcher's `Ref: Copy + Ord` requirement and —
/// crucially — matches the read item of the feedback arrangement's
/// `ColumnationStack<UpsertKey>` key container, so the paged `ValRow` builder
/// can reconcile the `Column` input against the spine (its `BuilderInput` bound
/// is `ReadItem: PartialEq<Ref<UpsertKey>>`). A derived impl would yield a
/// generated `UpsertKeyReference` (and route `[u8; 32]` through the generic,
/// non-fixed-stride array container), breaking that reconciliation.
mod columnar_upsert_key {
    use super::UpsertKey;
    use columnar::Columnar;
    use mz_ore::cast::CastFrom;
    use std::ops::Range;

    /// A newtype wrapper for a vector of `UpsertKey` values.
    #[derive(Clone, Copy, Default, Debug)]
    pub struct UpsertKeys<T>(T);
    impl<D, T: columnar::Push<D>> columnar::Push<D> for UpsertKeys<T> {
        #[inline(always)]
        fn push(&mut self, item: D) {
            self.0.push(item)
        }
    }
    impl<T: columnar::Clear> columnar::Clear for UpsertKeys<T> {
        #[inline(always)]
        fn clear(&mut self) {
            self.0.clear()
        }
    }
    impl<T: columnar::Len> columnar::Len for UpsertKeys<T> {
        #[inline(always)]
        fn len(&self) -> usize {
            self.0.len()
        }
    }
    impl<'a> columnar::Index for UpsertKeys<&'a [UpsertKey]> {
        type Ref = &'a UpsertKey;

        #[inline(always)]
        fn get(&self, index: usize) -> Self::Ref {
            &self.0[index]
        }
    }

    impl Columnar for UpsertKey {
        #[inline(always)]
        fn into_owned<'a>(other: columnar::Ref<'a, Self>) -> Self {
            *other
        }
        type Container = UpsertKeys<Vec<UpsertKey>>;
        #[inline(always)]
        fn reborrow<'b, 'a: 'b>(thing: columnar::Ref<'a, Self>) -> columnar::Ref<'b, Self>
        where
            Self: 'a,
        {
            thing
        }
    }

    impl columnar::Borrow for UpsertKeys<Vec<UpsertKey>> {
        type Ref<'a> = &'a UpsertKey;
        type Borrowed<'a>
            = UpsertKeys<&'a [UpsertKey]>
        where
            Self: 'a;
        #[inline(always)]
        fn borrow<'a>(&'a self) -> Self::Borrowed<'a> {
            UpsertKeys(self.0.as_slice())
        }
        #[inline(always)]
        fn reborrow<'b, 'a: 'b>(item: Self::Borrowed<'a>) -> Self::Borrowed<'b>
        where
            Self: 'a,
        {
            UpsertKeys(item.0)
        }
        #[inline(always)]
        fn reborrow_ref<'b, 'a: 'b>(item: Self::Ref<'a>) -> Self::Ref<'b>
        where
            Self: 'a,
        {
            item
        }
    }

    impl columnar::Container for UpsertKeys<Vec<UpsertKey>> {
        #[inline(always)]
        fn extend_from_self(&mut self, other: Self::Borrowed<'_>, range: Range<usize>) {
            self.0.extend_from_self(other.0, range)
        }
        #[inline(always)]
        fn reserve_for<'a, I>(&mut self, selves: I)
        where
            Self: 'a,
            I: Iterator<Item = Self::Borrowed<'a>> + Clone,
        {
            self.0.reserve_for(selves.map(|s| s.0));
        }
    }

    impl<'a> columnar::AsBytes<'a> for UpsertKeys<&'a [UpsertKey]> {
        const SLICE_COUNT: usize = 1;
        #[inline(always)]
        fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
            mz_ore::soft_assert_no_log!(index < Self::SLICE_COUNT);
            (
                u64::cast_from(align_of::<UpsertKey>()),
                bytemuck::cast_slice(self.0),
            )
        }
        #[inline(always)]
        fn as_bytes(&self) -> impl Iterator<Item = (u64, &'a [u8])> {
            std::iter::once((
                u64::cast_from(align_of::<UpsertKey>()),
                bytemuck::cast_slice(self.0),
            ))
        }
    }
    impl<'a> columnar::FromBytes<'a> for UpsertKeys<&'a [UpsertKey]> {
        const SLICE_COUNT: usize = 1;
        #[inline(always)]
        fn from_bytes(bytes: &mut impl Iterator<Item = &'a [u8]>) -> Self {
            UpsertKeys(bytemuck::cast_slice(
                bytes.next().expect("Iterator exhausted prematurely"),
            ))
        }
    }
}

/// Projects a source's native `FromTime` to a columnar, totally-ordered key
/// used by the upsert source stash to keep the latest update per `(key, time)`.
///
/// The upsert stash is a paged columnar merge batcher; its diff carries this
/// projection rather than the raw `FromTime`, so the only columnar type the
/// stash needs is `Order` — never the (possibly structurally complex) source
/// timestamp. This is what keeps the columnar requirement off the generic
/// source-render path: that path only ever needs `FromTime: UpsertSourceTime`.
/// Only the relative order matters; the value is never read back.
///
/// The upsert envelope is rendered for Kafka and the KEY VALUE load generator,
/// so those source times (`KafkaTimestamp`, `MzOffset`) project to a real order
/// key. The remaining source times implement the trait only for coherence on
/// the generic render path — their sources never render upsert — so their
/// projection is a panicking guard rather than a real key.
pub trait UpsertSourceTime
where
    for<'a> columnar::Ref<'a, Self::Order>: Ord,
{
    /// Columnar order key. Must order consistently with the source time it is
    /// projected from.
    type Order: columnar::Columnar + Clone + Default + Ord + Send + Sync + 'static;
    /// Project the source time onto its order key.
    fn upsert_order(&self) -> Self::Order;
}

impl UpsertSourceTime for KafkaTimestamp {
    /// Per-record Kafka source times are exact singletons (a single partition
    /// at a single offset; see the source reader), and `KafkaTimestamp`'s
    /// derived `Ord` is lexicographic on `(partition, offset)`, so this flat
    /// projection is order-preserving. `RangeBound`'s infinities map to the
    /// `i64` extrema to remain order-consistent for any non-singleton bound.
    type Order = (i64, u64);
    fn upsert_order(&self) -> (i64, u64) {
        let partition = match self.interval().lower {
            RangeBound::NegInfinity => i64::MIN,
            RangeBound::Elem(p, _) => i64::from(p),
            RangeBound::PosInfinity => i64::MAX,
        };
        (partition, self.timestamp().offset)
    }
}

/// Load-generator (and Postgres) source time. The KEY VALUE load generator is
/// the one non-Kafka source that renders the upsert envelope (see
/// `apply_source_envelope_encoding` in the planner), so this projects to the
/// record offset: offsets increase with each update, so "max order wins" is
/// exactly "latest update wins" for dedup.
impl UpsertSourceTime for MzOffset {
    type Order = u64;
    fn upsert_order(&self) -> u64 {
        self.offset
    }
}

/// Source times whose sources never render the upsert envelope (MySQL and SQL
/// Server CDC). `Order = ()` keeps the generic render path free of any columnar
/// requirement on the source time. The projection panics rather than returning:
/// `()` would collapse every `from_time` to equal, silently breaking "latest
/// offset wins" dedup, so if such a source ever reaches the upsert path we want
/// a loud failure, not arbitrary per-key output.
macro_rules! upsert_source_time_unit {
    ($($ty:ty),+ $(,)?) => {$(
        impl UpsertSourceTime for $ty {
            type Order = ();
            fn upsert_order(&self) {
                unreachable!(
                    "upsert source stash is not rendered for this source, but \
                     {} reached the projection",
                    std::any::type_name::<Self>(),
                )
            }
        }
    )+};
}
upsert_source_time_unit!(GtidPartition, Lsn);

/// The storage spill gate ([`mz_timely_util::columnar::chunk::StorageSpill`]),
/// used by the chunked upsert-v2 stash flavor.
///
/// In that flavor the source stash and feedback arrangement spill through the
/// process buffer pool: committed chunk bodies land in the pool once compute's
/// config handler has installed and budgeted it (storage and compute run in
/// the same `clusterd` process). The gate is consulted at every chunk commit,
/// so flips apply to running dataflows.
pub mod upsert_stash_spill {
    /// Enable or disable spilling of upsert chunk bodies to the buffer pool.
    pub fn set_enabled(enabled: bool) {
        mz_timely_util::columnar::chunk::set_storage_spill_enabled(enabled);
    }
}

/// Pager for the paged upsert-v2 stash flavor.
///
/// This draws from the same process-wide [`TieredPolicy`] budget pool as the
/// compute column-paged batcher — there is one budget and one underlying
/// `mz_ore::pager` — but whether the stash *uses* it is gated by storage's own
/// `enable_upsert_paged_spill` flag, independently of compute's
/// `enable_column_paged_batcher_spill`. The shared pool's budget / backend /
/// codec are configured by compute's `apply_tiered_config` (storage and compute
/// run in the same `clusterd` process).
///
/// Flipping the flag takes effect on dataflows created after the change: the
/// paged flavor captures the pager once at operator construction.
///
/// [`TieredPolicy`]: mz_timely_util::column_pager::policy::TieredPolicy
pub mod upsert_stash_pager {
    use std::sync::{LazyLock, RwLock};

    use mz_timely_util::column_pager::{ColumnPager, shared_pager};

    /// Active pager handed to upsert source-stash batchers. Defaults to
    /// disabled (every chunk resident) until [`set_enabled`] turns it on.
    static PAGER: LazyLock<RwLock<ColumnPager>> =
        LazyLock::new(|| RwLock::new(ColumnPager::disabled()));

    /// Enable or disable the stash's use of the shared column pager. When
    /// enabled, the stash spills through the shared budget pool; when disabled
    /// it keeps every chunk resident.
    pub fn set_enabled(enabled: bool) {
        *PAGER.write().expect("upsert stash pager poisoned") = shared_pager(enabled);
    }

    /// The current upsert-stash pager. Cheap: clones the inner `Arc`.
    pub fn pager() -> ColumnPager {
        PAGER.read().expect("upsert stash pager poisoned").clone()
    }
}

impl Debug for UpsertKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "0x")?;
        for byte in self.0 {
            write!(f, "{:02x}", byte)?;
        }
        Ok(())
    }
}

impl AsRef<[u8]> for UpsertKey {
    #[inline(always)]
    // Note we do 1 `multi_get` and 1 `multi_put` while processing a _batch of updates_. Within the
    // batch, we effectively consolidate each key, before persisting that consolidated value.
    // Easy!!
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

impl From<&[u8]> for UpsertKey {
    fn from(bytes: &[u8]) -> Self {
        UpsertKey(bytes.try_into().expect("invalid key length"))
    }
}

/// The hash function used to map upsert keys. It is important that this hash is a cryptographic
/// hash so that there is no risk of collisions. Collisions on SHA256 have a probability of 2^128
/// which is many orders of magnitude smaller than many other events that we don't even think about
/// (e.g bit flips). In short, we can safely assume that sha256(a) == sha256(b) iff a == b.
type KeyHash = Sha256;

impl UpsertKey {
    pub fn from_key(key: Result<&Row, &UpsertError>) -> Self {
        Self::from_iter(key.map(|r| r.iter()))
    }

    pub fn from_value(value: Result<&Row, &UpsertError>, key_indices: &[usize]) -> Self {
        thread_local! {
            /// A thread-local datum cache used to calculate hashes
            static VALUE_DATUMS: RefCell<DatumVec> = RefCell::new(DatumVec::new());
        }
        VALUE_DATUMS.with(|value_datums| {
            let mut value_datums = value_datums.borrow_mut();
            let value = value.map(|v| value_datums.borrow_with(v));
            let key = match value {
                Ok(ref datums) => Ok(key_indices.iter().map(|&idx| datums[idx])),
                Err(err) => Err(err),
            };
            Self::from_iter(key)
        })
    }

    pub fn from_iter<'a, 'b>(
        key: Result<impl Iterator<Item = Datum<'a>> + 'b, &UpsertError>,
    ) -> Self {
        thread_local! {
            /// A thread-local datum cache used to calculate hashes
            static KEY_DATUMS: RefCell<DatumVec> = RefCell::new(DatumVec::new());
        }
        KEY_DATUMS.with(|key_datums| {
            let mut key_datums = key_datums.borrow_mut();
            // Borrowing the DatumVec gives us a temporary buffer to store datums in that will be
            // automatically cleared on Drop. See the DatumVec docs for more details.
            let mut key_datums = key_datums.borrow();
            let key: Result<&[Datum], Datum> = match key {
                Ok(key) => {
                    for datum in key {
                        key_datums.push(datum);
                    }
                    Ok(&*key_datums)
                }
                Err(UpsertError::Value(err)) => {
                    key_datums.extend(err.for_key.iter());
                    Ok(&*key_datums)
                }
                Err(UpsertError::KeyDecode(err)) => Err(Datum::Bytes(&err.raw)),
                Err(UpsertError::NullKey(_)) => Err(Datum::Null),
            };
            let mut hasher = DigestHasher(KeyHash::new());
            key.hash(&mut hasher);
            Self(hasher.0.finalize().into())
        })
    }
}

struct DigestHasher<H: Digest>(H);

impl<H: Digest> Hasher for DigestHasher<H> {
    fn write(&mut self, bytes: &[u8]) {
        self.0.update(bytes);
    }

    fn finish(&self) -> u64 {
        panic!("digest wrapper used to produce a hash");
    }
}

use std::convert::Infallible;
use timely::container::CapacityContainerBuilder;
use timely::dataflow::channels::pact::Pipeline;

use self::types::ValueMetadata;

/// This leaf operator drops `token` after the input reaches the `resume_upper`.
/// This is useful to take coordinated actions across all workers, after the `upsert`
/// operator has rehydrated.
pub fn rehydration_finished<'scope, T: Timestamp>(
    scope: Scope<'scope, T>,
    source_config: &crate::source::RawSourceCreationConfig,
    // A token that we can drop to signal we are finished rehydrating.
    token: impl std::any::Any + 'static,
    resume_upper: Antichain<T>,
    input: StreamVec<'scope, T, Infallible>,
) {
    let worker_id = source_config.worker_id;
    let id = source_config.id;
    let mut builder = AsyncOperatorBuilder::new(format!("rehydration_finished({id}"), scope);
    let mut input = builder.new_disconnected_input(input, Pipeline);

    builder.build(move |_capabilities| async move {
        let mut input_upper = Antichain::from_elem(Timestamp::minimum());
        // Ensure this operator finishes if the resume upper is `[0]`
        while !PartialOrder::less_equal(&resume_upper, &input_upper) {
            let Some(event) = input.next().await else {
                break;
            };
            if let AsyncEvent::Progress(upper) = event {
                input_upper = upper;
            }
        }
        tracing::info!(
            %worker_id,
            source_id = %id,
            "upsert source has downgraded past the resume upper ({resume_upper:?}) across all workers",
        );
        drop(token);
    });
}

/// Keys this operator's previous output, read back from persist, for retraction.
///
/// Only [`UpsertError`] survives the error side. It is the one error this operator can have
/// written, and therefore the one it can retract; anything else entered the shard from
/// elsewhere and is not ours to take back.
pub(crate) fn key_persist_feedback<'scope, T>(
    ok: VecCollection<'scope, T, Row, Diff>,
    err: VecCollection<'scope, T, DataflowError, Diff>,
    key_indices: Vec<usize>,
) -> VecCollection<'scope, T, (UpsertKey, UpsertValue), Diff>
where
    T: Timestamp,
{
    let keyed_ok = {
        let key_indices = key_indices.clone();
        ok.map(move |row| {
            let key = UpsertKey::from_value(Ok(&row), &key_indices);
            (key, Ok(row))
        })
    };
    let keyed_err = err.flat_map(move |err| {
        let err = match err {
            DataflowError::EnvelopeError(err) => match *err {
                EnvelopeError::Upsert(err) => Box::new(err),
                EnvelopeError::Flat(_) => return None,
            },
            _ => return None,
        };
        let key = UpsertKey::from_value(Err(&err), &key_indices);
        Some((key, Err(err)))
    });
    keyed_ok.concat(keyed_err)
}

/// Resumes an upsert computation at `resume_upper` given as inputs a collection of upsert commands
/// and the collection of the previous output of this operator.
/// Returns a tuple of
/// - A collection of the computed upsert operator and,
/// - A health update stream to propagate errors
pub(crate) fn upsert<'scope, T, FromTime>(
    input: VecCollection<'scope, T, (UpsertKey, Option<UpsertValue>, FromTime), Diff>,
    upsert_envelope: UpsertEnvelope,
    resume_upper: Antichain<T>,
    previous_ok: VecCollection<'scope, T, Row, Diff>,
    previous_err: VecCollection<'scope, T, DataflowError, Diff>,
    previous_token: Option<Vec<PressOnDropButton>>,
    source_config: crate::source::SourceExportCreationConfig,
    instance_context: &StorageInstanceContext,
    storage_configuration: &StorageConfiguration,
    dataflow_paramters: &crate::internal_control::DataflowParameters,
    backpressure_metrics: Option<UpsertBackpressureMetrics>,
) -> (
    VecCollection<'scope, T, Result<Row, DataflowError>, Diff>,
    StreamVec<'scope, T, (Option<GlobalId>, HealthStatusUpdate)>,
    StreamVec<'scope, T, Infallible>,
    PressOnDropButton,
)
where
    T: Timestamp + TotalOrder + Sync,
    T: Refines<mz_repr::Timestamp> + TotalOrder + Sync,
    FromTime: Timestamp + Clone + Sync,
{
    let upsert_metrics = source_config.metrics.get_upsert_metrics(
        source_config.id,
        source_config.worker_id,
        backpressure_metrics,
    );

    let rocksdb_cleanup_tries =
        dyncfgs::STORAGE_ROCKSDB_CLEANUP_TRIES.get(storage_configuration.config_set());

    // Whether or not to partially drain the input buffer
    // to prevent buffering of the _upstream_ snapshot.
    let prevent_snapshot_buffering =
        dyncfgs::STORAGE_UPSERT_PREVENT_SNAPSHOT_BUFFERING.get(storage_configuration.config_set());
    // If the above is true, the number of timely batches to process at once.
    let snapshot_buffering_max = dyncfgs::STORAGE_UPSERT_MAX_SNAPSHOT_BATCH_BUFFERING
        .get(storage_configuration.config_set());

    // Whether we should provide the upsert state merge operator to the RocksDB instance
    // (for faster performance during snapshot hydration).
    let rocksdb_use_native_merge_operator =
        dyncfgs::STORAGE_ROCKSDB_USE_MERGE_OPERATOR.get(storage_configuration.config_set());

    let upsert_config = UpsertConfig {
        shrink_upsert_unused_buffers_by_ratio: storage_configuration
            .parameters
            .shrink_upsert_unused_buffers_by_ratio,
    };

    let thin_input = upsert_thinning(input);

    let tuning = dataflow_paramters.upsert_rocksdb_tuning_config.clone();

    // When running RocksDB in memory, the file system is emulated. However, we still need to
    // pick a path that exists because RocksDB will attempt to create the working directory
    // (see https://github.com/rust-rocksdb/rust-rocksdb/issues/1015) and write a lock file,
    // so we need to ensure the directory is unique per worker.
    let rocksdb_dir = instance_context
        .scratch_directory
        .clone()
        .unwrap_or_else(|| PathBuf::from("/tmp"))
        .join("storage")
        .join("upsert")
        .join(source_config.id.to_string())
        .join(source_config.worker_id.to_string());

    tracing::info!(
        worker_id = %source_config.worker_id,
        source_id = %source_config.id,
        ?rocksdb_dir,
        ?tuning,
        ?rocksdb_use_native_merge_operator,
        "rendering upsert source"
    );

    let rocksdb_shared_metrics = Arc::clone(&upsert_metrics.rocksdb_shared);
    let rocksdb_instance_metrics = Arc::clone(&upsert_metrics.rocksdb_instance_metrics);

    let env = instance_context
        .rocksdb_env()
        .expect("failed to create rocksdb env");

    // A closure that will initialize and return a configured RocksDB instance
    let rocksdb_init_fn = move || async move {
        let merge_operator = if rocksdb_use_native_merge_operator {
            Some((
                "upsert_state_snapshot_merge_v1".to_string(),
                |a: &[u8], b: ValueIterator<BincodeOpts, StateValue<T, FromTime>>| {
                    consolidating_merge_function::<T, FromTime>(a.into(), b)
                },
            ))
        } else {
            None
        };
        rocksdb::RocksDB::new(
            mz_rocksdb::RocksDBInstance::new(
                &rocksdb_dir,
                mz_rocksdb::InstanceOptions::new(
                    env,
                    rocksdb_cleanup_tries,
                    merge_operator,
                    // For now, just use the same config as the one used for
                    // merging snapshots.
                    upsert_bincode_opts(),
                ),
                tuning,
                rocksdb_shared_metrics,
                rocksdb_instance_metrics,
            )
            .unwrap(),
        )
    };

    upsert_operator(
        thin_input,
        upsert_envelope.key_indices,
        resume_upper,
        previous_ok,
        previous_err,
        previous_token,
        upsert_metrics,
        source_config,
        rocksdb_init_fn,
        upsert_config,
        storage_configuration,
        prevent_snapshot_buffering,
        snapshot_buffering_max,
    )
}

/// An experimental upsert implementation loosely described in this doc:
/// [Upsert V2 Much Simpler Boogaloo](https://www.notion.so/materialize/Upsert-V2-Much-Simpler-Boogaloo-31913f48d37b807fa88bdeafc27c02d9?source=copy_link)
///
/// Instead of using rocksdb as a state backend, this implementation uses a differential dataflow collection to hold the key state,
/// and performs consolidation of updates with matching keys and MZ timestamps, using max FromTime to choose winners,
/// resulting in only one record per key per time.
pub(crate) fn upsert_v2<'scope, T, FromTime>(
    input: VecCollection<'scope, T, (UpsertKey, Option<UpsertValue>, FromTime), Diff>,
    upsert_envelope: UpsertEnvelope,
    resume_upper: Antichain<T>,
    previous_ok: VecCollection<'scope, T, Row, Diff>,
    previous_err: VecCollection<'scope, T, DataflowError, Diff>,
    previous_token: Option<Vec<PressOnDropButton>>,
    source_config: crate::source::SourceExportCreationConfig,
    backpressure_metrics: Option<UpsertBackpressureMetrics>,
    stash_flavor: upsert_continual_feedback_v2::UpsertStashFlavor,
) -> (
    VecCollection<'scope, T, Result<Row, DataflowError>, Diff>,
    StreamVec<'scope, T, (Option<GlobalId>, HealthStatusUpdate)>,
    StreamVec<'scope, T, Infallible>,
    PressOnDropButton,
)
where
    T: Timestamp + TotalOrder + Sync,
    T: Refines<mz_repr::Timestamp> + differential_dataflow::lattice::Lattice,
    T: columnation::Columnation,
    T: columnar::Columnar + Default,
    for<'a> columnar::Ref<'a, T>: Copy + Ord,
    FromTime: Timestamp + Clone + Sync,
    FromTime: UpsertSourceTime,
{
    let upsert_metrics = source_config.metrics.get_upsert_metrics(
        source_config.id,
        source_config.worker_id,
        backpressure_metrics,
    );

    let thin_input = upsert_thinning(input);

    tracing::info!(
        worker_id = %source_config.worker_id,
        source_id = %source_config.id,
        ?stash_flavor,
        "rendering upsert source (btreemap backend)"
    );

    upsert_continual_feedback_v2::upsert_inner(
        stash_flavor,
        thin_input,
        upsert_envelope.key_indices,
        resume_upper,
        previous_ok,
        previous_err,
        previous_token,
        upsert_metrics,
        source_config,
    )
}

// A shim so we can dispatch based on the dyncfg that tells us which upsert
// operator to use.
fn upsert_operator<'scope, T, FromTime, F, Fut, US>(
    input: VecCollection<'scope, T, (UpsertKey, Option<UpsertValue>, FromTime), Diff>,
    key_indices: Vec<usize>,
    resume_upper: Antichain<T>,
    persist_ok: VecCollection<'scope, T, Row, Diff>,
    persist_err: VecCollection<'scope, T, DataflowError, Diff>,
    persist_token: Option<Vec<PressOnDropButton>>,
    upsert_metrics: UpsertMetrics,
    source_config: crate::source::SourceExportCreationConfig,
    state: F,
    upsert_config: UpsertConfig,
    _storage_configuration: &StorageConfiguration,
    prevent_snapshot_buffering: bool,
    snapshot_buffering_max: Option<usize>,
) -> (
    VecCollection<'scope, T, Result<Row, DataflowError>, Diff>,
    StreamVec<'scope, T, (Option<GlobalId>, HealthStatusUpdate)>,
    StreamVec<'scope, T, Infallible>,
    PressOnDropButton,
)
where
    T: Timestamp + TotalOrder + Sync,
    T: Refines<mz_repr::Timestamp> + TotalOrder + Sync,
    F: FnOnce() -> Fut + 'static,
    Fut: std::future::Future<Output = US>,
    US: UpsertStateBackend<T, FromTime>,
    FromTime: Debug + timely::ExchangeData + Clone + Ord + Sync,
{
    // Hard-coded to true because classic UPSERT cannot be used safely with
    // concurrent ingestions, which we need for both 0dt upgrades and
    // multi-replica ingestions.
    let use_continual_feedback_upsert = true;

    tracing::info!(id = %source_config.id, %use_continual_feedback_upsert, "upsert operator implementation");

    if use_continual_feedback_upsert {
        upsert_continual_feedback::upsert_inner(
            input,
            key_indices,
            resume_upper,
            persist_ok,
            persist_err,
            persist_token,
            upsert_metrics,
            source_config,
            state,
            upsert_config,
            prevent_snapshot_buffering,
            snapshot_buffering_max,
        )
    } else {
        upsert_classic(
            input,
            key_indices,
            resume_upper,
            persist_ok,
            persist_err,
            persist_token,
            upsert_metrics,
            source_config,
            state,
            upsert_config,
            prevent_snapshot_buffering,
            snapshot_buffering_max,
        )
    }
}

/// Renders an operator that discards updates that are known to not affect the outcome of upsert in
/// a streaming fashion. For each distinct (key, time) in the input it emits the value with the
/// highest from_time. Its purpose is to thin out data as much as possible before exchanging them
/// across workers.
fn upsert_thinning<'scope, T, K, V, FromTime>(
    input: VecCollection<'scope, T, (K, V, FromTime), Diff>,
) -> VecCollection<'scope, T, (K, V, FromTime), Diff>
where
    T: Timestamp + TotalOrder,
    K: timely::ExchangeData + Clone + Eq + Ord,
    V: timely::ExchangeData + Clone,
    FromTime: Timestamp,
{
    input
        .inner
        .unary(Pipeline, "UpsertThinning", |_, _| {
            // A capability suitable to emit all updates in `updates`, if any.
            let mut capability: Option<InputCapability<T>> = None;
            // A batch of received updates
            let mut updates = Vec::new();
            move |input, output| {
                input.for_each(|cap, data| {
                    assert!(
                        data.iter().all(|(_, _, diff)| diff.is_positive()),
                        "invalid upsert input"
                    );
                    updates.append(data);
                    match capability.as_mut() {
                        Some(capability) => {
                            if cap.time() <= capability.time() {
                                *capability = cap;
                            }
                        }
                        None => capability = Some(cap),
                    }
                });
                if let Some(capability) = capability.take() {
                    // Sort by (key, time, Reverse(from_time)) so that deduping by (key, time) gives
                    // the latest change for that key.
                    updates.sort_unstable_by(|a, b| {
                        let ((key1, _, from_time1), time1, _) = a;
                        let ((key2, _, from_time2), time2, _) = b;
                        Ord::cmp(
                            &(key1, time1, Reverse(from_time1)),
                            &(key2, time2, Reverse(from_time2)),
                        )
                    });
                    let mut session = output.session(&capability);
                    session.give_iterator(updates.drain(..).dedup_by(|a, b| {
                        let ((key1, _, _), time1, _) = a;
                        let ((key2, _, _), time2, _) = b;
                        (key1, time1) == (key2, time2)
                    }))
                }
            }
        })
        .as_collection()
}

/// Helper method for `upsert_classic` used to stage `data` updates
/// from the input/source timely edge.
fn stage_input<T, FromTime>(
    stash: &mut Vec<(T, UpsertKey, Reverse<FromTime>, Option<UpsertValue>)>,
    data: &mut Vec<((UpsertKey, Option<UpsertValue>, FromTime), T, Diff)>,
    input_upper: &Antichain<T>,
    resume_upper: &Antichain<T>,
    storage_shrink_upsert_unused_buffers_by_ratio: usize,
) where
    T: PartialOrder,
    FromTime: Ord,
{
    if PartialOrder::less_equal(input_upper, resume_upper) {
        data.retain(|(_, ts, _)| resume_upper.less_equal(ts));
    }

    stash.extend(data.drain(..).map(|((key, value, order), time, diff)| {
        assert!(diff.is_positive(), "invalid upsert input");
        (time, key, Reverse(order), value)
    }));

    if storage_shrink_upsert_unused_buffers_by_ratio > 0 {
        let reduced_capacity = stash.capacity() / storage_shrink_upsert_unused_buffers_by_ratio;
        if reduced_capacity > stash.len() {
            stash.shrink_to(reduced_capacity);
        }
    }
}

/// The style of drain we are performing on the stash. `AtTime`-drains cannot
/// assume that all values have been seen, and must leave tombstones behind for deleted values.
#[derive(Debug)]
enum DrainStyle<'a, T> {
    ToUpper(&'a Antichain<T>),
    AtTime(T),
}

/// Helper method for `upsert_inner` used to stage `data` updates
/// from the input timely edge.
async fn drain_staged_input<S, T, FromTime, E>(
    stash: &mut Vec<(T, UpsertKey, Reverse<FromTime>, Option<UpsertValue>)>,
    commands_state: &mut indexmap::IndexMap<UpsertKey, types::UpsertValueAndSize<T, FromTime>>,
    output_updates: &mut Vec<(UpsertValue, T, Diff)>,
    multi_get_scratch: &mut Vec<UpsertKey>,
    drain_style: DrainStyle<'_, T>,
    error_emitter: &mut E,
    state: &mut UpsertState<'_, S, T, FromTime>,
    source_config: &crate::source::SourceExportCreationConfig,
) where
    S: UpsertStateBackend<T, FromTime>,
    T: PartialOrder + Ord + Clone + Send + Sync + Serialize + Debug + 'static,
    FromTime: timely::ExchangeData + Clone + Ord + Sync,
    E: UpsertErrorEmitter<T>,
{
    stash.sort_unstable();

    // Find the prefix that we can emit
    let idx = stash.partition_point(|(ts, _, _, _)| match &drain_style {
        DrainStyle::ToUpper(upper) => !upper.less_equal(ts),
        DrainStyle::AtTime(time) => ts <= time,
    });

    tracing::trace!(?drain_style, updates = idx, "draining stash in upsert");

    // Read the previous values _per key_ out of `state`, recording it
    // along with the value with the _latest timestamp for that key_.
    commands_state.clear();
    for (_, key, _, _) in stash.iter().take(idx) {
        commands_state.entry(*key).or_default();
    }

    // These iterators iterate in the same order because `commands_state`
    // is an `IndexMap`.
    multi_get_scratch.clear();
    multi_get_scratch.extend(commands_state.iter().map(|(k, _)| *k));
    match state
        .multi_get(multi_get_scratch.drain(..), commands_state.values_mut())
        .await
    {
        Ok(_) => {}
        Err(e) => {
            error_emitter
                .emit("Failed to fetch records from state".to_string(), e)
                .await;
        }
    }

    // From the prefix that can be emitted we can deduplicate based on (ts, key) in
    // order to only process the command with the maximum order within the (ts,
    // key) group. This is achieved by wrapping order in `Reverse(FromTime)` above.;
    let mut commands = stash.drain(..idx).dedup_by(|a, b| {
        let ((a_ts, a_key, _, _), (b_ts, b_key, _, _)) = (a, b);
        a_ts == b_ts && a_key == b_key
    });

    let bincode_opts = types::upsert_bincode_opts();
    // Upsert the values into `commands_state`, by recording the latest
    // value (or deletion). These will be synced at the end to the `state`.
    //
    // Note that we are effectively doing "mini-upsert" here, using
    // `command_state`. This "mini-upsert" is seeded with data from `state`, using
    // a single `multi_get` above, and the final state is written out into
    // `state` using a single `multi_put`. This simplifies `UpsertStateBackend`
    // implementations, and reduces the number of reads and write we need to do.
    //
    // This "mini-upsert" technique is actually useful in `UpsertState`'s
    // `consolidate_snapshot_read_write_inner` implementation, minimizing gets and puts on
    // the `UpsertStateBackend` implementations. In some sense, its "upsert all the way down".
    while let Some((ts, key, from_time, value)) = commands.next() {
        let mut command_state = if let Entry::Occupied(command_state) = commands_state.entry(key) {
            command_state
        } else {
            panic!("key missing from commands_state");
        };

        let existing_value = &mut command_state.get_mut().value;

        if let Some(cs) = existing_value.as_mut() {
            cs.ensure_decoded(bincode_opts, source_config.id, Some(&key));
        }

        // Skip this command if its order key is below the one in the upsert state.
        // Note that the existing order key may be `None` if the existing value
        // is from snapshotting, which always sorts below new values/deletes.
        let existing_order = existing_value
            .as_ref()
            .and_then(|cs| cs.provisional_order(&ts));
        if existing_order >= Some(&from_time.0) {
            // Skip this update. If no later updates adjust this key, then we just
            // end up writing the same value back to state. If there
            // is nothing in the state, `existing_order` is `None`, and this
            // does not occur.
            continue;
        }

        match value {
            Some(value) => {
                if let Some(old_value) =
                    existing_value.replace(StateValue::finalized_value(value.clone()))
                {
                    if let Some(old_value) = old_value.into_decoded().finalized {
                        output_updates.push((old_value, ts.clone(), Diff::MINUS_ONE));
                    }
                }
                output_updates.push((value, ts, Diff::ONE));
            }
            None => {
                if let Some(old_value) = existing_value.take() {
                    if let Some(old_value) = old_value.into_decoded().finalized {
                        output_updates.push((old_value, ts, Diff::MINUS_ONE));
                    }
                }

                // Record a tombstone for deletes.
                *existing_value = Some(StateValue::tombstone());
            }
        }
    }

    match state
        .multi_put(
            true, // Do update per-update stats.
            commands_state.drain(..).map(|(k, cv)| {
                (
                    k,
                    types::PutValue {
                        value: cv.value.map(|cv| cv.into_decoded()),
                        previous_value_metadata: cv.metadata.map(|v| ValueMetadata {
                            size: v.size.try_into().expect("less than i64 size"),
                            is_tombstone: v.is_tombstone,
                        }),
                    },
                )
            }),
        )
        .await
    {
        Ok(_) => {}
        Err(e) => {
            error_emitter
                .emit("Failed to update records in state".to_string(), e)
                .await;
        }
    }
}

/// A no-op-ish error emitter for the fuzzing hook. With the in-memory backend
/// and the well-formed inputs the fuzzer builds, `multi_get`/`multi_put` never
/// error, so reaching this is itself a finding.
#[cfg(feature = "fuzzing")]
struct PanicErrorEmitter;

#[cfg(feature = "fuzzing")]
#[async_trait::async_trait(?Send)]
impl<T> UpsertErrorEmitter<T> for PanicErrorEmitter {
    async fn emit(&mut self, context: String, e: anyhow::Error) {
        panic!("unexpected upsert state error during fuzzing: {context}: {e}");
    }
}

/// Fuzzing hook: run a single `drain_staged_input` over `commands` (each a
/// `(timestamp, key, order, value)`, where `value == None` is a delete) against
/// a fresh empty in-memory state, draining everything strictly below
/// `drain_to`. Returns the emitted output updates and the final finalized value
/// of each key in `all_keys`. Exposed only for fuzzing. Not a stable public
/// API.
#[cfg(feature = "fuzzing")]
pub async fn fuzz_drain_staged_input(
    parts: &types::FuzzUpsertParts,
    source_config: &crate::source::SourceExportCreationConfig,
    commands: Vec<(u64, UpsertKey, u64, Option<UpsertValue>)>,
    drain_to: u64,
    all_keys: &[UpsertKey],
) -> (Vec<(UpsertValue, u64, Diff)>, Vec<Option<UpsertValue>>) {
    let mut state = parts.state();
    let mut stash: Vec<(u64, UpsertKey, Reverse<u64>, Option<UpsertValue>)> = commands
        .into_iter()
        .map(|(ts, key, order, value)| (ts, key, Reverse(order), value))
        .collect();
    let mut commands_state = indexmap::IndexMap::new();
    let mut output = Vec::new();
    let mut multi_get_scratch = Vec::new();
    let mut emitter = PanicErrorEmitter;

    drain_staged_input(
        &mut stash,
        &mut commands_state,
        &mut output,
        &mut multi_get_scratch,
        DrainStyle::ToUpper(&Antichain::from_elem(drain_to)),
        &mut emitter,
        &mut state,
        source_config,
    )
    .await;

    let bincode_opts = types::upsert_bincode_opts();
    let mut results = vec![types::UpsertValueAndSize::default(); all_keys.len()];
    state
        .multi_get(all_keys.iter().copied(), results.iter_mut())
        .await
        .expect("multi_get in fuzz hook should not error");
    let final_state = results
        .into_iter()
        .map(|r| match r.value {
            None => None,
            Some(mut sv) => {
                sv.ensure_decoded(bincode_opts, GlobalId::User(0), None);
                sv.into_decoded().finalized
            }
        })
        .collect();

    (output, final_state)
}

// Created a struct to hold the configs for upserts.
// So that new configs don't require a new method parameter.
pub(crate) struct UpsertConfig {
    pub shrink_upsert_unused_buffers_by_ratio: usize,
}

fn upsert_classic<'scope, T, FromTime, F, Fut, US>(
    input: VecCollection<'scope, T, (UpsertKey, Option<UpsertValue>, FromTime), Diff>,
    key_indices: Vec<usize>,
    resume_upper: Antichain<T>,
    previous_ok: VecCollection<'scope, T, Row, Diff>,
    previous_err: VecCollection<'scope, T, DataflowError, Diff>,
    previous_token: Option<Vec<PressOnDropButton>>,
    upsert_metrics: UpsertMetrics,
    source_config: crate::source::SourceExportCreationConfig,
    state: F,
    upsert_config: UpsertConfig,
    prevent_snapshot_buffering: bool,
    snapshot_buffering_max: Option<usize>,
) -> (
    VecCollection<'scope, T, Result<Row, DataflowError>, Diff>,
    StreamVec<'scope, T, (Option<GlobalId>, HealthStatusUpdate)>,
    StreamVec<'scope, T, Infallible>,
    PressOnDropButton,
)
where
    T: Timestamp + TotalOrder + Sync,
    F: FnOnce() -> Fut + 'static,
    Fut: std::future::Future<Output = US>,
    US: UpsertStateBackend<T, FromTime>,
    FromTime: timely::ExchangeData + Clone + Ord + Sync,
{
    let mut builder = AsyncOperatorBuilder::new("Upsert".to_string(), input.scope());

    let previous = key_persist_feedback(previous_ok, previous_err, key_indices);
    let (output_handle, output) = builder.new_output();

    // An output that just reports progress of the snapshot consolidation process upstream to the
    // persist source to ensure that backpressure is applied
    let (_snapshot_handle, snapshot_stream) =
        builder.new_output::<CapacityContainerBuilder<Vec<Infallible>>>();

    let (mut health_output, health_stream) = builder.new_output();
    let mut input = builder.new_input_for(
        input.inner,
        Exchange::new(move |((key, _, _), _, _)| UpsertKey::hashed(key)),
        &output_handle,
    );

    let mut previous = builder.new_input_for(
        previous.inner,
        Exchange::new(|((key, _), _, _)| UpsertKey::hashed(key)),
        &output_handle,
    );

    let upsert_shared_metrics = Arc::clone(&upsert_metrics.shared);
    let shutdown_button = builder.build(move |caps| async move {
        let [mut output_cap, mut snapshot_cap, health_cap]: [_; 3] = caps.try_into().unwrap();

        let mut state = UpsertState::<_, _, FromTime>::new(
            state().await,
            upsert_shared_metrics,
            &upsert_metrics,
            source_config.source_statistics.clone(),
            upsert_config.shrink_upsert_unused_buffers_by_ratio,
        );
        let mut events = vec![];
        let mut snapshot_upper = Antichain::from_elem(Timestamp::minimum());

        let mut stash = vec![];

        let mut error_emitter = (&mut health_output, &health_cap);

        tracing::info!(
            ?resume_upper,
            ?snapshot_upper,
            "timely-{} upsert source {} starting rehydration",
            source_config.worker_id,
            source_config.id
        );
        // Read and consolidate the snapshot from the 'previous' input until it
        // reaches the `resume_upper`.
        while !PartialOrder::less_equal(&resume_upper, &snapshot_upper) {
            previous.ready().await;
            while let Some(event) = previous.next_sync() {
                match event {
                    AsyncEvent::Data(_cap, data) => {
                        events.extend(data.into_iter().filter_map(|((key, value), ts, diff)| {
                            if !resume_upper.less_equal(&ts) {
                                Some((key, value, diff))
                            } else {
                                None
                            }
                        }))
                    }
                    AsyncEvent::Progress(upper) => {
                        snapshot_upper = upper;
                    }
                };
            }

            match state
                .consolidate_chunk(
                    events.drain(..),
                    PartialOrder::less_equal(&resume_upper, &snapshot_upper),
                )
                .await
            {
                Ok(_) => {
                    if let Some(ts) = snapshot_upper.clone().into_option() {
                        // As we shutdown, we could ostensibly get data from later than the
                        // `resume_upper`, which we ignore above. We don't want our output capability to make
                        // it further than the `resume_upper`.
                        if !resume_upper.less_equal(&ts) {
                            snapshot_cap.downgrade(&ts);
                            output_cap.downgrade(&ts);
                        }
                    }
                }
                Err(e) => {
                    UpsertErrorEmitter::<T>::emit(
                        &mut error_emitter,
                        "Failed to rehydrate state".to_string(),
                        e,
                    )
                    .await;
                }
            }
        }

        drop(events);
        drop(previous_token);
        drop(snapshot_cap);

        // Exchaust the previous input. It is expected to immediately reach the empty
        // antichain since we have dropped its token.
        //
        // Note that we do not need to also process the `input` during this, as the dropped token
        // will shutdown the `backpressure` operator
        while let Some(_event) = previous.next().await {}

        // After snapshotting, our output frontier is exactly the `resume_upper`
        if let Some(ts) = resume_upper.as_option() {
            output_cap.downgrade(ts);
        }

        tracing::info!(
            "timely-{} upsert source {} finished rehydration",
            source_config.worker_id,
            source_config.id
        );

        // A re-usable buffer of changes, per key. This is an `IndexMap` because it has to be `drain`-able
        // and have a consistent iteration order.
        let mut commands_state: indexmap::IndexMap<_, types::UpsertValueAndSize<T, FromTime>> =
            indexmap::IndexMap::new();
        let mut multi_get_scratch = Vec::new();

        // Now can can resume consuming the collection
        let mut output_updates = vec![];
        let mut input_upper = Antichain::from_elem(Timestamp::minimum());

        while let Some(event) = input.next().await {
            // Buffer as many events as possible. This should be bounded, as new data can't be
            // produced in this worker until we yield to timely.
            let events = [event]
                .into_iter()
                .chain(std::iter::from_fn(|| input.next().now_or_never().flatten()))
                .enumerate();

            let mut partial_drain_time = None;
            for (i, event) in events {
                match event {
                    AsyncEvent::Data(cap, mut data) => {
                        tracing::trace!(
                            time=?cap.time(),
                            updates=%data.len(),
                            "received data in upsert"
                        );
                        stage_input(
                            &mut stash,
                            &mut data,
                            &input_upper,
                            &resume_upper,
                            upsert_config.shrink_upsert_unused_buffers_by_ratio,
                        );

                        let event_time = cap.time();
                        // If the data is at _exactly_ the output frontier, we can preemptively drain it into the state.
                        // Data within this set events strictly beyond this time are staged as
                        // normal.
                        //
                        // This is a load-bearing optimization, as it is required to avoid buffering
                        // the entire source snapshot in the `stash`.
                        if prevent_snapshot_buffering && output_cap.time() == event_time {
                            partial_drain_time = Some(event_time.clone());
                        }
                    }
                    AsyncEvent::Progress(upper) => {
                        tracing::trace!(?upper, "received progress in upsert");
                        // Ignore progress updates before the `resume_upper`, which is our initial
                        // capability post-snapshotting.
                        if PartialOrder::less_than(&upper, &resume_upper) {
                            continue;
                        }

                        // Disable the partial drain as this progress event covers
                        // the `output_cap` time.
                        partial_drain_time = None;
                        drain_staged_input::<_, _, _, _>(
                            &mut stash,
                            &mut commands_state,
                            &mut output_updates,
                            &mut multi_get_scratch,
                            DrainStyle::ToUpper(&upper),
                            &mut error_emitter,
                            &mut state,
                            &source_config,
                        )
                        .await;

                        output_handle.give_container(&output_cap, &mut output_updates);

                        if let Some(ts) = upper.as_option() {
                            output_cap.downgrade(ts);
                        }
                        input_upper = upper;
                    }
                }
                let events_processed = i + 1;
                if let Some(max) = snapshot_buffering_max {
                    if events_processed >= max {
                        break;
                    }
                }
            }

            // If there were staged events that occurred at the capability time, drain
            // them. This is safe because out-of-order updates to the same key that are
            // drained in separate calls to `drain_staged_input` are correctly ordered by
            // their `FromTime` in `drain_staged_input`.
            //
            // Note also that this may result in more updates in the output collection than
            // the minimum. However, because the frontier only advances on `Progress` updates,
            // the collection always accumulates correctly for all keys.
            if let Some(partial_drain_time) = partial_drain_time {
                drain_staged_input::<_, _, _, _>(
                    &mut stash,
                    &mut commands_state,
                    &mut output_updates,
                    &mut multi_get_scratch,
                    DrainStyle::AtTime(partial_drain_time),
                    &mut error_emitter,
                    &mut state,
                    &source_config,
                )
                .await;

                output_handle.give_container(&output_cap, &mut output_updates);
            }
        }
    });

    (
        output.as_collection().map(|result| match result {
            Ok(ok) => Ok(ok),
            Err(err) => Err(DataflowError::from(EnvelopeError::Upsert(*err))),
        }),
        health_stream,
        snapshot_stream,
        shutdown_button.press_on_drop(),
    )
}

#[async_trait::async_trait(?Send)]
pub(crate) trait UpsertErrorEmitter<T> {
    async fn emit(&mut self, context: String, e: anyhow::Error);
}

#[async_trait::async_trait(?Send)]
impl<T: Timestamp> UpsertErrorEmitter<T>
    for (
        &mut AsyncOutputHandle<
            T,
            CapacityContainerBuilder<Vec<(Option<GlobalId>, HealthStatusUpdate)>>,
        >,
        &Capability<T>,
    )
{
    async fn emit(&mut self, context: String, e: anyhow::Error) {
        process_upsert_state_error::<T>(context, e, self.0, self.1).await
    }
}

/// Emit the given error, and stall till the dataflow is restarted.
async fn process_upsert_state_error<T: Timestamp>(
    context: String,
    e: anyhow::Error,
    health_output: &AsyncOutputHandle<
        T,
        CapacityContainerBuilder<Vec<(Option<GlobalId>, HealthStatusUpdate)>>,
    >,
    health_cap: &Capability<T>,
) {
    let update = HealthStatusUpdate::halting(e.context(context).to_string_with_causes(), None);
    health_output.give(health_cap, (None, update));
    std::future::pending::<()>().await;
    unreachable!("pending future never returns");
}