1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
//! Attempt work leases and stream-lease derivation for execution contexts.
use std::sync::Arc;
use cudf::Runtime as CudfRuntime;
use query_planning::error::{QueryError, QueryErrorCode, Result};
use super::{
super::{AttemptWorkLease, StreamLease},
AttemptAuxiliaryPermit, ExecutionContext, QueryScope, StreamBinding, validate_service_stream_pool_device,
};
use crate::observability::{RuntimeTaskId, RuntimeWorkerId};
impl ExecutionContext {
/// Acquires an attempt work lease (an attempt auxiliary permit, a service
/// stream lease, and the derived context owning them) when an auxiliary slot
/// is available.
///
/// Combines the resources required for auxiliary attempt work across execution
/// lanes and scheduled source reads. The derived context's alias holds an
/// auxiliary fork of the attempt pool for the lease lifetime, preventing
/// concurrent work from sharing stream dispatches. The acquired permit enforces
/// pool width bounds.
///
/// Work within an attempt avoids blocking on auxiliary permits, preventing
/// deadlocks when root execution jobs hold permits across multiple lanes.
/// Returns `Ok(None)` when the envelope is fully claimed; callers either
/// narrow concurrency or execute work directly on their existing context.
/// `source_worker` executes after permit acquisition to name the source-read
/// worker, ensuring worker registration occurs only for granted leases.
#[doc(hidden)]
pub fn try_acquire_attempt_work_lease(
&self,
source_worker: impl FnOnce() -> Option<(RuntimeTaskId, RuntimeWorkerId)>,
) -> Result<Option<AttemptWorkLease>> {
let attempt_auxiliary_permit = match self.admitted() {
Some(attempt) => {
match attempt
.auxiliary_slots
.try_acquire("query_engine.runtime.attempt_work_lease.attempt_auxiliary_try_acquire")?
{
Some(permit) => Some(permit),
None => return Ok(None),
}
},
None => None,
};
self.derive_attempt_work_lease(source_worker(), attempt_auxiliary_permit).map(Some)
}
/// Leases the service stream and derives the work context once the
/// auxiliary permit is held. The permit comes first so an attempt waiting
/// on its own envelope never holds a global pool stream while it waits.
/// The lease carries the plan's stream scope, extended with the source
/// worker when one is named, so stream observations correlate to the job
/// that issued them.
fn derive_attempt_work_lease(
&self,
source_worker: Option<(RuntimeTaskId, RuntimeWorkerId)>,
attempt_auxiliary_permit: Option<AttemptAuxiliaryPermit>,
) -> Result<AttemptWorkLease> {
let mut stream_scope = self.plan.runtime_stream_lease_scope.clone();
if let Some((task_id, worker_id)) = source_worker {
stream_scope = stream_scope.with_source_worker(task_id, worker_id);
}
let stream_lease = self.service_stream_pool().acquire_observed_with_cancellation(
self.cancellation(),
"query_engine.runtime.attempt_work_lease.acquire",
self.plan.runtime_observation.clone(),
stream_scope,
)?;
let context = self.with_stream_lease(&stream_lease)?;
Ok(AttemptWorkLease::new(context, stream_lease, attempt_auxiliary_permit))
}
/// Whether this context runs on a stream leased for attempt work rather
/// than on the attempt's own root stream.
///
/// Work issued from such a context already holds an attempt auxiliary
/// permit, so it must run inline and must not lease a second one.
#[doc(hidden)]
pub fn is_attempt_work_lease_context(&self) -> bool {
matches!(*self.stream, StreamBinding::AttemptWork { .. })
}
/// Which auxiliary fork of the attempt pool this context's alias holds
/// alone. A context with no alias (the attempt root, or a direct context)
/// reports `None`; the root runs on the route's reserved fork instead.
#[doc(hidden)]
pub fn attempt_claimed_fork_index(&self) -> Option<usize> {
match &*self.stream {
StreamBinding::AttemptWork { alias, .. } => Some(alias.claimed_fork_index()),
StreamBinding::Standalone { .. } | StreamBinding::AttemptRoot { .. } => None,
}
}
/// Builds a short-lived execution context from a borrowed stream lease.
///
/// The returned context shares this context's query policies, caches, and
/// cancellation flag, but executes GPU work on the leased runtime. The
/// caller must keep `stream_lease` alive until all work submitted through
/// the returned context has completed.
pub(crate) fn with_stream_lease(&self, stream_lease: &StreamLease) -> Result<Self> {
self.check_service_epoch("query_engine.runtime.context.with_stream_lease")?;
validate_service_stream_pool_device(stream_lease.pool(), self.device())?;
let runtime = Arc::new(stream_lease.runtime()?);
Ok(self.with_stream(self.work_stream(runtime, None)?))
}
/// Builds an auxiliary fragment context that owns `stream_lease` and the
/// attempt `permit` paired with it.
///
/// Both live in the context's stream binding, so cloning an outward
/// context cannot retain the stream while releasing its attempt envelope
/// share early.
pub(crate) fn with_owned_auxiliary_stream(
&self,
stream_lease: StreamLease,
permit: AttemptAuxiliaryPermit,
) -> Result<Self> {
let operation = "query_engine.runtime.context.with_owned_auxiliary_stream";
self.check_service_epoch(operation)?;
validate_service_stream_pool_device(stream_lease.pool(), self.device())?;
if self.admitted().is_none() {
return Err(QueryError::new(
QueryErrorCode::Internal,
operation,
"an auxiliary fragment stream belongs to an admitted attempt",
)
.with_fact("reason", "auxiliary_stream_without_admission"));
}
let runtime = Arc::new(stream_lease.runtime()?);
Ok(self.with_stream(self.work_stream(runtime, Some((stream_lease, permit)))?))
}
/// Binds `runtime`'s stream as work of this context's query. For an
/// attempt: a member of its allocation domain, then an alias on the route
/// of its root stream. Standalone work attaches a member of the same query
/// domain for the lifetime of the work context.
fn work_stream(
&self,
runtime: Arc<CudfRuntime>,
owned: Option<(StreamLease, AttemptAuxiliaryPermit)>,
) -> Result<StreamBinding> {
if let QueryScope::Standalone(query) = &self.query {
return query.bind_work(runtime);
}
let attempt = self.admitted().expect("an admitted query scope has an admitted owner");
let route = self.stream.route().expect("an admitted stream owns its attempt route");
let member = attempt.allocation_domain.create_member(runtime.stream())?;
let alias = Arc::new(route.attach(runtime.stream())?);
Ok(StreamBinding::AttemptWork { runtime, alias, member, route: Arc::clone(route), owned })
}
}