delta_kernel 0.28.0

//! Cooperative cancellation for long-running Kernel reads.
//!
//! Kernel never does I/O itself and owns no async runtime, so cancellation is *cooperative* and
//! *runtime-agnostic*: a caller supplies a [`CancellationToken`] (via
//! [`ScanBuilder::with_cancellation_token`](crate::scan::ScanBuilder::with_cancellation_token)),
//! Kernel polls it at action-batch boundaries, and cancellation-aware [`Engine`](crate::Engine)
//! reads may race their I/O against it. Cancellation is always surfaced as
//! [`Error::Cancelled`] -- never as normal iterator exhaustion -- so a partial listing can never be
//! mistaken for a complete one.

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use crate::{AsAny, DeltaResult, Error};

/// A shared, thread-safe cancellation token. Held as an `Arc` because the lazy scan iterator and
/// the engine reads it drives can outlive the builder call and run on other threads.
pub type CancellationTokenRef = Arc<dyn CancellationToken>;

/// A future that resolves when a [`CancellationToken`] is cancelled. Runtime-neutral: it is a
/// plain boxed [`Future`], so an engine can `select!` it against its own async reads without
/// Kernel taking on any async-runtime dependency.
pub type CancelledFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;

/// Returns `Err(Error::Cancelled)` if `token` is present and already cancelled, else `Ok(())`.
///
/// Used to fail fast before starting a setup/read operation (e.g. a footer read or a sidecar
/// listing) so cancelled work is not begun.
pub(crate) fn check_cancelled(token: Option<&CancellationTokenRef>) -> DeltaResult<()> {
    match token {
        Some(t) if t.is_cancelled() => Err(Error::Cancelled),
        _ => Ok(()),
    }
}

/// A cooperative cancellation signal supplied by a caller.
///
/// Implementors wrap whatever their runtime provides (e.g. `tokio_util::sync::CancellationToken`).
/// Kernel and cancellation-aware engines only *consume* it: Kernel polls [`is_cancelled`] between
/// action batches, while an async engine may await [`cancelled_future`] to wake blocked I/O.
///
/// # Recovering the underlying token
///
/// The `Arc` kernel hands to each `*_with_cancellation` [`Engine`] method is the same one the
/// caller supplied, so an engine that supplied its own implementation can downcast it back to the
/// concrete type through [`AsAny`]. (Kernel may poll a composed or derived token internally; the
/// guarantee is only about what the engine receives.) This lets an engine reach a native
/// cancellation handle it wrapped, for code that cannot accept a Rust trait object.
///
/// Borrow with `as_ref().any_ref()`, or take an owned handle with [`AsAny::as_any`] -- see
/// [`AsAny::any_ref`] for why the borrow must go through `as_ref()` first.
///
/// ```
/// # use delta_kernel::cancellation::{CancellationToken, CancellationTokenRef, CancelledFuture};
/// # use delta_kernel::AsAny;
/// # use std::sync::Arc;
/// # struct MyToken;
/// # impl CancellationToken for MyToken {
/// #     fn is_cancelled(&self) -> bool { false }
/// #     fn cancelled_future(&self) -> CancelledFuture<'_> { Box::pin(std::future::pending()) }
/// # }
/// # let token: CancellationTokenRef = Arc::new(MyToken);
/// // In an engine's `read_*_with_cancellation`, given `token: CancellationTokenRef`:
/// if let Some(mine) = token.as_ref().any_ref().downcast_ref::<MyToken>() {
///     let _ = mine.is_cancelled(); // recovered the caller's concrete token
/// }
/// ```
///
/// [`is_cancelled`]: CancellationToken::is_cancelled
/// [`cancelled_future`]: CancellationToken::cancelled_future
/// [`Engine`]: crate::Engine
/// [`AsAny`]: crate::AsAny
pub trait CancellationToken: AsAny {
    /// Returns `true` once cancellation has been requested. Cheap, synchronous, and monotonic:
    /// once it returns `true` it must never return `false` again.
    fn is_cancelled(&self) -> bool;

