mz_persist_client/internal/

watch.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.

//! Notifications for state changes.

use differential_dataflow::lattice::Lattice;
use mz_persist::location::SeqNo;
use std::fmt::{Debug, Formatter};
use std::ops::Deref;
use std::sync::{Arc, RwLock};
use timely::PartialOrder;
use timely::progress::{Antichain, Timestamp};
use tokio::sync::{Notify, broadcast, watch};
use tracing::debug;

use crate::cache::LockingTypedState;
use crate::internal::metrics::Metrics;

#[derive(Debug)]
pub struct StateWatchNotifier<T> {
    metrics: Arc<Metrics>,
    /// Fires on every strict seqno advance. Used by waiters that care about any
    /// state change (e.g. lease/seqno machinery via [StateWatch::wait_for_seqno_ge]).
    seqno_tx: broadcast::Sender<SeqNo>,
    /// Holds the current shard upper (write frontier), advanced only when the upper
    /// strictly moves forward. Used by [StateWatch::wait_for_upper_past] so upper
    /// waiters (Listen/Subscribe/snapshot) don't wake on GC, rollups, or since-
    /// downgrades that bump the seqno without moving the upper. The watch channel
    /// coalesces: a waiter always reads the latest upper, never a stale intermediate.
    upper_tx: watch::Sender<Antichain<T>>,
}

impl<T: Timestamp + Lattice> StateWatchNotifier<T> {
    pub(crate) fn new(metrics: Arc<Metrics>, initial_upper: Antichain<T>) -> Self {
        let (seqno_tx, _rx) = broadcast::channel(1);
        let (upper_tx, _rx) = watch::channel(initial_upper);
        StateWatchNotifier {
            metrics,
            seqno_tx,
            upper_tx,
        }
    }

    /// Wake up any watchers of this state.
    ///
    /// `seqno` is the state's seqno after the modification, and `upper` is the
    /// shard upper after the modification.
    ///
    /// This must be called while under the same lock that modified the state to
    /// avoid any potential for out of order SeqNos in the seqno broadcast channel,
    /// and so the watched upper never regresses.
    ///
    /// This restriction can be lifted (i.e. we could notify after releasing the
    /// write lock), but we'd have to reason about out of order SeqNos in the
    /// broadcast channel. In particular, if we see `RecvError::Lagged` then
    /// it's possible we lost X+1 and got X, so if X isn't sufficient to return,
    /// we'd need to grab the read lock and verify the real SeqNo.
    pub(crate) fn notify(&self, seqno: SeqNo, upper: &Antichain<T>) {
        match self.seqno_tx.send(seqno) {
            // Someone got woken up.
            Ok(_) => {
                self.metrics.watch.notify_sent.inc();
            }
            // No one is listening, that's also fine.
            Err(_) => {
                self.metrics.watch.notify_noop.inc();
            }
        }
        // Advance the watched upper only on a strict advance. send_if_modified does
        // not require live receivers and never errors, so an advance is recorded
        // even when no one waits; a later subscriber reads the up-to-date value.
        let upper_advanced = self.upper_tx.send_if_modified(|current| {
            if PartialOrder::less_than(current, upper) {
                current.clone_from(upper);
                true
            } else {
                false
            }
        });
        if upper_advanced {
            self.metrics.watch.notify_upper_sent.inc();
        }
    }
}

