mz_storage_operators/

persist.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::collections::BTreeMap;
use std::fmt::Debug;
use std::sync::{Arc, Weak};

use mz_ore::task::AbortOnDropHandle;
use mz_persist_client::batch::Batch;
use mz_persist_client::{PersistClient, Schemas};
use mz_persist_types::ShardId;
use mz_persist_types::part::Part;
use mz_repr::Timestamp;
use mz_storage_types::StorageDiff;
use mz_storage_types::sources::SourceData;
use serde::{Deserialize, Serialize};
use timely::progress::Antichain;
use uuid::Uuid;

/// A unique identifier for a [SharedBatchBuilder].
#[derive(
    Debug,
    Copy,
    Clone,
    Serialize,
    Deserialize,
    Ord,
    PartialOrd,
    Eq,
    PartialEq
)]
#[serde(transparent)]
pub struct SharedBatchId(Uuid);

impl SharedBatchId {
    /// Returns a fresh, random id.
    pub fn new() -> Self {
        Self(Uuid::new_v4())
    }
}

/// A struct from which to obtain a [SharedBatchBuilder].
#[derive(Debug, Clone)]
pub struct SharedBatches {
    data: Arc<std::sync::Mutex<Inner>>,
}

#[derive(Debug, Default)]
struct Inner {
    /// The size of `states` after the last sweep of dead entries.
    last_retained_len: usize,
    states: BTreeMap<SharedBatchId, Weak<BatchState>>,
}

fn upgrade_or_init<T>(weak: &mut Weak<T>, init: impl FnOnce() -> T) -> Arc<T> {
    if let Some(owned) = weak.upgrade() {
        owned
    } else {
        let owned = Arc::new(init());
        *weak = Arc::downgrade(&owned);
        owned
    }
}

impl SharedBatches {
    /// Returns an empty registry. Clones share the registry.
    pub fn new() -> Self {
        Self {
            data: Arc::default(),
        }
    }

    /// Get a builder for the specified batch. If there is an existing, live builder for the same
    /// batch id managed by [Self], any parts added to this builder will be included in the same,
    /// shared batch; if not, a new batch will be created.
    ///
    /// This is designed so that, in a multi-worker dataflow, workers that are building the same batch
    /// at the same time can write fewer, larger parts instead of many small ones. It does not guarantee
    /// that there will be _precisely_ one batch across all workers, however.
    ///
    /// The batch's description is fixed by whichever caller creates it, so every caller sharing a
    /// `shared_batch_id` must pass the same `shard_id`, `lower`, and `upper`, or this panics.
    pub fn builder(
        &self,
        shared_batch_id: SharedBatchId,
        client: PersistClient,
        shard_id: ShardId,
        shard_name: String,
        schemas: Schemas<SourceData, ()>,
        lower: Antichain<Timestamp>,
        upper: Antichain<Timestamp>,
    ) -> SharedBatchBuilder {
        let mut guard = self.data.lock().unwrap();
        let inner = &mut *guard;
        // We clean out entries from the map where all shared batch handles have been dropped.
        // Amortize by only scanning the map once it's doubled in size.
        if inner.states.len() > inner.last_retained_len * 2 {
            // Probe the count rather than `upgrade()`: the temporary strong reference an upgrade
            // creates can make a concurrent `SharedBatchBuilder::finish` see a strong count above
            // one and hand the batch to nobody, which seals the interval with no data. A count of
            // zero cannot become live again, and a stale non-zero count only keeps a dead entry
            // until the next sweep.
            inner.states.retain(|_, weak| weak.strong_count() > 0);
            inner.last_retained_len = inner.states.len();
        }
        let weak = inner.states.entry(shared_batch_id).or_default();
        let desc = BatchDesc {
            shard_id,
            lower: lower.clone(),
            upper: upper.clone(),
        };
        let state = upgrade_or_init(weak, || {
            let (tx, mut rx) = tokio::sync::mpsc::channel(PARTS_IN_FLIGHT);
            let task_id = format!("shared-batch-{shared_batch_id:?}-{shard_id}");
            BatchState {
                desc: desc.clone(),
                tx,
                handle: mz_ore::task::spawn(|| task_id, async move {
                    let mut builder = None;
                    while let Some(cmd) = rx.recv().await {
                        match cmd {
                            BatchCommand::Push(part) => {
                                let builder = builder.get_or_insert_with(|| {
                                    client.batch_builder(
                                        shard_id,
                                        &shard_name,
                                        schemas.clone(),
                                        lower.clone(),
                                        None,
                                    )
                                });
                                builder.add_part(part).await.expect("valid timestamps");
                            }
                        }
                    }
                    if let Some(builder) = builder {
                        Some(builder.finish(upper).await.expect("valid upper bound"))
                    } else {
                        None
                    }
                })
                .abort_on_drop(),
            }
        });

        // A caller with a different description would otherwise write its updates under another
        // caller's bounds.
        assert_eq!(
            state.desc, desc,
            "shared batch {shared_batch_id:?} requested with conflicting descriptions",
        );

        SharedBatchBuilder {
            batch_id: shared_batch_id,
            shard_id,
            state,
        }
    }
}