    /// Returns a future that resolves when the token is cancelled (immediately if it already is).
    ///
    /// There is no default implementation: a correct notification cannot be synthesized from
    /// [`is_cancelled`](Self::is_cancelled) alone without either busy-polling or a runtime, and
    /// Kernel has neither. Implementors back this with their own notification primitive.
    fn cancelled_future(&self) -> CancelledFuture<'_>;
}

/// Wraps a fallible iterator so that cancellation terminates it with a single
/// [`Error::Cancelled`] rather than silent truncation.
///
/// Before each pull, the token is polled: if cancelled, one `Err(Error::Cancelled)` is yielded
/// and every subsequent call returns `None` (the iterator is fused). An `Err(Error::Cancelled)`
/// arriving from the inner iterator (e.g. a cancellation-aware engine interrupting a read) fuses
/// it the same way, so a token shared with the engine still yields exactly one terminal error.
/// With no token, or before cancellation, items pass through unchanged. This is deliberately
/// **not** `take_while`, which would end the iterator with `None` and make a cancelled listing
/// look complete.
pub(crate) struct CancellableIterator<I> {
    inner: I,
    token: Option<CancellationTokenRef>,
    /// Set once cancellation has been observed and the terminal error emitted, fusing the
    /// iterator to `None` thereafter.
    done: bool,
}

impl<I> CancellableIterator<I> {
    pub(crate) fn new(inner: I, token: Option<CancellationTokenRef>) -> Self {
        Self {
            inner,
            token,
            done: false,
        }
    }
}

impl<I, T> Iterator for CancellableIterator<I>
where
    I: Iterator<Item = DeltaResult<T>>,
{
    type Item = DeltaResult<T>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }
        if self.token.as_ref().is_some_and(|t| t.is_cancelled()) {
            self.done = true;
            return Some(Err(Error::Cancelled));
        }
        let item = self.inner.next();
        // A cancellation-aware engine can itself surface `Err(Cancelled)` from an interrupted
        // read. Fuse on it so the composed pipeline still yields exactly one terminal error
        // rather than this layer re-injecting a second one on the next poll.
        if matches!(item, Some(Err(Error::Cancelled))) {
            self.done = true;
        }
        item
    }
}

#[cfg(test)]
mod tests {
    use std::future::ready;
    use std::sync::atomic::{AtomicBool, Ordering};

    use super::*;

    /// Minimal [`CancellationToken`] backed by an [`AtomicBool`], for tests.
    #[derive(Default)]
    struct TestToken(AtomicBool);

    impl TestToken {
        fn cancel(&self) {
            self.0.store(true, Ordering::SeqCst);
        }
    }

