moraine 0.6.3

//! The rendezvous concurrent commits meet at, so several become one batch.
//!
//! A store admits one batch at a time. A commit that finds the store free
//! opens a batch and stages onto it; a commit that finds one forming joins
//! it; a commit that finds one in flight waits for the next. A batch seals
//! the moment no caller is on its way into it; nothing waits on a timer.

use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};

use futures::FutureExt;
use slatedb::{Db, DbTransaction, IsolationLevel};
use tokio::sync::{Mutex, MutexGuard, watch};

use super::{
    CommitDurability, HeadTransition, HeadViewUpdate, Landed, Prepared, StagedWrite, commit_batch,
    fold, head_view_for, prepare_and_stage,
};
use crate::{
    catalog::{CatalogSnapshot, SnapshotId, projection::ProjectionCache},
    error::{Error, Result},
    store::StagedBytes,
    transaction::{operations::ChangeSet, verbs::Transaction},
};

/// The most commits one batch carries; a full batch seals even while
/// callers are still arriving.
const MAX_BATCH_MEMBERS: usize = 64;

/// What a batch did, told to every member that rode it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum Outcome {
    /// Committed. Every member's allocated ids stand.
    Committed,
    /// Lost the head race. Every member re-runs, classifying its own
    /// change set against the commits that won.
    LostRace,
    /// Abandoned before submission; every member may re-attempt.
    Abandoned,
    /// Submitted, but durability was not acknowledged. Never re-attempt.
    Unknown(String),
    /// The writer lost its epoch and cannot submit further work.
    Fenced(String),
}

/// Submission distinguishes an abandoned batch from a lost acknowledgement.
#[derive(Clone)]
pub(crate) enum BatchState {
    Forming,
    Submitted,
    Finished(Outcome),
}

/// What one caller staged onto a batch.
pub(crate) struct Staged {
    /// One snapshot id per member of that caller, in member order.
    pub(crate) ids: Vec<SnapshotId>,
    /// One change set per member that staged writes, for conflict
    /// classification if the batch loses its race.
    pub(crate) ours: Vec<ChangeSet>,
    /// The head the batch's premise was read at.
    pub(crate) head_before: u64,
    /// Whether this caller put anything in the batch.
    pub(crate) contributed: bool,
    /// The batch's outcome, once it lands.
    pub(crate) outcome: watch::Receiver<BatchState>,
}

/// A batch being formed: one open transaction carrying every member's
/// writes, and the premise the next member stages against.
struct Batch {
    db_tx: DbTransaction,
    /// The head view the batch opened at.
    base: Arc<CatalogSnapshot>,
    /// `base` folded forward through the members staged so far; `None`
    /// until a second member needs it.
    premise: Option<CatalogSnapshot>,
    /// How much of `writes` the premise already reflects.
    folded: usize,
    head_before: u64,
    /// The id the batch will leave at the head pointer: its last
    /// snapshot-minting member's.
    head: u64,
    /// Commits staged onto the batch, against [`MAX_BATCH_MEMBERS`].
    members: usize,
    writes: Vec<StagedWrite>,
    /// What the members have staged onto `db_tx`, index entries included.
    staged_bytes: StagedBytes,
    outcome: watch::Sender<BatchState>,
}

impl Batch {
    async fn open(db: &Db, projections: &std::sync::RwLock<ProjectionCache>) -> Result<Self> {
        let db_tx = db
            .begin(IsolationLevel::Snapshot)
            .await
            .map_err(Error::from)?;
        let base = match head_view_for(&db_tx, projections).await {
            Ok(base) => base,
            Err(err) => {
                db_tx.rollback();
                return Err(err);
            }
        };
        let head_before = base.snapshot.snapshot_id;

        Ok(Self {
            db_tx,
            base,
            premise: None,
            folded: 0,
            head_before,
            head: head_before,
            members: 0,
            writes: Vec::new(),
            staged_bytes: StagedBytes::default(),
            outcome: watch::Sender::new(BatchState::Forming),
        })
    }

    /// Brings the premise up to date with everything staged since it was
    /// last folded.
    fn refold(&mut self) -> Result<()> {
        if self.folded == self.writes.len() {
            return Ok(());
        }
        let mut premise = self.premise.take().unwrap_or_else(|| (*self.base).clone());
        fold::fold_batch(&mut premise, &self.writes[self.folded..])?;
        self.folded = self.writes.len();
        self.premise = Some(premise);
        Ok(())
    }