/// How many pushed parts a shared batch buffers before [SharedBatchBuilder::push] waits for the
/// builder task to take them.
const PARTS_IN_FLIGHT: usize = 4;

/// The bounds a shared batch is built under, fixed by the caller that creates it.
#[derive(Debug, Clone, PartialEq, Eq)]
struct BatchDesc {
    shard_id: ShardId,
    lower: Antichain<Timestamp>,
    upper: Antichain<Timestamp>,
}

enum BatchCommand {
    Push(Part),
}

#[derive(Debug)]
struct BatchState {
    desc: BatchDesc,
    tx: tokio::sync::mpsc::Sender<BatchCommand>,
    /// Aborts the builder task when the last handle goes away without finishing, so a torn-down
    /// sink stops writing parts.
    handle: AbortOnDropHandle<Option<Batch<SourceData, (), Timestamp, StorageDiff>>>,
}

/// A handle for a shared batch builder. Everyone with a handle for the same builder can
/// [Self::push] to that builder, but the last handle to call [Self::finish] will receive a batch
/// containing all the data.
///
/// Dropping the last handle instead aborts the builder task. Parts it already uploaded stay in blob
/// until persist's leaked-blob cleanup collects them.
pub struct SharedBatchBuilder {
    batch_id: SharedBatchId,
    shard_id: ShardId,
    state: Arc<BatchState>,
}

impl Debug for SharedBatchBuilder {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SharedBatchBuilder")
            .field("batch_id", &self.batch_id)
            .field("shard_id", &self.shard_id)
            .finish_non_exhaustive()
    }
}

impl SharedBatchBuilder {
    /// Include the provided part in the batch. If there are a large number of updates that have
    /// not been flushed to S3, this call may wait.
    pub async fn push(&self, part: Part) {
        if part.len() == 0 {
            return;
        }
        self.state
            .tx
            .send(BatchCommand::Push(part))
            .await
            .expect("shared batch builder task exited");
    }

