Line data Source code
1 : //! Unit tests of the execution state: input events, rollups, output events, snapshots and resets.
2 :
3 : use super::{EntityState, ProfileInputEvent, ProfileUpdate, QueryProfile, SlotRole, SlotSpec, Stats, Status};
4 : use dashql_buffers::profile as fb;
5 :
6 : /// A join in two pipelines below an execution target:
7 : /// operators 0 (scan build), 1 (scan probe), 2 (join), 3 (execution target),
8 : /// pipeline 0 = slots 0 (scan build, source), 1 (join, sink),
9 : /// pipeline 1 = slots 2 (scan probe, source), 3 (join, through), 4 (execution target, sink)
10 : const JOIN_PLAN: &str = r#"{
11 : "operator":"executiontarget","operatorId":1,
12 : "input":{"operator":"join","operatorId":2,
13 : "left":{"operator":"tablescan","operatorId":3},
14 : "right":{"operator":"tablescan","operatorId":4}},
15 : "pipelines":[{"pipelineId":1,"operators":[3,2]},{"pipelineId":2,"operators":[4,2,1]}]
16 : }"#;
17 :
18 12 : fn load(plan: &str) -> QueryProfile {
19 12 : let mut model = QueryProfile::new();
20 12 : model.parse_hyper_plan(plan.as_bytes()).expect("plan must parse");
21 12 : model
22 12 : }
23 :
24 15 : fn slot(time: u64, slot: u32, status: Status, stats: Stats) -> ProfileInputEvent {
25 15 : ProfileInputEvent::UpdateSlot { time, slot, status: Some(status), stats }
26 15 : }
27 :
28 10 : fn tuples(tuples_in: u64, tuples_out: u64) -> Stats {
29 10 : Stats { tuples_in: Some(tuples_in), tuples_out: Some(tuples_out), ..Default::default() }
30 10 : }
31 :
32 : /// The state of an entity in a list of updates
33 29 : fn state_of(updates: &[ProfileUpdate], entity: fb::ProfileEntity, index: u32) -> Option<EntityState> {
34 89 : updates.iter().find(|u| u.entity == entity && u.index == index).map(|u| u.state)
35 29 : }
36 :
37 11 : fn snapshot_state(model: &QueryProfile, entity: fb::ProfileEntity, index: u32) -> EntityState {
38 11 : state_of(&model.snapshot_events(), entity, index).unwrap_or_default()
39 11 : }
40 :
41 : #[test]
42 1 : fn an_operator_in_two_pipelines_rolls_up_its_slots() {
43 1 : let mut model = load(JOIN_PLAN);
44 1 : let updates = model.apply_events(&[
45 1 : slot(10, 1, Status::Succeeded, tuples(100, 0)),
46 1 : slot(20, 3, Status::Running, Stats { cpu_time_us: Some(7), ..tuples(40, 30) }),
47 1 : ]);
48 : // The join consumed both inputs and produced the output of its probe side, not of its build side
49 1 : let join = state_of(&updates, fb::ProfileEntity::OPERATOR, 2).unwrap();
50 1 : assert_eq!(join.stats.tuples_in, Some(140));
51 1 : assert_eq!(join.stats.tuples_out, Some(30));
52 1 : assert_eq!(join.stats.cpu_time_us, Some(7));
53 1 : assert_eq!(join.status, Status::Running);
54 : // Every update of a batch carries the time of the batch
55 7 : assert!(updates.iter().all(|u| u.time == 20));
56 : // The build pipeline ends at the join, its output is what reached the sink
57 1 : let build = state_of(&updates, fb::ProfileEntity::PIPELINE, 0).unwrap();
58 1 : assert_eq!((build.status, build.stats.tuples_out), (Status::Succeeded, Some(0)));
59 1 : let probe = state_of(&updates, fb::ProfileEntity::PIPELINE, 1).unwrap();
60 1 : assert_eq!((probe.status, probe.stats.tuples_in, probe.stats.cpu_time_us), (Status::Running, None, Some(7)));
61 : // The fragment and the query see both pipelines
62 1 : let fragment = state_of(&updates, fb::ProfileEntity::FRAGMENT, 0).unwrap();
63 1 : assert_eq!((fragment.status, fragment.stats.cpu_time_us), (Status::Running, Some(7)));
64 1 : let query = state_of(&updates, fb::ProfileEntity::QUERY, 0).unwrap();
65 1 : assert_eq!(query.status, Status::Running);
66 1 : }
67 :
68 : #[test]
69 1 : fn reported_values_win_over_rollups() {
70 1 : let mut model = load(JOIN_PLAN);
71 1 : model.apply_events(&[
72 1 : slot(1, 3, Status::Running, tuples(40, 30)),
73 1 : ProfileInputEvent::UpdateOperator {
74 1 : time: 2,
75 1 : operator: 2,
76 1 : status: Some(Status::Succeeded),
77 1 : stats: Stats { tuples_out: Some(99), ..Default::default() },
78 1 : },
79 1 : ]);
80 1 : let join = snapshot_state(&model, fb::ProfileEntity::OPERATOR, 2);
81 1 : assert_eq!(join.status, Status::Succeeded);
82 1 : assert_eq!(join.stats.tuples_out, Some(99));
83 : // Fields the operator did not report still come from its slots
84 1 : assert_eq!(join.stats.tuples_in, Some(40));
85 1 : }
86 :
87 : #[test]
88 1 : fn unknown_stays_unknown() {
89 1 : let mut model = load(JOIN_PLAN);
90 1 : model.apply_events(&[slot(1, 0, Status::Running, Stats { tuples_out: Some(5), ..Default::default() })]);
91 1 : let scan = snapshot_state(&model, fb::ProfileEntity::OPERATOR, 0);
92 1 : assert_eq!(scan.stats.tuples_in, None);
93 1 : assert_eq!(scan.stats.reported(), fb::ProfileStatsField::TUPLES_OUT.0);
94 : // Entities without any state have no snapshot event
95 1 : assert!(state_of(&model.snapshot_events(), fb::ProfileEntity::OPERATOR, 1).is_none());
96 1 : }
97 :
98 : #[test]
99 1 : fn finished_when_all_children_finished() {
100 1 : let mut model = load(JOIN_PLAN);
101 3 : let at = |started, finished| Stats { started_at: Some(started), finished_at: finished, ..Default::default() };
102 1 : model.apply_events(&[slot(1, 1, Status::Succeeded, at(5, Some(9))), slot(1, 3, Status::Running, at(3, None))]);
103 1 : let join = snapshot_state(&model, fb::ProfileEntity::OPERATOR, 2);
104 1 : assert_eq!((join.stats.started_at, join.stats.finished_at), (Some(3), None));
105 1 : model.apply_events(&[slot(2, 3, Status::Succeeded, at(3, Some(12)))]);
106 1 : let join = snapshot_state(&model, fb::ProfileEntity::OPERATOR, 2);
107 1 : assert_eq!((join.status, join.stats.finished_at), (Status::Succeeded, Some(12)));
108 1 : }
109 :
110 : #[test]
111 1 : fn applying_events_twice_changes_nothing() {
112 1 : let mut model = load(JOIN_PLAN);
113 1 : let events = [slot(1, 1, Status::Running, tuples(10, 0)), slot(1, 3, Status::Running, tuples(4, 3))];
114 1 : assert!(!model.apply_events(&events).is_empty());
115 1 : assert!(model.apply_events(&events).is_empty());
116 1 : }
117 :
118 : #[test]
119 1 : fn a_partial_report_keeps_the_other_fields() {
120 1 : let mut model = load(JOIN_PLAN);
121 1 : model.apply_events(&[slot(1, 0, Status::Running, tuples(10, 8))]);
122 1 : let updates = model.apply_events(&[ProfileInputEvent::UpdateSlot {
123 1 : time: 2,
124 1 : slot: 0,
125 1 : status: None,
126 1 : stats: Stats { tuples_out: Some(9), ..Default::default() },
127 1 : }]);
128 1 : let scan = state_of(&updates, fb::ProfileEntity::SLOT, 0).unwrap();
129 1 : assert_eq!((scan.status, scan.stats.tuples_in, scan.stats.tuples_out), (Status::Running, Some(10), Some(9)));
130 1 : }
131 :
132 : #[test]
133 1 : fn nested_fragments_count_operators_once() {
134 1 : let mut model = load(
135 1 : r#"{"operator":"federate","inputs":[{"operator":"map","input":
136 1 : {"operator":"federate","inputs":[{"operator":"scan"}]}}]}"#,
137 : );
138 1 : let memory = |operator, bytes| ProfileInputEvent::UpdateOperator {
139 : time: 1,
140 4 : operator,
141 4 : status: Some(Status::Succeeded),
142 4 : stats: Stats { memory_bytes: Some(bytes), ..Default::default() },
143 4 : };
144 : // Operators 0 (scan) and 1 (inner federate) belong to the inner fragment, 2 (map) and 3 to the outer one.
145 : // A fragment is its own work, the query all of it.
146 1 : model.apply_events(&[memory(0, 1), memory(1, 2), memory(2, 4), memory(3, 8)]);
147 1 : assert_eq!(snapshot_state(&model, fb::ProfileEntity::FRAGMENT, 1).stats.memory_bytes, Some(3));
148 1 : assert_eq!(snapshot_state(&model, fb::ProfileEntity::FRAGMENT, 0).stats.memory_bytes, Some(12));
149 1 : assert_eq!(snapshot_state(&model, fb::ProfileEntity::QUERY, 0).stats.memory_bytes, Some(15));
150 1 : assert_eq!(snapshot_state(&model, fb::ProfileEntity::QUERY, 0).status, Status::Succeeded);
151 1 : }
152 :
153 : #[test]
154 1 : fn analyzed_plans_are_an_event_batch() {
155 1 : let model = load(
156 1 : r#"{"operator":"executiontarget","statistics":{"output-rows":3},
157 1 : "input":{"operator":"tablescan","statistics":{"processed-rows":10,"output-rows":3}}}"#,
158 : );
159 1 : let scan = snapshot_state(&model, fb::ProfileEntity::OPERATOR, 0);
160 1 : assert_eq!((scan.stats.tuples_in, scan.stats.tuples_out), (Some(10), Some(3)));
161 : // The hidden root fragment takes its output from its root operator
162 1 : assert_eq!(snapshot_state(&model, fb::ProfileEntity::FRAGMENT, 0).stats.tuples_out, Some(3));
163 1 : }
164 :
165 : #[test]
166 1 : fn pipelines_and_fragments_register_at_runtime() {
167 1 : let mut model = load(r#"{"operator":"output","input":{"operator":"filter","input":{"operator":"scan"}}}"#);
168 1 : let version = model.structure_version();
169 1 : let updates = model.apply_events(&[
170 1 : ProfileInputEvent::RegisterFragment { time: 1, engine_id: Some("1".into()), parent_fragment: None },
171 1 : ProfileInputEvent::RegisterPipeline {
172 1 : time: 1,
173 1 : engine_id: Some("0".into()),
174 1 : fragment: Some(0),
175 1 : slots: vec![
176 1 : SlotSpec { operator: 0, role: SlotRole::Source },
177 1 : SlotSpec { operator: 1, role: SlotRole::Through },
178 1 : SlotSpec { operator: 2, role: SlotRole::Sink },
179 1 : ],
180 1 : },
181 1 : slot(2, 1, Status::Running, tuples(10, 4)),
182 1 : ]);
183 1 : assert_eq!(model.structure_version(), version + 2);
184 : // Registered entities without state are no update, the slot and what rolls up from it are
185 1 : assert!(state_of(&updates, fb::ProfileEntity::SLOT, 0).is_none());
186 1 : assert_eq!(state_of(&updates, fb::ProfileEntity::OPERATOR, 1).unwrap().stats.tuples_out, Some(4));
187 1 : assert_eq!(state_of(&updates, fb::ProfileEntity::PIPELINE, 0).unwrap().status, Status::Running);
188 1 : assert_eq!(state_of(&updates, fb::ProfileEntity::FRAGMENT, 0).unwrap().status, Status::Running);
189 :
190 : // The edges inside the pipeline are carried by the slots of their children
191 1 : let buffer = model.pack_buffer();
192 1 : let plan = flatbuffers::root::<fb::QueryProfile>(&buffer).unwrap();
193 1 : assert_eq!(plan.structure_version(), version + 2);
194 2 : let edge_slots: Vec<u32> = plan.operator_edges().unwrap().iter().map(|e| e.slot()).collect();
195 1 : assert_eq!(edge_slots, vec![0, 1]);
196 :
197 : // Events that reference unknown entities are ignored
198 1 : let ignored = model.apply_events(&[
199 1 : ProfileInputEvent::RegisterFragment { time: 3, engine_id: None, parent_fragment: Some(7) },
200 1 : ProfileInputEvent::RegisterPipeline {
201 1 : time: 3,
202 1 : engine_id: None,
203 1 : fragment: None,
204 1 : slots: vec![SlotSpec { operator: 9, role: SlotRole::Source }],
205 1 : },
206 1 : slot(3, 42, Status::Failed, Stats::default()),
207 1 : ]);
208 1 : assert!(ignored.is_empty());
209 1 : assert_eq!(model.structure_version(), version + 2);
210 1 : }
211 :
212 : #[test]
213 1 : fn reset_and_replay_seek_backwards() {
214 1 : let mut model = load(JOIN_PLAN);
215 1 : let version = model.structure_version();
216 1 : let log = vec![
217 1 : ProfileInputEvent::UpdateQuery { time: 1, status: Some(Status::Running), stats: Stats::default() },
218 1 : ProfileInputEvent::RegisterPipeline {
219 1 : time: 2,
220 1 : engine_id: None,
221 1 : fragment: Some(0),
222 1 : slots: vec![SlotSpec { operator: 0, role: SlotRole::Source }],
223 1 : },
224 1 : slot(3, 5, Status::Running, tuples(1, 1)),
225 1 : slot(4, 1, Status::Succeeded, tuples(100, 0)),
226 : ];
227 : // The states, the change times depend on the batches
228 2 : let states = |model: &QueryProfile| -> Vec<_> {
229 16 : model.snapshot_events().into_iter().map(|u| (u.entity, u.index, u.state)).collect()
230 2 : };
231 1 : model.apply_events(&log);
232 1 : let at_end = states(&model);
233 :
234 1 : model.reset_execution();
235 1 : assert!(model.snapshot_events().is_empty());
236 1 : assert_eq!(model.structure_version(), version + 2);
237 1 : assert_eq!(flatbuffers::root::<fb::QueryProfile>(&model.pack_buffer()).unwrap().slots().unwrap().len(), 5);
238 :
239 : // Replaying the log up to a point shows the state of that point, replaying all of it the end state
240 1 : model.apply_events(&log[..1]);
241 1 : assert_eq!(snapshot_state(&model, fb::ProfileEntity::QUERY, 0).status, Status::Running);
242 1 : assert!(state_of(&model.snapshot_events(), fb::ProfileEntity::OPERATOR, 2).is_none());
243 1 : model.apply_events(&log[1..]);
244 1 : assert_eq!(states(&model), at_end);
245 1 : }
246 :
247 : #[test]
248 1 : fn input_events_have_a_json_form() {
249 1 : let json = r#"[
250 1 : {"type":"updateQuery","time":1,"status":"running","stats":{"workDone":3,"workTotal":10}},
251 1 : {"type":"registerPipeline","time":2,"engineId":"7","fragment":0,
252 1 : "slots":[{"operator":0,"role":"source"},{"operator":2,"role":"sink"}]},
253 1 : {"type":"updateSlot","time":3,"slot":5,"stats":{"tuplesOut":4,"cpuTimeUs":12}}
254 1 : ]"#;
255 1 : let events: Vec<ProfileInputEvent> = serde_json::from_str(json).unwrap();
256 1 : assert_eq!(events.len(), 3);
257 1 : let reserialized = serde_json::to_string(&events).unwrap();
258 1 : assert_eq!(serde_json::from_str::<Vec<ProfileInputEvent>>(&reserialized).unwrap(), events);
259 :
260 1 : let mut model = load(JOIN_PLAN);
261 1 : let updates = model.apply_events_json(json.as_bytes()).unwrap();
262 1 : let query = state_of(&updates, fb::ProfileEntity::QUERY, 0).unwrap();
263 1 : assert_eq!((query.status, query.stats.work_done, query.stats.work_total), (Status::Running, Some(3), Some(10)));
264 1 : assert_eq!(state_of(&updates, fb::ProfileEntity::SLOT, 5).unwrap().stats.cpu_time_us, Some(12));
265 :
266 : // Unknown fields are rejected, a typo must not silently drop a counter
267 1 : let typo = r#"[{"type":"updateSlot","time":1,"slot":0,"stats":{"tuples":1}}]"#;
268 1 : assert!(serde_json::from_str::<Vec<ProfileInputEvent>>(typo).is_err());
269 1 : }
270 :
271 : #[test]
272 1 : fn packed_events_carry_the_reported_fields() {
273 1 : let mut model = load(JOIN_PLAN);
274 1 : let updates = model.apply_events(&[slot(5, 0, Status::Running, tuples(10, 8))]);
275 1 : let buffer = model.pack_events(&updates);
276 1 : let events = flatbuffers::root::<fb::QueryProfileEvents>(&buffer).unwrap();
277 1 : assert_eq!(events.structure_version(), model.structure_version());
278 4 : let event = events.events().unwrap().iter().find(|e| e.entity() == fb::ProfileEntity::SLOT).unwrap();
279 1 : assert_eq!((event.index(), event.event_time(), event.status()), (0, 5, fb::ProfileExecutionStatus::RUNNING));
280 1 : let stats = event.stats();
281 1 : assert_eq!(stats.reported(), fb::ProfileStatsField::TUPLES_IN.0 | fb::ProfileStatsField::TUPLES_OUT.0);
282 1 : assert_eq!((stats.tuples_in(), stats.tuples_out()), (10, 8));
283 1 : }
284 :
285 : #[test]
286 1 : fn trino_stages_register_once_and_tolerate_missing_fields() {
287 1 : let mut model = QueryProfile::new();
288 1 : model.start_query(fb::ProfileEngine::TRINO);
289 2 : let poll = |state: &str, completed: u64| {
290 2 : format!(
291 : r#"{{"state":"RUNNING","completedSplits":{completed},"totalSplits":4,
292 : "rootStage":{{"stageId":"0","state":"{state}",
293 : "subStages":[{{"stageId":"1","state":"FINISHED","cpuTimeMillis":2}}]}}}}"#
294 : )
295 2 : };
296 1 : let version = model.structure_version();
297 1 : let updates = model.apply_trino_statement_stats(poll("RUNNING", 1).as_bytes(), 10).unwrap();
298 1 : assert_eq!(model.structure_version(), version + 2);
299 1 : let query = state_of(&updates, fb::ProfileEntity::QUERY, 0).unwrap();
300 1 : assert_eq!((query.status, query.stats.work_done, query.stats.work_total), (Status::Running, Some(1), Some(4)));
301 : // Fields the payload lacks stay unknown
302 1 : assert_eq!(query.stats.memory_bytes, None);
303 1 : let stage = state_of(&updates, fb::ProfileEntity::FRAGMENT, 1).unwrap();
304 1 : assert_eq!((stage.status, stage.stats.cpu_time_us), (Status::Succeeded, Some(2000)));
305 :
306 : // The next poll knows the stages already
307 1 : let updates = model.apply_trino_statement_stats(poll("FINISHED", 4).as_bytes(), 20).unwrap();
308 1 : assert_eq!(model.structure_version(), version + 2);
309 1 : assert_eq!(state_of(&updates, fb::ProfileEntity::FRAGMENT, 0).unwrap().status, Status::Succeeded);
310 1 : assert!(state_of(&updates, fb::ProfileEntity::FRAGMENT, 1).is_none());
311 :
312 : // A poll without stage tree updates the query only, an unknown state keeps the status
313 1 : let updates = model.apply_trino_statement_stats(br#"{"state":"SOMETHING_NEW","completedSplits":4}"#, 30).unwrap();
314 1 : assert!(updates.is_empty());
315 1 : assert!(model.apply_trino_statement_stats(b"not json", 40).is_err());
316 :
317 1 : let buffer = model.pack_buffer();
318 1 : let plan = flatbuffers::root::<fb::QueryProfile>(&buffer).unwrap();
319 1 : assert_eq!(plan.engine(), fb::ProfileEngine::TRINO);
320 1 : let child = plan.fragments().unwrap().get(1);
321 1 : let engine_id = plan.string_dictionary().unwrap().get(child.engine_id() as usize);
322 1 : assert_eq!((engine_id, child.parent_fragment()), ("1", 0));
323 1 : }
|
