hawdb-executor 0.11.0

// Copyright 2026 Nowledge
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Query-result admission, consumer delivery, and memory-lease lifetime.

use crate::binding::{map_memory_bytes, map_payload_bytes, Binding};
use crate::pipeline::runtime_checkpoint;
use crate::{
    QueryMemoryAccount, QueryMemoryClass, QueryMemoryLease, QueryMemoryLedger, QueryRows,
    QueryRowsBuilder, Row,
};
use hawdb_core::{HawDBError, ReadBudgetResource, Result, RuntimeTaskContext};
use std::num::NonZeroUsize;

#[derive(Clone, Copy)]
pub enum ConsumerMemoryMode {
    /// The consumer keeps emitted rows alive through the final memory snapshot.
    Retained,
    /// Each row can be uncharged as soon as the consumer call returns.
    ReleasedAfterCall,
    /// Rows remain query-owned until all output budgets have been validated.
    DeferredUntilValidated,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamDelivery {
    /// Keep bounded rows query-owned until every output limit is validated.
    Validated,
    /// Deliver rows as produced and let the caller surface terminal failures.
    Incremental,
}

impl StreamDelivery {
    pub fn consumer_memory_mode(self, bounded: bool) -> ConsumerMemoryMode {
        match (self, bounded) {
            (Self::Validated, true) => ConsumerMemoryMode::DeferredUntilValidated,
            (Self::Validated, false) | (Self::Incremental, _) => {
                ConsumerMemoryMode::ReleasedAfterCall
            }
        }
    }
}

#[derive(Clone, Copy, Default)]
pub struct OutputLimits {
    pub max_rows: Option<usize>,
    pub max_payload_bytes: Option<usize>,
}

#[derive(Clone, Copy, Default)]
pub struct OutputMetrics {
    pub rows: usize,
    pub payload_bytes: usize,
}

pub struct QueryOutputAccumulator<'a> {
    consumer: Option<&'a mut dyn FnMut(Row) -> Result<()>>,
    owned_rows: Option<QueryRowsBuilder>,
    memory_mode: ConsumerMemoryMode,
    max_rows: Option<usize>,
    max_payload_bytes: Option<usize>,
    result_account: QueryMemoryAccount,
    retained_result_lease: QueryMemoryLease,
    deferred_rows: Vec<Row>,
    metrics: OutputMetrics,
}

impl<'a> QueryOutputAccumulator<'a> {
    pub fn new(
        limits: OutputLimits,
        result_memory_budget: NonZeroUsize,
        memory_ledger: &QueryMemoryLedger,
        memory_mode: ConsumerMemoryMode,
        consumer: &'a mut dyn FnMut(Row) -> Result<()>,
    ) -> Result<Self> {
        Self::new_with_budget_option(
            limits,
            result_memory_budget,
            memory_ledger,
            "QueryMemoryLedger::account(budget_bytes)",
            memory_mode,
            consumer,
        )
    }

    pub fn new_with_budget_option(
        limits: OutputLimits,
        result_memory_budget: NonZeroUsize,
        memory_ledger: &QueryMemoryLedger,
        budget_option: &'static str,
        memory_mode: ConsumerMemoryMode,
        consumer: &'a mut dyn FnMut(Row) -> Result<()>,
    ) -> Result<Self> {
        let result_account = memory_ledger.account_with_budget_option(
            QueryMemoryClass::ResultMaterialization,
            "query result",
            result_memory_budget,
            budget_option,
        );
        let retained_result_lease = result_account.reserve(0)?;
        Ok(Self {
            consumer: Some(consumer),
            owned_rows: None,
            memory_mode,
            max_rows: limits.max_rows,
            max_payload_bytes: limits.max_payload_bytes,
            result_account,
            retained_result_lease,
            deferred_rows: Vec::new(),
            metrics: OutputMetrics::default(),
        })
    }

    pub fn new_owned(
        limits: OutputLimits,
        result_memory_budget: NonZeroUsize,
        memory_ledger: &QueryMemoryLedger,
    ) -> Result<Self> {
        Self::new_owned_with_budget_option(
            limits,
            result_memory_budget,
            memory_ledger,
            "QueryMemoryLedger::account(budget_bytes)",
        )
    }

