query_runtime/observability/
schedule_projection.rs1use 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)]
27pub struct AttemptId([u8; 16]);
31
32impl AttemptId {
33 pub fn as_bytes(&self) -> &[u8; 16] {
35 &self.0
36 }
37
38 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]
85pub struct AttemptComposedScheduleOccurrence {
87 pub occurrence_id: String,
89 pub execution_scope_fingerprint: String,
91 pub isolated_peak_bytes: u64,
93 pub retained_output_bytes: u64,
95 pub stage_peak_bytes: u64,
97}
98
99#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
100#[non_exhaustive]
101pub struct AttemptComposedScheduleProjection {
106 pub schema: String,
108 pub attempt_id: AttemptId,
110 pub device_ordinal: usize,
112 pub grant_id: u64,
114 pub grant_bytes: u64,
116 pub attempt_wide_fixed_bytes: u64,
118 pub producer_occurrences: Vec<AttemptComposedScheduleOccurrence>,
120 pub consumer_scope_fingerprint: String,
122 pub consumer_additional_peak_bytes: u64,
124 pub consumer_stage_peak_bytes: u64,
126 pub complete_peak_bytes: u64,
128 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 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}