nodedb 0.4.0

// SPDX-License-Identifier: BUSL-1.1

//! Drain coordination: `begin_drain`, `wait_until_drained`, `clear_drain`.

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use super::refcount::CollectionQuiesce;

/// Exclusive per-name lifecycle hold used by no-Raft DDL.
///
/// Dropping the guard releases one drain hold and wakes CREATE waiters. Call
/// [`disarm`](LifecycleDrainGuard::disarm) after ownership is transferred to a
/// durable pending-reclaim record.
pub struct LifecycleDrainGuard {
    registry: Arc<CollectionQuiesce>,
    database_id: u64,
    tenant_id: u64,
    collection: String,
    active: bool,
}

impl LifecycleDrainGuard {
    pub fn disarm(mut self) {
        self.active = false;
    }
}

impl Drop for LifecycleDrainGuard {
    fn drop(&mut self) {
        if self.active {
            self.registry
                .clear_drain(self.database_id, self.tenant_id, &self.collection);
        }
    }
}

impl CollectionQuiesce {
    /// Acquire one lifecycle drain hold for `(database, tenant, collection)`.
    /// New scans and same-name CREATE operations remain blocked until every
    /// matching hold is released by `clear_drain` or `forget`.
    pub fn begin_drain(&self, database_id: u64, tenant_id: u64, collection: &str) {
        let mut inner = self.inner_mut();
        let entry = inner
            .states
            .entry((database_id, tenant_id, collection.to_string()))
            .or_default();
        entry.drain_holders = entry.drain_holders.saturating_add(1);
    }

    /// Stop the drain marker, allowing new scans again. Only called when
    /// the purge is aborted or a recreate happens — on a normal purge
    /// the collection metadata is gone so new scans naturally return
    /// `collection_not_found` from that point on, and the drain entry
    /// is garbage-collected via `forget`.
    pub fn clear_drain(&self, database_id: u64, tenant_id: u64, collection: &str) {
        let mut inner = self.inner_mut();
        let key = (database_id, tenant_id, collection.to_string());
        let remove = if let Some(state) = inner.states.get_mut(&key) {
            state.drain_holders = state.drain_holders.saturating_sub(1);
            state.drain_holders == 0 && state.open_scans == 0
        } else {
            false
        };
        if remove {
            inner.states.remove(&key);
        }
        drop(inner);
        self.notify.notify_waiters();
    }

    /// Drop the entry entirely once reclaim has completed. After this,
    /// `is_draining` returns false and `open_scans` is 0. Called by
    /// the purge handler right before emitting the reclaim ack.
    pub fn forget(&self, database_id: u64, tenant_id: u64, collection: &str) {
        let mut inner = self.inner_mut();
        let key = (database_id, tenant_id, collection.to_string());
        let remove = if let Some(state) = inner.states.get_mut(&key) {
            state.drain_holders = state.drain_holders.saturating_sub(1);
            state.drain_holders == 0
        } else {
            false
        };
        if remove {
            inner.states.remove(&key);
        }
        drop(inner);
        self.notify.notify_waiters();
    }

    /// Try to acquire the exclusive lifecycle drain without waiting.
    /// Synchronous DDL uses this form because it may run on a current-thread
    /// Tokio runtime where blocking on an async waiter would panic.
    pub fn try_acquire_lifecycle(
        self: &Arc<Self>,
        database_id: u64,
        tenant_id: u64,
        collection: &str,
    ) -> Option<LifecycleDrainGuard> {
        let mut inner = self.inner_mut();
        let entry = inner
            .states
            .entry((database_id, tenant_id, collection.to_string()))
            .or_default();
        if entry.drain_holders > 0 {
            return None;
        }
        entry.drain_holders = 1;
        Some(LifecycleDrainGuard {
            registry: Arc::clone(self),
            database_id,
            tenant_id,
            collection: collection.to_string(),
            active: true,
        })
    }

    /// Exclusively acquire the lifecycle drain for one collection name.
    ///
    /// The check-and-acquire is performed under the registry mutex, so two
    /// concurrent local DDL operations cannot both enter the destructive
    /// lifecycle section.
    pub async fn acquire_lifecycle(
        self: &Arc<Self>,
        database_id: u64,
        tenant_id: u64,
        collection: &str,
    ) -> LifecycleDrainGuard {
        loop {
            let notified = self.notify.notified();
            tokio::pin!(notified);
            notified.as_mut().enable();
            {
                let mut inner = self.inner_mut();
                let entry = inner
                    .states
                    .entry((database_id, tenant_id, collection.to_string()))
                    .or_default();
                if entry.drain_holders == 0 {
                    entry.drain_holders = 1;
                    return LifecycleDrainGuard {
                        registry: Arc::clone(self),
                        database_id,
                        tenant_id,
                        collection: collection.to_string(),
                        active: true,
                    };
                }
            }
            notified.await;
        }
    }

    /// Returns a future that resolves once every open scan against
    /// `(tenant_id, collection)` has completed. Safe to await from the
    /// Control Plane (tokio) — internally uses [`tokio::sync::Notify`]
    /// for wake-up; no polling.
    ///
    /// `begin_drain` must be called before awaiting this future, or
    /// new scans could continue to bump the counter and the future
    /// would never resolve.
    pub fn wait_until_drained(
        self: &Arc<Self>,
        database_id: u64,
        tenant_id: u64,
        collection: &str,
    ) -> WaitDrain {
        WaitDrain {
            registry: Arc::clone(self),
            database_id,
            tenant_id,
            collection: collection.to_string(),
            notified: None,
        }
    }