    pub fn new_owned_with_budget_option(
        limits: OutputLimits,
        result_memory_budget: NonZeroUsize,
        memory_ledger: &QueryMemoryLedger,
        budget_option: &'static str,
    ) -> Result<Self> {
        let result_account = memory_ledger.account_with_budget_option(
            QueryMemoryClass::ResultMaterialization,
            "query result",
            result_memory_budget,
            budget_option,
        );
        let retained_result_lease =
            result_account.reserve(QueryRowsBuilder::owner_memory_bytes())?;
        Ok(Self {
            consumer: None,
            owned_rows: Some(QueryRowsBuilder::new()),
            memory_mode: ConsumerMemoryMode::Retained,
            max_rows: limits.max_rows,
            max_payload_bytes: limits.max_payload_bytes,
            result_account,
            retained_result_lease,
            deferred_rows: Vec::new(),
            metrics: OutputMetrics::default(),
        })
    }

    /// Finalize only after execution, output validation, and profile capture.
    /// The accumulator keeps its retained charge until the execution call ends.
    pub fn take_owned_rows(&mut self) -> Option<QueryRows> {
        self.owned_rows.take().map(QueryRowsBuilder::finish)
    }

    pub fn emit(&mut self, binding: Binding) -> Result<()> {
        if let Some(max_rows) = self.max_rows
            && self.metrics.rows >= max_rows
        {
            return Err(HawDBError::read_budget_exceeded(ReadBudgetResource::Rows, max_rows, format!(
                "read query returned more than {max_rows} rows, exceeding max_read_result_rows {max_rows}"
            )));
        }

        let row = binding.values;
        let row_payload_bytes = map_payload_bytes(&row);
        let next_payload_bytes = self.metrics.payload_bytes.saturating_add(row_payload_bytes);
        if let Some(max_payload_bytes) = self.max_payload_bytes
            && next_payload_bytes > max_payload_bytes
        {
            return Err(HawDBError::read_budget_exceeded(ReadBudgetResource::PayloadBytes, max_payload_bytes, format!(
                "read query payload would exceed max_payload_bytes {max_payload_bytes} (max_read_result_payload_bytes {max_payload_bytes}; next total {next_payload_bytes})"
            )));
        }
        if let Some(rows) = &mut self.owned_rows {
            rows.push_budgeted_named_row(
                row,
                &self.result_account,
                &mut self.retained_result_lease,
            )?;
            self.metrics.rows = self.metrics.rows.saturating_add(1);
            self.metrics.payload_bytes = next_payload_bytes;
            return Ok(());
        }
        let row_memory_bytes = map_memory_bytes(&row);
        let transient_result_lease = match self.memory_mode {
            ConsumerMemoryMode::Retained | ConsumerMemoryMode::DeferredUntilValidated => {
                self.retained_result_lease.grow(row_memory_bytes)?;
                None
            }
            ConsumerMemoryMode::ReleasedAfterCall => {
                Some(self.result_account.reserve(row_memory_bytes)?)
            }
        };

        match self.memory_mode {
            ConsumerMemoryMode::Retained | ConsumerMemoryMode::ReleasedAfterCall => {
                (self.consumer.as_mut().expect("row consumer is present"))(row)?;
            }
            ConsumerMemoryMode::DeferredUntilValidated => self.deferred_rows.push(row),
        }
        drop(transient_result_lease);
        self.metrics.rows = self.metrics.rows.saturating_add(1);
        self.metrics.payload_bytes = next_payload_bytes;
        Ok(())
    }

    pub fn metrics(&self) -> OutputMetrics {
        self.metrics
    }

    pub fn finish_delivery(&mut self, task_context: Option<&RuntimeTaskContext>) -> Result<()> {
        if !matches!(self.memory_mode, ConsumerMemoryMode::DeferredUntilValidated) {
            return Ok(());
        }
        for row in self.deferred_rows.drain(..) {
            runtime_checkpoint(task_context)?;
            (self.consumer.as_mut().expect("row consumer is present"))(row)?;
            runtime_checkpoint(task_context)?;
        }
        self.retained_result_lease.reset();
        Ok(())
    }
}

#[cfg(test)]
mod tests;