/// A reactive subscription to changes in [LockingTypedState].
///
/// This exposes two independent signals over the same state:
/// - [Self::wait_for_seqno_ge] wakes on any strict seqno advance.
/// - [Self::wait_for_upper_past] wakes only when the shard upper advances.
///
/// Invariants:
/// - The `state.seqno` only advances (never regresses). This is guaranteed by
///   LockingTypedState.
/// - `seqno_high_water` is always <= `state.seqno`.
/// - If `seqno_high_water` is < `state.seqno`, then we'll get a notification on
///   `seqno_rx`. This is maintained by notifying new seqnos under the same lock
///   which adds them.
/// - `seqno_high_water` always holds the highest value received on `seqno_rx`.
///   This is maintained by the wait method taking an exclusive reference to self.
/// - `upper_rx` holds the shard upper, advanced under the same lock that advances
///   the upper. It coalesces, so a waiter always observes the latest upper.
#[derive(Debug)]
pub struct StateWatch<K, V, T, D> {
    metrics: Arc<Metrics>,
    state: Arc<LockingTypedState<K, V, T, D>>,
    seqno_high_water: SeqNo,
    seqno_rx: broadcast::Receiver<SeqNo>,
    upper_rx: watch::Receiver<Antichain<T>>,
}

impl<K, V, T, D> StateWatch<K, V, T, D> {
    pub(crate) fn new(state: Arc<LockingTypedState<K, V, T, D>>, metrics: Arc<Metrics>) -> Self {
        // Important! We have to subscribe to the seqno broadcast channel _before_
        // we grab the current seqno. Otherwise, we could race with a write to
        // state and miss a notification. Tokio guarantees that "the returned
        // Receiver will receive values sent after the call to subscribe", and
        // the read_lock linearizes the subscribe to be _before_ whatever
        // seqno we get here.
        //
        // The upper watch needs no such care: it retains the latest upper, and
        // `subscribe` hands the receiver that value, so a waiter reads the current
        // upper directly and parks only for the next advance.
        let seqno_rx = state.notifier().seqno_tx.subscribe();
        let upper_rx = state.notifier().upper_tx.subscribe();
        let seqno_high_water = state.read_lock(&metrics.locks.watch, |x| x.seqno);
        StateWatch {
            metrics,
            state,
            seqno_high_water,
            seqno_rx,
            upper_rx,
        }
    }

    /// Blocks until the State has a SeqNo >= the requested one.
    ///
    /// This method is cancel-safe.
    pub async fn wait_for_seqno_ge(&mut self, requested: SeqNo) -> &mut Self {
        self.metrics.watch.notify_wait_started.inc();
        debug!("wait_for_seqno_ge {} {}", self.state.shard_id(), requested);
        loop {
            if self.seqno_high_water >= requested {
                break;
            }
            match self.seqno_rx.recv().await {
                Ok(x) => {
                    self.metrics.watch.notify_recv.inc();
                    assert!(x >= self.seqno_high_water);
                    self.seqno_high_water = x;
                }
                Err(broadcast::error::RecvError::Closed) => {
                    unreachable!("we're holding on to a reference to the sender")
                }
                Err(broadcast::error::RecvError::Lagged(_)) => {
                    self.metrics.watch.notify_lagged.inc();
                    // This is just a hint that our buffer (of size 1) filled
                    // up, which is totally fine. The broadcast channel
                    // guarantees that the most recent N (again, =1 here) are
                    // kept, so just loop around. This branch means we should be
                    // able to read a new value immediately.
                    continue;
                }
            }
        }
        self.metrics.watch.notify_wait_finished.inc();
        debug!(
            "wait_for_seqno_ge {} {} returning",
            self.state.shard_id(),
            requested
        );
        self
    }

    /// Blocks until the shard upper has advanced past `frontier`.
    ///
    /// Unlike [Self::wait_for_seqno_ge], this only wakes on upper advances, not on
    /// seqno bumps caused by GC, rollups, since-downgrades, or no-op CaAs.
    ///
    /// This method is cancel-safe.
    pub async fn wait_for_upper_past(&mut self, frontier: &Antichain<T>) -> &mut Self
    where
        T: Timestamp,
    {
        self.metrics.watch.notify_wait_started.inc();
        debug!(
            "wait_for_upper_past {} {:?}",
            self.state.shard_id(),
            frontier.elements()
        );
        loop {
            // Bind the comparison so the borrow guard drops before the await.
            let past = PartialOrder::less_than(frontier, &*self.upper_rx.borrow_and_update());
            if past {
                break;
            }
            // `changed` only errors once every sender has dropped. We hold the
            // state Arc that owns the notifier's sender, so it outlives us.
            self.upper_rx
                .changed()
                .await
                .expect("notifier sender outlives the watch");
            self.metrics.watch.notify_recv.inc();
        }
        self.metrics.watch.notify_wait_finished.inc();
        debug!(
            "wait_for_upper_past {} {:?} returning",
            self.state.shard_id(),
            frontier.elements()
        );
        self
    }
}