    impl CancellationToken for TestToken {
        fn is_cancelled(&self) -> bool {
            self.0.load(Ordering::SeqCst)
        }
        fn cancelled_future(&self) -> CancelledFuture<'_> {
            // Tests only drive `is_cancelled`; a resolved/pending future is enough here.
            Box::pin(ready(()))
        }
    }

    fn ok_iter(n: usize) -> impl Iterator<Item = DeltaResult<usize>> {
        (0..n).map(Ok)
    }

    #[test]
    fn no_token_passes_through_unchanged() {
        let out: Vec<_> = CancellableIterator::new(ok_iter(3), None)
            .map(Result::unwrap)
            .collect();
        assert_eq!(out, vec![0, 1, 2]);
    }

    #[test]
    fn uncancelled_token_passes_through_unchanged() {
        let token: CancellationTokenRef = Arc::new(TestToken::default());
        let out: Vec<_> = CancellableIterator::new(ok_iter(3), Some(token))
            .map(Result::unwrap)
            .collect();
        assert_eq!(out, vec![0, 1, 2]);
    }

    #[test]
    fn pre_cancelled_yields_one_error_then_ends() {
        let token = Arc::new(TestToken::default());
        token.cancel();
        let mut iter = CancellableIterator::new(ok_iter(3), Some(token as CancellationTokenRef));
        assert!(matches!(iter.next(), Some(Err(Error::Cancelled))));
        // Fused: never a `Some(Ok(..))` after cancellation, and no infinite error stream.
        assert!(iter.next().is_none());
        assert!(iter.next().is_none());
    }

    #[test]
    fn mid_stream_cancellation_yields_error_not_silent_truncation() {
        let token = Arc::new(TestToken::default());
        let ct: CancellationTokenRef = token.clone();
        let mut iter = CancellableIterator::new(ok_iter(5), Some(ct));
        assert!(matches!(iter.next(), Some(Ok(0))));
        assert!(matches!(iter.next(), Some(Ok(1))));
        token.cancel();
        // The terminal item is an error, so a cancelled listing can't look complete (which a
        // bare `None` / `take_while` would).
        assert!(matches!(iter.next(), Some(Err(Error::Cancelled))));
        assert!(iter.next().is_none());
    }

    // An `Err(Cancelled)` from the inner iterator (as a cancellation-aware engine emits) must
    // fuse this layer, so a token shared between engine and kernel yields exactly ONE terminal
    // error, not two. Regression guard for the double-emit the layered pipeline would otherwise
    // produce. The token is left uncancelled so the fuse comes solely from the inner error.
    #[test]
    fn inner_cancelled_error_fuses_without_double_emit() {
        let token: CancellationTokenRef = Arc::new(TestToken::default());
        let inner = vec![Ok(0), Err(Error::Cancelled), Ok(99)].into_iter();
        let mut iter = CancellableIterator::new(inner, Some(token));
        assert!(matches!(iter.next(), Some(Ok(0))));
        assert!(matches!(iter.next(), Some(Err(Error::Cancelled))));
        // Fused on the inner error: the trailing Ok is never yielded, and no second error.
        assert!(iter.next().is_none());
    }

    #[test]
    fn check_cancelled_reports_state() {
        let token = Arc::new(TestToken::default());
        let ct: CancellationTokenRef = token.clone();
        assert!(check_cancelled(Some(&ct)).is_ok());
        assert!(check_cancelled(None).is_ok());
        token.cancel();
        assert!(matches!(check_cancelled(Some(&ct)), Err(Error::Cancelled)));
    }

    /// A second token type, to check that a downcast discriminates rather than always succeeding.
    #[derive(Default)]
    struct OtherToken;

    impl CancellationToken for OtherToken {
        fn is_cancelled(&self) -> bool {
            false
        }
        fn cancelled_future(&self) -> CancelledFuture<'_> {
            Box::pin(ready(()))
        }
    }

    // Downcasting an erased token recovers the original value, not a copy: cancelling through the
    // recovered handle is observable through the erased one.
    #[test]
    fn downcast_recovers_the_same_token() {
        let erased: CancellationTokenRef = Arc::new(TestToken::default());

        let recovered = erased
            .clone()
            .as_any()
            .downcast::<TestToken>()
            .expect("erased token should downcast to its concrete type");
        recovered.cancel();

        assert!(erased.is_cancelled());
    }

    #[test]
    fn downcast_to_the_wrong_type_fails() {
        let erased: CancellationTokenRef = Arc::new(TestToken::default());
        assert!(erased.clone().as_any().downcast::<OtherToken>().is_err());
        assert!(erased
            .as_ref()
            .any_ref()
            .downcast_ref::<OtherToken>()
            .is_none());
    }

    // `Arc<dyn CancellationToken>` satisfies the blanket `AsAny` impl in its own right, so
    // `arc.any_ref()` resolves to the *`Arc`* rather than the token inside it and downcasts to the
    // concrete type fail. Borrowing goes through the trait object: `arc.as_ref().any_ref()`.
    #[test]
    fn any_ref_borrows_the_token_through_the_trait_object() {
        let token = Arc::new(TestToken::default());
        let erased: CancellationTokenRef = token.clone();

        assert!(erased.any_ref().downcast_ref::<TestToken>().is_none());

        let borrowed = erased
            .as_ref()
            .any_ref()
            .downcast_ref::<TestToken>()
            .expect("erased token should downcast to its concrete type");
        assert!(!borrowed.is_cancelled());
        token.cancel();
        assert!(borrowed.is_cancelled());
    }
}