    /// Runs every member's closure in turn and stages the result. An error
    /// leaves the batch poisoned: the caller must discard it.
    async fn stage<F>(
        &mut self,
        members: &[F],
        projections: &std::sync::RwLock<ProjectionCache>,
    ) -> Result<(Vec<SnapshotId>, Vec<ChangeSet>)>
    where
        F: Fn(&mut Transaction) -> Result<()>,
    {
        let mut ids = Vec::with_capacity(members.len());
        let mut ours = Vec::new();

        for member in members {
            self.members += 1;
            self.refold()?;
            let premise = self.premise.as_ref().unwrap_or(&self.base);
            let prepared =
                prepare_and_stage(&self.db_tx, projections, member, premise, self.members == 1)
                    .await?;
            match prepared {
                Prepared::Nothing { head } => ids.push(SnapshotId::new(head)),
                Prepared::Staged {
                    ours: theirs,
                    commits,
                    writes,
                    staged_bytes,
                } => {
                    ids.push(SnapshotId::new(commits));
                    ours.push(*theirs);
                    self.head = commits;
                    self.writes.extend(writes);
                    self.staged_bytes.0 = self.staged_bytes.0.saturating_add(staged_bytes.0);
                }
            }
        }

        Ok((ids, ours))
    }
}

/// The forming batch and whether one is in flight, under one lock so a
/// batch leaves `forming` and enters flight atomically.
struct Shared {
    forming: Option<Batch>,
    in_flight: bool,
}

/// Where concurrent commits meet. One per store handle, shared by every
/// clone of the catalog.
pub(crate) struct Coalescer {
    shared: Mutex<Shared>,
    /// Callers that have arrived and not yet staged; a batch seals when
    /// this reaches zero.
    arriving: AtomicUsize,
    /// Bumped under `shared` whenever a batch leaves flight.
    flights: watch::Sender<u64>,
    projections: Arc<std::sync::RwLock<ProjectionCache>>,
    /// How a sealed batch reaches object storage.
    durability: CommitDurability,
}

impl Coalescer {
    pub(crate) fn new(
        projections: Arc<std::sync::RwLock<ProjectionCache>>,
        durability: CommitDurability,
    ) -> Self {
        Self {
            shared: Mutex::new(Shared {
                forming: None,
                in_flight: false,
            }),
            arriving: AtomicUsize::new(0),
            flights: watch::Sender::new(0),
            projections,
            durability,
        }
    }

    /// Registers a caller as on its way into the batch now forming.
    fn arrive(self: &Arc<Self>) -> Arrival {
        self.arriving.fetch_add(1, Ordering::AcqRel);
        Arrival {
            coalescer: Arc::clone(self),
            settled: false,
        }
    }

    /// Waits until no batch is in flight and returns the lock on the
    /// forming one. Callers are admitted one at a time.
    async fn admit(&self) -> MutexGuard<'_, Shared> {
        loop {
            let shared = self.shared.lock().await;
            if !shared.in_flight {
                return shared;
            }
            // Subscribed under the lock: the generation only moves under it,
            // so the flight cannot end unnoticed in between.
            let mut flights = self.flights.subscribe();
            drop(shared);
            let _ = flights.changed().await;
        }
    }

    /// Stages `members` onto whichever batch is forming, opening one if
    /// none is, and seals the batch if this caller is the last one in.
    pub(crate) async fn stage<F>(self: &Arc<Self>, db: &Db, members: &[F]) -> Result<Staged>
    where
        F: Fn(&mut Transaction) -> Result<()>,
    {
        let arrival = self.arrive();
        let mut shared = self.admit().await;

        let mut batch = match shared.forming.take() {
            Some(batch) => batch,
            None => Batch::open(db, &self.projections).await?,
        };
        let staged_before = batch.writes.len();

        let staged = match batch.stage(members, &self.projections).await {
            Ok(staged) => staged,
            Err(err) => {
                // The batch is poisoned; its members learn that nothing
                // landed when the outcome sender drops.
                batch.db_tx.rollback();
                return Err(err);
            }
        };
        let (ids, ours) = staged;

        let outcome = batch.outcome.subscribe();
        let head_before = batch.head_before;
        let contributed = batch.writes.len() > staged_before;

        let alone = arrival.settle() == 0;
        let full = batch.members >= MAX_BATCH_MEMBERS;
        if (alone || full) && !batch.writes.is_empty() {
            shared.in_flight = true;
            drop(shared);
            self.launch(batch);
        } else if batch.writes.is_empty() {
            batch.db_tx.rollback();
        } else {
            shared.forming = Some(batch);
        }

        Ok(Staged {
            ids,
            ours,
            head_before,
            contributed,
            outcome,
        })
    }

    /// Commits `batch` on a task of its own, so the caller that sealed it
    /// cannot take the batch down with it.
    fn launch(self: &Arc<Self>, batch: Batch) {
        let coalescer = Arc::clone(self);
        drop(tokio::spawn(async move { coalescer.land(batch).await }));
    }

    /// Commits one batch, tells its members, and reopens the store.
    async fn land(self: Arc<Self>, batch: Batch) {
        let Batch {
            db_tx,
            base,
            head_before,
            head,
            writes,
            staged_bytes,
            outcome,
            ..
        } = batch;

        outcome.send_replace(BatchState::Submitted);
        let landed = match std::panic::AssertUnwindSafe(commit_batch(
            db_tx,
            HeadTransition {
                before: head_before,
                after: head,
            },
            &writes,
            staged_bytes,
            HeadViewUpdate::Rebuild(base),
            &self.projections,
            &self.durability,
        ))
        .catch_unwind()
        .await
        {
            Ok(Ok(Landed::Committed(_))) => Outcome::Committed,
            Ok(Ok(Landed::LostRace)) => Outcome::LostRace,
            Ok(Err(Error::Fenced(reason))) => Outcome::Fenced(reason),
            Ok(Err(err)) => {
                tracing::warn!(error = %err, "commit batch outcome is unknown; its members must not retry");
                Outcome::Unknown(err.to_string())
            }
            Err(_) => {
                super::invalidate_head_view(&self.projections);
                Outcome::Unknown("the submitted batch panicked before acknowledgement".into())
            }
        };
        // Every member may already have gone; the batch landed regardless.
        let _ = outcome.send(BatchState::Finished(landed));

        let mut shared = self.shared.lock().await;
        shared.in_flight = false;
        self.flights
            .send_modify(|flight| *flight = flight.wrapping_add(1));
    }

    /// Seals whatever batch is forming, if nothing else will: the path a
    /// caller that vanished before staging leaves behind.
    async fn seal_abandoned(self: Arc<Self>) {
        let mut shared = self.shared.lock().await;
        if shared.in_flight || self.arriving.load(Ordering::Acquire) > 0 {
            return;
        }
        let Some(batch) = shared.forming.take() else {
            return;
        };
        if batch.writes.is_empty() {
            batch.db_tx.rollback();
            return;
        }
        shared.in_flight = true;
        drop(shared);
        self.land(batch).await;
    }
}