    fn inner_mut(&self) -> std::sync::MutexGuard<'_, super::refcount::Inner> {
        self.inner.lock().expect("CollectionQuiesce mutex poisoned")
    }
}

/// Future returned by [`CollectionQuiesce::wait_until_drained`].
///
/// Completes when the `(tenant, collection)` open-scan count reaches 0.
/// Implementation detail: each poll takes a fresh `Notify::notified()`
/// future so we don't race against a notification that fires between
/// check and await.
pub struct WaitDrain {
    registry: Arc<CollectionQuiesce>,
    database_id: u64,
    tenant_id: u64,
    collection: String,
    notified: Option<Pin<Box<tokio::sync::futures::Notified<'static>>>>,
}

// Safety: Notified borrows from the Notify inside `registry` (Arc).
// We transmute the lifetime to `'static` because we own the Arc for the
// future's lifetime, guaranteeing the Notify outlives the Notified.
unsafe impl Send for WaitDrain {}

impl Future for WaitDrain {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        loop {
            if self
                .registry
                .open_scans(self.database_id, self.tenant_id, &self.collection)
                == 0
            {
                return Poll::Ready(());
            }
            // Arm a notification, then re-check. If a release fires
            // between the check and the arm we handled it on the next
            // iteration (open_scans would be 0 then).
            if self.notified.is_none() {
                let notify: &tokio::sync::Notify = &self.registry.notify;
                // SAFETY: we hold an Arc<CollectionQuiesce>; the Notify
                // inside it outlives `self`.
                let notified: tokio::sync::futures::Notified<'_> = notify.notified();
                let notified: tokio::sync::futures::Notified<'static> =
                    unsafe { std::mem::transmute(notified) };
                self.notified = Some(Box::pin(notified));
            }
            let fut = self.notified.as_mut().expect("just set");
            match fut.as_mut().poll(cx) {
                Poll::Ready(()) => {
                    self.notified = None;
                    // Loop: re-check open_scans.
                    continue;
                }
                Poll::Pending => {
                    // Re-check once before sleeping to close the race
                    // between arming the notified future and a release
                    // that just happened.
                    if self
                        .registry
                        .open_scans(self.database_id, self.tenant_id, &self.collection)
                        == 0
                    {
                        return Poll::Ready(());
                    }
                    return Poll::Pending;
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const DB: u64 = 0;

    #[tokio::test]
    async fn drain_resolves_immediately_when_no_open_scans() {
        let q = CollectionQuiesce::new();
        q.begin_drain(DB, 1, "c");
        q.wait_until_drained(DB, 1, "c").await;
    }

    #[tokio::test]
    async fn drain_waits_for_last_scan_to_release() {
        let q = CollectionQuiesce::new();
        let g1 = q.try_start_scan(DB, 1, "c").unwrap();
        let g2 = q.try_start_scan(DB, 1, "c").unwrap();
        q.begin_drain(DB, 1, "c");

        let q_clone = Arc::clone(&q);
        let drain_task = tokio::spawn(async move {
            q_clone.wait_until_drained(DB, 1, "c").await;
        });

        // Briefly yield so the drain task parks.
        tokio::task::yield_now().await;
        assert!(
            !drain_task.is_finished(),
            "drain must not resolve while scans open"
        );

        drop(g1);
        tokio::task::yield_now().await;
        assert!(
            !drain_task.is_finished(),
            "drain must not resolve with 1 scan still open"
        );

        drop(g2);
        drain_task.await.unwrap();
    }

    #[tokio::test]
    async fn single_lifecycle_holder_clears_on_forget() {
        let q = CollectionQuiesce::new();
        q.begin_drain(DB, 1, "c");
        assert!(q.is_draining(DB, 1, "c"));

        q.forget(DB, 1, "c");
        assert!(!q.is_draining(DB, 1, "c"));
    }

    #[tokio::test]
    async fn lifecycle_acquisition_is_exclusive() {
        let q = CollectionQuiesce::new();
        let first = q.acquire_lifecycle(DB, 1, "c").await;

        let q_clone = Arc::clone(&q);
        let second = tokio::spawn(async move { q_clone.acquire_lifecycle(DB, 1, "c").await });
        tokio::task::yield_now().await;
        assert!(!second.is_finished());

        drop(first);
        let second = second.await.unwrap();
        drop(second);
        assert!(!q.is_draining(DB, 1, "c"));
    }

    #[tokio::test]
    async fn is_draining_until_every_lifecycle_holder_forgets() {
        let q = CollectionQuiesce::new();
        q.begin_drain(DB, 1, "c");
        q.begin_drain(DB, 1, "c");

        // One holder released — still draining while the other holds.
        q.forget(DB, 1, "c");
        assert!(q.is_draining(DB, 1, "c"));

        // Last holder released — drain clears.
        q.forget(DB, 1, "c");
        assert!(!q.is_draining(DB, 1, "c"));
    }

    #[tokio::test]
    async fn forget_clears_state() {
        let q = CollectionQuiesce::new();
        q.begin_drain(DB, 1, "c");
        q.wait_until_drained(DB, 1, "c").await;
        q.forget(DB, 1, "c");
        assert!(!q.is_draining(DB, 1, "c"));
        assert!(q.try_start_scan(DB, 1, "c").is_ok());
    }
}