polyc-query 2026.9.0

//! Releases exact query results through a consumer-polled bounded stream.

use std::collections::BTreeSet;
use std::error::Error;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;

use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
#[cfg(test)]
use datafusion::catalog::TableProvider;
use datafusion::execution::memory_pool::MemoryReservation;
use datafusion::physical_plan::SendableRecordBatchStream;
use futures::{Stream, StreamExt};
use polyc_state::journal::GetJournalSource;
use polyc_state::query_audit::{ErrorClass, QueryOutcome};
use tokio::sync::{OwnedSemaphorePermit, mpsc};
use tokio_util::sync::CancellationToken;

use super::latch::{BatchRelease, DeliveredCounts, TerminalLatch};
use super::{
    CoreExecutionError, CoreMetadataAuthority, CoreOperationContext, CurrentCredentialAuthority,
    EffectiveCoreBounds, PermitGuardian, QueryScope, classify_error, operation_refusal,
    operation_wait,
};
#[cfg(test)]
use crate::core_resolution::CoreTable;

/// Owns a physical plan whose permit, exact pins, and limits cannot change.
pub(crate) struct BoundCoreQuery {
    pub(super) guardian: PermitGuardian,
    pub(super) dataframe: datafusion::dataframe::DataFrame,
    pub(super) manifests: Vec<polyc_state::projection::ProjectionManifest>,
    pub(super) original_scope: QueryScope,
    pub(super) metadata: Arc<dyn CoreMetadataAuthority>,
    pub(super) scope_revalidator: Arc<dyn CurrentCredentialAuthority>,
    pub(super) operation: Arc<CoreOperationContext>,
    pub(super) cancellation: CancellationToken,
    pub(super) bounds: EffectiveCoreBounds,
    pub(super) revalidation_interval: Duration,
    pub(super) _explain_enabled: bool,
    pub(super) execution_admission: OwnedSemaphorePermit,
    pub(super) source_decode_reservation: MemoryReservation,
    /// Deschedules the producer's terminal report by this much.
    ///
    /// Zero in production. A test uses it to hold the producer away from the
    /// guardian after rows were already delivered, which is the exact race a
    /// scheduling grace used to resolve wrongly.
    #[cfg(test)]
    pub(super) report_delay: Duration,
    #[cfg(test)]
    pub(super) providers: std::collections::BTreeMap<CoreTable, Arc<dyn TableProvider>>,
}

impl BoundCoreQuery {
    #[cfg(test)]
    pub(super) const fn delay_report_for_test(&mut self, delay: Duration) {
        self.report_delay = delay;
    }

    #[cfg(test)]
    pub(super) fn provider(&self, table: CoreTable) -> Arc<dyn TableProvider> {
        Arc::clone(
            self.providers
                .get(&table)
                .expect("bound dependency provider"),
        )
    }