/// Waits for acknowledgement, retaining the submission boundary if the task
/// exits.
pub(crate) async fn await_outcome(mut outcome: watch::Receiver<BatchState>) -> Outcome {
    loop {
        if let BatchState::Finished(landed) = &*outcome.borrow_and_update() {
            return landed.clone();
        }
        if outcome.changed().await.is_err() {
            return match &*outcome.borrow() {
                BatchState::Forming => Outcome::Abandoned,
                BatchState::Submitted => {
                    Outcome::Unknown("the submitted batch stopped reporting".into())
                }
                BatchState::Finished(landed) => landed.clone(),
            };
        }
    }
}

/// Counts one caller from the moment it arrives until it has staged.
/// Dropping it without settling both drops the count and seals what is
/// left, so a vanished caller cannot leave the batch waiting.
struct Arrival {
    coalescer: Arc<Coalescer>,
    settled: bool,
}

impl Arrival {
    /// Marks this caller as staged, reporting how many are still on their
    /// way in.
    fn settle(mut self) -> usize {
        self.settled = true;
        self.coalescer.arriving.fetch_sub(1, Ordering::AcqRel) - 1
    }
}

impl Drop for Arrival {
    fn drop(&mut self) {
        if self.settled {
            return;
        }
        if self.coalescer.arriving.fetch_sub(1, Ordering::AcqRel) != 1 {
            return;
        }
        // Outside a runtime, the next commit finds the batch and seals it.
        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
            let coalescer = Arc::clone(&self.coalescer);
            drop(runtime.spawn(async move { coalescer.seal_abandoned().await }));
        }
    }
}

#[cfg(test)]
mod tests {
    use object_store::memory::InMemory;

    use super::*;
    use crate::{Catalog, CatalogOptions};

    #[tokio::test]
    async fn a_lost_sender_preserves_the_submission_boundary() {
        for (state, expected) in [
            (BatchState::Forming, Outcome::Abandoned),
            (
                BatchState::Submitted,
                Outcome::Unknown("the submitted batch stopped reporting".into()),
            ),
            (BatchState::Finished(Outcome::Committed), Outcome::Committed),
        ] {
            let (sender, receiver) = watch::channel(state);
            drop(sender);
            assert_eq!(await_outcome(receiver).await, expected);
        }
    }

    #[tokio::test]
    async fn a_write_error_is_not_reported_as_safe_to_retry() {
        let catalog = Catalog::open(Arc::new(InMemory::new()), CatalogOptions::default())
            .await
            .unwrap();
        let store = catalog.store();
        let db = store.writer_db().unwrap();
        let mut batch = Batch::open(db, catalog.projections()).await.unwrap();
        batch
            .stage(
                &[|tx: &mut Transaction| {
                    tx.create_schema("pending")?;
                    Ok(())
                }],
                catalog.projections(),
            )
            .await
            .unwrap();
        let outcome = batch.outcome.subscribe();
        catalog.close().await.unwrap();
        let coalescer = Arc::new(Coalescer::new(
            Arc::clone(catalog.projections()),
            CommitDurability::OnFlushInterval,
        ));
        coalescer.land(batch).await;
        assert!(matches!(await_outcome(outcome).await, Outcome::Unknown(_)));
    }
}