/// A concurrent state - one which allows reading, writing, and waiting for changes made by
/// another concurrent writer.
///
/// This is morally similar to a mutex with a condvar, but allowing asynchronous waits and with
/// access methods that make it a little trickier to accidentally hold a lock across a yield point.
pub(crate) struct AwaitableState<T> {
    state: Arc<RwLock<T>>,
    /// NB: we can't wrap the [Notify] in the lock since the signature of [Notify::notified]
    /// doesn't allow it, but this is only accessed while holding the lock.
    notify: Arc<Notify>,
}

impl<T: Debug> Debug for AwaitableState<T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        self.state.read().fmt(f)
    }
}

impl<T> Clone for AwaitableState<T> {
    fn clone(&self) -> Self {
        Self {
            state: Arc::clone(&self.state),
            notify: Arc::clone(&self.notify),
        }
    }
}

/// A wrapper around a mutable ref that tracks whether it's ever accessed mutably. See
/// [AwaitableState::maybe_modify] for usage.
pub struct ModifyGuard<'a, T> {
    mut_ref: &'a mut T,
    modified: bool,
}

impl<'a, T> Deref for ModifyGuard<'a, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &*self.mut_ref
    }
}

impl<'a, T> ModifyGuard<'a, T> {
    pub fn get_mut(&mut self) -> &mut T {
        self.modified = true;
        &mut *self.mut_ref
    }
}

impl<T> AwaitableState<T> {
    pub fn new(value: T) -> Self {
        Self {
            state: Arc::new(RwLock::new(value)),
            notify: Arc::new(Notify::new()),
        }
    }

    #[allow(dead_code)]
    pub fn read<A>(&self, read_fn: impl FnOnce(&T) -> A) -> A {
        let guard = self.state.read().expect("not poisoned");
        let state = &*guard;
        read_fn(state)
    }

    /// Conditionally modify the state. This method passes a guard to the provided function,
    /// which only allows mutable access to the data via [ModifyGuard::get_mut]. If that method
    /// is not called, waiters will not be woken up.
    pub fn maybe_modify<A>(&self, write_fn: impl FnOnce(&mut ModifyGuard<T>) -> A) -> A {
        let mut guard = self.state.write().expect("not poisoned");
        let mut state = ModifyGuard {
            mut_ref: &mut *guard,
            modified: false,
        };
        let result = write_fn(&mut state);
        // Notify everyone while holding the guard. This guarantees that all waiters will observe
        // the just-updated state, assuming the state was accessed mutably.
        if state.modified {
            self.notify.notify_waiters();
        }
        drop(guard);
        result
    }

    pub fn modify<A>(&self, write_fn: impl FnOnce(&mut T) -> A) -> A {
        self.maybe_modify(|guard| write_fn(guard.get_mut()))
    }

    pub async fn wait_for<A>(&self, mut wait_fn: impl FnMut(&T) -> Option<A>) -> A {
        loop {
            let notified = {
                let guard = self.state.read().expect("not poisoned");
                let state = &*guard;
                if let Some(result) = wait_fn(state) {
                    return result;
                }
                // Grab the notified future while holding the guard. This ensures that we will see any
                // future modifications to this state, even if they happen before the first poll.
                let notified = self.notify.notified();
                drop(guard);
                notified
            };

            notified.await;
        }
    }

    pub async fn wait_while(&self, mut wait_fn: impl FnMut(&T) -> bool) {
        self.wait_for(|s| (!wait_fn(s)).then_some(())).await
    }
}