    /// Starts the producer task and hands back the consumer-polled stream.
    ///
    /// Nothing is awaited here. The producer waits for the consumer's first
    /// readiness signal. No batch is produced before a poll asks for one.
    pub(crate) fn execute(self) -> CoreResultStream {
        // Captured before the plan is consumed. A caller must be able to
        // announce the result schema before the first batch exists, and the
        // durable audit's own source vector is what the terminal reports.
        let released_schema = self.dataframe.schema().inner().clone();
        let released_source = self.guardian.source().clone();
        let started = tokio::time::Instant::now();
        let Self {
            guardian,
            dataframe,
            manifests,
            original_scope,
            metadata,
            scope_revalidator,
            operation,
            cancellation,
            bounds,
            revalidation_interval,
            _explain_enabled,
            execution_admission,
            source_decode_reservation,
            #[cfg(test)]
            report_delay,
            #[cfg(test)]
                providers: _,
        } = self;
        // The latch is the one place the producer's report, the consumer's
        // withdrawal, and the delivered counts meet. The guardian reads it, so
        // an exact row count never depends on which task ran last.
        let latch = guardian.latch();
        let consumer_latch = Arc::clone(&latch);
        let consumer_cancellation = cancellation.clone();
        let terminal_cancellation = cancellation.clone();
        let (sender, receiver) = mpsc::channel(1);
        let (readiness_signal, readiness_requests) = mpsc::unbounded_channel();
        tokio::spawn(async move {
            let stream = operation_wait(&operation, &cancellation, dataframe.execute_stream())
                .await
                .and_then(|result| result.map_err(CoreExecutionError::from));
            let end = match stream {
                Ok(upstream) => {
                    let state = CoreStreamState {
                        upstream,
                        execution_admission,
                        _source_decode_reservation: source_decode_reservation,
                        revalidation: CoreRevalidationWitness {
                            manifests,
                            original_scope,
                            metadata,
                            scope_revalidator,
                            operation: Arc::clone(&operation),
                        },
                        operation,
                        cancellation,
                        bounds,
                        revalidation_interval,
                        last_revalidation: None,
                        released_rows: 0,
                        released_bytes: 0,
                        known_extra_row: false,
                        latch: Arc::clone(&latch),
                    };
                    release_from_producer(state, &sender, readiness_requests).await
                }
                Err(error) => ProducerEnd::Failed(error),
            };
            // A consumer that withdrew after the last row still withdrew.
            let end = match end {
                ProducerEnd::Ended
                    if terminal_cancellation.is_cancelled() || sender.is_closed() =>
                {
                    ProducerEnd::Cancelled
                }
                end => end,
            };
            #[cfg(test)]
            if !report_delay.is_zero() {
                tokio::time::sleep(report_delay).await;
            }
            // The outcome only. The counts belong to the terminal, and the
            // terminal is built inside the latch's own critical section, so
            // nothing can be released between reading them and recording it.
            let settlement = guardian.finish(producer_outcome(&end)).await;
            let frame = match (settlement, end) {
                (Err(error), _) | (Ok(()), ProducerEnd::Failed(error)) => {
                    CoreStreamFrame::Failure(error)
                }
                (Ok(()), ProducerEnd::Cancelled) => {
                    CoreStreamFrame::Failure(CoreExecutionError::Cancelled)
                }
                (Ok(()), ProducerEnd::Ended) => CoreStreamFrame::Complete,
            };
            let _ = sender.send(frame).await;
        });
        CoreResultStream {
            receiver,
            readiness_signal,
            poll_outstanding: false,
            terminal_seen: false,
            cancellation: consumer_cancellation,
            latch: consumer_latch,
            schema: released_schema,
            source: released_source,
            started,
        }
    }
}

/// How the producer's own loop ended.
enum ProducerEnd {
    /// The upstream stream ended after every produced row was delivered.
    Ended,
    /// The consumer withdrew.
    Cancelled,
    /// The producer measured this exact failure.
    Failed(CoreExecutionError),
}

struct CoreStreamState {
    upstream: SendableRecordBatchStream,
    execution_admission: OwnedSemaphorePermit,
    _source_decode_reservation: MemoryReservation,
    revalidation: CoreRevalidationWitness,
    operation: Arc<CoreOperationContext>,
    cancellation: CancellationToken,
    bounds: EffectiveCoreBounds,
    revalidation_interval: Duration,
    last_revalidation: Option<tokio::time::Instant>,
    released_rows: u64,
    released_bytes: u64,
    known_extra_row: bool,
    latch: Arc<TerminalLatch>,
}

async fn release_from_producer(
    mut state: CoreStreamState,
    sender: &mpsc::Sender<CoreStreamFrame>,
    mut readiness: mpsc::UnboundedReceiver<()>,
) -> ProducerEnd {
    loop {
        // The original deadline also bounds the wait for a readiness token. A
        // consumer that stops polling without dropping would otherwise pin the
        // permit, the admission slot, and the memory reservation forever, and
        // leave the durable intent unmatched past its declared deadline.
        let Ok(remaining) = state.operation.remaining() else {
            state.cancellation.cancel();
            return ProducerEnd::Failed(CoreExecutionError::Deadline);
        };
        tokio::select! {
            ready = readiness.recv() => {
                if ready.is_none() {
                    state.cancellation.cancel();
                    return ProducerEnd::Cancelled;
                }
            }
            () = state.cancellation.cancelled() => return ProducerEnd::Cancelled,
            () = tokio::time::sleep(remaining) => {
                state.cancellation.cancel();
                return ProducerEnd::Failed(CoreExecutionError::Deadline);
            }
        }
        match state.next().await {
            Ok(Some((batch, next))) => {
                state = next;
                if sender.send(CoreStreamFrame::Batch(batch)).await.is_err() {
                    state.cancellation.cancel();
                    return ProducerEnd::Cancelled;
                }
            }
            Ok(None) => return ProducerEnd::Ended,
            Err(error) => return ProducerEnd::Failed(error),
        }
    }
}

