1use std::convert::Infallible;
17use std::time::Duration;
18use std::time::Instant;
19
20use async_stream::stream;
21use axum::Router;
22use axum::extract::{Path, State};
23use axum::response::IntoResponse;
24use axum::response::sse::{Event, KeepAlive, Sse};
25use axum::routing::get;
26use bhtune_db::models::{TuneOutcome, TuneRunRow, TuneSampleRow};
27use serde::Serialize;
28use utoipa::ToSchema;
29
30use crate::error::ApiError;
31use crate::routes::history::{InitialReadingsResponse, SampleResponse};
32use crate::state::AppState;
33
34const POLL_INTERVAL: Duration = Duration::from_millis(300);
40
41#[derive(Debug, Serialize, ToSchema)]
47pub(crate) struct RunStreamDone {
48 outcome: TuneOutcome,
49}
50
51fn ok_event(event: Event) -> Result<Event, Infallible> {
58 Ok(event)
59}
60
61#[utoipa::path(
73 get,
74 path = "/api/runs/{id}/stream",
75 tag = "runs",
76 params(
77 ("id" = i64, Path, description = "Run id"),
78 ),
79 responses(
80 (
81 status = 200,
82 description = "A `text/event-stream` with an optional `initial` event \
83 (data: InitialReadingsResponse), `sample` events (data: SampleResponse), \
84 and one final `done` event (data: RunStreamDone).",
85 content_type = "text/event-stream",
86 body = SampleResponse,
87 ),
88 (status = 404, description = "No run with that id.", body = crate::error::ErrorBody),
89 ),
90)]
91pub(crate) async fn stream_run(
92 State(state): State<AppState>,
93 Path(run_id): Path<i64>,
94) -> Result<impl IntoResponse, ApiError> {
95 stream_run_with_permit(state, run_id, None, None).await
96}
97
98pub(crate) async fn stream_run_with_permit(
99 state: AppState,
100 run_id: i64,
101 permit: Option<Vec<tokio::sync::OwnedSemaphorePermit>>,
102 timeout: Option<Duration>,
103) -> Result<impl IntoResponse, ApiError> {
104 if TuneRunRow::get(&state.pool, run_id).await?.is_none() {
105 return Err(ApiError::NotFound(format!("no run with id {run_id}")));
106 }
107
108 let pool = state.pool.clone();
109 let events = stream! {
110 let _permit = permit;
111 let started = Instant::now();
112 let mut last_tick: i64 = -1;
116 let mut sent_initial = false;
117 loop {
118 if timeout.is_some_and(|limit| started.elapsed() >= limit) {
119 yield ok_event(Event::default().event("error").data("SSE demo time limit exceeded"));
120 break;
121 }
122 let run_outcome = match TuneRunRow::get(&pool, run_id).await {
123 Ok(Some(run)) => {
124 if !sent_initial && let Some(initial) = run.initial_readings.as_ref() {
125 let response = InitialReadingsResponse::from(initial.clone());
126 match Event::default().event("initial").json_data(response) {
127 Ok(event) => {
128 sent_initial = true;
129 yield ok_event(event);
130 }
131 Err(err) => tracing::error!(
132 run_id,
133 error = %err,
134 "failed to encode initial readings as an SSE event"
135 ),
136 }
137 }
138 Some(run.outcome)
139 }
140 Ok(_) => None,
144 Err(err) => {
145 tracing::error!(
146 run_id,
147 error = %err,
148 "failed to poll tune_runs for the run stream; ending the stream"
149 );
150 yield ok_event(Event::default().event("error").data(err.to_string()));
151 break;
152 }
153 };
154
155 match TuneSampleRow::list_for_run_since(&pool, run_id, last_tick).await {
156 Ok(samples) => {
157 for sample in &samples {
158 last_tick = sample.tick_index;
159 let response = SampleResponse::from(sample);
160 match Event::default().event("sample").json_data(response) {
161 Ok(event) => yield ok_event(event),
162 Err(err) => tracing::error!(
163 run_id,
164 error = %err,
165 "failed to encode a tune sample as an SSE event"
166 ),
167 }
168 }
169 }
170 Err(err) => {
171 tracing::error!(
172 run_id,
173 error = %err,
174 "failed to poll tune_samples for the run stream; ending the stream"
175 );
176 yield ok_event(Event::default().event("error").data(err.to_string()));
177 break;
178 }
179 }
180
181 if let Some(outcome) = run_outcome
182 && outcome != TuneOutcome::Running
183 {
184 let done = RunStreamDone { outcome };
185 if let Ok(event) = Event::default().event("done").json_data(done) {
186 yield ok_event(event);
187 }
188 break;
189 }
190
191 tokio::time::sleep(POLL_INTERVAL).await;
192 }
193 };
194
195 Ok(Sse::new(events).keep_alive(KeepAlive::default()))
196}
197
198pub fn router() -> Router<AppState> {
199 Router::new().route("/api/runs/{id}/stream", get(stream_run))
200}
201
202#[cfg(test)]
203mod tests {
204 use super::*;
205 use axum::body::{Body, to_bytes};
206 use axum::http::{Request, StatusCode};
207 use bhtune_core::{ControllerType, LoopConfig, LoopTags, MrftState, ProcessType, Tick};
208 use bhtune_db::models::{SampleQuality, TuneDriver, TuneRunInitialReadings};
209 use chrono::Utc;
210 use tower::ServiceExt;
211
212 async fn seed_running_run(state: &AppState, name: &str) -> i64 {
213 let template_row =
214 bhtune_db::models::DcsTemplateRow::get_by_name(&state.pool, "Yokogawa CentumVP")
215 .await
216 .unwrap()
217 .unwrap();
218 let template = template_row.template;
219 let config = LoopConfig {
220 process_type: ProcessType::Flow,
221 controller_type: ControllerType::Pi,
222 relay_amp_percent: 5.0,
223 num_cycles_skip: 1,
224 num_cycles_count: 3,
225 noise_protection_secs: 0,
226 mrft_delay_secs: 0,
227 };
228 let tags = LoopTags::derive_from_pv_tag(&format!("{name}.PV"), &template);
229 let run = TuneRunRow::start(
230 &state.pool,
231 None,
232 name,
233 TuneDriver::Simulator,
234 config,
235 template_row.origin,
236 &template,
237 &tags,
238 Utc::now(),
239 )
240 .await
241 .unwrap();
242 run.id
243 }
244
245 async fn insert_sample(state: &AppState, run_id: i64, tick: i64) {
246 let sample = Tick {
247 time: Utc::now(),
248 pv: 50.0 + tick as f32,
249 };
250 let mrft_state = MrftState {
251 hysteresis: 1.0,
252 mv_value_current: 55.0,
253 mv_sign_next_step: 1,
254 counter_all_switches: tick as u32,
255 cycles_completed: 0,
256 cycles_remaining: 2,
257 };
258 TuneSampleRow::insert(
259 &state.pool,
260 run_id,
261 tick,
262 sample,
263 mrft_state,
264 SampleQuality::Good,
265 )
266 .await
267 .unwrap();
268 }
269
270 async fn record_initial_readings(state: &AppState, run_id: i64) {
271 TuneRunRow::record_initial_readings(
272 &state.pool,
273 run_id,
274 TuneRunInitialReadings {
275 pv_ini: 48.0,
276 mv_ini: 42.0,
277 mv_range_low: 0.0,
278 mv_range_high: 100.0,
279 pv_range_high: 100.0,
280 pv_range_low: 0.0,
281 controller_direction: bhtune_core::ControllerDirection::Reverse,
282 mode_raw: Some("AUTO".to_string()),
283 mode_attribute_raw: None,
284 setpoint_ini: Some(50.0),
285 },
286 )
287 .await
288 .unwrap();
289 }
290
291 fn parse_sse(body: &str) -> Vec<(String, String)> {
297 body.split("\n\n")
298 .filter(|chunk| !chunk.trim().is_empty())
299 .map(|chunk| {
300 let mut event = String::new();
301 let mut data = String::new();
302 for line in chunk.lines() {
303 if let Some(rest) = line.strip_prefix("event:") {
304 event = rest.trim().to_string();
305 } else if let Some(rest) = line.strip_prefix("data:") {
306 data = rest.trim().to_string();
307 }
308 }
309 (event, data)
310 })
311 .collect()
312 }
313
314 #[test]
315 fn sse_parser_ignores_unrecognized_lines_between_event_and_data() {
316 assert_eq!(
317 parse_sse("event: ping\nid: 42\ndata: payload\n\n"),
318 vec![("ping".to_string(), "payload".to_string())]
319 );
320 }
321
322 #[tokio::test]
323 async fn streaming_an_unknown_run_returns_404() {
324 let state = crate::test_support::in_memory_state().await;
325 let app = crate::build_router(state);
326 let response = app
327 .oneshot(
328 Request::get("/api/runs/999/stream")
329 .body(Body::empty())
330 .unwrap(),
331 )
332 .await
333 .unwrap();
334 assert_eq!(response.status(), StatusCode::NOT_FOUND);
335 }
336
337 #[tokio::test]
338 async fn streaming_after_its_deadline_emits_an_error_event_and_closes() {
339 let state = crate::test_support::in_memory_state().await;
340 let run_id = seed_running_run(&state, "LIC-STREAM-TIMEOUT").await;
341 let response = stream_run_with_permit(state, run_id, None, Some(Duration::ZERO))
342 .await
343 .unwrap()
344 .into_response();
345
346 let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
347 let events = parse_sse(core::str::from_utf8(&bytes).unwrap());
348
349 assert_eq!(
350 events,
351 vec![(
352 "error".to_string(),
353 "SSE demo time limit exceeded".to_string()
354 )]
355 );
356 }
357
358 #[tokio::test]
359 async fn streaming_a_completed_run_replays_every_sample_then_emits_done() {
360 let state = crate::test_support::in_memory_state().await;
361 let run_id = seed_running_run(&state, "LIC-STREAM-1").await;
362 insert_sample(&state, run_id, 0).await;
363 insert_sample(&state, run_id, 1).await;
364 TuneRunRow::complete(&state.pool, run_id, Utc::now())
365 .await
366 .unwrap();
367
368 let app = crate::build_router(state);
369 let response = app
370 .oneshot(
371 Request::get(format!("/api/runs/{run_id}/stream"))
372 .body(Body::empty())
373 .unwrap(),
374 )
375 .await
376 .unwrap();
377 assert_eq!(response.status(), StatusCode::OK);
378 assert_eq!(
379 response
380 .headers()
381 .get(axum::http::header::CONTENT_TYPE)
382 .unwrap(),
383 "text/event-stream"
384 );
385
386 let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
387 let body = String::from_utf8(bytes.to_vec()).unwrap();
388 let events = parse_sse(&body);
389
390 let sample_events: Vec<_> = events.iter().filter(|(e, _)| e == "sample").collect();
391 assert_eq!(
392 sample_events.len(),
393 2,
394 "both pre-recorded samples must be replayed"
395 );
396 let first: serde_json::Value = serde_json::from_str(&sample_events[0].1).unwrap();
397 assert_eq!(first["tick_index"], 0);
398 let second: serde_json::Value = serde_json::from_str(&sample_events[1].1).unwrap();
399 assert_eq!(second["tick_index"], 1);
400
401 let (last_event, last_data) = events.last().expect("stream must emit at least `done`");
402 assert_eq!(last_event, "done");
403 let done: serde_json::Value = serde_json::from_str(last_data).unwrap();
404 assert_eq!(done["outcome"], "completed");
405 }
406
407 #[tokio::test]
408 async fn streaming_emits_initial_readings_before_replayed_samples() {
409 let state = crate::test_support::in_memory_state().await;
410 let run_id = seed_running_run(&state, "LIC-STREAM-INITIAL").await;
411 record_initial_readings(&state, run_id).await;
412 insert_sample(&state, run_id, 0).await;
413 TuneRunRow::complete(&state.pool, run_id, Utc::now())
414 .await
415 .unwrap();
416
417 let app = crate::build_router(state);
418 let response = app
419 .oneshot(
420 Request::get(format!("/api/runs/{run_id}/stream"))
421 .body(Body::empty())
422 .unwrap(),
423 )
424 .await
425 .unwrap();
426 let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
427 let events = parse_sse(core::str::from_utf8(&bytes).unwrap());
428
429 assert_eq!(events[0].0, "initial");
430 let initial: serde_json::Value = serde_json::from_str(&events[0].1).unwrap();
431 assert_eq!(initial["pv_ini"], 48.0);
432 assert_eq!(initial["mv_ini"], 42.0);
433 assert_eq!(events[1].0, "sample");
434 assert_eq!(events[2].0, "done");
435 }
436
437 #[tokio::test]
438 async fn streaming_a_run_with_no_samples_still_terminates_with_done() {
439 let state = crate::test_support::in_memory_state().await;
440 let run_id = seed_running_run(&state, "LIC-STREAM-2").await;
441 TuneRunRow::abort(&state.pool, run_id, Utc::now())
442 .await
443 .unwrap();
444
445 let app = crate::build_router(state);
446 let response = app
447 .oneshot(
448 Request::get(format!("/api/runs/{run_id}/stream"))
449 .body(Body::empty())
450 .unwrap(),
451 )
452 .await
453 .unwrap();
454 let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
455 let events = parse_sse(core::str::from_utf8(&bytes).unwrap());
456 assert_eq!(
457 events.len(),
458 1,
459 "no samples were recorded, so only `done` is emitted"
460 );
461 assert_eq!(events[0].0, "done");
462 let done: serde_json::Value = serde_json::from_str(&events[0].1).unwrap();
463 assert_eq!(done["outcome"], "aborted");
464 }
465
466 #[tokio::test]
467 async fn streaming_a_still_running_run_waits_for_it_to_finish_before_closing() {
468 let state = crate::test_support::in_memory_state().await;
469 let run_id = seed_running_run(&state, "LIC-STREAM-3").await;
470 insert_sample(&state, run_id, 0).await;
471 let pool = state.pool.clone();
478 tokio::spawn(async move {
479 tokio::time::sleep(Duration::from_millis(400)).await;
480 TuneRunRow::complete(&pool, run_id, Utc::now())
481 .await
482 .unwrap();
483 });
484
485 let app = crate::build_router(state);
486 let response = tokio::time::timeout(
487 Duration::from_secs(5),
488 app.oneshot(
489 Request::get(format!("/api/runs/{run_id}/stream"))
490 .body(Body::empty())
491 .unwrap(),
492 ),
493 )
494 .await
495 .expect("the stream must eventually close on its own")
496 .unwrap();
497
498 let bytes = tokio::time::timeout(
499 Duration::from_secs(5),
500 to_bytes(response.into_body(), usize::MAX),
501 )
502 .await
503 .expect("reading the full (finite) SSE body must not hang")
504 .unwrap();
505 let events = parse_sse(core::str::from_utf8(&bytes).unwrap());
506
507 assert!(
508 events
509 .iter()
510 .any(|(e, d)| e == "sample" && d.contains("\"tick_index\":0")),
511 "the pre-recorded sample must have been replayed: {events:?}"
512 );
513 let (last_event, last_data) = events.last().unwrap();
514 assert_eq!(last_event, "done");
515 let done: serde_json::Value = serde_json::from_str(last_data).unwrap();
516 assert_eq!(done["outcome"], "completed");
517 }
518}