#[cfg(test)]
mod tests {
    use std::future::Future;
    use std::pin::Pin;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::task::{Context, Poll, Wake, Waker};
    use std::time::Duration;

    use futures::FutureExt;
    use futures_task::noop_waker;
    use itertools::Itertools;
    use mz_build_info::DUMMY_BUILD_INFO;
    use mz_dyncfg::ConfigUpdates;
    use mz_ore::assert_none;
    use mz_ore::cast::CastFrom;
    use mz_ore::metrics::MetricsRegistry;
    use rand::prelude::SliceRandom;
    use timely::progress::Antichain;
    use tokio::task::JoinSet;

    use crate::cache::StateCache;
    use crate::cfg::PersistConfig;
    use crate::internal::machine::{
        NEXT_LISTEN_BATCH_RETRYER_CLAMP, NEXT_LISTEN_BATCH_RETRYER_INITIAL_BACKOFF,
        NEXT_LISTEN_BATCH_RETRYER_MULTIPLIER,
    };
    use crate::internal::state::TypedState;
    use crate::tests::new_test_client;
    use crate::{Diagnostics, ShardId};

    use super::*;

    #[mz_ore::test(tokio::test)]
    async fn state_watch() {
        mz_ore::test::init_logging();
        let metrics = Arc::new(Metrics::new(
            &PersistConfig::new_for_tests(),
            &MetricsRegistry::new(),
        ));
        let cache = StateCache::new_no_metrics();
        let shard_id = ShardId::new();
        let state = cache
            .get::<(), (), u64, i64, _, _>(
                shard_id,
                || async {
                    Ok(TypedState::new(
                        DUMMY_BUILD_INFO.semver_version(),
                        shard_id,
                        "host".to_owned(),
                        0u64,
                    ))
                },
                &Diagnostics::for_tests(),
            )
            .await
            .unwrap();
        assert_eq!(state.read_lock(&metrics.locks.watch, |x| x.seqno), SeqNo(0));

        // A watch for 0 resolves immediately.
        let mut w0 = StateWatch::new(Arc::clone(&state), Arc::clone(&metrics));
        let _ = w0.wait_for_seqno_ge(SeqNo(0)).await;

        // A watch for 1 does not yet resolve.
        let w0s1 = w0.wait_for_seqno_ge(SeqNo(1)).map(|_| ()).shared();
        assert_eq!(w0s1.clone().now_or_never(), None);

        // After mutating state, the watch for 1 does resolve.
        state.write_lock(&metrics.locks.applier_write, |state| {
            state.seqno = state.seqno.next()
        });
        let () = w0s1.await;

        // A watch for an old seqno immediately resolves.
        let _ = w0.wait_for_seqno_ge(SeqNo(0)).await;

        // We can create a new watch and it also behaves.
        let mut w1 = StateWatch::new(Arc::clone(&state), Arc::clone(&metrics));
        let _ = w1.wait_for_seqno_ge(SeqNo(0)).await;
        let _ = w1.wait_for_seqno_ge(SeqNo(1)).await;
        assert_none!(w1.wait_for_seqno_ge(SeqNo(2)).now_or_never());
    }

    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
    #[cfg_attr(miri, ignore)] // error: unsupported operation: integer-to-pointer casts and `ptr::from_exposed_addr` are not supported with `-Zmiri-strict-provenance`
    async fn state_watch_concurrency() {
        mz_ore::test::init_logging();
        let metrics = Arc::new(Metrics::new(
            &PersistConfig::new_for_tests(),
            &MetricsRegistry::new(),
        ));
        let cache = StateCache::new_no_metrics();
        let shard_id = ShardId::new();
        let state = cache
            .get::<(), (), u64, i64, _, _>(
                shard_id,
                || async {
                    Ok(TypedState::new(
                        DUMMY_BUILD_INFO.semver_version(),
                        shard_id,
                        "host".to_owned(),
                        0u64,
                    ))
                },
                &Diagnostics::for_tests(),
            )
            .await
            .unwrap();
        assert_eq!(state.read_lock(&metrics.locks.watch, |x| x.seqno), SeqNo(0));

        const NUM_WATCHES: usize = 100;
        const NUM_WRITES: usize = 20;

        let watches = (0..NUM_WATCHES)
            .map(|idx| {
                let state = Arc::clone(&state);
                let metrics = Arc::clone(&metrics);
                mz_ore::task::spawn(|| "watch", async move {
                    let mut watch = StateWatch::new(Arc::clone(&state), Arc::clone(&metrics));
                    // We stared at 0, so N writes means N+1 seqnos.
                    let wait_seqno = SeqNo(u64::cast_from(idx % NUM_WRITES + 1));
                    let _ = watch.wait_for_seqno_ge(wait_seqno).await;
                    let observed_seqno =
                        state.read_lock(&metrics.locks.applier_read_noncacheable, |x| x.seqno);
                    assert!(
                        wait_seqno <= observed_seqno,
                        "{} vs {}",
                        wait_seqno,
                        observed_seqno
                    );
                })
            })
            .collect::<Vec<_>>();
        let writes = (0..NUM_WRITES)
            .map(|_| {
                let state = Arc::clone(&state);
                let metrics = Arc::clone(&metrics);
                mz_ore::task::spawn(|| "write", async move {
                    state.write_lock(&metrics.locks.applier_write, |x| {
                        x.seqno = x.seqno.next();
                    });
                })
            })
            .collect::<Vec<_>>();
        for watch in watches {
            watch.await;
        }
        for write in writes {
            write.await;
        }
    }

