Line data Source code
1 : //! The Trino adapters: raw Trino API payloads to input events.
2 : //!
3 : //! `trino_statement_events` reads the `stats` of a `/v1/statement` response. Trino sends them with every poll:
4 : //! the query (state, splits, rows, cpu, memory) and the stage tree (`rootStage`). Stages are the fragments of
5 : //! the plan, their ids are the fragment ids. A stage that the model does not know yet is registered.
6 : //!
7 : //! `trino_query_info_events` reads a QueryInfo (`/v1/query/{id}`): the operator summaries of the stages give the
8 : //! pipelines of a stage and the statistics of their slots (a plan node in a pipeline). The pipelines are
9 : //! registered when they appear.
10 : //!
11 : //! Every field is optional, a missing field is not reported (proxies like the Trino Gateway and older servers
12 : //! may leave fields out).
13 :
14 : use super::{ProfileInputEvent, SlotRole, SlotSpec, Stats, Status};
15 : use crate::query_profile::QueryProfile;
16 : use serde_json::Value;
17 : use std::collections::HashMap;
18 :
19 : /// The status of a Trino query state (`QueryState`)
20 90 : fn query_status(state: &str) -> Option<Status> {
21 90 : match state {
22 90 : "QUEUED" | "WAITING_FOR_RESOURCES" | "DISPATCHING" | "PLANNING" | "STARTING" => Some(Status::Pending),
23 54 : "RUNNING" | "FINISHING" => Some(Status::Running),
24 12 : "FINISHED" => Some(Status::Succeeded),
25 4 : "FAILED" => Some(Status::Failed),
26 1 : _ => None,
27 : }
28 90 : }
29 :
30 : /// The status of a Trino stage state (`StageState`). A stage is `PENDING` after `RUNNING`: its tasks are done and
31 : /// it waits for the query to finish, its work is done.
32 179 : fn stage_status(state: &str) -> Option<Status> {
33 179 : match state {
34 179 : "PLANNED" | "SCHEDULING" => Some(Status::Pending),
35 178 : "SCHEDULED" | "RUNNING" => Some(Status::Running),
36 87 : "PENDING" | "FINISHED" => Some(Status::Succeeded),
37 6 : "FAILED" | "ABORTED" => Some(Status::Failed),
38 2 : "CANCELED" => Some(Status::Skipped),
39 1 : _ => None,
40 : }
41 179 : }
42 :
43 : /// A counter of a JSON object
44 2064 : fn count(object: &Value, name: &str) -> Option<u64> {
45 2064 : object.get(name).and_then(Value::as_u64)
46 2064 : }
47 :
48 : /// Milliseconds as microseconds
49 444 : fn micros(object: &Value, name: &str) -> Option<u64> {
50 444 : count(object, name).map(|ms| ms.saturating_mul(1000))
51 444 : }
52 :
53 : /// The statistics that the query and its stages share
54 222 : fn common_stats(object: &Value) -> Stats {
55 222 : Stats {
56 222 : tuples_in: count(object, "processedRows"),
57 222 : bytes_in: count(object, "processedBytes"),
58 222 : cpu_time_us: micros(object, "cpuTimeMillis"),
59 222 : wall_time_us: micros(object, "wallTimeMillis"),
60 222 : work_done: count(object, "completedSplits"),
61 222 : work_total: count(object, "totalSplits"),
62 222 : ..Default::default()
63 222 : }
64 222 : }
65 :
66 : /// The input events of the `stats` of a `/v1/statement` response
67 89 : pub(crate) fn trino_statement_events(model: &QueryProfile, stats: &Value, time: u64) -> Vec<ProfileInputEvent> {
68 89 : let mut events = Vec::new();
69 89 : let status = stats.get("state").and_then(Value::as_str).and_then(query_status);
70 89 : let query = Stats {
71 89 : memory_bytes: count(stats, "peakMemoryBytes"),
72 89 : spilled_bytes: count(stats, "spilledBytes"),
73 89 : ..common_stats(stats)
74 89 : };
75 89 : events.push(ProfileInputEvent::UpdateQuery { time, status, stats: query });
76 :
77 : // The known fragments by engine id
78 89 : let mut fragments: HashMap<String, u32> = model
79 89 : .fragments
80 89 : .iter()
81 89 : .enumerate()
82 176 : .filter_map(|(i, f)| Some((f.engine_id.clone()?, i as u32)))
83 89 : .collect();
84 89 : let mut next_fragment = model.fragments.len() as u32;
85 :
86 : // Walk the stage tree, parents first
87 89 : let mut pending: Vec<(&Value, Option<u32>)> = stats.get("rootStage").map(|s| vec![(s, None)]).unwrap_or_default();
88 222 : while let Some((stage, parent)) = pending.pop() {
89 133 : let Some(stage_id) = stage.get("stageId").and_then(Value::as_str) else {
90 0 : continue;
91 : };
92 133 : let fragment = match fragments.get(stage_id) {
93 120 : Some(fragment) => *fragment,
94 : None => {
95 13 : let fragment = next_fragment;
96 13 : next_fragment += 1;
97 13 : fragments.insert(stage_id.to_string(), fragment);
98 13 : events.push(ProfileInputEvent::RegisterFragment {
99 13 : time,
100 13 : engine_id: Some(stage_id.to_string()),
101 13 : parent_fragment: parent,
102 13 : });
103 13 : fragment
104 : }
105 : };
106 133 : let status = stage.get("state").and_then(Value::as_str).and_then(stage_status);
107 133 : events.push(ProfileInputEvent::UpdateFragment { time, fragment, status, stats: common_stats(stage) });
108 133 : let children = stage.get("subStages").and_then(Value::as_array).into_iter().flatten();
109 133 : pending.extend(children.rev().map(|child| (child, Some(fragment))));
110 : }
111 89 : events
112 89 : }
113 :
114 : /// A Trino duration (`"956.59us"`, `"2.15s"`) in microseconds
115 880 : fn duration_us(value: Option<&Value>) -> Option<u64> {
116 880 : let text = value?.as_str()?.trim();
117 4829 : let split = text.find(|c: char| c.is_ascii_alphabetic())?;
118 859 : let (number, unit) = text.split_at(split);
119 859 : let number: f64 = number.trim().parse().ok()?;
120 858 : let factor = match unit {
121 858 : "ns" => 1e-3,
122 499 : "us" => 1.0,
123 206 : "ms" => 1e3,
124 51 : "s" => 1e6,
125 4 : "m" => 60e6,
126 3 : "h" => 3600e6,
127 2 : "d" => 86400e6,
128 1 : _ => return None,
129 : };
130 857 : Some((number * factor).round() as u64)
131 880 : }
132 :
133 : /// A Trino data size (`"271278B"`, `"1.5MB"`) in bytes
134 456 : fn data_size_bytes(value: Option<&Value>) -> Option<u64> {
135 456 : let text = value?.as_str()?.trim();
136 1934 : let split = text.find(|c: char| c.is_ascii_alphabetic())?;
137 450 : let (number, unit) = text.split_at(split);
138 450 : let number: f64 = number.trim().parse().ok()?;
139 450 : let factor = match unit {
140 450 : "B" => 1.0,
141 17 : "kB" => 1024.0,
142 12 : "MB" => 1024.0 * 1024.0,
143 5 : "GB" => 1024.0 * 1024.0 * 1024.0,
144 4 : "TB" => 1024.0 * 1024.0 * 1024.0 * 1024.0,
145 3 : "PB" => 1024.0 * 1024.0 * 1024.0 * 1024.0 * 1024.0,
146 2 : _ => return None,
147 : };
148 448 : Some((number * factor).round() as u64)
149 456 : }
150 :
151 : /// Sum the durations of an operator summary
152 248 : fn sum_durations(summary: &Value, names: &[&str]) -> Option<u64> {
153 744 : names.iter().filter_map(|name| duration_us(summary.get(*name))).reduce(|a, b| a.saturating_add(b))
154 248 : }
155 :
156 : /// The statistics of a slot: the Trino operators of one plan node in one pipeline, ordered by operator id.
157 : /// The input is the input of the first operator, the output the output of the last.
158 100 : fn slot_stats(operators: &[&Value]) -> Stats {
159 100 : let (first, last) = (operators[0], operators[operators.len() - 1]);
160 100 : let mut stats = Stats {
161 100 : tuples_in: count(first, "inputPositions"),
162 100 : tuples_out: count(last, "outputPositions"),
163 100 : bytes_in: data_size_bytes(first.get("inputDataSize")),
164 100 : bytes_out: data_size_bytes(last.get("outputDataSize")),
165 100 : ..Default::default()
166 100 : };
167 620 : let sum = |values: &mut Option<u64>, value: Option<u64>| {
168 620 : if let Some(value) = value {
169 608 : *values = Some(values.unwrap_or(0).saturating_add(value));
170 608 : }
171 620 : };
172 124 : for op in operators {
173 124 : sum(&mut stats.cpu_time_us, sum_durations(op, &["addInputCpu", "getOutputCpu", "finishCpu"]));
174 124 : sum(&mut stats.wall_time_us, sum_durations(op, &["addInputWall", "getOutputWall", "finishWall"]));
175 124 : sum(&mut stats.blocked_time_us, duration_us(op.get("blockedWall")));
176 124 : sum(&mut stats.memory_bytes, data_size_bytes(op.get("peakUserMemoryReservation")));
177 124 : sum(&mut stats.spilled_bytes, data_size_bytes(op.get("spilledDataSize")));
178 124 : }
179 100 : stats
180 100 : }
181 :
182 : /// The input events of a QueryInfo (`/v1/query/{id}`, pruned or not): the pipelines of the stages and the
183 : /// statistics of their slots from the operator summaries. Trino reports the operator summaries of a stage once its
184 : /// tasks are done (the stage is `PENDING` until the query finishes, then `FINISHED`), and of every stage for
185 : /// EXPLAIN ANALYZE. A pruned QueryInfo keeps the output stage only. Operators that are not in the plan are skipped.
186 15 : pub(crate) fn trino_query_info_events(model: &QueryProfile, info: &Value, time: u64) -> Vec<ProfileInputEvent> {
187 15 : let mut events = Vec::new();
188 15 : let fragments: HashMap<&str, u32> = model
189 15 : .fragments
190 15 : .iter()
191 15 : .enumerate()
192 46 : .filter_map(|(i, f)| Some((f.engine_id.as_deref()?, i as u32)))
193 15 : .collect();
194 15 : let operators: HashMap<&str, u32> = model
195 15 : .operators
196 15 : .iter()
197 15 : .enumerate()
198 178 : .filter_map(|(i, op)| Some((op.engine_id.as_deref()?, i as u32)))
199 15 : .collect();
200 15 : let mut pipelines: HashMap<String, u32> = model
201 15 : .pipelines
202 15 : .iter()
203 15 : .enumerate()
204 15 : .filter_map(|(i, p)| Some((p.engine_id.clone()?, i as u32)))
205 15 : .collect();
206 15 : let mut next_pipeline = model.pipelines.len() as u32;
207 15 : let mut next_slot = model.slots.len() as u32;
208 15 : let mut slot_of: HashMap<(u32, u32), u32> =
209 18 : model.slots.iter().enumerate().map(|(i, s)| ((s.pipeline, s.operator), i as u32)).collect();
210 :
211 15 : let stages = info.get("stages").and_then(|s| s.get("stages")).and_then(Value::as_array).into_iter().flatten();
212 40 : for stage in stages {
213 40 : let Some(stage_id) = stage.get("stageId").and_then(Value::as_str) else { continue };
214 40 : let fragment_id = stage_id.rsplit('.').next().unwrap_or(stage_id);
215 40 : let Some(&fragment) = fragments.get(fragment_id) else { continue };
216 40 : let status = stage.get("state").and_then(Value::as_str).and_then(stage_status);
217 40 : let summaries = stage.get("stageStats").and_then(|s| s.get("operatorSummaries")).and_then(Value::as_array);
218 :
219 : // The operators of every pipeline, ordered by operator id
220 40 : let mut by_pipeline: std::collections::BTreeMap<u64, Vec<&Value>> = Default::default();
221 126 : for summary in summaries.into_iter().flatten() {
222 126 : if let Some(pipeline) = count(summary, "pipelineId") {
223 126 : by_pipeline.entry(pipeline).or_default().push(summary);
224 126 : }
225 : }
226 46 : for (pipeline_id, mut summaries) in by_pipeline {
227 228 : summaries.sort_by_key(|s| count(s, "operatorId").unwrap_or(0));
228 : // The slots: the operators of a plan node, in the order of their first operator
229 46 : let mut slots: Vec<(u32, Vec<&Value>)> = Vec::new();
230 126 : for summary in &summaries {
231 126 : let Some(&op) = summary.get("planNodeId").and_then(Value::as_str).and_then(|id| operators.get(id))
232 : else {
233 2 : continue;
234 : };
235 124 : match slots.iter_mut().find(|(o, _)| *o == op) {
236 24 : Some((_, ops)) => ops.push(summary),
237 100 : None => slots.push((op, vec![summary])),
238 : }
239 : }
240 46 : if slots.is_empty() {
241 0 : continue;
242 46 : }
243 46 : let engine_id = format!("{fragment_id}.{pipeline_id}");
244 46 : let pipeline = match pipelines.get(&engine_id) {
245 11 : Some(pipeline) => *pipeline,
246 : None => {
247 35 : let last = slots.len() - 1;
248 35 : let specs = slots
249 35 : .iter()
250 35 : .enumerate()
251 35 : .map(|(i, (operator, _))| SlotSpec {
252 82 : operator: *operator,
253 82 : role: if i == 0 {
254 35 : SlotRole::Source
255 47 : } else if i == last {
256 31 : SlotRole::Sink
257 : } else {
258 16 : SlotRole::Through
259 : },
260 82 : })
261 35 : .collect();
262 35 : events.push(ProfileInputEvent::RegisterPipeline {
263 35 : time,
264 35 : engine_id: Some(engine_id.clone()),
265 35 : fragment: Some(fragment),
266 35 : slots: specs,
267 35 : });
268 35 : let pipeline = next_pipeline;
269 35 : next_pipeline += 1;
270 82 : for (operator, _) in &slots {
271 82 : slot_of.insert((pipeline, *operator), next_slot);
272 82 : next_slot += 1;
273 82 : }
274 35 : pipelines.insert(engine_id, pipeline);
275 35 : pipeline
276 : }
277 : };
278 100 : for (operator, ops) in &slots {
279 100 : let Some(&slot) = slot_of.get(&(pipeline, *operator)) else { continue };
280 100 : events.push(ProfileInputEvent::UpdateSlot { time, slot, status, stats: slot_stats(ops) });
281 : }
282 : }
283 : }
284 15 : events
285 15 : }
286 :
287 : impl QueryProfile {
288 : /// Apply a Trino QueryInfo (JSON, pruned or not), received at a time in milliseconds: the pipelines and the
289 : /// slot statistics of the stages of the plan
290 15 : pub fn apply_trino_query_info(
291 15 : &mut self,
292 15 : info: &[u8],
293 15 : time: u64,
294 15 : ) -> Result<Vec<super::ProfileUpdate>, serde_json::Error> {
295 15 : let info: Value = serde_json::from_slice(info)?;
296 15 : let events = trino_query_info_events(self, &info, time);
297 15 : Ok(self.apply_events(&events))
298 15 : }
299 :
300 : /// Apply the `stats` of a Trino `/v1/statement` response (JSON), received at a time in milliseconds
301 90 : pub fn apply_trino_statement_stats(
302 90 : &mut self,
303 90 : stats: &[u8],
304 90 : time: u64,
305 90 : ) -> Result<Vec<super::ProfileUpdate>, serde_json::Error> {
306 90 : let stats: Value = serde_json::from_slice(stats)?;
307 89 : let events = trino_statement_events(self, &stats, time);
308 89 : Ok(self.apply_events(&events))
309 90 : }
310 : }
311 :
312 : #[cfg(test)]
313 : mod tests {
314 : use super::*;
315 : use serde_json::json;
316 :
317 : #[test]
318 1 : fn durations_of_every_unit() {
319 10 : let us = |text: &str| duration_us(Some(&json!(text)));
320 1 : assert_eq!(us("1500ns"), Some(2));
321 1 : assert_eq!(us("956.59us"), Some(957));
322 1 : assert_eq!(us("2.15ms"), Some(2150));
323 1 : assert_eq!(us(" 2.15s "), Some(2_150_000));
324 1 : assert_eq!(us("1.5m"), Some(90_000_000));
325 1 : assert_eq!(us("2h"), Some(7_200_000_000));
326 1 : assert_eq!(us("1d"), Some(86_400_000_000));
327 : // Unknown units, missing numbers and non-strings are no durations
328 1 : assert_eq!(us("3w"), None);
329 1 : assert_eq!(us("s"), None);
330 1 : assert_eq!(us("12"), None);
331 1 : assert_eq!(duration_us(Some(&json!(12))), None);
332 1 : assert_eq!(duration_us(None), None);
333 1 : }
334 :
335 : #[test]
336 1 : fn data_sizes_of_every_unit() {
337 8 : let bytes = |text: &str| data_size_bytes(Some(&json!(text)));
338 1 : assert_eq!(bytes("271278B"), Some(271_278));
339 1 : assert_eq!(bytes("2kB"), Some(2048));
340 1 : assert_eq!(bytes("1.5MB"), Some(1_572_864));
341 1 : assert_eq!(bytes("2GB"), Some(2 << 30));
342 1 : assert_eq!(bytes("1TB"), Some(1 << 40));
343 1 : assert_eq!(bytes("1PB"), Some(1 << 50));
344 : // Data sizes are binary and case sensitive
345 1 : assert_eq!(bytes("1KB"), None);
346 1 : assert_eq!(bytes("1EB"), None);
347 1 : }
348 :
349 : #[test]
350 1 : fn stage_and_query_states() {
351 1 : assert_eq!(stage_status("PLANNED"), Some(Status::Pending));
352 1 : assert_eq!(stage_status("SCHEDULED"), Some(Status::Running));
353 : // A pending stage after running has done its work
354 1 : assert_eq!(stage_status("PENDING"), Some(Status::Succeeded));
355 1 : assert_eq!(stage_status("ABORTED"), Some(Status::Failed));
356 1 : assert_eq!(stage_status("CANCELED"), Some(Status::Skipped));
357 1 : assert_eq!(stage_status("SOMETHING_NEW"), None);
358 1 : assert_eq!(query_status("FAILED"), Some(Status::Failed));
359 1 : }
360 : }
|
