query_runtime/runtime/query_service/service/
api.rs

1//! Public request/snapshot DTOs for query-service admission.
2
3#[cfg(test)]
4use std::time::Duration;
5use std::{fmt, time::Instant};
6
7#[cfg(test)]
8use query_planning::error::QueryErrorCode;
9use query_planning::{
10    admission::AttemptCapabilityEvidence,
11    error::{QueryError, Result},
12};
13use serde::Serialize;
14
15#[cfg(any(test, feature = "test-hooks"))]
16use crate::attempt::AttemptOwner;
17#[cfg(test)]
18use crate::runtime::query_service::config::QueueBounds;
19use crate::{
20    attempt::{AttemptAdmission, AttemptAdmissionBinding, AttemptLifecycleError},
21    observability::{AttemptId, RuntimeObservationHandle},
22    runtime::{
23        ExecutionCancellation,
24        cache::{CacheSourceEventSnapshot, DeviceCacheReport},
25        gpu::{DeviceLedgerReport, DeviceMemoryFailureReport},
26        gpu_execution_pool::{GpuExecutionPoolSnapshot, GpuExecutionQuerySnapshot},
27    },
28};
29
30/// Result of a controlled device-quarantine injection.
31///
32/// The injection changes only the engine's admission ledger. It deliberately
33/// does not manufacture a CUDA error or mutate cuDF's process-wide device
34/// health; controlled validation uses it to exercise the same device-local
35/// scheduler transition that a proven native fault would perform.
36#[cfg(feature = "qualification")]
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub struct DeviceQuarantineInjection {
39    device_ordinal: usize,
40    device_health_generation: u64,
41    newly_quarantined: bool,
42}
43
44#[cfg(feature = "qualification")]
45impl DeviceQuarantineInjection {
46    pub(crate) fn new(device_ordinal: usize, device_health_generation: u64, newly_quarantined: bool) -> Self {
47        Self { device_ordinal, device_health_generation, newly_quarantined }
48    }
49
50    /// Ordinal of the admission ledger that was targeted.
51    pub fn device_ordinal(self) -> usize {
52        self.device_ordinal
53    }
54
55    /// Generation of the targeted ledger's closed admission state.
56    pub fn device_health_generation(self) -> u64 {
57        self.device_health_generation
58    }
59
60    /// Whether this call performed the open-to-closed transition.
61    pub fn newly_quarantined(self) -> bool {
62        self.newly_quarantined
63    }
64}
65
66/// Point-in-time service admission state for bounded-queue and overload reports.
67#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
68pub struct QueryServiceAdmissionSnapshot {
69    /// Number of currently admitted query handles.
70    pub active_queries: usize,
71    /// Number of admissions currently waiting in the service queue.
72    pub queued_queries: usize,
73    /// Largest number of admissions queued concurrently since this service
74    /// was constructed. This is maintained by queue transitions and is not
75    /// inferred from snapshot polling. It therefore remains available after
76    /// the queue drains.
77    pub queue_high_watermark: usize,
78}
79
80/// One read of a live service's state, each part taken under its owner's
81/// lock. The device ledgers are the admission authority; cache, memory
82/// failure, and execution-pool reports are diagnostic cross-checks.
83#[derive(Debug, Clone, PartialEq, Eq)]
84pub struct BackendDiagnostics {
85    /// Shared service admission queue state.
86    pub admission: QueryServiceAdmissionSnapshot,
87    /// Authoritative device admission balances sampled under each ledger owner.
88    pub ledgers: Vec<DeviceLedgerReport>,
89    /// Physical allocation-domain usage and activity of each cache-enabled
90    /// device.
91    pub caches: Vec<DeviceCacheReport>,
92    /// Bounded source-cache events from the cache owners, for correlation.
93    pub cache_source_events: Vec<CacheSourceEventSnapshot>,
94    /// Device allocator failure counters used for diagnosis, independent of admission.
95    pub memory_failures: Vec<DeviceMemoryFailureReport>,
96    /// Device execution-pool concurrency snapshots.
97    pub execution_pools: Vec<GpuExecutionPoolSnapshot>,
98    /// Service-wide aggregate query-execution concurrency snapshot.
99    pub execution_queries: GpuExecutionQuerySnapshot,
100}
101
102/// Caller-facing classification for a failed native admission or execution
103/// attempt. This projection is intentionally policy-neutral: applications
104/// decide whether a capacity result should be queued, shed, or surfaced.
105///
106/// Admission and execution are separate caller boundaries. Use
107/// [`Self::from_error`] for an error returned while an admission is being
108/// attempted, and [`Self::from_execution_error`] after a handle has been
109/// published. The engine does not infer that boundary from a diagnostic
110/// reason or message.
111///
112/// A frontend can keep its own protocol policy at the edge without teaching
113/// the engine about HTTP:
114///
115/// ```text
116/// match QueryAdmissionOutcome::from_error_with_queue_bounds(&error, &queue_bounds) {
117///     Capacity { queued_queries, queue_capacity, max_queue_wait } => {
118///         // Enqueue when the application's wait budget allows it. If its
119///         // own bounded queue is full, return 429. `max_queue_wait` is only
120///         // the engine's rejection bound; it is not a relief-time ETA.
121///         if should_enqueue(queued_queries, queue_capacity, max_queue_wait) {
122///             AppAction::Enqueue
123///         } else {
124///             AppAction::Http429
125///         }
126///     }
127///     BudgetUnavailable { .. } => AppAction::ShedWithNewCapOrDevice,
128///     CapabilityRejected | InvalidInput => AppAction::Shed,
129///     Cancelled => AppAction::Shed,
130///     // Returned both for pre-admission errors outside the categories
131///     // above and for every post-admission error except cancellation.
132///     ExecutionFailure => AppAction::Http503,
133/// }
134/// ```
135///
136/// `AppAction` and `should_enqueue` belong to the application. The original
137/// [`QueryError`] remains the structured diagnostic source; this enum is only
138/// its stable policy-selection projection.
139#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
140#[cfg(test)]
141pub(crate) enum QueryAdmissionOutcome {
142    /// Bounded capacity is temporarily unavailable.  Queue facts are present
143    /// when the service attached them to the structured error; the configured
144    /// wait bound is supplied by [`Self::from_error_with_queue_bounds`].
145    Capacity { queued_queries: Option<usize>, queue_capacity: Option<usize>, max_queue_wait: Option<Duration> },
146    /// The requested query cap does not fit any eligible device after
147    /// resident ownership is deducted.
148    ///
149    /// This is deliberately not [`Self::Capacity`]: queueing and backoff
150    /// cannot make an impossible cap fit, so the caller must release resident
151    /// owners or change the cap, device, or query.
152    BudgetUnavailable { requested_cap_bytes: Option<u64>, largest_eligible_cap_bytes: Option<u64> },
153    /// The plan or operator is permanently outside the supported capability
154    /// catalogue.
155    CapabilityRejected,
156    /// The request, plan, or materialized input is invalid for this operation.
157    InvalidInput,
158    /// Caller cancellation or deadline stopped the attempt before completion.
159    Cancelled,
160    /// A failure that is not admission capacity, an impossible cap, a
161    /// capability rejection, invalid input, or cancellation. [`Self::from_error`]
162    /// returns it for any such pre-admission error; [`Self::from_execution_error`]
163    /// returns it for every post-admission error except cancellation. The
164    /// label therefore cannot prove execution started, and no retry policy
165    /// is implied by it.
166    ExecutionFailure,
167}
168
169#[cfg(test)]
170impl QueryAdmissionOutcome {
171    /// Classifies an error returned while admission is being attempted.
172    ///
173    /// This method does not inspect diagnostic reason strings. A caller that
174    /// has already published a handle must use [`Self::from_execution_error`]
175    /// instead, even when the execution error has a resource-exhaustion code.
176    #[cfg(test)]
177    pub(crate) fn from_error(error: &QueryError) -> Self {
178        Self::from_error_with_optional_queue_bounds(error, None)
179    }
180
181    /// Classifies a structured native error and attaches the service-owned
182    /// queue wait bound when the caller has the corresponding configuration.
183    #[cfg(test)]
184    pub(crate) fn from_error_with_queue_bounds(error: &QueryError, queue_bounds: &QueueBounds) -> Self {
185        Self::from_error_with_optional_queue_bounds(error, Some(queue_bounds))
186    }
187
188    /// Classifies an error observed after admission published a handle.
189    ///
190    /// Execution-time resource errors are deliberately not projected as
191    /// admission capacity: the service cannot ask the caller to queue an
192    /// attempt that has already started. Cancellation remains cancellation;
193    /// every other post-admission error is an execution failure.
194    #[cfg(test)]
195    pub(crate) fn from_execution_error(error: &QueryError) -> Self {
196        if error.code() == QueryErrorCode::Cancelled { Self::Cancelled } else { Self::ExecutionFailure }
197    }
198
199    #[cfg(test)]
200    fn from_error_with_optional_queue_bounds(error: &QueryError, queue_bounds: Option<&QueueBounds>) -> Self {
201        match error.code() {
202            QueryErrorCode::QueryBudgetUnavailable => Self::BudgetUnavailable {
203                requested_cap_bytes: fact_u64(error, "requested_cap_bytes"),
204                largest_eligible_cap_bytes: fact_u64(error, "largest_eligible_cap_bytes"),
205            },
206            QueryErrorCode::ServiceOverloaded | QueryErrorCode::MemoryReservationFailed => Self::Capacity {
207                queued_queries: fact_usize(error, "queued_queries"),
208                queue_capacity: fact_usize(error, "max_waiters")
209                    .or_else(|| queue_bounds.map(|bounds| bounds.max_waiters)),
210                max_queue_wait: queue_bounds.map(|bounds| bounds.max_wait),
211            },
212            QueryErrorCode::UnsupportedPipelineContract
213            | QueryErrorCode::UnsupportedOperator
214            | QueryErrorCode::UnsupportedExpression
215            | QueryErrorCode::UnsupportedType
216            | QueryErrorCode::CudfUnsupported => Self::CapabilityRejected,
217            QueryErrorCode::InvalidPlan
218            | QueryErrorCode::InvalidAttemptCorrelation
219            | QueryErrorCode::InvalidConfiguration
220            | QueryErrorCode::GraphInputInvalid
221            | QueryErrorCode::CudfInvalidInput => Self::InvalidInput,
222            QueryErrorCode::Cancelled => Self::Cancelled,
223            QueryErrorCode::CudfTemporaryFailure
224            | QueryErrorCode::CudfPersistentFailure
225            | QueryErrorCode::AlgorithmDidNotConverge
226            | QueryErrorCode::CudfDependencyFailure
227            | QueryErrorCode::SourceAccessDenied
228            | QueryErrorCode::SourceUnauthenticated
229            | QueryErrorCode::SourceObjectNotFound
230            | QueryErrorCode::SourceTemporaryFailure
231            | QueryErrorCode::CudfInvariantViolation
232            | QueryErrorCode::QueryMemoryLimitExceeded
233            | QueryErrorCode::DeviceOutOfMemory
234            | QueryErrorCode::DeviceAllocationFailed
235            | QueryErrorCode::DeviceFault
236            | QueryErrorCode::CleanupFailed
237            | QueryErrorCode::InternalInvariantViolation
238            | QueryErrorCode::SizeOverflow
239            | QueryErrorCode::HostMaterializationLimitExceeded
240            | QueryErrorCode::Internal => Self::ExecutionFailure,
241        }
242    }
243
244    /// Returns the queue depth attached to a capacity result, if known.
245    #[cfg(test)]
246    pub(crate) fn queued_queries(&self) -> Option<usize> {
247        match self {
248            Self::Capacity { queued_queries, .. } => *queued_queries,
249            _ => None,
250        }
251    }
252
253    /// Returns the configured queue capacity attached to a capacity result,
254    /// if the rejection carried it.
255    #[cfg(test)]
256    pub(crate) fn queue_capacity(&self) -> Option<usize> {
257        match self {
258            Self::Capacity { queue_capacity, .. } => *queue_capacity,
259            _ => None,
260        }
261    }
262
263    /// Returns the configured maximum queue wait when supplied by the caller.
264    #[cfg(test)]
265    pub(crate) fn max_queue_wait(&self) -> Option<Duration> {
266        match self {
267            Self::Capacity { max_queue_wait, .. } => *max_queue_wait,
268            _ => None,
269        }
270    }
271}
272
273#[cfg(test)]
274fn fact_usize(error: &QueryError, key: &str) -> Option<usize> {
275    error.fact_value(key)?.parse().ok()
276}
277
278#[cfg(test)]
279fn fact_u64(error: &QueryError, key: &str) -> Option<u64> {
280    error.fact_value(key)?.parse().ok()
281}
282
283#[cfg(test)]
284mod outcome_tests {
285    use super::*;
286
287    fn error(code: QueryErrorCode) -> QueryError {
288        QueryError::new(code, "query_engine.test", "test error")
289    }
290
291    #[test]
292    fn capacity_projection_preserves_queue_facts_and_configured_wait_bound() {
293        let error = error(QueryErrorCode::ServiceOverloaded)
294            .with_fact("reason", "admission_queue_full")
295            .with_fact("queued_queries", 7)
296            .with_fact("max_waiters", 8);
297        let bounds = QueueBounds { max_waiters: 8, max_wait: Duration::from_secs(11) };
298
299        let outcome = QueryAdmissionOutcome::from_error_with_queue_bounds(&error, &bounds);
300        assert_eq!(outcome.queued_queries(), Some(7));
301        assert_eq!(outcome.queue_capacity(), Some(8));
302        assert_eq!(outcome.max_queue_wait(), Some(Duration::from_secs(11)));
303        assert!(matches!(outcome, QueryAdmissionOutcome::Capacity { .. }));
304    }
305
306    #[test]
307    fn admission_projection_is_exhaustive_and_does_not_parse_reason_strings() {
308        let cases = [
309            (QueryErrorCode::MemoryReservationFailed, "capacity"),
310            // An impossible cap is not queue pressure: the caller must change
311            // the cap or release resident owners, never back off.
312            (QueryErrorCode::QueryBudgetUnavailable, "budget_unavailable"),
313            (QueryErrorCode::UnsupportedPipelineContract, "capability"),
314            (QueryErrorCode::InvalidPlan, "invalid"),
315            (QueryErrorCode::InvalidConfiguration, "invalid"),
316            (QueryErrorCode::Cancelled, "cancelled"),
317            (QueryErrorCode::CudfInvariantViolation, "execution"),
318        ];
319        for (code, expected) in cases {
320            let outcome = QueryAdmissionOutcome::from_error(&error(code).with_fact("reason", "unrelated_reason_name"));
321            let actual = match outcome {
322                QueryAdmissionOutcome::Capacity { .. } => "capacity",
323                QueryAdmissionOutcome::BudgetUnavailable { .. } => "budget_unavailable",
324                QueryAdmissionOutcome::CapabilityRejected => "capability",
325                QueryAdmissionOutcome::InvalidInput => "invalid",
326                QueryAdmissionOutcome::Cancelled => "cancelled",
327                QueryAdmissionOutcome::ExecutionFailure => "execution",
328            };
329            assert_eq!(actual, expected, "classification for {code:?}");
330        }
331        // The budget projection never reports queue metrics, so a backend
332        // cannot mistake it for queue pressure with a wait bound.
333        let bounds = QueueBounds { max_waiters: 8, max_wait: Duration::from_secs(11) };
334        let budget = QueryAdmissionOutcome::from_error_with_queue_bounds(
335            &error(QueryErrorCode::QueryBudgetUnavailable).with_fact("queued_queries", 7).with_fact("max_waiters", 8),
336            &bounds,
337        );
338        assert_eq!(budget.queued_queries(), None);
339        assert_eq!(budget.queue_capacity(), None);
340        assert_eq!(budget.max_queue_wait(), None);
341    }
342
343    #[test]
344    fn capacity_projection_leaves_unknown_queue_metrics_unknown() {
345        let outcome = QueryAdmissionOutcome::from_error(&error(QueryErrorCode::MemoryReservationFailed));
346        assert!(matches!(
347            outcome,
348            QueryAdmissionOutcome::Capacity { queued_queries: None, queue_capacity: None, max_queue_wait: None }
349        ));
350
351        let bounds = QueueBounds { max_waiters: 4, max_wait: Duration::from_secs(3) };
352        let outcome = QueryAdmissionOutcome::from_error_with_queue_bounds(
353            &error(QueryErrorCode::MemoryReservationFailed),
354            &bounds,
355        );
356        assert_eq!(outcome.queued_queries(), None);
357        assert_eq!(outcome.queue_capacity(), Some(4));
358        assert_eq!(outcome.max_queue_wait(), Some(Duration::from_secs(3)));
359    }
360
361    #[test]
362    fn execution_projection_does_not_reclassify_post_admission_resource_errors() {
363        for code in [QueryErrorCode::MemoryReservationFailed, QueryErrorCode::DeviceAllocationFailed] {
364            assert_eq!(
365                QueryAdmissionOutcome::from_execution_error(&error(code)),
366                QueryAdmissionOutcome::ExecutionFailure,
367                "post-admission {code:?} is an execution failure"
368            );
369        }
370        assert_eq!(
371            QueryAdmissionOutcome::from_execution_error(&error(QueryErrorCode::Cancelled)),
372            QueryAdmissionOutcome::Cancelled
373        );
374    }
375}
376
377/// Complete request-scoped metadata for one query-service admission attempt.
378#[derive(Debug)]
379#[non_exhaustive]
380pub struct QueryAdmissionRequest {
381    /// Explicit immutable query cap. When absent, the selected device resolves
382    /// its configured default inside the admission transaction.
383    requested_query_cap_bytes: Option<u64>,
384    /// Opaque capability evidence emitted by planning.
385    ///
386    /// Request callers provide lifecycle and correlation metadata only;
387    /// raw floors, working-set hints, placement choices, and capability sets are not a public
388    /// admission input.  When absent, admission uses empty evidence only when
389    /// every installed profile accepts every modeled capability.
390    pub(in crate::runtime::query_service) attempt_requirement: Option<AttemptCapabilityEvidence>,
391    /// Public evidence retained for diagnostic inspection until submission.
392    planned_capability_evidence: Option<AttemptCapabilityEvidence>,
393    /// Caller-provided bound for queueing this attempt. The service combines it
394    /// with its configured queue bound, choosing the earlier instant.
395    pub(crate) queue_deadline: Option<Instant>,
396    /// Attempt-wide cancellation shared by queueing and transactional resource
397    /// acquisition. The same token must remain attached until the admission
398    /// either publishes a complete envelope or rolls every provisional
399    /// resource back.
400    pub(crate) execution_cancellation: ExecutionCancellation,
401    /// The admitted attempt's observation and cleanup attachment.
402    attempt: AttemptAdmissionBinding,
403    runtime_observation_activity: Option<Result<QueryActivityLease>>,
404    /// The owner of a [`Self::test_attempt`] request, handed to the admitted
405    /// query handle.
406    #[cfg(any(test, feature = "test-hooks"))]
407    pub(in crate::runtime::query_service) test_attempt: Option<AttemptOwner>,
408}
409
410pub(in crate::runtime::query_service) struct AdmissionObservation {
411    pub(in crate::runtime::query_service) handle: RuntimeObservationHandle,
412    pub(in crate::runtime::query_service) activity: Option<QueryActivityLease>,
413    pub(in crate::runtime::query_service) activity_error: Option<QueryError>,
414    /// The engine-minted attempt identity, which may receive a post-callback
415    /// allocator correlation after a physical domain has been admitted.
416    pub(in crate::runtime::query_service) attempt_id: AttemptId,
417}
418
419/// Opaque RAII token retained until query-service cleanup is complete.
420///
421/// The native engine never invokes the value. Dropping it returns activity
422/// ownership to the adapter coordinator after queued admission, contexts,
423/// grants, device slots, and cleanup work have all been released.
424pub struct QueryActivityLease {
425    _lease: Box<dyn Send + Sync>,
426}
427
428impl QueryActivityLease {
429    pub fn new(lease: impl Send + Sync + 'static) -> Self {
430        Self { _lease: Box::new(lease) }
431    }
432}
433
434impl fmt::Debug for QueryActivityLease {
435    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
436        formatter.debug_struct("QueryActivityLease").finish_non_exhaustive()
437    }
438}
439
440impl QueryAdmissionRequest {
441    /// Creates the admission request of one execution attempt, with no
442    /// explicit capability requirement.
443    ///
444    /// The binding is the request's only observation and cleanup attachment.
445    /// [`crate::attempt::AttemptOwner`] mints it as one unit, so a request
446    /// cannot pair one attempt's observation with another attempt's cleanup.
447    pub fn new(attempt: AttemptAdmissionBinding) -> Self {
448        Self {
449            requested_query_cap_bytes: None,
450            attempt_requirement: None,
451            planned_capability_evidence: None,
452            queue_deadline: None,
453            execution_cancellation: ExecutionCancellation::default(),
454            attempt,
455            runtime_observation_activity: None,
456            #[cfg(any(test, feature = "test-hooks"))]
457            test_attempt: None,
458        }
459    }
460
461    /// Creates the request of a fresh attempt classified as native, for
462    /// tests that admit without a planning frontend.
463    ///
464    /// The admitted query handle owns the attempt: only an attempt's terminal
465    /// commit releases its grant, so dropping that handle ends the attempt as
466    /// `Dropped` and the release follows.
467    #[cfg(any(test, feature = "test-hooks"))]
468    #[doc(hidden)]
469    pub fn test_attempt() -> Self {
470        let owner = AttemptOwner::native_for_test();
471        let mut request = Self::new(owner.admission_binding());
472        request.test_attempt = Some(owner);
473        request
474    }
475
476    /// Separates a test request's attempt owner, which the admitted handle
477    /// receives, from the request that enters admission.
478    #[cfg(any(test, feature = "test-hooks"))]
479    pub(in crate::runtime::query_service) fn split_test_attempt(mut self) -> (Self, Option<AttemptOwner>) {
480        let owner = self.test_attempt.take();
481        (self, owner)
482    }
483
484    /// Requests an exact immutable query cap. Admission never clamps it.
485    pub fn with_query_cap_bytes(mut self, query_cap_bytes: u64) -> Self {
486        self.requested_query_cap_bytes = Some(query_cap_bytes);
487        self
488    }
489
490    pub(in crate::runtime::query_service) const fn requested_query_cap_bytes(&self) -> Option<u64> {
491        self.requested_query_cap_bytes
492    }
493
494    /// Attaches an opaque evidence computed by capability-first planning.
495    ///
496    /// A public request may prove required device capabilities and plan
497    /// scope, but it cannot carry a model byte target into grant selection.
498    /// The installed device profile therefore remains the sole authority for
499    /// the immutable grant.
500    pub fn with_query_plan_capability_evidence(mut self, evidence: AttemptCapabilityEvidence) -> Self {
501        self.planned_capability_evidence = Some(evidence);
502        self
503    }
504
505    /// Internal controller/test constructor for synthetic capability evidence.
506    ///
507    /// Kept as crate-private API surface: the public-API boundary UI tests
508    /// pin that callers outside the crate cannot attach their own requirement.
509    #[cfg(test)]
510    pub(crate) fn with_attempt_requirement(mut self, attempt_requirement: AttemptCapabilityEvidence) -> Self {
511        self.attempt_requirement = Some(attempt_requirement);
512        self
513    }
514
515    pub(in crate::runtime::query_service) fn resolved_attempt_requirement(&self) -> Option<AttemptCapabilityEvidence> {
516        self.planned_capability_evidence.clone().or_else(|| self.attempt_requirement.clone())
517    }
518
519    /// Returns the evidence that was submitted by planning for the
520    /// execution-time coverage check.  Internal requirements still size and
521    /// place an admission, but must not serve as evidence for an arbitrary
522    /// later-built plan.
523    pub(in crate::runtime::query_service) fn submitted_plan_capability_evidence(
524        &self,
525    ) -> Option<AttemptCapabilityEvidence> {
526        self.planned_capability_evidence.clone()
527    }
528
529    /// Sets an attempt-bound queue deadline.
530    pub fn with_queue_deadline(mut self, queue_deadline: Instant) -> Self {
531        self.queue_deadline = Some(queue_deadline);
532        self
533    }
534
535    /// Carries the attempt-wide cancellation token into admission.
536    pub fn with_execution_cancellation(mut self, execution_cancellation: ExecutionCancellation) -> Self {
537        self.execution_cancellation = execution_cancellation;
538        self
539    }
540
541    /// Retains an adapter-owned activity token through delayed native cleanup.
542    #[cfg(test)]
543    pub(crate) fn with_runtime_observation_activity(mut self, activity: QueryActivityLease) -> Self {
544        self.runtime_observation_activity = Some(Ok(activity));
545        self
546    }
547
548    /// Carries fallible adapter activity acquisition into native admission.
549    ///
550    /// Coordinated callers must pass acquisition failure through this field.
551    /// Admission then rejects with that error, and no attempt runs without its
552    /// lease.
553    #[cfg(test)]
554    pub(crate) fn with_runtime_observation_activity_result(mut self, activity: Result<QueryActivityLease>) -> Self {
555        self.runtime_observation_activity = Some(activity);
556        self
557    }
558
559    pub(in crate::runtime::query_service) fn into_admission_observation(self) -> AdmissionObservation {
560        let (activity, activity_error) = match self.runtime_observation_activity {
561            Some(Ok(activity)) => (Some(activity), None),
562            Some(Err(error)) => (None, Some(error)),
563            None => (None, None),
564        };
565        AdmissionObservation {
566            handle: self.attempt.observation_handle(),
567            activity,
568            activity_error,
569            attempt_id: self.attempt.attempt_id(),
570        }
571    }
572
573    /// Claims the bound attempt's only admission right.
574    pub(in crate::runtime::query_service) fn claim_attempt_admission(
575        &self,
576    ) -> std::result::Result<AttemptAdmission, AttemptLifecycleError> {
577        self.attempt.claim()
578    }
579}
580
581#[cfg(test)]
582mod tests {
583    use std::sync::Arc;
584
585    use super::*;
586    use crate::{
587        attempt::{AttemptBackendCounters, BeginAttemptRequest},
588        observability::{AttemptDisposition, RuntimeObservationConfig},
589        runtime::BackendLifecycle,
590    };
591
592    #[test]
593    fn attempt_binding_attaches_observation_and_release_together() {
594        let lifecycle = BackendLifecycle::new();
595        let attempt = AttemptOwner::begin(
596            BeginAttemptRequest {
597                observation: RuntimeObservationConfig::bounded(8, 8 * 1024),
598                ..BeginAttemptRequest::default()
599            },
600            lifecycle.try_begin_observation().expect("attempt observation permit"),
601            Arc::new(AttemptBackendCounters::default()),
602            None,
603        );
604        let attempt_id = attempt.owner.attempt_id();
605        attempt.owner.handle().set_disposition(AttemptDisposition::Native).expect("native disposition");
606        let request = QueryAdmissionRequest::new(attempt.owner.admission_binding());
607
608        assert!(request.claim_attempt_admission().is_ok());
609        assert!(matches!(request.claim_attempt_admission(), Err(AttemptLifecycleError::AdmissionAlreadyClaimed)));
610        let observation = request.into_admission_observation();
611        assert_eq!(observation.attempt_id, attempt_id);
612    }
613}
query_runtime/runtime/query_service/service/api.rs

query_runtime/runtime/query_service/service/
api.rs