    /// A [Waker] that counts how many times it has been woken.
    struct CountingWaker(AtomicUsize);

    impl Wake for CountingWaker {
        fn wake(self: Arc<Self>) {
            self.0.fetch_add(1, Ordering::SeqCst);
        }
        fn wake_by_ref(self: &Arc<Self>) {
            self.0.fetch_add(1, Ordering::SeqCst);
        }
    }

    /// A shard upper waiter must not be woken by seqno bumps that don't advance
    /// the upper (GC, rollups, since-downgrades, other writers' no-op CaAs). It
    /// must still be woken when the upper actually advances.
    #[mz_persist_proc::test(tokio::test)]
    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
    async fn wait_for_upper_past_ignores_non_upper_seqno_bumps(dyncfgs: ConfigUpdates) {
        mz_ore::test::init_logging();

        let client = new_test_client(&dyncfgs).await;
        // Disable the listen poll so the only wakeups are watch notifications.
        client.cfg.set_config(
            &NEXT_LISTEN_BATCH_RETRYER_INITIAL_BACKOFF,
            Duration::from_secs(1_000_000),
        );
        client
            .cfg
            .set_config(&NEXT_LISTEN_BATCH_RETRYER_MULTIPLIER, 1);
        client.cfg.set_config(
            &NEXT_LISTEN_BATCH_RETRYER_CLAMP,
            Duration::from_secs(1_000_000),
        );

        let shard_id = ShardId::new();
        let (mut write, mut read) = client.expect_open::<(), (), u64, i64>(shard_id).await;
        // A second writer so we can advance the upper while `write` is borrowed
        // by the in-flight wait future.
        let (mut writer2, _) = client.expect_open::<(), (), u64, i64>(shard_id).await;

        // Move the upper to 10 so there's room to downgrade `since` underneath it.
        write
            .expect_compare_and_append(&[(((), ()), 0, 1)], 0, 10)
            .await;

        // Arm a waiter for an upper past 100. It stays pending.
        let counter = Arc::new(CountingWaker(AtomicUsize::new(0)));
        let waker = Waker::from(Arc::clone(&counter));
        let mut cx = Context::from_waker(&waker);
        let frontier = Antichain::from_elem(100);
        let mut fut = Box::pin(write.wait_for_upper_past(&frontier));
        assert!(matches!(Pin::new(&mut fut).poll(&mut cx), Poll::Pending));
        let wakes_after_arm = counter.0.load(Ordering::SeqCst);

        // Bump the seqno WITHOUT advancing the upper (a since-downgrade). This is
        // the regression: before the fix this woke the upper waiter.
        read.downgrade_since(&Antichain::from_elem(5)).await;
        assert!(matches!(Pin::new(&mut fut).poll(&mut cx), Poll::Pending));
        assert_eq!(
            counter.0.load(Ordering::SeqCst),
            wakes_after_arm,
            "a non-upper seqno bump must not wake an upper waiter"
        );

        // Advancing the upper must wake the waiter and let it resolve.
        writer2
            .expect_compare_and_append(&[(((), ()), 10, 1)], 10, 200)
            .await;
        assert!(
            counter.0.load(Ordering::SeqCst) > wakes_after_arm,
            "an upper advance must wake the upper waiter"
        );
        fut.await;
    }

