gracefull shutdown of SSE loop

This commit is contained in:
Djeeberjr 2025-05-27 14:06:38 +02:00
parent e527d41cb7
commit 2150fff6a5

View File

@ -2,7 +2,7 @@ use log::{error, info, warn};
use rocket::http::Status; use rocket::http::Status;
use rocket::response::stream::{Event, EventStream}; use rocket::response::stream::{Event, EventStream};
use rocket::serde::json::Json; use rocket::serde::json::Json;
use rocket::{Config, State, post}; use rocket::{Config, Shutdown, State, post};
use rocket::{get, http::ContentType, response::content::RawHtml, routes}; use rocket::{get, http::ContentType, response::content::RawHtml, routes};
use rust_embed::Embed; use rust_embed::Embed;
use serde::Deserialize; use serde::Deserialize;
@ -10,6 +10,7 @@ use std::borrow::Cow;
use std::env; use std::env;
use std::ffi::OsStr; use std::ffi::OsStr;
use std::sync::Arc; use std::sync::Arc;
use tokio::select;
use tokio::sync::Mutex; use tokio::sync::Mutex;
use tokio::sync::broadcast::Sender; use tokio::sync::broadcast::Sender;
@ -84,15 +85,22 @@ fn static_files(file: std::path::PathBuf) -> Option<(ContentType, Vec<u8>)> {
} }
#[get("/api/idevent")] #[get("/api/idevent")]
fn id_event(sse_broadcaster: &State<Sender<String>>) -> EventStream![] { fn id_event(sse_broadcaster: &State<Sender<String>>, shutdown: Shutdown) -> EventStream![] {
let mut rx = sse_broadcaster.subscribe(); let mut rx = sse_broadcaster.subscribe();
EventStream! { EventStream! {
loop { loop {
let msg = rx.recv().await; select! {
msg = rx.recv() => {
if let Ok(id) = msg { if let Ok(id) = msg {
yield Event::data(id); yield Event::data(id);
} }
} }
_ = &mut shutdown.clone() => {
// Shutdown signal received, exit the loop
break;
}
}
}
} }
} }