Skip to main content

query_runtime/observability/
schedule_projection.rs

1use std::{
2    fmt,
3    sync::atomic::{AtomicU64, Ordering},
4};
5
6use query_planning::util::StableFnv1a64;
7use serde::{Deserialize, Serialize, de};
8
9pub const COMPOSED_SCHEDULE_PROJECTION_SCHEMA: &str = "query_engine.composed_schedule_projection";
10
11static NEXT_ATTEMPT_ID: AtomicU64 = AtomicU64::new(1);
12
13fn attempt_id_seed() -> u64 {
14    static SEED: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
15    *SEED.get_or_init(|| {
16        let mut random = [0_u8; std::mem::size_of::<u64>()];
17        let nanos = std::time::SystemTime::now()
18            .duration_since(std::time::UNIX_EPOCH)
19            .map(|duration| duration.as_nanos() as u64)
20            .unwrap_or(0);
21        random.copy_from_slice(&(nanos ^ ((std::process::id() as u64) << 32) ^ 0x9e37_79b9_7f4a_7c15).to_le_bytes());
22        u64::from_le_bytes(random)
23    })
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
27/// Opaque execution-attempt identity minted by the runtime.
28/// Display and serialization use `attempt-` followed by 32 hexadecimal digits.
29/// The identity is independent of caller query and trace correlations.
30pub struct AttemptId([u8; 16]);
31
32impl AttemptId {
33    /// Returns the opaque 16-byte identity without its display prefix.
34    pub fn as_bytes(&self) -> &[u8; 16] {
35        &self.0
36    }
37
38    /// Mints a process-local execution-attempt correlation identity independent
39    /// of caller-supplied query and trace identifiers.
40    pub fn mint() -> Self {
41        let counter = NEXT_ATTEMPT_ID.fetch_add(1, Ordering::Relaxed).to_be_bytes();
42        let seed = attempt_id_seed().to_be_bytes();
43        let mut bytes = [0; 16];
44        bytes[..8].copy_from_slice(&counter);
45        bytes[8..].copy_from_slice(&seed);
46        Self(bytes)
47    }
48}
49
50impl fmt::Display for AttemptId {
51    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
52        write!(formatter, "attempt-")?;
53        for byte in self.0 {
54            write!(formatter, "{byte:02x}")?;
55        }
56        Ok(())
57    }
58}
59
60impl Serialize for AttemptId {
61    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
62    where S: serde::Serializer {
63        serializer.collect_str(self)
64    }
65}
66
67impl<'de> Deserialize<'de> for AttemptId {
68    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
69    where D: serde::Deserializer<'de> {
70        let value = String::deserialize(deserializer)?;
71        let hex = value.strip_prefix("attempt-").ok_or_else(|| de::Error::custom("attempt id prefix"))?;
72        if hex.len() != 32 {
73            return Err(de::Error::custom("attempt id length"));
74        }
75        let mut bytes = [0; 16];
76        for (index, byte) in bytes.iter_mut().enumerate() {
77            *byte = u8::from_str_radix(&hex[index * 2..index * 2 + 2], 16).map_err(de::Error::custom)?;
78        }
79        Ok(Self(bytes))
80    }
81}
82
83#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
84#[non_exhaustive]
85/// One producer occurrence in a composed execution memory schedule.
86pub struct AttemptComposedScheduleOccurrence {
87    /// Producer occurrence identity within the composed schedule.
88    pub occurrence_id: String,
89    /// Fingerprint of the producer's isolated execution scope.
90    pub execution_scope_fingerprint: String,
91    /// Peak capacity required by this producer's isolated scope, in bytes.
92    pub isolated_peak_bytes: u64,
93    /// Producer output bytes that remain owned after its scope completes.
94    pub retained_output_bytes: u64,
95    /// Peak scheduled capacity for this occurrence, including previously retained outputs.
96    pub stage_peak_bytes: u64,
97}
98
99#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
100#[non_exhaustive]
101/// Immutable composed memory schedule bound to an admitted attempt.
102///
103/// The projection records producer order, retained outputs and the complete
104/// capacity requirement. It cannot grant capacity or change the selected device.
105pub struct AttemptComposedScheduleProjection {
106    /// Composed-schedule schema label checked by the content fingerprint validator.
107    pub schema: String,
108    /// Attempt that owns this immutable composed schedule.
109    pub attempt_id: AttemptId,
110    /// CUDA device selected by the admission grant.
111    pub device_ordinal: usize,
112    /// Immutable ledger grant identity to which the schedule is bound.
113    pub grant_id: u64,
114    /// Immutable admitted query capacity in bytes.
115    pub grant_bytes: u64,
116    /// Capacity required for the entire attempt outside individual scope peaks, in bytes.
117    pub attempt_wide_fixed_bytes: u64,
118    /// Producer occurrences in their committed execution order.
119    pub producer_occurrences: Vec<AttemptComposedScheduleOccurrence>,
120    /// Fingerprint of the final consumer's execution scope.
121    pub consumer_scope_fingerprint: String,
122    /// Consumer capacity beyond retained producer inputs, in bytes.
123    pub consumer_additional_peak_bytes: u64,
124    /// Consumer stage peak including retained producer outputs, in bytes.
125    pub consumer_stage_peak_bytes: u64,
126    /// Largest stage capacity required by the complete composed schedule, in bytes.
127    pub complete_peak_bytes: u64,
128    /// Deterministic content fingerprint of memory quantities and scope order.
129    pub fingerprint: String,
130}
131
132impl AttemptComposedScheduleProjection {
133    #[cfg(test)]
134    pub(crate) fn for_test(device_ordinal: usize, grant_id: u64, grant_bytes: u64) -> Self {
135        let mut projection = Self {
136            schema: COMPOSED_SCHEDULE_PROJECTION_SCHEMA.to_owned(),
137            attempt_id: AttemptId::mint(),
138            device_ordinal,
139            grant_id,
140            grant_bytes,
141            attempt_wide_fixed_bytes: 0,
142            producer_occurrences: Vec::new(),
143            consumer_scope_fingerprint: "consumer".to_owned(),
144            consumer_additional_peak_bytes: 1,
145            consumer_stage_peak_bytes: 1,
146            complete_peak_bytes: 1,
147            fingerprint: String::new(),
148        };
149        projection.fingerprint = Self::fingerprint_for(
150            projection.attempt_wide_fixed_bytes,
151            &projection.producer_occurrences,
152            &projection.consumer_scope_fingerprint,
153            projection.consumer_additional_peak_bytes,
154            projection.consumer_stage_peak_bytes,
155            projection.complete_peak_bytes,
156        );
157        projection
158    }
159
160    /// Checks the schema label and deterministic memory-schedule fingerprint.
161    /// Attempt identity, device ordinal and grant binding are excluded from the
162    /// fingerprint; this check does not validate the admission grant.
163    pub fn fingerprint_matches_contents(&self) -> bool {
164        self.schema == COMPOSED_SCHEDULE_PROJECTION_SCHEMA
165            && self.fingerprint
166                == Self::fingerprint_for(
167                    self.attempt_wide_fixed_bytes,
168                    &self.producer_occurrences,
169                    &self.consumer_scope_fingerprint,
170                    self.consumer_additional_peak_bytes,
171                    self.consumer_stage_peak_bytes,
172                    self.complete_peak_bytes,
173                )
174    }
175
176    pub(crate) fn fingerprint_for(
177        attempt_wide_fixed_bytes: u64,
178        producer_occurrences: &[AttemptComposedScheduleOccurrence],
179        consumer_scope_fingerprint: &str,
180        consumer_additional_peak_bytes: u64,
181        consumer_stage_peak_bytes: u64,
182        complete_peak_bytes: u64,
183    ) -> String {
184        let mut hasher = StableFnv1a64::new();
185        hasher.update_field(COMPOSED_SCHEDULE_PROJECTION_SCHEMA);
186        hasher.update_field(&attempt_wide_fixed_bytes.to_string());
187        hasher.update_field(&producer_occurrences.len().to_string());
188        for occurrence in producer_occurrences {
189            hasher.update_field(&occurrence.occurrence_id);
190            hasher.update_field(&occurrence.execution_scope_fingerprint);
191            hasher.update_field(&occurrence.isolated_peak_bytes.to_string());
192            hasher.update_field(&occurrence.retained_output_bytes.to_string());
193            hasher.update_field(&occurrence.stage_peak_bytes.to_string());
194        }
195        hasher.update_field(consumer_scope_fingerprint);
196        hasher.update_field(&consumer_additional_peak_bytes.to_string());
197        hasher.update_field(&consumer_stage_peak_bytes.to_string());
198        hasher.update_field(&complete_peak_bytes.to_string());
199        format!("composed-schedule-fnv1a64:{:016x}", hasher.finish())
200    }
201}