algeon_datafusion/native_exec/config/
admission.rs

1//! The admission authority a config carries.
2//!
3//! At planning the config holds a `QueryAdmissionSource`, the backend that
4//! will admit the plan. Once an invocation minted its ticket, or a composed
5//! root bound one, the config can hold that `QueryAdmissionTicket` too. Lookup
6//! prefers the ticket of the executing task context over the config's own, so
7//! every fragment of one execution admits under the same ticket whatever its
8//! config carried at planning.
9
10use std::sync::Arc;
11
12use query_runtime::runtime::{ExecutionCancellation, GpuBackendSession};
13
14use super::{
15    AlgeonNativeExecutionConfig, QueryAdmissionSource, QueryAdmissionTicket, TaskContext, invocation_admission_ticket,
16};
17use crate::{error::Result, native_exec::admission::attempt_execution_required_error, planner::native::Invocation};
18
19impl AlgeonNativeExecutionConfig {
20    pub(crate) fn with_query_admission_source(mut self, source: QueryAdmissionSource) -> Self {
21        self.query_admission_source = Some(source);
22        self.query_admission_ticket = None;
23        self
24    }
25
26    pub(crate) fn query_admission_source(&self) -> Option<QueryAdmissionSource> {
27        self.query_admission_source.clone()
28    }
29
30    pub(crate) fn with_existing_query_admission_ticket(mut self, ticket: Arc<QueryAdmissionTicket>) -> Self {
31        self.query_admission_ticket = Some(ticket);
32        self
33    }
34
35    /// Drops a planner-time ticket before a frontend resolves metadata that
36    /// participates in its immutable admission evidence.
37    #[cfg(feature = "cugraph")]
38    pub(crate) fn without_query_admission_ticket(mut self) -> Self {
39        self.query_admission_ticket = None;
40        self
41    }
42
43    #[cfg(test)]
44    pub(crate) fn query_admission_ticket_id(&self) -> Option<u64> {
45        self.query_admission_ticket.as_ref().map(|ticket| ticket.id())
46    }
47
48    pub(crate) fn query_admission_ticket_for_task_context(
49        &self,
50        task_context: &TaskContext,
51    ) -> Option<Arc<QueryAdmissionTicket>> {
52        invocation_admission_ticket(task_context).or_else(|| self.query_admission_ticket.as_ref().map(Arc::clone))
53    }
54
55    #[cfg(test)]
56    pub(crate) fn query_admission_ticket_id_for_task_context(&self, task_context: &TaskContext) -> Option<u64> {
57        self.query_admission_ticket_for_task_context(task_context).map(|ticket| ticket.id())
58    }
59
60    /// The cancellation of the attempt that native work in `task_context`
61    /// runs under: the invocation's attempt token, else the token of the
62    /// attempt its admission ticket belongs to.
63    pub(crate) fn attempt_cancellation(
64        &self,
65        task_context: &TaskContext,
66        operation: &'static str,
67    ) -> Result<ExecutionCancellation> {
68        Invocation::of_task(task_context)
69            .and_then(Invocation::attempt)
70            .map(|attempt| attempt.cancellation().execution_cancellation())
71            .or_else(|| {
72                self.query_admission_ticket_for_task_context(task_context).map(|ticket| ticket.attempt_cancellation())
73            })
74            .ok_or_else(|| attempt_execution_required_error(operation))
75    }
76
77    pub(crate) fn query_admission_ticket(&self) -> Option<Arc<QueryAdmissionTicket>> {
78        self.query_admission_ticket.as_ref().map(Arc::clone)
79    }
80
81    /// Returns whether this config carries a backend admission source or an attempt-bound ticket.
82    ///
83    /// This presence check does not claim capacity or establish that the backend is open.
84    pub fn has_query_admission_authority(&self) -> bool {
85        self.query_admission_source.is_some() || self.query_admission_ticket.is_some()
86    }
87
88    pub(crate) fn query_admission_source_matches(&self, session: &GpuBackendSession) -> bool {
89        self.query_admission_source.as_ref().is_some_and(|source| source.backend_matches(session))
90    }
91}
algeon_datafusion/native_exec/config/admission.rs

algeon_datafusion/native_exec/config/
admission.rs