    #[mz_persist_proc::test(tokio::test)]
    #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
    async fn state_watch_listen_snapshot(dyncfgs: ConfigUpdates) {
        mz_ore::test::init_logging();
        let waker = noop_waker();
        let mut cx = Context::from_waker(&waker);

        let client = new_test_client(&dyncfgs).await;
        // Override the listen poll so that it's useless.
        client.cfg.set_config(
            &NEXT_LISTEN_BATCH_RETRYER_INITIAL_BACKOFF,
            Duration::from_secs(1_000_000),
        );
        client
            .cfg
            .set_config(&NEXT_LISTEN_BATCH_RETRYER_MULTIPLIER, 1);
        client.cfg.set_config(
            &NEXT_LISTEN_BATCH_RETRYER_CLAMP,
            Duration::from_secs(1_000_000),
        );

        let (mut write, mut read) = client.expect_open::<(), (), u64, i64>(ShardId::new()).await;

        // Grab a snapshot for 1, which doesn't resolve yet. Also grab a listen
        // for 0, which resolves but doesn't yet resolve the next batch.
        let mut listen = read
            .clone("test")
            .await
            .listen(Antichain::from_elem(0))
            .await
            .unwrap();
        let mut snapshot = Box::pin(read.snapshot(Antichain::from_elem(0)));
        assert!(Pin::new(&mut snapshot).poll(&mut cx).is_pending());
        let mut listen_next_batch = Box::pin(listen.next(None));
        assert!(Pin::new(&mut listen_next_batch).poll(&mut cx).is_pending());

        // Now update the frontier, which should allow the snapshot to resolve
        // and the listen to resolve its next batch. Because we disabled the
        // polling, the listen_next_batch future will block forever and timeout
        // the test if the watch doesn't work.
        write.expect_compare_and_append(&[], 0, 1).await;
        let _ = listen_next_batch.await;

        // For good measure, also resolve the snapshot, though we haven't broken
        // the polling on this.
        let _ = snapshot.await;
    }

    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
    #[allow(clippy::disallowed_methods)] // For JoinSet.
    #[cfg_attr(miri, ignore)]
    async fn wait_on_awaitable_state() {
        const TASKS: usize = 1000;
        // Launch a bunch of tasks, have them all wait for a specific number, then increment it
        // by one. Lost notifications would cause this test to time out.
        let mut set = JoinSet::new();
        let state = AwaitableState::new(0);
        let mut tasks = (0..TASKS).collect_vec();
        let mut rng = rand::rng();
        tasks.shuffle(&mut rng);
        for i in tasks {
            set.spawn({
                let state = state.clone();
                async move {
                    state.wait_while(|v| *v != i).await;
                    state.modify(|v| *v += 1);
                }
            });
        }
        set.join_all().await;
        assert_eq!(state.read(|i| *i), TASKS);
    }
}