impl CoreStreamState {
    async fn next(mut self) -> Result<Option<(RecordBatch, Self)>, CoreExecutionError> {
        loop {
            if self.cancellation.is_cancelled() {
                return Err(CoreExecutionError::Cancelled);
            }
            if self.known_extra_row {
                return Ok(None);
            }
            let next =
                match operation_wait(&self.operation, &self.cancellation, self.upstream.next())
                    .await
                {
                    Ok(next) => next,
                    Err(error) => return Err(error),
                };
            let Some(batch) = next else {
                return Ok(None);
            };
            let batch = match batch {
                Ok(batch) => batch,
                Err(error) => return Err(stream_error(error)),
            };
            if batch.num_rows() == 0 {
                continue;
            }
            if self.released_rows >= self.bounds.rows() {
                self.latch.record_truncation();
                return Ok(None);
            }
            let remaining_rows = self.bounds.rows() - self.released_rows;
            let release_rows = usize::try_from(remaining_rows)
                .unwrap_or(usize::MAX)
                .min(batch.num_rows());
            let release = batch.slice(0, release_rows);
            if release_rows < batch.num_rows() {
                self.latch.record_truncation();
                self.known_extra_row = true;
            }
            let bytes = u64::try_from(release.get_array_memory_size()).unwrap_or(u64::MAX);
            let total = match self.released_bytes.checked_add(bytes) {
                Some(total) if total <= self.bounds.result_release_bytes() => total,
                Some(total) => {
                    return Err(CoreExecutionError::ReleaseBound {
                        observed: total,
                        limit: self.bounds.result_release_bytes(),
                    });
                }
                None => {
                    return Err(CoreExecutionError::ReleaseBound {
                        observed: u64::MAX,
                        limit: self.bounds.result_release_bytes(),
                    });
                }
            };
            if self
                .last_revalidation
                .is_none_or(|last| last.elapsed() >= self.revalidation_interval)
            {
                self.revalidation.revalidate(&self.cancellation).await?;
                self.last_revalidation = Some(tokio::time::Instant::now());
            }
            self.released_rows += u64::try_from(release_rows).unwrap_or(u64::MAX);
            self.released_bytes = total;
            return Ok(Some((release, self)));
        }
    }
}