    /// Fetch the results of the batch builder from the shared state, if any.
    ///
    /// Only the last builder to call this method will obtain the resulting batch, and that batch
    /// will include all data written by all workers. Every handle must call this, also when it
    /// pushed nothing, since any handle may be the last one.
    pub async fn finish(self) -> Option<Batch<SourceData, (), Timestamp, StorageDiff>> {
        let state = Arc::into_inner(self.state)?;
        drop(state.tx); // The task only finishes the batch once the channel is closed.
        state.handle.await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use mz_persist_types::codec_impls::UnitSchema;
    use mz_persist_types::part::PartBuilder;
    use mz_repr::{Datum, RelationDesc, Row, SqlScalarType};

    fn bool_schemas() -> Schemas<SourceData, ()> {
        Schemas {
            id: None,
            key: Arc::new(
                RelationDesc::builder()
                    .with_column("test", SqlScalarType::Bool.nullable(true))
                    .finish(),
            ),
            val: Arc::new(UnitSchema),
        }
    }

    fn bool_part(schemas: &Schemas<SourceData, ()>) -> Part {
        let mut builder = PartBuilder::new(&*schemas.key, &*schemas.val);
        builder.push(
            &SourceData(Ok(Row::pack_slice(&[Datum::True]))),
            &(),
            Timestamp::new(0),
            1i64,
        );
        builder.finish()
    }

    #[mz_ore::test(tokio::test)]
    async fn test_shared_batch() {
        let client = PersistClient::new_for_tests().await;
        let shared = SharedBatches::new();
        let batch_id = SharedBatchId::new();
        let shard_id = ShardId::new();
        let schemas = bool_schemas();
        let lower = Antichain::from_elem(0.into());
        let upper = Antichain::from_elem(1.into());
        let first = shared.builder(
            batch_id,
            client.clone(),
            shard_id,
            "test".to_string(),
            schemas.clone(),
            lower.clone(),
            upper.clone(),
        );
        let second = shared.builder(
            batch_id,
            client,
            shard_id,
            "test".to_string(),
            schemas.clone(),
            lower.clone(),
            upper.clone(),
        );
        let part = bool_part(&schemas);
        first.push(part.clone()).await;
        second.push(part.clone()).await;
        assert!(
            second.finish().await.is_none(),
            "a handle that is not the last to finish receives no batch"
        );
        let batch = first.finish().await.unwrap();
        assert_eq!(
            batch.shard_id(),
            shard_id,
            "batch should be for the expected shard"
        );
        assert_eq!(
            batch.into_hollow_batch().len,
            2,
            "batch should include updates from both pushes"
        )
    }

    fn test_builder(
        shared: &SharedBatches,
        client: &PersistClient,
        batch_id: SharedBatchId,
        shard_id: ShardId,
    ) -> SharedBatchBuilder {
        shared.builder(
            batch_id,
            client.clone(),
            shard_id,
            "test".to_string(),
            bool_schemas(),
            Antichain::from_elem(0.into()),
            Antichain::from_elem(1.into()),
        )
    }

    #[mz_ore::test(tokio::test)]
    #[should_panic(expected = "requested with conflicting descriptions")]
    async fn test_shared_batch_conflicting_descriptions() {
        let client = PersistClient::new_for_tests().await;
        let shared = SharedBatches::new();
        let batch_id = SharedBatchId::new();
        let shard_id = ShardId::new();
        let _first = test_builder(&shared, &client, batch_id, shard_id);
        let _second = shared.builder(
            batch_id,
            client,
            shard_id,
            "test".to_string(),
            bool_schemas(),
            Antichain::from_elem(0.into()),
            Antichain::from_elem(2.into()),
        );
    }

    // A worker that goes away after pushing, without finishing, leaves its data to whichever
    // handle finishes last.
    #[mz_ore::test(tokio::test)]
    async fn test_shared_batch_dropped_handle() {
        let client = PersistClient::new_for_tests().await;
        let shared = SharedBatches::new();
        let batch_id = SharedBatchId::new();
        let shard_id = ShardId::new();
        let schemas = bool_schemas();
        let first = test_builder(&shared, &client, batch_id, shard_id);
        let second = test_builder(&shared, &client, batch_id, shard_id);
        first.push(bool_part(&schemas)).await;
        second.push(bool_part(&schemas)).await;
        drop(second);
        let batch = first
            .finish()
            .await
            .expect("last finisher receives the batch");
        assert_eq!(
            batch.into_hollow_batch().len,
            2,
            "batch should include the dropped handle's updates"
        );
    }

    // When every handle goes away without finishing, the batch is abandoned: it is delivered to
    // nobody, and a later builder for the same id starts a fresh batch without its data.
    #[mz_ore::test(tokio::test)]
    async fn test_shared_batch_abandoned() {
        let client = PersistClient::new_for_tests().await;
        let shared = SharedBatches::new();
        let batch_id = SharedBatchId::new();
        let shard_id = ShardId::new();
        let schemas = bool_schemas();
        let first = test_builder(&shared, &client, batch_id, shard_id);
        let second = test_builder(&shared, &client, batch_id, shard_id);
        first.push(bool_part(&schemas)).await;
        second.push(bool_part(&schemas)).await;
        drop(first);
        drop(second);
        let third = test_builder(&shared, &client, batch_id, shard_id);
        third.push(bool_part(&schemas)).await;
        let batch = third
            .finish()
            .await
            .expect("sole handle receives the batch");
        assert_eq!(
            batch.into_hollow_batch().len,
            1,
            "batch should exclude the abandoned batch's updates"
        );
    }

    // Sweeping dead entries while handles for a live batch finish concurrently must not lose the
    // batch: exactly one of the finishing handles receives it.
    #[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 4))]
    #[cfg_attr(miri, ignore)] // too slow
    async fn test_shared_batch_sweep_races_finish() {
        let client = PersistClient::new_for_tests().await;
        let shared = SharedBatches::new();
        let shard_id = ShardId::new();
        let schemas = bool_schemas();
        let lower = Antichain::from_elem(0.into());
        let upper = Antichain::from_elem(1.into());
        let builder = {
            let schemas = schemas.clone();
            move |batch_id| {
                shared.builder(
                    batch_id,
                    client.clone(),
                    shard_id,
                    "test".to_string(),
                    schemas.clone(),
                    lower.clone(),
                    upper.clone(),
                )
            }
        };

        // Each call registers an entry whose handle drops at once, so the map keeps doubling and
        // nearly every call sweeps.
        let churn: Vec<_> = (0..3)
            .map(|_| {
                let builder = builder.clone();
                mz_ore::task::spawn(|| "churn", async move {
                    loop {
                        drop(builder(SharedBatchId::new()));
                        tokio::task::yield_now().await;
                    }
                })
                .abort_on_drop()
            })
            .collect();

        for i in 0..5000 {
            let batch_id = SharedBatchId::new();
            let first = builder(batch_id);
            let second = builder(batch_id);
            first.push(bool_part(&schemas)).await;
            // When both handles push, a sweep that evicted the live entry between the two
            // `builder` calls would split the updates across two batches.
            let pushes = if i % 2 == 0 {
                second.push(bool_part(&schemas)).await;
                2
            } else {
                1
            };
            let first = mz_ore::task::spawn(|| "finish-first", first.finish());
            let second = mz_ore::task::spawn(|| "finish-second", second.finish());
            let batches: Vec<_> = [first.await, second.await].into_iter().flatten().collect();
            assert_eq!(
                batches.len(),
                1,
                "exactly one finishing handle receives the batch"
            );
            for batch in batches {
                assert_eq!(
                    batch.into_hollow_batch().len,
                    pushes,
                    "the batch holds every handle's updates"
                );
            }
        }
        drop(churn);
    }
}