fn stream_error(error: datafusion::error::DataFusionError) -> CoreExecutionError {
    let mut source: Option<&(dyn Error + 'static)> = Some(&error);
    while let Some(current) = source {
        if matches!(
            current.downcast_ref::<CoreExecutionError>(),
            Some(CoreExecutionError::Deadline)
        ) {
            return CoreExecutionError::Deadline;
        }
        if matches!(
            current.downcast_ref::<CoreExecutionError>(),
            Some(CoreExecutionError::Cancelled)
        ) {
            return CoreExecutionError::Cancelled;
        }
        source = current.source();
    }
    CoreExecutionError::DataFusion(error)
}

struct CoreRevalidationWitness {
    manifests: Vec<polyc_state::projection::ProjectionManifest>,
    original_scope: QueryScope,
    metadata: Arc<dyn CoreMetadataAuthority>,
    scope_revalidator: Arc<dyn CurrentCredentialAuthority>,
    operation: Arc<CoreOperationContext>,
}

impl CoreRevalidationWitness {
    async fn revalidate(&self, cancellation: &CancellationToken) -> Result<(), CoreExecutionError> {
        operation_wait(&self.operation, cancellation, self.revalidate_inner()).await?
    }

    async fn revalidate_inner(&self) -> Result<(), CoreExecutionError> {
        let current = self
            .scope_revalidator
            .current_scope(&self.operation)
            .await?;
        if !scope_contains(&current, &self.original_scope) {
            return Err(CoreExecutionError::AuthorityNarrowed);
        }
        for manifest in &self.manifests {
            self.operation.check().map_err(operation_refusal)?;
            let expected = manifest.checkpoint().source();
            let observed = self
                .metadata
                .source_head(
                    &self.operation,
                    GetJournalSource::new(expected.partition().clone()),
                )
                .await?
                .ok_or_else(|| CoreExecutionError::SourceChanged(expected.partition().clone()))?;
            if observed.source() != expected {
                return Err(CoreExecutionError::SourceChanged(
                    expected.partition().clone(),
                ));
            }
        }
        Ok(())
    }
}

fn scope_contains(current: &QueryScope, original: &QueryScope) -> bool {
    match (current, original) {
        (QueryScope::Fleet, _) => true,
        (QueryScope::Conversations(_), QueryScope::Fleet) => false,
        (QueryScope::Conversations(current), QueryScope::Conversations(original)) => {
            let current = current.iter().collect::<BTreeSet<_>>();
            original
                .iter()
                .all(|conversation| current.contains(conversation))
        }
    }
}

/// Returns the outcome the producer measured.
///
/// The outcome only. The rows and truncation that go with it belong to the
/// latch, which builds the completion from the counts it holds, in the same
/// critical section that records the terminal.
fn producer_outcome(end: &ProducerEnd) -> QueryOutcome {
    match end {
        ProducerEnd::Ended => QueryOutcome::Succeeded,
        ProducerEnd::Cancelled => QueryOutcome::Failed(ErrorClass::Cancelled),
        ProducerEnd::Failed(error) => QueryOutcome::Failed(classify_error(error)),
    }
}

/// Carries the private bounded Arrow stream.
///
/// End-of-stream becomes visible only after the producer durably settles its
/// exact completion. Dropping this value signals cancellation and nothing
/// more. The producer still owns the one-shot permit. It settles under its own
/// terminal budget.
pub(crate) struct CoreResultStream {
    receiver: mpsc::Receiver<CoreStreamFrame>,
    readiness_signal: mpsc::UnboundedSender<()>,
    poll_outstanding: bool,
    terminal_seen: bool,
    cancellation: CancellationToken,
    latch: Arc<TerminalLatch>,
    /// The plan's own output schema, captured before the plan was consumed.
    schema: SchemaRef,
    /// The durable audit's source vector for this execution.
    source: polyc_state::query_audit::SourceSnapshot,
    started: tokio::time::Instant,
}

impl CoreResultStream {
    /// Returns the result schema, available before the first batch.
    pub(crate) const fn schema(&self) -> &SchemaRef {
        &self.schema
    }

    /// Returns the exact source vector this execution read.
    pub(crate) const fn source(&self) -> &polyc_state::query_audit::SourceSnapshot {
        &self.source
    }

    /// Returns what the consumer has actually received so far.
    ///
    /// Read from the latch, not from a producer-side tally. What "received"
    /// means follows the live accounting: batches that left this stream, or —
    /// once a consumer re-frames them — the frames that left that consumer.
    pub(crate) fn delivered(&self) -> DeliveredCounts {
        self.latch.delivered()
    }

    /// Returns how long this execution has run.
    pub(crate) fn elapsed(&self) -> Duration {
        self.started.elapsed()
    }

    /// Hands the delivery count to a consumer that re-frames each batch.
    ///
    /// After this, a batch leaving the stream counts nothing. The consumer
    /// admits what it releases, through [`Self::admit_release`], so the
    /// durable record counts rows a caller actually received rather than rows
    /// this stream handed to a framer.
    pub(crate) fn account_at_consumer(&self) {
        self.latch.account_at_consumer();
    }

    /// Admits one release into the record this query will settle.
    ///
    /// Returns whether it was admitted. A refusal means a terminal has
    /// already been chosen, so the frame must not reach the caller: the
    /// record it would belong to is already fixed without it.
    pub(crate) fn admit_release(&self, rows: u64, bytes: u64) -> bool {
        self.latch.admit_release(rows, bytes)
    }

    /// Settles the terminal the consumer measured at its own admitted bound.
    ///
    /// The consumer stopped because it reached a ceiling the caller asked
    /// for. It then drops this stream to stop the producer, and that drop
    /// must not turn a bounded success into a withdrawal. The stop is always
    /// a truncation: release ended because the next frame did not fit.
    ///
    /// Returns the counts the record carries, or `None` when another party
    /// had already ended the query.
    pub(crate) fn settle_consumer_bound(&self) -> Option<DeliveredCounts> {
        self.latch
            .settle_consumer_bound(&self.source, self.started.elapsed())
    }

    /// Records a failure the consumer measured after the batch left this
    /// stream.
    ///
    /// A consumer that cannot use a delivered batch ends the query, and
    /// dropping this stream would then record a withdrawal. That is the wrong
    /// word: the caller withdrew nothing. The latch keeps a measured failure
    /// over a later cancellation, so reporting here is what makes the durable
    /// terminal say what actually happened.
    /// The record's counts are the admitted releases, which is what the
    /// caller received: a batch that left this stream but could not be framed
    /// reached nobody, and was never admitted.
    ///
    /// Returns whether this failure became the selected terminal. A caller
    /// must not publish its local terminal when another terminal won first.
    pub(crate) fn report_failure(&self, class: ErrorClass) -> bool {
        self.latch.report(
            QueryOutcome::Failed(class),
            self.started.elapsed(),
            &self.source,
        )
    }

    /// Requests one producer batch without dequeuing it.
    ///
    /// A case uses this to hold a batch in front of a producer terminal. It
    /// can then prove that a later framing failure does not replace the
    /// terminal that already became the durable record.
    #[cfg(test)]
    pub(crate) fn request_buffered_batch(&mut self) {
        if !self.poll_outstanding {
            let _ = self.readiness_signal.send(());
            self.poll_outstanding = true;
        }
    }

    /// Returns how many producer frames wait in front of this consumer.
    #[cfg(test)]
    pub(crate) fn buffered_frames(&self) -> usize {
        self.receiver.len()
    }

    /// Returns whether a terminal has won the latch.
    #[cfg(test)]
    pub(crate) fn terminal_selected(&self) -> bool {
        self.latch.is_settled()
    }
}

enum CoreStreamFrame {
    Batch(RecordBatch),
    Failure(CoreExecutionError),
    Complete,
}

impl Stream for CoreResultStream {
    type Item = Result<RecordBatch, CoreExecutionError>;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            if !self.poll_outstanding {
                // The producer may close readiness while it is durably
                // settling a terminal result. The result channel, not this
                // hint, owns EOF.
                let _ = self.readiness_signal.send(());
                self.poll_outstanding = true;
            }
            let frame = self.receiver.poll_recv(context);
            return match frame {
                Poll::Ready(Some(CoreStreamFrame::Batch(batch))) => {
                    self.poll_outstanding = false;
                    // Offered at the exact point the batch leaves this stream,
                    // through the same admission a frame goes through.
                    let rows = u64::try_from(batch.num_rows()).unwrap_or(u64::MAX);
                    let bytes = u64::try_from(batch.get_array_memory_size()).unwrap_or(u64::MAX);
                    match self.latch.record_batch_delivery(rows, bytes) {
                        // A terminal was already chosen and this batch is not
                        // in the record. Handing it over would give the caller
                        // rows the record does not carry, so it is dropped and
                        // the producer's own terminal is taken instead. The
                        // producer sends nothing after its terminal, so this
                        // skips at most what the channel already buffered.
                        BatchRelease::Refused => continue,
                        BatchRelease::Admitted | BatchRelease::CountedByConsumer => {
                            Poll::Ready(Some(Ok(batch)))
                        }
                    }
                }
                Poll::Ready(Some(CoreStreamFrame::Failure(error))) => {
                    self.poll_outstanding = false;
                    self.terminal_seen = true;
                    Poll::Ready(Some(Err(error)))
                }
                Poll::Ready(Some(CoreStreamFrame::Complete)) => {
                    self.poll_outstanding = false;
                    self.terminal_seen = true;
                    Poll::Ready(None)
                }
                Poll::Ready(None) if self.terminal_seen => Poll::Ready(None),
                Poll::Ready(None) => {
                    self.poll_outstanding = false;
                    self.terminal_seen = true;
                    Poll::Ready(Some(Err(CoreExecutionError::AuditCompletionUnavailable)))
                }
                Poll::Pending => Poll::Pending,
            };
        }
    }
}

impl Drop for CoreResultStream {
    fn drop(&mut self) {
        // Withdraw through the latch first, so the transition is recorded
        // before the producer can observe the cancellation token and report a
        // success the consumer will never see.
        self.latch.cancel();
        self.cancellation.cancel();
    }
}