diff --git a/CHANGELOG.md b/CHANGELOG.md index 403e5cf23..b150bd9c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -166,6 +166,9 @@ active URL of the specified provider. item now stores retry/cooldown state for `resolve`, `probe`, and `tmdb`. - **TMDB No-Match Cooldown**: Added explicit TMDB cooldown handling. When TMDB resolve completes successfully but returns no match, TMDB reasons are suppressed for that item during cooldown to prevent endless requeue loops. +- **Custom Stream Response Timeout**: Added support to limit how long custom fallback stream responses are served. + Set `config.custom_stream_response_timeout_secs` to a value `> 0` to auto-stop these streams after N seconds. + If unset or `0`, custom responses are streamed without timeout. ## ⚙️ New Settings @@ -173,6 +176,8 @@ active URL of the specified provider. - Added `priority` (`i8`, default `0`) to user credentials. Lower value = higher priority (nice-style). Configurable via Web UI user editor. Negative values are valid and represent higher-than-default priority. - **config.yml**: + - Added `custom_stream_response_timeout_secs` (`u32`, default `0`): maximum duration in seconds for custom stream response videos. + `0` disables the timeout and keeps existing behavior. - Added `metadata_update.probe.user_priority` (`i8`, default `127`): priority assigned to probe connections. Probe tasks run at the lowest priority by default; reduce this value to give probes more connection access. - Added `metadata_update` (optional) with grouped sections: `log`, `resolve`, `probe`, `ffprobe`, `tmdb`. diff --git a/README.md b/README.md index 74a5aaad1..3cb7d8d00 100644 --- a/README.md +++ b/README.md @@ -131,6 +131,8 @@ Top level entries in the config files are: - `user_access_control` _optional_ - `connect_timeout_secs`: _optional_ and used for provider requests connection timeout, for only the connect phase. - `custom_stream_response_path` _optional_ +- `custom_stream_response_timeout_secs` _optional_, max duration in seconds for serving custom stream response videos. + If unset or `0`, custom stream responses are served without timeout. - `hdhomerun` _optional_ - `proxy` _optional_ - `ipcheck` _optional_ @@ -1080,6 +1082,16 @@ and add it to the `config.yml`. custom_stream_response_path: /home/tuliprox/resources ``` +Optional timeout for custom stream responses: + +- `custom_stream_response_timeout_secs`: if set to `> 0`, custom response videos are stopped after the given number of seconds. +- If unset or set to `0`, custom response videos are served as-is (no timeout). + +```yaml +custom_stream_response_path: /home/tuliprox/resources +custom_stream_response_timeout_secs: 20 +``` + ### 1.15 `user_config_dir` It is the storage path for user configurations (f.e. bouquets). diff --git a/backend/Cargo.toml b/backend/Cargo.toml index dca732b11..4a20c0b08 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -45,7 +45,7 @@ flate2 = "1.1.9" blake3 = "1.8.3" bytes = "1.11.1" tokio-stream = { version = "0.1.18", features = ["sync"] } -tokio = { version = "1.49.0", features = ["rt-multi-thread", "parking_lot", "fs", "process"] } +tokio = { version = "1.49.0", features = ["rt-multi-thread", "parking_lot", "fs", "process", "signal"] } #console-subscriber = "0" #tracing = "0.1" #tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } diff --git a/backend/src/api/api_utils.rs b/backend/src/api/api_utils.rs index 50894b3e5..7f139afcd 100644 --- a/backend/src/api/api_utils.rs +++ b/backend/src/api/api_utils.rs @@ -488,6 +488,10 @@ async fn create_stream_response_details( &streaming_strategy.provider_stream_state, grace_period_options.period_millis, ); + let provider_grace_active = matches!( + streaming_strategy.provider_stream_state, + ProviderStreamState::GracePeriod(_, _) + ); let guard_provider_name = streaming_strategy.provider_handle.as_ref().and_then(|guard| guard.allocation.get_provider_name()); @@ -522,6 +526,7 @@ async fn create_stream_response_details( provider_name: guard_provider_name.clone(), request_url: None, grace_period: grace_period_options, + provider_grace_active: false, disable_provider_grace: false, reconnect_flag: None, provider_handle: streaming_strategy.provider_handle.clone(), @@ -534,39 +539,57 @@ async fn create_stream_response_details( sanitize_sensitive_info(guard_provider_name.as_deref().unwrap_or("?")), sanitize_sensitive_info(resolve_request_url_for_logging(input, request_url.as_ref()).as_ref()) ); - let parsed_url = Url::parse(&request_url); - let ((stream, stream_info), reconnect_flag) = if let Ok(url) = parsed_url { - let default_user_agent = app_state.app_config.config.load().default_user_agent.clone(); - let disabled_headers = app_state.get_disabled_headers(); - let mut provider_stream_factory_options = ProviderStreamFactoryOptions::new( - fingerprint.addr, - item_type, - share_stream, - stream_options, - &url, - req_headers, - streaming_strategy.input_headers.as_ref(), - disabled_headers.as_ref(), - default_user_agent.as_deref(), - ); - - let provider_config = input.get_resolve_provider(url.as_ref()); - provider_stream_factory_options.set_provider(provider_config); - - let reconnect_flag = provider_stream_factory_options.get_reconnect_flag_clone(); - let provider_stream = match create_provider_stream( - app_state, - &app_state.http_client.load(), - provider_stream_factory_options, - ) - .await - { - None => (None, None), - Some((stream, info)) => (Some(stream), info), - }; - (provider_stream, Some(reconnect_flag)) + let defer_provider_stream_until_grace_check = if provider_grace_active && grace_period_options.hold_stream { + if let Some(provider_name) = guard_provider_name.as_ref() { + app_state.active_provider.is_over_limit(provider_name).await + } else { + false + } } else { - ((None, None), None) + false + }; + let (stream, stream_info, reconnect_flag) = if defer_provider_stream_until_grace_check { + debug_if_enabled!( + "Deferring provider stream open until grace check completes for {}", + sanitize_sensitive_info(resolve_request_url_for_logging(input, request_url.as_ref()).as_ref()) + ); + (None, None, None) + } else { + let parsed_url = Url::parse(&request_url); + let ((stream, stream_info), reconnect_flag) = if let Ok(url) = parsed_url { + let default_user_agent = app_state.app_config.config.load().default_user_agent.clone(); + let disabled_headers = app_state.get_disabled_headers(); + let mut provider_stream_factory_options = ProviderStreamFactoryOptions::new( + fingerprint.addr, + item_type, + share_stream, + stream_options, + &url, + req_headers, + streaming_strategy.input_headers.as_ref(), + disabled_headers.as_ref(), + default_user_agent.as_deref(), + ); + + let provider_config = input.get_resolve_provider(url.as_ref()); + provider_stream_factory_options.set_provider(provider_config); + + let reconnect_flag = provider_stream_factory_options.get_reconnect_flag_clone(); + let provider_stream = match create_provider_stream( + app_state, + &app_state.http_client.load(), + provider_stream_factory_options, + ) + .await + { + None => (None, None), + Some((stream, info)) => (Some(stream), info), + }; + (provider_stream, Some(reconnect_flag)) + } else { + ((None, None), None) + }; + (stream, stream_info, reconnect_flag) }; if log_enabled!(log::Level::Debug) { @@ -580,8 +603,9 @@ async fn create_stream_response_details( } } - // if we have no stream, we should release the provider - let provider_handle = if stream.is_none() { + // if we have no stream, we should release the provider unless we intentionally defer + // opening it for provider-grace hold_stream behavior. + let provider_handle = if stream.is_none() && !defer_provider_stream_until_grace_check { let provider_handle = streaming_strategy.provider_handle.take(); app_state.connection_manager.release_provider_handle(provider_handle).await; error!("Can't open stream {}", sanitize_sensitive_info(&request_url)); @@ -596,6 +620,7 @@ async fn create_stream_response_details( provider_name: guard_provider_name.clone(), request_url: Some(request_url.clone()), grace_period: grace_period_options, + provider_grace_active, disable_provider_grace: false, reconnect_flag, provider_handle, @@ -810,7 +835,10 @@ pub async fn force_provider_stream_response( } }; - if stream_details.has_stream() { + let deferred_grace_hold_stream = + !stream_details.has_stream() && stream_details.provider_grace_active && stream_details.grace_period.hold_stream; + + if stream_details.has_stream() || deferred_grace_hold_stream { let provider_response = stream_details.stream_info.as_ref().map(|(h, sc, url, cvt)| (h.clone(), *sc, url.clone(), *cvt)); app_state.active_users.update_session_addr(&user.username, &user_session.token, &fingerprint.addr).await; @@ -883,14 +911,18 @@ pub async fn stream_response( &fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted, ) - .await .into_response(); } let virtual_id = stream_channel.virtual_id; let item_type = stream_channel.item_type; - let share_stream = is_stream_share_enabled(item_type, target); + // Grace candidates are intentionally not shared: + // they are transient and may be terminated after the grace check. + // Keeping them out of the shared-stream lifecycle avoids cross-effects on + // concurrently running streams of the same user. + let share_stream = + is_stream_share_enabled(item_type, target) && connection_permission != UserConnectionPermission::GracePeriod; let _shared_lock = if share_stream { let write_lock = app_state.app_config.file_locks.write_lock_str(stream_url).await; @@ -960,7 +992,10 @@ pub async fn stream_response( } }; - if stream_details.has_stream() { + let deferred_grace_hold_stream = + !stream_details.has_stream() && stream_details.provider_grace_active && stream_details.grace_period.hold_stream; + + if stream_details.has_stream() || deferred_grace_hold_stream { // let content_length = get_stream_content_length(provider_response.as_ref()); let provider_response = stream_details .stream_info @@ -982,14 +1017,16 @@ pub async fn stream_response( debug_if_enabled!("panel_api provisioning response to client: status={} headers={:?}", status, headers); } - let provider_handle = if share_stream { stream_details.provider_handle.take() } else { None }; - let mut is_stream_shared = share_stream; + if stream_details.provider_grace_active { + is_stream_shared = false; + } if let Some((_header, _status_code, _url, Some(_custom_video))) = stream_details.stream_info.as_ref() { if stream_details.stream.is_some() { is_stream_shared = false; } } + let provider_handle = if is_stream_shared { stream_details.provider_handle.take() } else { None }; stream_channel.shared = is_stream_shared; let stream = create_active_client_stream( @@ -1137,6 +1174,10 @@ async fn try_shared_stream_response_if_any( session_token: &str, req_headers: &HeaderMap, ) -> Option { + if connect_permission == UserConnectionPermission::GracePeriod { + return None; + } + if let Some((stream, provider)) = SharedStreamManager::subscribe_shared_stream(app_state, stream_url, &fingerprint.addr).await { @@ -1195,7 +1236,6 @@ pub async fn local_stream_response( &fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted, ) - .await .into_response(); } diff --git a/backend/src/api/endpoints/custom_video_stream_api.rs b/backend/src/api/endpoints/custom_video_stream_api.rs index 1705d74e6..6990aeda6 100644 --- a/backend/src/api/endpoints/custom_video_stream_api.rs +++ b/backend/src/api/endpoints/custom_video_stream_api.rs @@ -24,7 +24,7 @@ async fn cvs_api( return axum::http::StatusCode::FORBIDDEN.into_response(); } - create_custom_video_stream_response(&app_state, &fingerprint.addr, custom_video_type).await.into_response() + create_custom_video_stream_response(&app_state, &fingerprint.addr, custom_video_type).into_response() } pub fn cvs_api_register() -> axum::Router> { diff --git a/backend/src/api/endpoints/hls_api.rs b/backend/src/api/endpoints/hls_api.rs index dc7761c51..1a078fe47 100644 --- a/backend/src/api/endpoints/hls_api.rs +++ b/backend/src/api/endpoints/hls_api.rs @@ -301,7 +301,6 @@ async fn hls_api_stream( &fingerprint.addr, CustomVideoStreamType::UserAccountExpired, ) - .await .into_response(); } @@ -325,7 +324,6 @@ async fn hls_api_stream( &fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted, ) - .await .into_response(); } @@ -335,7 +333,6 @@ async fn hls_api_stream( &fingerprint.addr, CustomVideoStreamType::ProviderConnectionsExhausted, ) - .await .into_response(); } @@ -374,7 +371,6 @@ async fn hls_api_stream( &fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted, ) - .await .into_response(); } diff --git a/backend/src/api/endpoints/m3u_api.rs b/backend/src/api/endpoints/m3u_api.rs index 1273b848e..cf2030333 100644 --- a/backend/src/api/endpoints/m3u_api.rs +++ b/backend/src/api/endpoints/m3u_api.rs @@ -92,7 +92,6 @@ async fn m3u_api_stream( &fingerprint.addr, CustomVideoStreamType::UserAccountExpired, ) - .await .into_response(); } @@ -153,7 +152,6 @@ async fn m3u_api_stream( &fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted, ) - .await .into_response(); } @@ -163,7 +161,6 @@ async fn m3u_api_stream( &fingerprint.addr, CustomVideoStreamType::ProviderConnectionsExhausted, ) - .await .into_response(); } if session.virtual_id == virtual_id && is_seek_request(cluster, req_headers).await { @@ -192,7 +189,6 @@ async fn m3u_api_stream( &fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted, ) - .await .into_response(); } diff --git a/backend/src/api/endpoints/v1_api_playlist.rs b/backend/src/api/endpoints/v1_api_playlist.rs index 627be92a2..9a518828d 100644 --- a/backend/src/api/endpoints/v1_api_playlist.rs +++ b/backend/src/api/endpoints/v1_api_playlist.rs @@ -11,7 +11,6 @@ use crate::{ }, auth::create_access_token, model::{parse_xmltv_for_web_ui_from_url, ConfigInput, ConfigInputFlags, ConfigInputFlagsSet, ConfigInputOptions}, - processing::processor::exec_processing, repository::xtream_get_item_for_stream_id, }; use axum::{response::IntoResponse, Router}; @@ -79,35 +78,21 @@ async fn playlist_update( let process_targets = app_state.app_config.sources.load().validate_targets(user_targets.as_ref()); match process_targets { Ok(valid_targets) => { - let http_client = app_state.http_client.load().as_ref().clone(); - let app_config = Arc::clone(&app_state.app_config); - let event_manager = Arc::clone(&app_state.event_manager); - let playlist_state = Arc::clone(&app_state.playlists); let valid_targets = Arc::new(valid_targets); - let provider_manager = Arc::clone(&app_state.active_provider); - let disabled_headers = app_state.get_disabled_headers(); - let metadata_manager = Arc::clone(&app_state.metadata_manager); - let update_guard = app_state.update_guard.clone(); - tokio::spawn({ - async move { - exec_processing( - &http_client, - app_config, - valid_targets, - Some(event_manager), - Some(app_state.clone()), - Some(playlist_state), - Some(update_guard), - disabled_headers, - Some(provider_manager), - Some(metadata_manager), - None, - None, - ) - .await; + // Deduplicate rapid clicks: the channel has capacity 1, so at most one + // update is queued at any time. Additional requests while the channel + // is full are silently dropped — the pending run already covers them. + match app_state.manual_update_sender.try_send(valid_targets) { + Ok(()) => axum::http::StatusCode::ACCEPTED.into_response(), + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + debug!("Manual playlist update deduplicated: an update is already pending or running"); + axum::http::StatusCode::ACCEPTED.into_response() } - }); - axum::http::StatusCode::ACCEPTED.into_response() + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + debug!("Manual playlist update rejected: worker channel closed (server shutting down)"); + axum::http::StatusCode::SERVICE_UNAVAILABLE.into_response() + } + } } Err(err) => { error!("Failed playlist update {}", sanitize_sensitive_info(err.to_string().as_str())); diff --git a/backend/src/api/endpoints/xtream_api.rs b/backend/src/api/endpoints/xtream_api.rs index e12c22f38..c45c91121 100644 --- a/backend/src/api/endpoints/xtream_api.rs +++ b/backend/src/api/endpoints/xtream_api.rs @@ -237,7 +237,6 @@ async fn xtream_player_api_stream( &fingerprint.addr, CustomVideoStreamType::UserAccountExpired, ) - .await .into_response(); } @@ -249,7 +248,6 @@ async fn xtream_player_api_stream( &fingerprint.addr, CustomVideoStreamType::ChannelUnavailable, ) - .await .into_response(); } @@ -310,7 +308,6 @@ async fn xtream_player_api_stream( &fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted, ) - .await .into_response(); } @@ -320,7 +317,6 @@ async fn xtream_player_api_stream( &fingerprint.addr, CustomVideoStreamType::ProviderConnectionsExhausted, ) - .await .into_response(); } @@ -353,7 +349,6 @@ async fn xtream_player_api_stream( &fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted, ) - .await .into_response(); } diff --git a/backend/src/api/main_api.rs b/backend/src/api/main_api.rs index 778a64ead..a162ada74 100644 --- a/backend/src/api/main_api.rs +++ b/backend/src/api/main_api.rs @@ -49,6 +49,7 @@ use std::{ path::PathBuf, sync::{atomic::AtomicI8, Arc}, }; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use tower_governor::key_extractor::SmartIpKeyExtractor; use tower_http::services::ServeDir; @@ -199,7 +200,7 @@ async fn healthcheck() -> impl axum::response::IntoResponse { axum::Json(create_ async fn create_shared_data( app_config: &Arc, forced_targets: &Arc, -) -> Result { +) -> Result<(AppState, mpsc::Receiver>), TuliproxError> { let config = app_config.config.load(); let use_geoip = config.is_geoip_enabled(); @@ -224,6 +225,7 @@ async fn create_shared_data( let event_manager = Arc::new(EventManager::new()); let active_provider = Arc::new(ActiveProviderManager::new(app_config, &event_manager)); let shared_stream_manager = Arc::new(SharedStreamManager::new(Arc::clone(&active_provider))); + active_provider.set_shared_stream_manager(Arc::clone(&shared_stream_manager)); let active_users = Arc::new(ActiveUserManager::new(&config, &geoip, &event_manager)); let connection_manager = Arc::new(ConnectionManager::new(&active_users, &active_provider, &shared_stream_manager, &event_manager)); @@ -231,27 +233,69 @@ async fn create_shared_data( let client = create_http_client(app_config)?; let client_no_redirect = create_http_client_no_redirect(app_config)?; - let cancel_tokens = Arc::new(ArcSwap::from_pointee(CancelTokens::default())); - let metadata_manager = Arc::new(MetadataUpdateManager::new(cancel_tokens.load().scheduler.clone())); + let tokens = CancelTokens::default(); + let metadata_manager = Arc::new(MetadataUpdateManager::new(tokens.metadata.clone())); + let cancel_tokens = Arc::new(ArcSwap::from_pointee(tokens)); - Ok(AppState { - forced_targets: Arc::new(ArcSwap::new(Arc::clone(forced_targets))), - app_config: Arc::clone(app_config), - http_client: Arc::new(ArcSwap::from_pointee(client)), - http_client_no_redirect: Arc::new(ArcSwap::from_pointee(client_no_redirect)), - downloads: Arc::new(DownloadQueue::new()), - cache: Arc::new(ArcSwapOption::from(cache)), - shared_stream_manager, - active_users, - active_provider, - connection_manager, - event_manager, - cancel_tokens, - playlists: Arc::new(PlaylistStorageState::new()), - geoip, - update_guard: UpdateGuard::new(), - metadata_manager, - }) + let (manual_update_sender, manual_update_rx) = mpsc::channel::>(1); + + Ok(( + AppState { + forced_targets: Arc::new(ArcSwap::new(Arc::clone(forced_targets))), + app_config: Arc::clone(app_config), + http_client: Arc::new(ArcSwap::from_pointee(client)), + http_client_no_redirect: Arc::new(ArcSwap::from_pointee(client_no_redirect)), + downloads: Arc::new(DownloadQueue::new()), + cache: Arc::new(ArcSwapOption::from(cache)), + shared_stream_manager, + active_users, + active_provider, + connection_manager, + event_manager, + cancel_tokens, + playlists: Arc::new(PlaylistStorageState::new()), + geoip, + update_guard: UpdateGuard::new(), + metadata_manager, + manual_update_sender, + }, + manual_update_rx, + )) +} + +async fn run_manual_update_worker( + client: reqwest::Client, + app_state: Arc, + mut rx: mpsc::Receiver>, +) { + while let Some(targets) = rx.recv().await { + exec_processing( + &client, + Arc::clone(&app_state.app_config), + targets, + Some(Arc::clone(&app_state.event_manager)), + Some(Arc::clone(&app_state)), + Some(Arc::clone(&app_state.playlists)), + Some(app_state.update_guard.clone()), + app_state.get_disabled_headers(), + Some(Arc::clone(&app_state.active_provider)), + Some(Arc::clone(&app_state.metadata_manager)), + None, + None, + ) + .await; + } +} + +fn cancel_all_service_tokens(app_state: &Arc) { + let cancel_tokens = app_state.cancel_tokens.load(); + cancel_tokens.scheduler.cancel(); + cancel_tokens.hdhomerun.cancel(); + cancel_tokens.file_watch.cancel(); + cancel_tokens.provider_dns.cancel(); + // Use the manager's shutdown() rather than cancelling the token directly so + // the is_shutdown flag is set and workers do not attempt to restart after cancellation. + app_state.metadata_manager.shutdown(); } fn exec_update_on_boot(client: &reqwest::Client, app_state: &Arc, targets: &Arc) -> bool { @@ -393,13 +437,21 @@ pub async fn start_server(app_config: Arc, targets: Arc, targets: Arc, targets: Arc { + info!("Received shutdown signal (Ctrl+C), cancelling all background services"); + cancel_all_service_tokens(&app_state_signal); + server_cancel_token_signal.cancel(); + } + Err(err) => { + error!("Failed to listen for Ctrl+C: {err}"); + } + } + }); + + serve(listener, router, Some(server_cancel_token), &shared_data.connection_manager).await; + + // Final shutdown safeguard for all background services (idempotent). + cancel_all_service_tokens(&shared_data); Ok(()) } diff --git a/backend/src/api/model/active_provider_manager.rs b/backend/src/api/model/active_provider_manager.rs index f0fee65dc..8213e45da 100644 --- a/backend/src/api/model/active_provider_manager.rs +++ b/backend/src/api/model/active_provider_manager.rs @@ -1,20 +1,20 @@ use crate::{ api::model::{ provider_lineup_manager::{ProviderAllocation, ProviderLineupManager}, - EventManager, ProviderConfig, + EventManager, ProviderConfig, SharedStreamManager, }, model::{AppConfig, ConfigInput, GracePeriodOptions}, utils::debug_if_enabled, }; -use log::error; +use log::{error, warn}; use shared::utils::sanitize_sensitive_info; -// trace_if_enabled removed -use std::collections::{HashMap, HashSet}; +use std::cmp::Reverse; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::{ net::SocketAddr, sync::{ atomic::{AtomicU64, Ordering}, - Arc, LazyLock, + Arc, LazyLock, OnceLock, }, time::{Duration, Instant}, }; @@ -27,6 +27,20 @@ static DUMMY_ADDR: LazyLock = LazyLock::new(|| "127.0.0.1:0".parse:: pub type ClientConnectionId = SocketAddr; type AllocationId = u64; +type SharedConnectionId = AllocationId; +// Key for BTreeMap priority index: (priority, Reverse, AllocationId) +// Semantics: lower numeric priority value = higher importance (0 = highest, 127 = lowest). +// `.last()` on the BTreeMap returns the entry with the highest priority value, which is +// the lowest-importance connection and therefore the best eviction victim. +// Ties are broken by `Reverse`: among equal priority values, the oldest connection +// (smallest `created_at`) sorts last and is evicted first. +type PriorityKey = (i8, Reverse, AllocationId); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PriorityOwner { + Single(ClientConnectionId), + Shared(SharedConnectionId), +} #[derive(Debug, Clone)] pub struct ProviderHandle { @@ -55,6 +69,7 @@ struct SharedAllocation { connections: HashSet, priority: i8, created_at: Instant, + cancel_token: Option, } #[derive(Debug, Clone)] @@ -82,6 +97,9 @@ struct Connections { // Index to quickly find connections by provider name for preemption // ProviderName -> Vec<(ClientConnectionId, AllocationId)> by_provider: HashMap, Vec<(ClientConnectionId, AllocationId)>>, + // Priority index per provider alias for O(log n) victim lookup + // ProviderName -> BTreeMap + priority_index: HashMap, BTreeMap>, } pub struct ActiveProviderManager { @@ -89,6 +107,7 @@ pub struct ActiveProviderManager { connections: RwLock, next_allocation_id: AtomicU64, preempted_grace_semaphore: Arc, + shared_stream_manager: OnceLock>, } impl ActiveProviderManager { @@ -100,9 +119,14 @@ impl ActiveProviderManager { connections: RwLock::new(Connections::default()), next_allocation_id: AtomicU64::new(1), preempted_grace_semaphore: Arc::new(Semaphore::new(PREEMPTED_GRACE_MAX_PENDING)), + shared_stream_manager: OnceLock::new(), } } + pub fn set_shared_stream_manager(&self, manager: Arc) { + let _ = self.shared_stream_manager.set(manager); + } + fn get_config_inputs(cfg: &AppConfig) -> Vec> { cfg.sources.load().inputs.iter().filter(|i| i.enabled).map(Arc::clone).collect() } @@ -164,6 +188,32 @@ impl ActiveProviderManager { }; if !matches!(allocation, ProviderAllocation::Exhausted) { + if matches!(&allocation, ProviderAllocation::GracePeriod(_)) && !force { + // Grace allocation received — try to evict a lower-priority victim + // across the entire input lineup (all aliases) to free capacity. + if self.evict_lower_priority_on_input(provider_or_input_name, priority).await { + // A victim was evicted. If it was on the same alias as our grace + // allocation, the provider is back at max (not over limit). + // If it was on a different alias, release the grace allocation + // and re-acquire on the input to land on the freed alias. + let evicted_on_same = !self.providers.is_over_limit( + &allocation.get_provider_name().unwrap_or_default(), + ).await; + if !evicted_on_same { + // Victim was on a different alias — first try to re-acquire. + // Keep the existing grace allocation as fallback if re-acquire fails. + let new_alloc = self.providers.acquire_connection_with_grace_override( + provider_or_input_name, allow_grace, + ).await; + if !matches!(new_alloc, ProviderAllocation::Exhausted) { + allocation.release().await; + return self.register_allocation(new_alloc, addr, priority, is_probe).await; + } + return self.register_allocation(allocation, addr, priority, is_probe).await; + } + } + // Eviction on same alias succeeded, or no victim found → keep grace allocation + } return self.register_allocation(allocation, addr, priority, is_probe).await; } @@ -189,6 +239,7 @@ impl ActiveProviderManager { let provider_name = allocation.get_provider_name().unwrap_or_default(); let allocation_id = self.next_allocation_id.fetch_add(1, Ordering::Relaxed); let cancel_token = CancellationToken::new(); + let now = Instant::now(); let mut connections = self.connections.write().await; let per_addr = connections.single.entry(*addr).or_default(); @@ -200,11 +251,13 @@ impl ActiveProviderManager { priority, is_probe, cancel_token: cancel_token.clone(), - created_at: Instant::now(), + created_at: now, }, ); connections.by_provider.entry(provider_name.clone()).or_default().push((*addr, allocation_id)); + connections.priority_index.entry(provider_name.clone()).or_default() + .insert((priority, Reverse(now), allocation_id), PriorityOwner::Single(*addr)); debug_if_enabled!( "Added provider connection {provider_name:?} for {} (prio={})", @@ -215,80 +268,69 @@ impl ActiveProviderManager { } #[allow(clippy::too_many_lines)] - async fn try_preempt_connection( + /// Evict a single lower-priority connection across the entire input lineup + /// (all provider aliases). Used when a `GracePeriod` allocation was granted. + /// Returns true if a victim was successfully evicted. + async fn evict_lower_priority_on_input( &self, input_name: &Arc, new_priority: i8, - allow_grace: bool, - ) -> Option { - // Victim: (addr, alloc_id, priority, created_at, is_shared) - let mut victim: Option<(ClientConnectionId, AllocationId, i8, Instant, bool)> = None; - - { + ) -> bool { + // Find best victim across all aliases under read lock + let victim: Option<(PriorityOwner, AllocationId, i8, Instant)> = { let connections = self.connections.read().await; - - for (prov_name, active_conns) in &connections.by_provider { - if self.providers.is_provider_for_input(prov_name, input_name) { - for (addr, alloc_id) in active_conns { - // Check single connections - if let Some(conns_map) = connections.single.get(addr) { - if let Some(info) = conns_map.get(alloc_id) { - if info.priority > new_priority { - let is_better = match victim { - None => true, - Some((_, _, v_prio, v_created, _)) => { - info.priority > v_prio - || (info.priority == v_prio && info.created_at < v_created) - } - }; - if is_better { - victim = Some((*addr, *alloc_id, info.priority, info.created_at, false)); - } - } - } - } - - // Check shared connections (only if exactly 1 listener) - if let Some(key) = connections.shared.key_by_addr.get(addr) { - if let Some(shared) = connections.shared.by_key.get(key) { - if shared.allocation_id == *alloc_id - && shared.connections.len() == 1 - && shared.priority > new_priority - { - let is_better = match victim { - None => true, - Some((_, _, v_prio, v_created, _)) => { - shared.priority > v_prio - || (shared.priority == v_prio && shared.created_at < v_created) - } - }; - if is_better { - victim = Some((*addr, *alloc_id, shared.priority, shared.created_at, true)); - } - } - } + let mut found: Option<(PriorityOwner, AllocationId, i8, Instant)> = None; + for (prov_name, tree) in &connections.priority_index { + if !self.providers.is_provider_for_input(prov_name, input_name) { + continue; + } + for ((prio, Reverse(created_at), alloc_id), owner) in tree.iter().rev() { + if *prio <= new_priority { + break; // No more victims on this alias + } + if let PriorityOwner::Shared(shared_id) = owner { + let evictable = connections + .shared + .shared_by_allocation_id + .get(shared_id) + .and_then(|key| connections.shared.by_key.get(key)) + .is_some_and(|s| s.connections.len() == 1); + if !evictable { + continue; } } + let is_better = match found { + None => true, + Some((_, _, v_prio, v_created)) => { + *prio > v_prio || (*prio == v_prio && *created_at < v_created) + } + }; + if is_better { + found = Some((*owner, *alloc_id, *prio, *created_at)); + } + break; // Only best candidate per alias } } - } + found + }; - if let Some((addr, alloc_id, v_prio, victim_created_at, is_shared)) = victim { - debug_if_enabled!( - "Preempting {} connection from {} (prio={}) for higher priority request (prio={})", - if is_shared { "shared" } else { "single" }, - sanitize_sensitive_info(&addr.to_string()), - v_prio, - new_priority - ); + let Some((owner, alloc_id, v_prio, victim_created_at)) = victim else { + return false; + }; + match owner { + PriorityOwner::Shared(shared_id) => { + debug_if_enabled!( + "Grace-evicting shared connection (allocation_id={shared_id}, prio={v_prio}) on input {} for higher priority request (prio={})", + sanitize_sensitive_info(input_name), + new_priority + ); - if is_shared { - let released_shared_allocation = { + let released = { let mut connections = self.connections.write().await; - let key = connections.shared.key_by_addr.get(&addr).cloned()?; + let Some(key) = connections.shared.shared_by_allocation_id.get(&shared_id).cloned() else { + return false; + }; - // Revalidate the selected shared victim under WRITE lock to avoid evicting - // a connection that gained additional listeners concurrently. let still_single = connections.shared.by_key.get(&key).is_some_and(|shared| { shared.allocation_id == alloc_id && shared.connections.len() == 1 @@ -296,39 +338,58 @@ impl ActiveProviderManager { && shared.created_at == victim_created_at }); if !still_single { - None - } else if let Some(shared) = connections.shared.by_key.remove(&key) { + return false; + } + + if let Some(shared) = connections.shared.by_key.remove(&key) { connections.shared.shared_by_allocation_id.remove(&shared.allocation_id); for shared_addr in &shared.connections { connections.shared.key_by_addr.remove(shared_addr); } - if let Some(name) = shared.allocation.get_provider_name() { if let Some(list) = connections.by_provider.get_mut(&name) { list.retain(|(_, i)| *i != shared.allocation_id); } + if let Some(tree) = connections.priority_index.get_mut(&name) { + tree.remove(&(v_prio, Reverse(victim_created_at), alloc_id)); + } } - Some(shared.allocation) + Some((key, shared.allocation, shared.cancel_token)) } else { - None + return false; } }; - let allocation = released_shared_allocation?; - allocation.release().await; - } else { - // Atomically remove the victim single connection and take ownership of the token. - // This guarantees only one concurrent preemptor can schedule delayed cancellation. + if let Some((stream_url, allocation, cancel_token)) = released { + if let Some(token) = cancel_token { + token.cancel(); + } + allocation.release().await; + // Stop the shared stream broadcast task to match the released capacity. + // Without this, the broadcast keeps running and consuming a provider slot + // that was already freed by allocation.release(). + if let Some(ssm) = self.shared_stream_manager.get() { + ssm.teardown_preempted_stream(&stream_url).await; + } + } + } + PriorityOwner::Single(addr) => { + debug_if_enabled!( + "Grace-evicting single connection from {} (prio={}) on input {} for higher priority request (prio={})", + sanitize_sensitive_info(&addr.to_string()), + v_prio, + sanitize_sensitive_info(input_name), + new_priority + ); + let removed_info = { let mut connections = self.connections.write().await; - let mut removed_info = None; let mut removed_provider_name = None; let mut remove_addr_entry = false; if let Some(per_addr) = connections.single.get_mut(&addr) { if let Some(info) = per_addr.get(&alloc_id) { - // Revalidate victim selection under write lock. if info.priority != v_prio || info.created_at != victim_created_at { - return None; + return false; } } if let Some(info) = per_addr.remove(&alloc_id) { @@ -337,29 +398,25 @@ impl ActiveProviderManager { removed_info = Some(info); } } - if remove_addr_entry { connections.single.remove(&addr); } - if let Some(name) = removed_provider_name { - if let Some(list) = connections.by_provider.get_mut(&name) { + if let Some(ref name) = removed_provider_name { + if let Some(list) = connections.by_provider.get_mut(name) { if let Some(idx) = list.iter().position(|(a, i)| *a == addr && *i == alloc_id) { list.remove(idx); } } + if let Some(tree) = connections.priority_index.get_mut(name) { + tree.remove(&(v_prio, Reverse(victim_created_at), alloc_id)); + } } removed_info }; - let Some(info) = removed_info else { - // Another preemptor already removed this victim. - return None; - }; - + let Some(info) = removed_info else { return false; }; let token = info.cancel_token; if info.is_probe { - // Probe preemption gets a short grace window, but cap detached sleep/cancel - // tasks so bursts cannot spawn unbounded background work. if let Ok(permit) = Arc::clone(&self.preempted_grace_semaphore).try_acquire_owned() { tokio::spawn(async move { let _permit = permit; @@ -367,14 +424,199 @@ impl ActiveProviderManager { token.cancel(); }); } else { + warn!( + "Preemption grace semaphore exhausted ({PREEMPTED_GRACE_MAX_PENDING} tasks pending); \ + falling back to immediate cancellation for preempted probe" + ); token.cancel(); } } else { token.cancel(); } - info.allocation.release().await; } + } + + true + } + + #[allow(clippy::too_many_lines)] + async fn try_preempt_connection( + &self, + input_name: &Arc, + new_priority: i8, + allow_grace: bool, + ) -> Option { + // Victim: (owner, alloc_id, priority, created_at) + let mut victim: Option<(PriorityOwner, AllocationId, i8, Instant)> = None; + + { + let connections = self.connections.read().await; + + // Search across ALL provider aliases of this input using the priority index + for (prov_name, tree) in &connections.priority_index { + if !self.providers.is_provider_for_input(prov_name, input_name) { + continue; + } + // Iterate from highest priority value (lowest importance) = best victim + for ((prio, Reverse(created_at), alloc_id), owner) in tree.iter().rev() { + if *prio <= new_priority { + break; // No more victims on this alias + } + if let PriorityOwner::Shared(shared_id) = owner { + let evictable = connections + .shared + .shared_by_allocation_id + .get(shared_id) + .and_then(|key| connections.shared.by_key.get(key)) + .is_some_and(|s| s.connections.len() == 1); + if !evictable { + continue; + } + } + let is_better = match victim { + None => true, + Some((_, _, v_prio, v_created)) => { + *prio > v_prio || (*prio == v_prio && *created_at < v_created) + } + }; + if is_better { + victim = Some((*owner, *alloc_id, *prio, *created_at)); + } + break; // Only need the best candidate per alias + } + } + } + + if let Some((owner, alloc_id, v_prio, victim_created_at)) = victim { + match owner { + PriorityOwner::Shared(shared_id) => { + debug_if_enabled!( + "Preempting shared connection (allocation_id={shared_id}, prio={v_prio}) for higher priority request (prio={new_priority})" + ); + let released_shared_allocation = { + let mut connections = self.connections.write().await; + let key = connections.shared.shared_by_allocation_id.get(&shared_id).cloned()?; + + // Revalidate the selected shared victim under WRITE lock to avoid evicting + // a connection that gained additional listeners concurrently. + let still_single = connections.shared.by_key.get(&key).is_some_and(|shared| { + shared.allocation_id == alloc_id + && shared.connections.len() == 1 + && shared.priority == v_prio + && shared.created_at == victim_created_at + }); + if !still_single { + None + } else if let Some(shared) = connections.shared.by_key.remove(&key) { + connections.shared.shared_by_allocation_id.remove(&shared.allocation_id); + for shared_addr in &shared.connections { + connections.shared.key_by_addr.remove(shared_addr); + } + + if let Some(name) = shared.allocation.get_provider_name() { + if let Some(list) = connections.by_provider.get_mut(&name) { + list.retain(|(_, i)| *i != shared.allocation_id); + } + if let Some(tree) = connections.priority_index.get_mut(&name) { + tree.remove(&(v_prio, Reverse(victim_created_at), alloc_id)); + } + } + Some((key, shared.allocation, shared.cancel_token)) + } else { + None + } + }; + let (stream_url, allocation, cancel_token) = released_shared_allocation?; + if let Some(token) = cancel_token { + token.cancel(); + } + allocation.release().await; + // Shared broadcast has to be torn down explicitly, otherwise the provider + // stream may continue after allocation counters were already released. + if let Some(ssm) = self.shared_stream_manager.get() { + ssm.teardown_preempted_stream(&stream_url).await; + } else { + error!( + "SharedStreamManager not initialised during preemption teardown for {}; \ + shared stream may linger after allocation release", + sanitize_sensitive_info(&stream_url) + ); + } + } + PriorityOwner::Single(addr) => { + debug_if_enabled!( + "Preempting single connection from {} (prio={v_prio}) for higher priority request (prio={new_priority})", + sanitize_sensitive_info(&addr.to_string()) + ); + + // Atomically remove the victim single connection and take ownership of the token. + // This guarantees only one concurrent preemptor can schedule delayed cancellation. + let removed_info = { + let mut connections = self.connections.write().await; + + let mut removed_info = None; + let mut removed_provider_name = None; + let mut remove_addr_entry = false; + if let Some(per_addr) = connections.single.get_mut(&addr) { + if let Some(info) = per_addr.get(&alloc_id) { + // Revalidate victim selection under write lock. + if info.priority != v_prio || info.created_at != victim_created_at { + return None; + } + } + if let Some(info) = per_addr.remove(&alloc_id) { + removed_provider_name = info.allocation.get_provider_name(); + remove_addr_entry = per_addr.is_empty(); + removed_info = Some(info); + } + } + + if remove_addr_entry { + connections.single.remove(&addr); + } + if let Some(name) = removed_provider_name { + if let Some(list) = connections.by_provider.get_mut(&name) { + if let Some(idx) = list.iter().position(|(a, i)| *a == addr && *i == alloc_id) { + list.remove(idx); + } + } + if let Some(tree) = connections.priority_index.get_mut(&name) { + tree.remove(&(v_prio, Reverse(victim_created_at), alloc_id)); + } + } + removed_info + }; + + let Some(info) = removed_info else { + // Another preemptor already removed this victim. + return None; + }; + + let token = info.cancel_token; + if info.is_probe { + // Probe preemption gets a short grace window, but cap detached sleep/cancel + // tasks so bursts cannot spawn unbounded background work. + if let Ok(permit) = Arc::clone(&self.preempted_grace_semaphore).try_acquire_owned() { + tokio::spawn(async move { + let _permit = permit; + tokio::time::sleep(PREEMPTED_PROBE_CANCEL_GRACE).await; + token.cancel(); + }); + } else { + warn!( + "Preemption grace semaphore exhausted ({PREEMPTED_GRACE_MAX_PENDING} tasks pending); \ + falling back to immediate cancellation for preempted probe" + ); + token.cancel(); + } + } else { + token.cancel(); + } + + info.allocation.release().await; + } + } // Now try acquire again preserving the original grace policy. let allocation = self.providers.acquire_connection_with_grace_override(input_name, allow_grace).await; @@ -454,7 +696,7 @@ impl ActiveProviderManager { let single_allocations = { let mut connections = self.connections.write().await; if let Some(allocations) = connections.single.remove(addr) { - // Remove from by_provider index while still holding the lock + // Remove from by_provider and priority_index while still holding the lock for (id, info) in &allocations { if let Some(name) = info.allocation.get_provider_name() { if let Some(list) = connections.by_provider.get_mut(&name) { @@ -462,6 +704,9 @@ impl ActiveProviderManager { list.remove(idx); } } + if let Some(tree) = connections.priority_index.get_mut(&name) { + tree.remove(&(info.priority, Reverse(info.created_at), *id)); + } } } Some(allocations) @@ -510,6 +755,9 @@ impl ActiveProviderManager { if let Some(list) = connections.by_provider.get_mut(&name) { list.retain(|(_, i)| *i != shared.allocation_id); } + if let Some(tree) = connections.priority_index.get_mut(&name) { + tree.remove(&(shared.priority, Reverse(shared.created_at), shared.allocation_id)); + } } Some(shared.allocation) } else { @@ -532,12 +780,14 @@ impl ActiveProviderManager { pub async fn release_handle(&self, handle: &ProviderHandle) { let mut released = None; + let mut released_priority_key: Option<(Arc, PriorityKey)> = None; { let mut connections = self.connections.write().await; // Try removing from Single if let Some(per_addr) = connections.single.get_mut(&handle.client_id) { if let Some(info) = per_addr.remove(&handle.allocation_id) { + let pkey = (info.priority, Reverse(info.created_at), handle.allocation_id); released = Some(info.allocation); if per_addr.is_empty() { connections.single.remove(&handle.client_id); @@ -552,6 +802,7 @@ impl ActiveProviderManager { list.remove(idx); } } + released_priority_key = Some((name, pkey)); } } } @@ -560,6 +811,7 @@ impl ActiveProviderManager { // Try removing from Shared if let Some(key) = connections.shared.shared_by_allocation_id.remove(&handle.allocation_id) { if let Some(shared) = connections.shared.by_key.remove(&key) { + let pkey = (shared.priority, Reverse(shared.created_at), handle.allocation_id); released = Some(shared.allocation); for addr in shared.connections { connections.shared.key_by_addr.remove(&addr); @@ -568,10 +820,18 @@ impl ActiveProviderManager { if let Some(list) = connections.by_provider.get_mut(&name) { list.retain(|(_, i)| *i != handle.allocation_id); } + released_priority_key = Some((name, pkey)); } } } } + + // Remove from priority_index + if let Some((name, pkey)) = &released_priority_key { + if let Some(tree) = connections.priority_index.get_mut(name) { + tree.remove(pkey); + } + } } if let Some(allocation) = released { @@ -594,15 +854,39 @@ impl ActiveProviderManager { let mut iter = m.drain(); if let Some((id, info)) = iter.next() { // Collect others as extras to release - for (_, extra_info) in iter { - extras.push(extra_info.allocation); - } + let extra_entries: Vec<_> = iter.collect(); // Cleanup indices if let Some(name) = info.allocation.get_provider_name() { if let Some(list) = connections.by_provider.get_mut(&name) { - list.retain(|(a, _)| *a != *addr); + if let Some(idx) = list.iter().position(|(a, alloc_id)| *a == *addr && *alloc_id == id) { + list.remove(idx); + } } + // Remove old single entry from priority_index + if let Some(tree) = connections.priority_index.get_mut(&name) { + tree.remove(&(info.priority, Reverse(info.created_at), id)); + } + } + // Remove extras from provider-specific indexes. + for (extra_id, extra_info) in &extra_entries { + if let Some(extra_provider_name) = extra_info.allocation.get_provider_name() { + if let Some(list) = connections.by_provider.get_mut(&extra_provider_name) { + if let Some(idx) = list + .iter() + .position(|(extra_addr, alloc_id)| *extra_addr == *addr && *alloc_id == *extra_id) + { + list.remove(idx); + } + } + if let Some(tree) = connections.priority_index.get_mut(&extra_provider_name) { + tree.remove(&(extra_info.priority, Reverse(extra_info.created_at), *extra_id)); + } + } + } + + for (_, extra_info) in extra_entries { + extras.push(extra_info.allocation); } connections.single.remove(addr); // Map is drained/empty now @@ -636,10 +920,18 @@ impl ActiveProviderManager { connections: HashSet::from([*addr]), priority: handle.1, created_at: handle.2, + cancel_token: handle.0.cancel_token.clone(), }, ); connections.shared.key_by_addr.insert(*addr, key.to_string()); connections.shared.shared_by_allocation_id.insert(handle.0.allocation_id, key.to_string()); + + // Insert new shared entry into priority_index + connections.priority_index.entry(provider_name.clone()).or_default() + .insert( + (handle.1, Reverse(handle.2), handle.0.allocation_id), + PriorityOwner::Shared(handle.0.allocation_id), + ); } extras }; @@ -977,4 +1269,139 @@ mod tests { manager.release_connection(&high_prio_addr).await; } + + #[tokio::test] + async fn test_grace_period_triggers_preemption_of_lower_priority() { + // Provider full, high-prio user arrives with grace allowed, + // low-prio victim should be evicted and provider should not be over limit. + let app_cfg = create_test_app_config_single_provider_pool(); + let event_manager = Arc::new(EventManager::new()); + let manager = ActiveProviderManager::new(&app_cfg, &event_manager); + + let input_name = "provider_1".intern(); + let low_prio_addr: SocketAddr = "127.0.0.1:47001".parse().unwrap(); + let high_prio_addr: SocketAddr = "127.0.0.1:47002".parse().unwrap(); + + // Low-priority user connects (priority 20 = low importance) + let low_alloc = manager + .acquire_connection(&input_name, &low_prio_addr, 20) + .await + .expect("low-priority user should get connection"); + assert_eq!(low_alloc.allocation.get_provider_name().as_deref(), Some(input_name.as_ref())); + let low_token = low_alloc.cancel_token.clone().expect("must have cancel token"); + + // Provider is now exhausted + assert!(manager.is_exhausted(&input_name).await); + + // High-priority user arrives WITH grace allowed (default streaming path) + // This should get a GracePeriod allocation and then evict the low-prio user + let high_alloc = manager + .acquire_connection(&input_name, &high_prio_addr, 0) + .await + .expect("high-priority user should get grace allocation and evict low-prio"); + assert_eq!(high_alloc.allocation.get_provider_name().as_deref(), Some(input_name.as_ref())); + + // Low-priority user's cancel token should be cancelled + assert!(low_token.is_cancelled(), "low-prio user should be cancelled after eviction"); + + // Provider should not be over limit (eviction freed a slot) + assert!(!manager.is_over_limit(&input_name).await, "provider should not be over limit after eviction"); + + manager.release_connection(&high_prio_addr).await; + } + + #[tokio::test] + async fn test_grace_period_no_victim_keeps_grace_behavior() { + // Provider full, same-prio user arrives with grace allowed, + // no victim available → normal grace behavior (allocation still returned). + let app_cfg = create_test_app_config_single_provider_pool(); + let event_manager = Arc::new(EventManager::new()); + let manager = ActiveProviderManager::new(&app_cfg, &event_manager); + + let input_name = "provider_1".intern(); + let user_1_addr: SocketAddr = "127.0.0.1:48001".parse().unwrap(); + let user_2_addr: SocketAddr = "127.0.0.1:48002".parse().unwrap(); + + // User 1 connects with priority 0 + let alloc1 = manager + .acquire_connection(&input_name, &user_1_addr, default_user_priority()) + .await + .expect("user1 should get connection"); + assert_eq!(alloc1.allocation.get_provider_name().as_deref(), Some(input_name.as_ref())); + let token1 = alloc1.cancel_token.clone().expect("must have cancel token"); + + // Provider is now exhausted + assert!(manager.is_exhausted(&input_name).await); + + // User 2 arrives with same priority and grace allowed + // Should get a grace allocation but NOT evict user 1 + let alloc2 = manager + .acquire_connection(&input_name, &user_2_addr, default_user_priority()) + .await; + + // Grace allocation should be returned (provider allows grace) + if let Some(alloc2) = &alloc2 { + assert_eq!(alloc2.allocation.get_provider_name().as_deref(), Some(input_name.as_ref())); + } + // User 1 should NOT be cancelled + assert!(!token1.is_cancelled(), "same-prio user should not be evicted"); + + manager.release_connection(&user_1_addr).await; + manager.release_connection(&user_2_addr).await; + } + + #[tokio::test] + async fn test_btree_index_consistent_after_lifecycle() { + // Verify that the priority_index stays consistent through add, evict, release. + let app_cfg = create_test_app_config_single_provider_pool(); + let event_manager = Arc::new(EventManager::new()); + let manager = ActiveProviderManager::new(&app_cfg, &event_manager); + + let input_name = "provider_1".intern(); + let addr_a: SocketAddr = "127.0.0.1:49001".parse().unwrap(); + let addr_b: SocketAddr = "127.0.0.1:49002".parse().unwrap(); + + // Add connection A (low priority) + let alloc_a = manager + .acquire_connection(&input_name, &addr_a, 10) + .await + .expect("alloc_a"); + + // Check index has 1 entry + { + let connections = manager.connections.read().await; + let tree = connections.priority_index.get(&input_name).expect("index for provider_1"); + assert_eq!(tree.len(), 1, "index should have 1 entry after first allocation"); + } + + // High-priority user evicts low-priority via grace path + let alloc_b = manager + .acquire_connection(&input_name, &addr_b, -5) + .await + .expect("alloc_b should evict alloc_a"); + + // Check index: should have 1 entry (alloc_a evicted, alloc_b added) + { + let connections = manager.connections.read().await; + let tree = connections.priority_index.get(&input_name).expect("index for provider_1"); + assert_eq!(tree.len(), 1, "index should have 1 entry after eviction + new allocation"); + // The remaining entry should be alloc_b + let ((prio, _, _), _) = tree.iter().next().expect("one entry"); + assert_eq!(*prio, -5, "remaining entry should be the high-prio connection"); + } + + // Release alloc_b + manager.release_handle(&alloc_b).await; + + // Check index: should be empty + { + let connections = manager.connections.read().await; + let tree = connections.priority_index.get(&input_name); + let is_empty = tree.is_none_or(|t| t.is_empty()); + assert!(is_empty, "index should be empty after releasing all connections"); + } + + // Verify alloc_a handle can be safely released (already evicted - no-op) + manager.release_handle(&alloc_a).await; + } } diff --git a/backend/src/api/model/active_user_manager.rs b/backend/src/api/model/active_user_manager.rs index 2d4da81fb..34fdf9e55 100644 --- a/backend/src/api/model/active_user_manager.rs +++ b/backend/src/api/model/active_user_manager.rs @@ -140,13 +140,58 @@ impl ActiveUserManager { } } - pub async fn release_connection(&self, addr: &SocketAddr) { - let (log_active_user, disconnected_user) = { + /// Releases an active stream for the given socket address without removing the + /// socket registration (`key_by_addr`). This is used when a stream ends while + /// the underlying HTTP connection may still remain open. + pub async fn release_stream(&self, addr: &SocketAddr) -> bool { + let (removed_stream, username) = { + let mut user_connections = self.connections.write().await; + + let Some(username) = user_connections.key_by_addr.get(addr).cloned() else { + return false; + }; + + let mut removed_stream = false; + if let Some(connection_data) = user_connections.by_key.get_mut(&username) { + let prev_len = connection_data.streams.len(); + connection_data.streams.retain(|c| c.addr != *addr); + removed_stream = connection_data.streams.len() != prev_len; + + if removed_stream { + if connection_data.connections > 0 { + connection_data.connections -= 1; + } + if connection_data.connections < connection_data.max_connections { + connection_data.granted_grace = false; + connection_data.grace_ts = 0; + } + } + } + (removed_stream, username) + }; + + if removed_stream { + if !username.is_empty() { + debug_if_enabled!( + "Released stream for user {username} at {}", + sanitize_sensitive_info(&addr.to_string()) + ); + } + self.log_active_user().await; + } + + removed_stream + } + + pub async fn release_connection(&self, addr: &SocketAddr) -> bool { + let (addr_removed, disconnected_user) = { let mut user_connections = self.connections.write().await; if let Some(username) = user_connections.key_by_addr.remove(addr) { if let Some(connection_data) = user_connections.by_key.get_mut(&username) { - if connection_data.connections > 0 { + let prev_len = connection_data.streams.len(); + connection_data.streams.retain(|c| c.addr != *addr); + if connection_data.streams.len() != prev_len && connection_data.connections > 0 { connection_data.connections -= 1; } @@ -154,7 +199,6 @@ impl ActiveUserManager { connection_data.granted_grace = false; connection_data.grace_ts = 0; } - connection_data.streams.retain(|c| c.addr != *addr); } (true, Some(username)) } else { @@ -171,9 +215,11 @@ impl ActiveUserManager { } } - if log_active_user { + if addr_removed { self.log_active_user().await; } + + addr_removed } pub fn update_config(&self, config: &Config) { @@ -223,7 +269,11 @@ impl ActiveUserManager { if self.grace_period_millis.load(Ordering::Relaxed) > 0 && current_connections == connection_data.max_connections { - // Allow a grace period once + // Intentional asymmetry: grace is granted when current == max (AT limit), while + // Exhausted is returned when current > max (OVER limit after the grace window). + // This allows exactly one extra connection during the grace window — the new + // connection is accepted now, and a background check evicts it if the count is + // still over max after the grace period elapses. connection_data.granted_grace = true; connection_data.grace_ts = now; debug!("Granted a grace period for user access: {username}"); diff --git a/backend/src/api/model/app_state.rs b/backend/src/api/model/app_state.rs index 819542a74..dedc4dcda 100644 --- a/backend/src/api/model/app_state.rs +++ b/backend/src/api/model/app_state.rs @@ -33,7 +33,8 @@ use std::{ sync::{atomic::AtomicI8, Arc}, time::Duration, }; -use tokio::{sync::Mutex, task}; +use tokio::sync::{mpsc, Mutex}; +use tokio::task; use tokio_util::sync::CancellationToken; use url::Url; @@ -70,7 +71,7 @@ struct TargetChanges { target: Arc, } -create_bitset!(u8, UpdateChangesFlags, Scheduler, Hdhomerun, FileWatch, Geoip, ProviderDns); +create_bitset!(u8, UpdateChangesFlags, Scheduler, Hdhomerun, FileWatch, Geoip, ProviderDns, Metadata); pub(in crate::api) struct UpdateChanges { flags: UpdateChangesFlagsSet, @@ -151,12 +152,20 @@ fn cancel_services(app_state: &Arc, changes: &UpdateChanges) { let hdhomerun = cancel_service!(hdhomerun, UpdateChangesFlags::Hdhomerun, changes, cancel_tokens); let file_watch = cancel_service!(file_watch, UpdateChangesFlags::FileWatch, changes, cancel_tokens); let provider_dns = cancel_service!(provider_dns, UpdateChangesFlags::ProviderDns, changes, cancel_tokens); + let metadata = if changes.flags.contains(UpdateChangesFlags::Metadata) { + let token = CancellationToken::new(); + app_state.metadata_manager.rotate_cancel_token(token.clone()); + token + } else { + cancel_tokens.metadata.clone() + }; let tokens = CancelTokens { scheduler, hdhomerun, file_watch, provider_dns, + metadata, }; app_state.cancel_tokens.store(Arc::new(tokens)); @@ -297,6 +306,7 @@ pub struct CancelTokens { pub(crate) hdhomerun: CancellationToken, pub(crate) file_watch: CancellationToken, pub(crate) provider_dns: CancellationToken, + pub(crate) metadata: CancellationToken, } impl Default for CancelTokens { fn default() -> Self { @@ -305,6 +315,7 @@ impl Default for CancelTokens { hdhomerun: CancellationToken::new(), file_watch: CancellationToken::new(), provider_dns: CancellationToken::new(), + metadata: CancellationToken::new(), } } } @@ -337,6 +348,11 @@ pub struct AppState { pub geoip: Arc>, pub update_guard: UpdateGuard, pub metadata_manager: Arc, + /// Bounded channel (capacity 1) for manual playlist update requests. + /// `try_send` deduplicates rapid clicks: if an update is already pending + /// or the channel is full, the request is silently dropped so at most one + /// update is queued at any time regardless of how many times the button is clicked. + pub manual_update_sender: mpsc::Sender>, } impl AppState { @@ -450,12 +466,14 @@ impl AppState { let geoip_enabled = config.is_geoip_enabled(); let geoip_enabled_old = old_config.is_geoip_enabled(); + let changed_storage_dir = old_config.storage_dir != config.storage_dir; let mut changes = UpdateChanges { flags: UpdateChangesFlagsSet::new(), targets: None }; changes.set_flag_if(changed_schedules || changed_library_enabled || geoip_enabled != geoip_enabled_old, UpdateChangesFlags::Scheduler); changes.set_flag_if(changed_hdhomerun, UpdateChangesFlags::Hdhomerun); changes.set_flag_if(changed_file_watch, UpdateChangesFlags::FileWatch); changes.set_flag_if(geoip_enabled != geoip_enabled_old, UpdateChangesFlags::Geoip); + changes.set_flag_if(changed_storage_dir, UpdateChangesFlags::Metadata); changes } diff --git a/backend/src/api/model/connection_manager.rs b/backend/src/api/model/connection_manager.rs index 64fd89d08..9bdbd5de8 100644 --- a/backend/src/api/model/connection_manager.rs +++ b/backend/src/api/model/connection_manager.rs @@ -6,12 +6,28 @@ use crate::{ auth::Fingerprint, utils::debug_if_enabled, }; -use log::warn; +use log::{debug, warn}; use shared::{ model::{ActiveUserConnectionChange, StreamChannel, VirtualId}, utils::sanitize_sensitive_info, }; use std::{borrow::Cow, net::SocketAddr, sync::Arc}; +use tokio::sync::mpsc; + +pub(crate) enum CleanupEvent { + ReleaseStream { addr: SocketAddr }, + ReleaseProviderHandle { handle: Option }, + ReleaseStreamAndProviderHandle { addr: SocketAddr, handle: Option }, + UpdateDetailAndReleaseProvider { + addr: SocketAddr, + video_type: CustomVideoStreamType, + handle: Option, + }, + UpdateDetailAndReleaseProviderConnection { + addr: SocketAddr, + video_type: CustomVideoStreamType, + }, +} pub struct ConnectionManager { pub user_manager: Arc, @@ -19,6 +35,7 @@ pub struct ConnectionManager { pub shared_stream_manager: Arc, event_manager: Arc, close_socket_signal_tx: tokio::sync::broadcast::Sender, + cleanup_tx: mpsc::UnboundedSender, } impl ConnectionManager { @@ -29,13 +46,91 @@ impl ConnectionManager { event_manager: &Arc, ) -> Self { let (close_socket_signal_tx, _) = tokio::sync::broadcast::channel(256); + let (cleanup_tx, cleanup_rx) = mpsc::unbounded_channel(); - Self { + let mgr = Self { user_manager: Arc::clone(user_manager), provider_manager: Arc::clone(provider_manager), shared_stream_manager: Arc::clone(shared_stream_manager), event_manager: Arc::clone(event_manager), close_socket_signal_tx, + cleanup_tx, + }; + + Self::spawn_cleanup_worker( + cleanup_rx, + Arc::clone(user_manager), + Arc::clone(provider_manager), + Arc::clone(shared_stream_manager), + Arc::clone(event_manager), + ); + + mgr + } + + fn spawn_cleanup_worker( + mut rx: mpsc::UnboundedReceiver, + user_manager: Arc, + provider_manager: Arc, + shared_stream_manager: Arc, + event_manager: Arc, + ) { + tokio::spawn(async move { + while let Some(event) = rx.recv().await { + match event { + CleanupEvent::ReleaseStream { addr } => { + if user_manager.release_stream(&addr).await { + event_manager.send_event(EventMessage::ActiveUser( + ActiveUserConnectionChange::Disconnected(addr), + )); + } + } + CleanupEvent::ReleaseProviderHandle { handle } => { + if let Some(h) = handle { + provider_manager.release_handle(&h).await; + } + } + CleanupEvent::ReleaseStreamAndProviderHandle { addr, handle } => { + // Release provider handle first to avoid a race window where the user + // connection count drops (making capacity appear available) before the + // provider slot is actually freed. + if let Some(h) = handle { + provider_manager.release_handle(&h).await; + } + if user_manager.release_stream(&addr).await { + event_manager.send_event(EventMessage::ActiveUser( + ActiveUserConnectionChange::Disconnected(addr), + )); + } + } + CleanupEvent::UpdateDetailAndReleaseProvider { addr, video_type, handle } => { + if let Some(stream_info) = user_manager.update_stream_detail(&addr, video_type).await { + event_manager.send_event(EventMessage::ActiveUser( + ActiveUserConnectionChange::Updated(stream_info), + )); + } + if let Some(h) = handle { + provider_manager.release_handle(&h).await; + } + } + CleanupEvent::UpdateDetailAndReleaseProviderConnection { addr, video_type } => { + if let Some(stream_info) = user_manager.update_stream_detail(&addr, video_type).await { + event_manager.send_event(EventMessage::ActiveUser( + ActiveUserConnectionChange::Updated(stream_info), + )); + } + provider_manager.release_connection(&addr).await; + shared_stream_manager.release_connection(&addr, false).await; + } + } + } + debug!("Cleanup worker exiting"); + }); + } + + pub(crate) fn send_cleanup(&self, event: CleanupEvent) { + if let Err(e) = self.cleanup_tx.send(event) { + debug!("Cleanup channel closed, dropping event: {e}"); } } @@ -63,10 +158,12 @@ impl ConnectionManager { } pub async fn release_connection(&self, addr: &SocketAddr) { - self.user_manager.release_connection(addr).await; + let removed = self.user_manager.release_connection(addr).await; self.provider_manager.release_connection(addr).await; self.shared_stream_manager.release_connection(addr, true).await; - self.event_manager.send_event(EventMessage::ActiveUser(ActiveUserConnectionChange::Disconnected(*addr))); + if removed { + self.event_manager.send_event(EventMessage::ActiveUser(ActiveUserConnectionChange::Disconnected(*addr))); + } } pub async fn release_provider_connection(&self, addr: &SocketAddr) { @@ -74,6 +171,12 @@ impl ConnectionManager { self.shared_stream_manager.release_connection(addr, false).await; } + pub async fn release_stream(&self, addr: &SocketAddr) { + if self.user_manager.release_stream(addr).await { + self.event_manager.send_event(EventMessage::ActiveUser(ActiveUserConnectionChange::Disconnected(*addr))); + } + } + pub async fn release_provider_handle(&self, provider_handle: Option) { if let Some(handle) = provider_handle { self.provider_manager.release_handle(&handle).await; diff --git a/backend/src/api/model/metadata_update_manager.rs b/backend/src/api/model/metadata_update_manager.rs index 10053de43..038861a74 100644 --- a/backend/src/api/model/metadata_update_manager.rs +++ b/backend/src/api/model/metadata_update_manager.rs @@ -12,6 +12,7 @@ use crate::{ }, utils::{debug_if_enabled, FileReadGuard}, }; +use arc_swap::ArcSwap; use dashmap::{mapref::entry::Entry, DashMap}; use log::{debug, error, info, warn}; use parking_lot::Mutex as ParkingMutex; @@ -44,6 +45,9 @@ const METADATA_RETRY_STATE_FILE: &str = "metadata_retry_state.db"; const TASK_ERR_NO_CONNECTION: &str = "No connection available"; const TASK_ERR_PREEMPTED: &str = "Task preempted"; const TASK_ERR_UPDATE_IN_PROGRESS: &str = "Playlist update in progress"; +// Per-task execution timeout. ffprobe defaults to 60 s; allow extra headroom for network +// fetches, TMDB lookups, and B+Tree writes before declaring the task stuck. +const PROBE_TASK_TIMEOUT_SECS: u64 = 30; const BLOCKING_DB_MIN_CONCURRENCY: usize = 4; const BLOCKING_DB_MAX_CONCURRENCY: usize = 32; const RETRY_STATE_PRUNE_INTERVAL_SECS: i64 = 300; @@ -534,6 +538,10 @@ impl PendingTask { pub struct MetadataUpdateManager { /// Per-input worker senders. Worker is spawned when entry is created. workers: DashMap, InputWorkerContext>, + /// Synchronizes cancellation token rotation with worker context creation. + worker_lifecycle_lock: ParkingMutex<()>, + /// Terminal shutdown flag; once set, token rotation must not reactivate workers. + is_shutdown_flag: AtomicBool, /// Global application state (weak reference to avoid cycles) app_state: tokio::sync::Mutex>>, /// Global gate: @@ -542,15 +550,11 @@ pub struct MetadataUpdateManager { /// This guarantees that no background task runs while an update is active. update_pause_gate: Arc>, /// Global cancellation token for shutdown - cancel_token: CancellationToken, + cancel_token: ArcSwap, /// Monotonic worker generation id used to avoid removing a newly spawned worker context. next_worker_id: AtomicU64, } -impl Default for MetadataUpdateManager { - fn default() -> Self { Self::new(CancellationToken::new()) } -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum SubmitTaskResult { QueuedOrMerged, @@ -559,16 +563,48 @@ enum SubmitTaskResult { } impl MetadataUpdateManager { + pub fn new(cancel_token: CancellationToken) -> Self { Self { workers: DashMap::new(), + worker_lifecycle_lock: ParkingMutex::new(()), + is_shutdown_flag: AtomicBool::new(false), app_state: tokio::sync::Mutex::new(None), update_pause_gate: Arc::new(RwLock::new(())), - cancel_token, + cancel_token: ArcSwap::from_pointee(cancel_token), next_worker_id: AtomicU64::new(1), } } + /// Requests graceful shutdown for all metadata workers and delayed requeue tasks. + /// This operation is idempotent. + pub fn shutdown(&self) { + let _guard = self.worker_lifecycle_lock.lock(); + self.is_shutdown_flag.store(true, Ordering::Release); + self.cancel_token.load_full().cancel(); + } + + /// Returns true once shutdown has been requested. + pub fn is_shutdown(&self) -> bool { self.is_shutdown_flag.load(Ordering::Acquire) } + + /// Rotates the cancellation token for metadata workers. + /// Existing workers are cancelled and removed so new tasks start with fresh runtime state. + pub fn rotate_cancel_token(&self, cancel_token: CancellationToken) { + let old_token = { + let _guard = self.worker_lifecycle_lock.lock(); + if self.is_shutdown_flag.load(Ordering::Acquire) { + self.workers.clear(); + cancel_token.cancel(); + return; + } + + let old_token = self.cancel_token.swap(Arc::new(cancel_token)); + self.workers.clear(); + old_token + }; + old_token.cancel(); + } + /// Acquire exclusive gate for a foreground playlist update. /// While this guard is held, background workers wait before starting heavy metadata/probe steps. pub async fn acquire_update_pause_guard(&self) -> tokio::sync::OwnedRwLockWriteGuard<()> { @@ -613,50 +649,63 @@ impl MetadataUpdateManager { loop { // Atomically ensure there is exactly one worker context per input. let mut worker_to_spawn: Option<(u64, InputWorker)> = None; - let ctx = match self.workers.entry(input_name.clone()) { - Entry::Occupied(entry) => entry.get().clone(), - Entry::Vacant(entry) => { - let (tx, rx) = mpsc::channel::(max_queue_size); - let pending_tasks = Arc::new(DashMap::new()); - let pending_task_count = Arc::new(AtomicUsize::new(0)); - let worker_id = self.next_worker_id.fetch_add(1, Ordering::Relaxed); - - let ctx = InputWorkerContext { - worker_id, - sender: tx.clone(), - pending_tasks: pending_tasks.clone(), - pending_task_count: pending_task_count.clone(), - }; - entry.insert(ctx.clone()); - - worker_to_spawn = Some(( - worker_id, - InputWorker { - input_name: input_name.clone(), - sender: tx, - receiver: rx, - pending_tasks, - pending_task_count, - app_state_weak: app_state_weak.clone(), - update_pause_gate: Arc::clone(&self.update_pause_gate), - cancel_token: self.cancel_token.clone(), - batch_buffer: BatchResultCollector::new(), - db_handles: HashMap::new(), - failed_clusters: HashSet::new(), - retry_states: HashMap::new(), - resolve_exhausted: HashMap::new(), - last_cycle_completed_at_ts: None, - metadata_retry_state_path: None, - metadata_retry_loaded: false, - metadata_retry_load_retry_at_ts: None, - last_retry_state_prune_at_ts: None, - scheduled_requeues: Arc::new(DashMap::new()), - recently_completed_no_change: HashMap::new(), - }, - )); - - ctx + let (cancel_token, ctx) = { + let _guard = self.worker_lifecycle_lock.lock(); + let cancel_token = self.cancel_token.load_full(); + if cancel_token.is_cancelled() { + debug_if_enabled!( + "Aborting metadata enqueue loop for input {input_name} because cancellation was requested: {task_to_queue}" + ); + return; } + + let ctx = match self.workers.entry(input_name.clone()) { + Entry::Occupied(entry) => entry.get().clone(), + Entry::Vacant(entry) => { + let (tx, rx) = mpsc::channel::(max_queue_size); + let pending_tasks = Arc::new(DashMap::new()); + let pending_task_count = Arc::new(AtomicUsize::new(0)); + let worker_id = self.next_worker_id.fetch_add(1, Ordering::Relaxed); + + let ctx = InputWorkerContext { + worker_id, + sender: tx.clone(), + pending_tasks: pending_tasks.clone(), + pending_task_count: pending_task_count.clone(), + }; + entry.insert(ctx.clone()); + + worker_to_spawn = Some(( + worker_id, + InputWorker { + input_name: input_name.clone(), + sender: tx, + receiver: rx, + pending_tasks, + pending_task_count, + app_state_weak: app_state_weak.clone(), + update_pause_gate: Arc::clone(&self.update_pause_gate), + cancel_token: (*cancel_token).clone(), + batch_buffer: BatchResultCollector::new(), + db_handles: HashMap::new(), + failed_clusters: HashSet::new(), + retry_states: HashMap::new(), + resolve_exhausted: HashMap::new(), + last_cycle_completed_at_ts: None, + metadata_retry_state_path: None, + metadata_retry_loaded: false, + metadata_retry_load_retry_at_ts: None, + last_retry_state_prune_at_ts: None, + scheduled_requeues: Arc::new(DashMap::new()), + recently_completed_no_change: HashMap::new(), + }, + )); + + ctx + } + }; + + (cancel_token, ctx) }; if let Some((worker_id, worker)) = worker_to_spawn { @@ -674,6 +723,17 @@ impl MetadataUpdateManager { }); } + let sender_still_current = { + let _guard = self.worker_lifecycle_lock.lock(); + let token_matches = Arc::ptr_eq(&self.cancel_token.load_full(), &cancel_token); + let worker_matches = + self.workers.get(&input_name).is_some_and(|current| current.worker_id == ctx.worker_id); + token_matches && worker_matches + }; + if !sender_still_current { + continue; + } + match Self::submit_task( ctx.sender.clone(), ctx.pending_tasks.clone(), @@ -686,7 +746,11 @@ impl MetadataUpdateManager { { SubmitTaskResult::QueuedOrMerged => return, SubmitTaskResult::QueueFull => { - warn!("Metadata queue full for input {input_name}, dropping task"); + error!( + "Metadata queue full for input {input_name} (max_queue_size={max_queue_size}), \ + dropping task: {task_to_queue}; consider increasing max_queue_size or reducing \ + probe frequency" + ); return; } SubmitTaskResult::ChannelClosed => { @@ -707,7 +771,7 @@ impl MetadataUpdateManager { let backoff_ms = 25_u64.saturating_mul(factor).min(2_000); let backoff = Duration::from_millis(backoff_ms); tokio::select! { - () = self.cancel_token.cancelled() => { + () = cancel_token.cancelled() => { warn!( "Aborting metadata task enqueue for input {input_name} because cancellation was requested: {task_to_queue}" ); @@ -996,6 +1060,9 @@ impl InputWorker { }; let Some((current_key, current_task, current_generation)) = task_data else { break }; + if self.cancel_token.is_cancelled() { + break; + } if !self.metadata_retry_loaded { self.ensure_metadata_retry_state_loaded(&input_name, app_state_weak.as_ref(), &runtime_settings).await; } @@ -1525,7 +1592,13 @@ impl InputWorker { // Try to get the next task immediately to keep locks open. // Ignore phantom signals (channel key without pending map entry). if next_task.is_none() { + if self.cancel_token.is_cancelled() { + break; + } while let Ok(key) = self.receiver.try_recv() { + if self.cancel_token.is_cancelled() { + break; + } if let Some(snapshot) = self.load_task_snapshot(key) { next_task = Some(snapshot); break; @@ -1737,10 +1810,20 @@ impl InputWorker { &mut self, runtime_settings: &MetadataUpdateRuntimeSettings, ) -> Option<(TaskKey, UpdateTask, u64)> { + if self.cancel_token.is_cancelled() { + return None; + } + // Fast path: drain immediate signals until we find a real pending task. loop { + if self.cancel_token.is_cancelled() { + return None; + } match self.receiver.try_recv() { Ok(key) => { + if self.cancel_token.is_cancelled() { + return None; + } if let Some(snapshot) = self.load_task_snapshot(key) { return Some(snapshot); } @@ -1764,6 +1847,9 @@ impl InputWorker { res = tokio::time::timeout(Duration::from_secs(runtime_settings.worker_idle_timeout_secs), self.receiver.recv()) => { match res { Ok(Some(key)) => { + if self.cancel_token.is_cancelled() { + return None; + } if let Some(snapshot) = self.load_task_snapshot(key) { return Some(snapshot); } @@ -1771,8 +1857,14 @@ impl InputWorker { Ok(None) => return None, Err(_) => { loop { + if self.cancel_token.is_cancelled() { + return None; + } match self.receiver.try_recv() { Ok(key) => { + if self.cancel_token.is_cancelled() { + return None; + } if let Some(snapshot) = self.load_task_snapshot(key) { return Some(snapshot); } @@ -2799,19 +2891,40 @@ impl InputWorker { }; // Execute task; probe tasks get a reserved provider handle, resolve tasks don't. - let res = if let Some(handle) = provider_handle.as_ref() { - if let Some(token) = &handle.cancel_token { - tokio::select! { - biased; + // The hard timeout is applied only to tasks that perform network probing — Probe* + // variants and Resolve tasks whose reason includes ResolveReason::Probe — because + // those are the only ones that can stall on an unresponsive provider. Pure metadata + // resolve tasks (Info / TMDB / Date only) run without a timeout so they are not + // subject to the probe hard limit. + let exec_fut = async { + if let Some(handle) = provider_handle.as_ref() { + if let Some(token) = &handle.cancel_token { + tokio::select! { + biased; - () = token.cancelled() => { - debug_if_enabled!("Metadata update task preempted by user request for input {}", input_name); - Err(shared::error::info_err!("{}", TASK_ERR_PREEMPTED)) - } + () = token.cancelled() => { + debug_if_enabled!("Metadata update task preempted by user request for input {}", input_name); + Err(shared::error::info_err!("{}", TASK_ERR_PREEMPTED)) + } - res = Self::execute_task_inner_static(&app_state, &client, &input_to_use, task, item_title.as_deref(), Some(handle), probe_priority, collector, db_handles, failed_clusters) => { - res + res = Self::execute_task_inner_static(&app_state, &client, &input_to_use, task, item_title.as_deref(), Some(handle), probe_priority, collector, db_handles, failed_clusters) => { + res + } } + } else { + Self::execute_task_inner_static( + &app_state, + &client, + &input_to_use, + task, + item_title.as_deref(), + Some(handle), + probe_priority, + collector, + db_handles, + failed_clusters, + ) + .await } } else { Self::execute_task_inner_static( @@ -2820,7 +2933,7 @@ impl InputWorker { &input_to_use, task, item_title.as_deref(), - Some(handle), + None, probe_priority, collector, db_handles, @@ -2828,20 +2941,24 @@ impl InputWorker { ) .await } + }; + let needs_probe_timeout = Self::is_probe_task(task) + || (Self::is_resolve_task(task) && Self::task_reason(task).contains(ResolveReason::Probe)); + let res = if needs_probe_timeout { + match tokio::time::timeout(Duration::from_secs(PROBE_TASK_TIMEOUT_SECS), exec_fut).await { + Ok(result) => result, + Err(_elapsed) => { + error!( + "Metadata probe task timed out after {PROBE_TASK_TIMEOUT_SECS}s for input {input_name}: {task}; \ + releasing provider handle and skipping task" + ); + Err(shared::error::info_err!( + "Task timed out after {PROBE_TASK_TIMEOUT_SECS}s for input {input_name}: {task}" + )) + } + } } else { - Self::execute_task_inner_static( - &app_state, - &client, - &input_to_use, - task, - item_title.as_deref(), - None, - probe_priority, - collector, - db_handles, - failed_clusters, - ) - .await + exec_fut.await }; if provider_handle.is_some() { @@ -3090,7 +3207,7 @@ mod tests { #[tokio::test] async fn queue_task_creates_single_worker_per_input_under_concurrency() { let cancel_token = CancellationToken::new(); - let manager = Arc::new(MetadataUpdateManager::new(cancel_token.clone())); + let manager = Arc::new(MetadataUpdateManager::new(cancel_token)); let input_name: Arc = Arc::from("race_input"); let mut joins = Vec::new(); @@ -3119,7 +3236,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(50)).await; assert_eq!(manager.active_worker_count(), 1); - cancel_token.cancel(); + manager.shutdown(); } #[tokio::test] diff --git a/backend/src/api/model/stream.rs b/backend/src/api/model/stream.rs index 40df403b7..0d8b3e870 100644 --- a/backend/src/api/model/stream.rs +++ b/backend/src/api/model/stream.rs @@ -1,12 +1,12 @@ use crate::{ api::model::{CustomVideoStreamType, ProviderHandle, StreamError}, model::GracePeriodOptions, - tools::atomic_once_flag::AtomicOnceFlag, }; use axum::http::StatusCode; use bytes::Bytes; use futures::stream::BoxStream; use std::{collections::HashMap, sync::Arc}; +use tokio_util::sync::CancellationToken; use url::Url; pub type BoxedProviderStream = BoxStream<'static, Result>; @@ -32,8 +32,9 @@ pub struct StreamDetails { pub provider_name: Option>, pub request_url: Option>, pub grace_period: GracePeriodOptions, + pub provider_grace_active: bool, pub disable_provider_grace: bool, - pub reconnect_flag: Option>, + pub reconnect_flag: Option, pub provider_handle: Option, } @@ -45,6 +46,7 @@ impl StreamDetails { provider_name: None, request_url: None, grace_period: grace_period_options, + provider_grace_active: false, disable_provider_grace: false, reconnect_flag: None, provider_handle: None, diff --git a/backend/src/api/model/streams/active_client_stream.rs b/backend/src/api/model/streams/active_client_stream.rs index aae9bee92..d14eda5ac 100644 --- a/backend/src/api/model/streams/active_client_stream.rs +++ b/backend/src/api/model/streams/active_client_stream.rs @@ -1,19 +1,18 @@ use crate::{ api::{ model::{ - AppState, BoxedProviderStream, ConnectionManager, CustomVideoStreamType, ProviderHandle, StreamDetails, - StreamError, TimedClientStream, TransportStreamBuffer, + AppState, BoxedProviderStream, CleanupEvent, ConnectionManager, CustomVideoStreamType, ProviderHandle, + StreamDetails, StreamError, TimedClientStream, TransportStreamBuffer, }, panel_api::{can_provision_on_exhausted, find_input_by_provider_name, run_panel_api_provisioning_probe}, }, auth::Fingerprint, model::{ConfigInput, ProxyUserCredentials}, - tools::atomic_once_flag::AtomicOnceFlag, utils::debug_if_enabled, }; use axum::http::{header::USER_AGENT, HeaderMap}; use bytes::Bytes; -use futures::{task::AtomicWaker, Stream, StreamExt}; +use futures::{task::AtomicWaker, Future, Stream, StreamExt}; use log::{debug, error, info}; use shared::{ model::{StreamChannel, UserConnectionPermission, VirtualId}, @@ -27,48 +26,156 @@ use std::{ }, task::{Context, Poll}, }; +use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned}; -const INNER_STREAM: u8 = 0_u8; -const USER_EXHAUSTED_STREAM: u8 = 1_u8; -const PROVIDER_EXHAUSTED_STREAM: u8 = 2_u8; -const CHANNEL_UNAVAILABLE_STREAM: u8 = 3_u8; -const PROVISIONING_STREAM: u8 = 4_u8; -const GRACE_PENDING: u8 = 255_u8; // Grace period check hasn't completed yet +/// Discriminates which byte-stream the client is consuming at any moment. +/// Stored as `u8` in an `AtomicU8` for lock-free access inside `poll_next`. +/// Lower numeric values correspond to a live or custom stream; `GracePending` +/// (255) is a transient sentinel that parks the poll until the grace task resolves. +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StreamMode { + /// Forward bytes directly from the upstream provider. + Inner = 0, + /// Show the "user connections exhausted" custom video. + UserExhausted = 1, + /// Show the "provider connections exhausted" custom video. + ProviderExhausted = 2, + /// Show the "channel unavailable" custom video. + ChannelUnavailable = 3, + /// Show the provisioning/placeholder custom video while probing for capacity. + Provisioning = 4, + /// Transient: grace-period check is still in progress; `poll_next` must park. + GracePending = 255, +} + +impl StreamMode { + fn from_u8(v: u8) -> Self { + match v { + 0 => Self::Inner, + 1 => Self::UserExhausted, + 2 => Self::ProviderExhausted, + 3 => Self::ChannelUnavailable, + 4 => Self::Provisioning, + _ => Self::GracePending, + } + } +} + +/// Holds the optional custom video buffers for each error/placeholder scenario. +/// Using named fields avoids the positional-indexing confusion of a 4-tuple. +struct CustomVideoBuffers { + user_exhausted: Option, + provider_exhausted: Option, + unavailable: Option, + provisioning: Option, +} struct GraceProvisioningInfo { input: Arc, - stop_signal: Arc, + stop_signal: CancellationToken, } +#[allow(clippy::struct_excessive_bools)] struct ActiveClientStreamState { inner: BoxedProviderStream, send_custom_stream_flag: Option>, provider_handle: Option, - custom_video: ( - Option, - Option, - Option, - Option, - ), + preempt_cancelled: Option>>, + grace_task_handle: Option>, + provisionable: bool, + custom_video: CustomVideoBuffers, waker: Option>, connection_manager: Arc, fingerprint: Arc, provider_stopped: bool, + user_stream_released: bool, + /// Mirrors `user_stream_released` for the provider handle to guard against double-release + /// when preemption and Drop race. + provider_handle_released: bool, } impl ActiveClientStreamState { - fn stop_provider_stream(&mut self, unavailable: bool) { + fn release_user_stream(&mut self) { + if self.user_stream_released { + return; + } + self.user_stream_released = true; + self.connection_manager.send_cleanup(CleanupEvent::ReleaseStream { + addr: self.fingerprint.addr, + }); + } + + fn stop_grace_task(&mut self) { + if let Some(task) = self.grace_task_handle.take() { + task.abort(); + } + } + + fn clear_finished_grace_task(&mut self) { + if self + .grace_task_handle + .as_ref() + .is_some_and(tokio::task::JoinHandle::is_finished) + { + self.grace_task_handle = None; + // If the task finished but the flag is still GRACE_PENDING (e.g. the task + // panicked or was cancelled before it could update the flag), reset the flag + // to INNER_STREAM so the client stream is not hung indefinitely. + if let Some(flag) = &self.send_custom_stream_flag { + let _ = flag.compare_exchange( + StreamMode::GracePending as u8, + StreamMode::Inner as u8, + Ordering::AcqRel, + Ordering::Relaxed, + ); + } + } + } + + fn stop_provider_stream_preempted(&mut self) { self.provider_stopped = true; + self.preempt_cancelled = None; + self.stop_grace_task(); + self.release_user_stream(); if self.provider_handle.is_some() { - let mgr = Arc::clone(&self.connection_manager); let handle = self.provider_handle.take(); + self.provider_handle_released = true; + if let Some(flag) = &self.send_custom_stream_flag { + flag.store(StreamMode::Inner as u8, Ordering::Release); + } + + if let Some(waker) = &self.waker { + waker.wake(); + } + + let addr = self.fingerprint.addr; + self.inner = futures::stream::empty::>().boxed(); + + debug_if_enabled!( + "Provider stream preempted for {}; stopping client stream", + sanitize_sensitive_info(&addr.to_string()) + ); + self.connection_manager.send_cleanup(CleanupEvent::ReleaseProviderHandle { handle }); + } + } + + fn stop_provider_stream(&mut self, unavailable: bool) { + self.provider_stopped = true; + self.preempt_cancelled = None; + self.stop_grace_task(); + self.release_user_stream(); + + if self.provider_handle.is_some() { + let handle = self.provider_handle.take(); + self.provider_handle_released = true; if unavailable { if let Some(flag) = &self.send_custom_stream_flag { let _ = flag.compare_exchange( - INNER_STREAM, - CHANNEL_UNAVAILABLE_STREAM, + StreamMode::Inner as u8, + StreamMode::ChannelUnavailable as u8, Ordering::AcqRel, Ordering::Relaxed, ); @@ -79,44 +186,53 @@ impl ActiveClientStreamState { waker.wake(); } - let con_man = Arc::clone(&self.connection_manager); let addr = self.fingerprint.addr; self.inner = futures::stream::empty::>().boxed(); - tokio::spawn(async move { - let stream_type = if unavailable { - CustomVideoStreamType::ChannelUnavailable - } else { - CustomVideoStreamType::UserConnectionsExhausted - }; - con_man.update_stream_detail(&addr, stream_type).await; - debug_if_enabled!( - "Provider stream stopped due to grace period or unavailable provider channel for {}", - sanitize_sensitive_info(&addr.to_string()) - ); - mgr.release_provider_handle(handle).await; + let video_type = if unavailable { + CustomVideoStreamType::ChannelUnavailable + } else { + CustomVideoStreamType::UserConnectionsExhausted + }; + debug_if_enabled!( + "Provider stream stopped due to grace period or unavailable provider channel for {}", + sanitize_sensitive_info(&addr.to_string()) + ); + self.connection_manager.send_cleanup(CleanupEvent::UpdateDetailAndReleaseProvider { + addr, + video_type, + handle, }); } } fn poll_next_base(&mut self, cx: &mut Context<'_>) -> Poll>> { + self.clear_finished_grace_task(); if let Some(waker) = &self.waker { waker.register(cx.waker()); } - let flag = match &self.send_custom_stream_flag { - Some(flag) => flag.load(Ordering::Acquire), - None => INNER_STREAM, + + if let Some(fut) = self.preempt_cancelled.as_mut() { + if fut.as_mut().poll(cx).is_ready() { + self.stop_provider_stream_preempted(); + return Poll::Ready(None); + } + } + + let mode = match &self.send_custom_stream_flag { + Some(flag) => StreamMode::from_u8(flag.load(Ordering::Acquire)), + None => StreamMode::Inner, }; - // When hold_stream_active is true and flag is GRACE_PENDING, we wait for the grace period + // When hold_stream_active is true and mode is GracePending, we wait for the grace period // check to complete before starting to stream. The grace period task will update the flag. - if flag == GRACE_PENDING { + if mode == StreamMode::GracePending { // Still waiting for grace period check to complete // The grace period task will wake us when done return Poll::Pending; } - if flag == INNER_STREAM { + if mode == StreamMode::Inner { match Pin::new(&mut self.inner).poll_next(cx) { Poll::Ready(Some(Err(e))) => { error!("Inner stream error: {e:?}"); @@ -150,74 +266,25 @@ impl Stream for ActiveClientStream { return res; } - let flag = match &self.state.send_custom_stream_flag { - Some(flag) => flag.load(Ordering::Acquire), - None => INNER_STREAM, + let mode = match &self.state.send_custom_stream_flag { + Some(flag) => StreamMode::from_u8(flag.load(Ordering::Acquire)), + None => StreamMode::Inner, }; - if flag == INNER_STREAM { - return Poll::Ready(None); + // Only stay pending if provisionable AND the grace task is still running. + // Once the grace task completes (handle is None), the decision is final — + // if the stream ended, we must close to avoid hanging the client forever. + let grace_active = self.state.provisionable && self.state.grace_task_handle.is_some(); + + if mode == StreamMode::Inner { + return if grace_active { Poll::Pending } else { Poll::Ready(None) }; } - let buffer_opt = match flag { - USER_EXHAUSTED_STREAM => self.state.custom_video.0.as_mut(), - PROVIDER_EXHAUSTED_STREAM => self.state.custom_video.1.as_mut(), - CHANNEL_UNAVAILABLE_STREAM => self.state.custom_video.2.as_mut(), - PROVISIONING_STREAM => self.state.custom_video.3.as_mut(), - _ => None, - }; - - if let Some(buffer) = buffer_opt { - buffer.register_waker(cx.waker()); - if let Some(chunk) = buffer.next_chunk() { - Poll::Ready(Some(Ok(chunk))) - } else { - // log::warn!("Custom video buffer empty/finished for {} (flag: {})", self.state.fingerprint.addr, flag); - Poll::Ready(None) - } - } else { - // log::warn!("No custom video buffer configured for flag {} for {}", flag, self.state.fingerprint.addr); - // Poll inner again (will return None) - Pin::new(&mut self.state.inner).poll_next(cx) - } - } -} - -impl Drop for ActiveClientStream { - fn drop(&mut self) { - let mgr = Arc::clone(&self.state.connection_manager); - let hndl = self.state.provider_handle.take(); - if let Ok(handle) = tokio::runtime::Handle::try_current() { - handle.spawn(async move { - mgr.release_provider_handle(hndl).await; - }); - } - } -} - -pub(in crate::api) struct ProvisionableActiveClientStream { - state: ActiveClientStreamState, -} - -impl Stream for ProvisionableActiveClientStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let flag = match &self.state.send_custom_stream_flag { - Some(flag) => flag.load(Ordering::Acquire), - None => INNER_STREAM, - }; - - let res = self.state.poll_next_base(cx); - if !matches!(res, Poll::Ready(None)) || flag == INNER_STREAM { - return res; - } - - let buffer_opt = match flag { - USER_EXHAUSTED_STREAM => self.state.custom_video.0.as_mut(), - PROVIDER_EXHAUSTED_STREAM => self.state.custom_video.1.as_mut(), - CHANNEL_UNAVAILABLE_STREAM => self.state.custom_video.2.as_mut(), - PROVISIONING_STREAM => self.state.custom_video.3.as_mut(), + let buffer_opt = match mode { + StreamMode::UserExhausted => self.state.custom_video.user_exhausted.as_mut(), + StreamMode::ProviderExhausted => self.state.custom_video.provider_exhausted.as_mut(), + StreamMode::ChannelUnavailable => self.state.custom_video.unavailable.as_mut(), + StreamMode::Provisioning => self.state.custom_video.provisioning.as_mut(), _ => None, }; @@ -225,7 +292,7 @@ impl Stream for ProvisionableActiveClientStream { buffer.register_waker(cx.waker()); match buffer.next_chunk() { Some(chunk) => Poll::Ready(Some(Ok(chunk))), - None => Poll::Pending, + None => if grace_active { Poll::Pending } else { Poll::Ready(None) }, } } else { Pin::new(&mut self.state.inner).poll_next(cx) @@ -233,13 +300,26 @@ impl Stream for ProvisionableActiveClientStream { } } -impl Drop for ProvisionableActiveClientStream { +impl Drop for ActiveClientStream { fn drop(&mut self) { - let mgr = Arc::clone(&self.state.connection_manager); - let hndl = self.state.provider_handle.take(); - tokio::spawn(async move { - mgr.release_provider_handle(hndl).await; - }); + self.state.stop_grace_task(); + let addr = self.state.fingerprint.addr; + let handle = self.state.provider_handle.take(); + // `provider_handle_released` mirrors `user_stream_released` for the provider slot. + // When preemption already released the handle, `provider_handle` is None and the + // flag is true — sending None here would be a no-op, but the explicit guard makes + // the invariant visible and safe against future call-site additions. + let handle_for_cleanup = if self.state.provider_handle_released { None } else { handle }; + if self.state.user_stream_released { + if !self.state.provider_handle_released { + self.state.provider_handle_released = true; + self.state.connection_manager.send_cleanup(CleanupEvent::ReleaseProviderHandle { handle: handle_for_cleanup }); + } + } else { + self.state.user_stream_released = true; + self.state.provider_handle_released = true; + self.state.connection_manager.send_cleanup(CleanupEvent::ReleaseStreamAndProviderHandle { addr, handle: handle_for_cleanup }); + } } } @@ -264,6 +344,7 @@ pub(crate) async fn create_active_client_stream( let user_agent = req_headers.get(USER_AGENT).map(|h| String::from_utf8_lossy(h.as_bytes())).unwrap_or_default(); let virtual_id = stream_channel.virtual_id; + let is_shared_source_stream = stream_channel.shared && stream_details.request_url.is_some(); app_state .connection_manager .update_connection( @@ -280,53 +361,75 @@ pub(crate) async fn create_active_client_stream( app_state.connection_manager.update_stream_detail(&fingerprint.addr, *cvt).await; } + // Shared broadcaster source (first subscriber path): feed provider bytes directly. + // Grace/custom handling is not needed here because this stream is only the fan-out source. + if is_shared_source_stream { + let stream = match stream_details.stream.take() { + None => { + let provider_handle = stream_details.provider_handle.take(); + app_state.connection_manager.release_provider_handle(provider_handle).await; + futures::stream::empty::>().boxed() + } + Some(stream) => { + let config = app_state.app_config.config.load(); + match config.sleep_timer_mins { + None => stream, + Some(mins) => { + let secs = u32::try_from((u64::from(mins) * 60).min(u64::from(u32::MAX))).unwrap_or(0); + if secs > 0 { + TimedClientStream::new(app_state, stream, secs, fingerprint.addr, virtual_id).boxed() + } else { + stream + } + } + } + } + }; + return stream; + } + let provisioning_info = resolve_grace_period_provisioning(app_state, &stream_details); let has_provisioning = provisioning_info.is_some(); let hold_stream = stream_details.grace_period.hold_stream; - let (grace_stop_flag, waker) = - if grant_user_grace_period || (stream_details.has_grace_period() && stream_details.provider_name.is_some()) { - let waker = Arc::new(AtomicWaker::new()); - let flag = stream_grace_period( - app_state, - &stream_details, - grant_user_grace_period, - user, - fingerprint, - virtual_id, - provisioning_info, - Some(Arc::clone(&waker)), - hold_stream, - ); - let maybe_waker = flag.as_ref().map(|_| waker); - (flag, maybe_waker) - } else { - ( - stream_grace_period( - app_state, - &stream_details, - grant_user_grace_period, - user, - fingerprint, - virtual_id, - provisioning_info, - None, - hold_stream, - ), - None, - ) - }; + let waker = Arc::new(AtomicWaker::new()); + let (grace_stop_flag, grace_task_handle) = if grant_user_grace_period || stream_details.provider_grace_active { + stream_grace_period( + app_state, + &stream_details, + grant_user_grace_period, + user, + fingerprint, + virtual_id, + provisioning_info, + Some(Arc::clone(&waker)), + hold_stream, + ) + } else { + stream_grace_period( + app_state, + &stream_details, + grant_user_grace_period, + user, + fingerprint, + virtual_id, + provisioning_info, + None, + hold_stream, + ) + }; let cfg = &app_state.app_config; let custom_response = cfg.custom_stream_response.load(); - let custom_video = custom_response.as_ref().map_or((None, None, None, None), |c| { - ( - c.user_connections_exhausted.clone(), - c.provider_connections_exhausted.clone(), - c.channel_unavailable.clone(), - c.panel_api_provisioning.clone(), - ) - }); + let custom_video = custom_response.as_ref().map_or( + CustomVideoBuffers { user_exhausted: None, provider_exhausted: None, unavailable: None, provisioning: None }, + |c| CustomVideoBuffers { + user_exhausted: c.user_connections_exhausted.clone(), + provider_exhausted: c.provider_connections_exhausted.clone(), + unavailable: c.channel_unavailable.clone(), + provisioning: c.panel_api_provisioning.clone(), + }, + ); let stream = match stream_details.stream.take() { None => { @@ -350,29 +453,35 @@ pub(crate) async fn create_active_client_stream( } }; + let preempt_cancelled = stream_details.provider_handle + .as_ref() + .and_then(|h| h.cancel_token.as_ref()) + .map(|token| Box::pin(token.clone().cancelled_owned())); + let state = ActiveClientStreamState { inner: stream, + preempt_cancelled, + grace_task_handle, provider_handle: stream_details.provider_handle, send_custom_stream_flag: grace_stop_flag, custom_video, - waker, + provisionable: has_provisioning, + waker: Some(waker), connection_manager: Arc::clone(&app_state.connection_manager), fingerprint: Arc::new(fingerprint.clone()), provider_stopped: false, + user_stream_released: false, + provider_handle_released: false, }; - if has_provisioning { - ProvisionableActiveClientStream { state }.boxed() - } else { - ActiveClientStream { state }.boxed() - } + ActiveClientStream { state }.boxed() } fn resolve_grace_period_provisioning( app_state: &Arc, stream_details: &StreamDetails, ) -> Option { - if stream_details.disable_provider_grace { + if stream_details.disable_provider_grace || !stream_details.provider_grace_active { return None; } let provider_name = stream_details.provider_name.as_deref(); @@ -381,7 +490,7 @@ fn resolve_grace_period_provisioning( return None; } - let stop_signal = Arc::new(AtomicOnceFlag::new()); + let stop_signal = CancellationToken::new(); Some(GraceProvisioningInfo { input, stop_signal }) } @@ -396,12 +505,12 @@ fn stream_grace_period( provisioning_info: Option, waker: Option>, hold_stream: bool, -) -> Option> { +) -> (Option>, Option>) { let active_users = Arc::clone(&app_state.active_users); let active_provider = Arc::clone(&app_state.active_provider); let connection_manager = Arc::clone(&app_state.connection_manager); - let provider_grace_check = if stream_details.has_grace_period() + let provider_grace_check = if stream_details.provider_grace_active && stream_details.provider_name.is_some() && !stream_details.disable_provider_grace { @@ -421,7 +530,9 @@ fn stream_grace_period( debug!("hold stream {hold_stream}"); if provider_grace_check.is_some() || user_grace_check.is_some() { - let stream_strategy_flag = Arc::new(AtomicU8::new(if hold_stream { GRACE_PENDING } else { INNER_STREAM })); + let stream_strategy_flag = Arc::new(AtomicU8::new( + if hold_stream { StreamMode::GracePending as u8 } else { StreamMode::Inner as u8 }, + )); let stream_strategy_flag_copy = Arc::clone(&stream_strategy_flag); let grace_period_millis = stream_details.grace_period.period_millis; @@ -431,79 +542,106 @@ fn stream_grace_period( let reconnect_flag = stream_details.reconnect_flag.clone(); let fingerprint = fingerprint.clone(); let app_state = Arc::clone(app_state); - tokio::spawn(async move { - tokio::time::sleep(tokio::time::Duration::from_millis(grace_period_millis)).await; + // Safety timeout: if async operations inside the grace task stall, force the flag + // out of GRACE_PENDING so the client stream is not hung indefinitely. + // Allow grace_period_millis for the intentional delay plus a 10-second buffer + // for the async connection checks that follow. + let grace_task_timeout = + tokio::time::Duration::from_millis(grace_period_millis.saturating_add(10_000)); + // Clone handles for use in the timeout fallback, in case the inner async block is cancelled. + let flag_for_fallback = Arc::clone(&stream_strategy_flag_copy); + let waker_for_fallback = waker.clone(); + let grace_task_handle = tokio::spawn(async move { + let timed_out = tokio::time::timeout(grace_task_timeout, async move { + tokio::time::sleep(tokio::time::Duration::from_millis(grace_period_millis)).await; - let mut updated = false; - if let Some((username, max_connections)) = user_grace_check { - let active_connections = user_manager.user_connections(&username).await; - if active_connections > max_connections { - stream_strategy_flag_copy.store(USER_EXHAUSTED_STREAM, Ordering::Release); - connection_manager - .update_stream_detail(&fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted) - .await; - // Release the shared stream subscription to stop the subscriber loop - connection_manager.shared_stream_manager.release_connection(&fingerprint.addr, true).await; - info!("User connections exhausted for active clients: {username}"); - updated = true; - } - } - - if !updated { - if let Some(provider_name) = provider_grace_check { - if provider_manager.is_over_limit(&provider_name).await { - if let Some(provisioning_info) = provisioning_info { - stream_strategy_flag_copy.store(PROVISIONING_STREAM, Ordering::Release); - connection_manager - .update_stream_detail(&fingerprint.addr, CustomVideoStreamType::Provisioning) - .await; - debug_if_enabled!( - "Provider grace period exhausted; provisioning for active clients: {provider_name}" - ); - let app_state = Arc::clone(&app_state); - let input = (*provisioning_info.input).clone(); - let stop_signal = provisioning_info.stop_signal; - let addr = fingerprint.addr; - tokio::spawn(async move { - if let Err(err) = - run_panel_api_provisioning_probe(app_state, input, stop_signal, addr, virtual_id) - .await - { - error!("Error running Probe: {err:?}"); - } - }); - } else { - stream_strategy_flag_copy.store(PROVIDER_EXHAUSTED_STREAM, Ordering::Release); - connection_manager - .update_stream_detail( - &fingerprint.addr, - CustomVideoStreamType::ProviderConnectionsExhausted, - ) - .await; - // Release the shared stream subscription to stop the subscriber loop - connection_manager.shared_stream_manager.release_connection(&fingerprint.addr, true).await; - info!("Provider connections exhausted for active clients: {provider_name}"); - } + let mut updated = false; + if let Some((username, max_connections)) = user_grace_check { + let active_connections = user_manager.user_connections(&username).await; + if active_connections > max_connections { + stream_strategy_flag_copy.store(StreamMode::UserExhausted as u8, Ordering::Release); + connection_manager + .update_stream_detail(&fingerprint.addr, CustomVideoStreamType::UserConnectionsExhausted) + .await; + // Release the shared stream subscription to stop the subscriber loop + connection_manager.shared_stream_manager.release_connection(&fingerprint.addr, true).await; + info!("User connections exhausted for active clients: {username}"); updated = true; } } - } - if !updated { - stream_strategy_flag_copy.store(INNER_STREAM, Ordering::Release); - } + if !updated { + if let Some(provider_name) = provider_grace_check { + if provider_manager.is_over_limit(&provider_name).await { + if let Some(provisioning_info) = provisioning_info { + stream_strategy_flag_copy.store(StreamMode::Provisioning as u8, Ordering::Release); + connection_manager + .update_stream_detail(&fingerprint.addr, CustomVideoStreamType::Provisioning) + .await; + debug_if_enabled!( + "Provider grace period exhausted; provisioning for active clients: {provider_name}" + ); + let app_state = Arc::clone(&app_state); + let input = (*provisioning_info.input).clone(); + let stop_signal = provisioning_info.stop_signal; + let addr = fingerprint.addr; + tokio::spawn(async move { + if let Err(err) = + run_panel_api_provisioning_probe(app_state, input, stop_signal, addr, virtual_id) + .await + { + error!("Error running Probe: {err:?}"); + } + }); + } else { + stream_strategy_flag_copy.store(StreamMode::ProviderExhausted as u8, Ordering::Release); + connection_manager + .update_stream_detail( + &fingerprint.addr, + CustomVideoStreamType::ProviderConnectionsExhausted, + ) + .await; + // Release the shared stream subscription to stop the subscriber loop + connection_manager.shared_stream_manager.release_connection(&fingerprint.addr, true).await; + info!("Provider connections exhausted for active clients: {provider_name}"); + } + updated = true; + } + } + } - if updated { - if let Some(flag) = reconnect_flag { - flag.notify(); + if !updated { + stream_strategy_flag_copy.store(StreamMode::Inner as u8, Ordering::Release); + } + + if updated { + if let Some(flag) = reconnect_flag { + flag.cancel(); + } + } + + if let Some(w) = waker.as_ref() { + w.wake(); + } + }) + .await; + + if timed_out.is_err() { + // Grace task exceeded its budget without updating the flag — reset GRACE_PENDING + // to INNER_STREAM so the client stream is not hung indefinitely. + error!("Grace period task timed out; resetting stream flag to prevent client hang"); + let _ = flag_for_fallback.compare_exchange( + StreamMode::GracePending as u8, + StreamMode::Inner as u8, + Ordering::AcqRel, + Ordering::Relaxed, + ); + if let Some(w) = waker_for_fallback.as_ref() { + w.wake(); } } - - if let Some(w) = waker.as_ref() { - w.wake(); - } }); - return Some(stream_strategy_flag); + return (Some(stream_strategy_flag), Some(grace_task_handle)); } - None + (None, None) } diff --git a/backend/src/api/model/streams/buffered_stream.rs b/backend/src/api/model/streams/buffered_stream.rs index e82731120..5715083cf 100644 --- a/backend/src/api/model/streams/buffered_stream.rs +++ b/backend/src/api/model/streams/buffered_stream.rs @@ -1,6 +1,5 @@ use crate::{ api::model::{BoxedProviderStream, StreamError, STREAM_IDLE_TIMEOUT}, - tools::atomic_once_flag::AtomicOnceFlag, }; use futures::{ stream::Stream, @@ -8,78 +7,145 @@ use futures::{ StreamExt, }; use log::debug; -use std::{cmp::max, pin::Pin, sync::Arc}; +use std::{cmp::max, future::Future, pin::Pin, sync::Arc}; use tokio::{ select, - sync::mpsc::{channel, Sender}, + sync::{ + mpsc::{error::TrySendError, channel, Sender}, + Semaphore, + }, time::{sleep, Duration, Instant}, }; use tokio_stream::wrappers::ReceiverStream; +use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned}; pub const CHANNEL_SIZE: usize = 1024; +pub const MAX_BUFFER_BYTES: usize = 5 * 1024 * 1024; pub(in crate::api::model) struct BufferedStream { stream: ReceiverStream>, - close_signal: Arc, + close_cancelled: Pin>, + semaphore: Arc, } impl BufferedStream { pub fn new( stream: BoxedProviderStream, buffer_size: usize, - client_close_signal: Arc, + client_close_signal: CancellationToken, _url: &str, ) -> Self { - // TODO make channel_size based on bytes not entries + // Item-count limit remains as a secondary cap; byte-level backpressure + // is enforced via `MAX_BUFFER_BYTES` and `Semaphore`. let (tx, rx) = channel(max(buffer_size, CHANNEL_SIZE)); - tokio::spawn(Self::buffer_stream(tx, stream, Arc::clone(&client_close_signal))); - Self { stream: ReceiverStream::new(rx), close_signal: client_close_signal } + let semaphore = Arc::new(Semaphore::new(MAX_BUFFER_BYTES)); + tokio::spawn(Self::buffer_stream( + tx, + stream, + client_close_signal.clone(), + Arc::clone(&semaphore), + )); + Self { + stream: ReceiverStream::new(rx), + close_cancelled: Box::pin(client_close_signal.cancelled_owned()), + semaphore, + } } async fn buffer_stream( tx: Sender>, mut stream: BoxedProviderStream, - client_close_signal: Arc, + client_close_signal: CancellationToken, + semaphore: Arc, ) { let idle_timeout = Duration::from_secs(STREAM_IDLE_TIMEOUT); let idle = sleep(idle_timeout); tokio::pin!(idle); - while client_close_signal.is_active() { + while !client_close_signal.is_cancelled() { select! { + biased; + () = client_close_signal.cancelled() => { + break; + } () = &mut idle => { debug!("Buffered stream idle for too long, closing"); - client_close_signal.notify(); + client_close_signal.cancel(); break; } chunk = stream.next() => { idle.as_mut().reset(Instant::now() + idle_timeout); match chunk { Some(Ok(chunk)) => { - let chunk_len = chunk.len(); - if tx.send(Ok(chunk)).await.is_err() { - debug!("Buffered stream channel closed before delivering {chunk_len} bytes to client"); - client_close_signal.notify(); + let chunk_len = chunk.len(); + // Cap permits at MAX_BUFFER_BYTES per chunk. A single chunk larger + // than the cap consumes fewer permits than its actual byte count, so + // the semaphore may temporarily allow more bytes in the channel than + // MAX_BUFFER_BYTES. This is an intentional trade-off: upstream + // providers are expected to emit chunks well below this limit; the + // inaccuracy is bounded to a single oversized chunk and is self- + // correcting once that chunk is delivered. + let permits = chunk_len.min(MAX_BUFFER_BYTES); + if permits > 0 { + let acquired = select! { + biased; + () = client_close_signal.cancelled() => None, + permit = Arc::clone(&semaphore).acquire_many_owned(u32::try_from(permits).unwrap_or(u32::MAX)) => permit.ok(), + }; + let Some(permit) = acquired else { + client_close_signal.cancel(); break; - } + }; + permit.forget(); } - Some(Err(err)) => { - let err_msg = err.to_string(); - if tx.send(Err(err)).await.is_err() { - debug!("Buffered stream dropped stream error due to closed receiver: {err_msg}"); - client_close_signal.notify(); + let send_res = match tx.try_send(Ok(chunk)) { + Ok(()) => Ok(()), + Err(TrySendError::Full(item)) => { + select! { + biased; + () = client_close_signal.cancelled() => Err(()), + res = tx.send(item) => res.map_err(|_| ()), + } } + Err(TrySendError::Closed(_)) => Err(()), + }; + if send_res.is_err() { + if permits > 0 { + semaphore.add_permits(permits); + } + debug!("Buffered stream channel closed before delivering {chunk_len} bytes to client"); + client_close_signal.cancel(); break; } - None => { - debug!("Upstream provider completed buffered stream"); - break; + } + Some(Err(err)) => { + let err_msg = err.to_string(); + let send_err_res = match tx.try_send(Err(err)) { + Ok(()) => Ok(()), + Err(TrySendError::Full(item)) => { + select! { + biased; + () = client_close_signal.cancelled() => Err(()), + res = tx.send(item) => res.map_err(|_| ()), + } + } + Err(TrySendError::Closed(_)) => Err(()), + }; + if send_err_res.is_err() { + debug!("Buffered stream dropped stream error due to closed receiver: {err_msg}"); + client_close_signal.cancel(); } + break; + } + None => { + debug!("Upstream provider completed buffered stream"); + break; + } } } } } - if !client_close_signal.is_active() { + if client_close_signal.is_cancelled() { debug!("Client close signal fired; buffered stream exiting"); } drop(tx); @@ -90,10 +156,17 @@ impl Stream for BufferedStream { type Item = Result; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - if self.close_signal.is_active() { - Pin::new(&mut self.get_mut().stream).poll_next(cx) - } else { + let this = self.get_mut(); + if this.close_cancelled.as_mut().poll(cx).is_ready() { Poll::Ready(None) + } else { + match Pin::new(&mut this.stream).poll_next(cx) { + Poll::Ready(Some(Ok(bytes))) => { + this.semaphore.add_permits(bytes.len().min(MAX_BUFFER_BYTES)); + Poll::Ready(Some(Ok(bytes))) + } + other => other, + } } } } diff --git a/backend/src/api/model/streams/client_stream.rs b/backend/src/api/model/streams/client_stream.rs index 32223ba1a..a90c8b3dc 100644 --- a/backend/src/api/model/streams/client_stream.rs +++ b/backend/src/api/model/streams/client_stream.rs @@ -1,6 +1,5 @@ use crate::{ api::model::{BoxedProviderStream, StreamError}, - tools::atomic_once_flag::AtomicOnceFlag, utils::trace_if_enabled, }; use bytes::Bytes; @@ -8,6 +7,7 @@ use futures::Stream; use log::trace; use shared::utils::sanitize_sensitive_info; use std::{ + future::Future, pin::Pin, sync::{ atomic::{AtomicUsize, Ordering}, @@ -15,57 +15,74 @@ use std::{ }, task::Poll, }; +use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned}; /// This stream counts the send bytes for reconnecting to the actual position and /// sets the `close_signal` if the client drops the connection. -#[repr(align(64))] pub(in crate::api::model) struct ClientStream { inner: BoxedProviderStream, - close_signal: Arc, - total_bytes: Arc>, + close_signal: CancellationToken, + close_cancelled: Pin>, + total_bytes: Option>, url: String, } impl ClientStream { pub(crate) fn new( inner: BoxedProviderStream, - close_signal: Arc, - total_bytes: Arc>, + close_signal: CancellationToken, + total_bytes: Option>, url: &str, ) -> Self { - Self { inner, close_signal, total_bytes, url: url.to_string() } + let close_cancelled = Box::pin(close_signal.clone().cancelled_owned()); + Self { inner, close_signal, close_cancelled, total_bytes, url: url.to_string() } } } + impl Stream for ClientStream { type Item = Result; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { - if self.close_signal.is_active() { - match Pin::as_mut(&mut self.inner).poll_next(cx) { + fn poll_next(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.close_cancelled.as_mut().poll(cx).is_ready() { + return Poll::Ready(None); + } + // Bound empty-chunk skips per poll invocation. A misbehaving provider + // that sends an endless run of empty keep-alive chunks must not spin + // the executor indefinitely: after 10 consecutive empty chunks we + // yield back to the runtime via wake_by_ref + Poll::Pending. + let mut empty_chunk_count = 0u32; + loop { + match Pin::new(&mut this.inner).poll_next(cx) { Poll::Ready(Some(Ok(bytes))) => { if bytes.is_empty() { - trace!("client stream empty bytes"); - // Empty payload signals upstream closure; notify and let consumer see final chunk - self.close_signal.notify(); - } else if let Some(counter) = self.total_bytes.as_ref() { - counter.fetch_add(bytes.len(), Ordering::AcqRel); + // Skip keep-alive empty chunks rather than treating them as EOF. + // Some providers send empty chunks as heartbeats; closing on them + // would prematurely terminate valid streams. + empty_chunk_count += 1; + if empty_chunk_count > 10 { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + trace!("client stream: skipping empty keep-alive chunk"); + continue; } - - Poll::Ready(Some(Ok(bytes))) + if let Some(counter) = &this.total_bytes { + counter.fetch_add(bytes.len(), Ordering::Relaxed); + } + return Poll::Ready(Some(Ok(bytes))); } Poll::Ready(None) => { - self.close_signal.notify(); - Poll::Ready(None) + this.close_signal.cancel(); + return Poll::Ready(None); } - Poll::Pending => Poll::Pending, Poll::Ready(Some(Err(err))) => { trace!("client stream error: {err}"); - self.close_signal.notify(); - Poll::Ready(Some(Err(err))) + this.close_signal.cancel(); + return Poll::Ready(Some(Err(err))); } + Poll::Pending => return Poll::Pending, } - } else { - Poll::Ready(None) } } } @@ -73,6 +90,6 @@ impl Stream for ClientStream { impl Drop for ClientStream { fn drop(&mut self) { trace_if_enabled!("Client disconnected {}", sanitize_sensitive_info(&self.url)); - self.close_signal.notify(); + self.close_signal.cancel(); } } diff --git a/backend/src/api/model/streams/provider_stream.rs b/backend/src/api/model/streams/provider_stream.rs index a386fa08f..8a49394af 100644 --- a/backend/src/api/model/streams/provider_stream.rs +++ b/backend/src/api/model/streams/provider_stream.rs @@ -2,12 +2,12 @@ use crate::{ api::{ api_utils::{try_unwrap_body, HeaderFilter}, model::{ - stream::ProviderStreamResponse, AppState, CustomVideoStream, ProvisioningStream, ThrottledStream, + stream::{BoxedProviderStream, ProviderStreamResponse}, + AppState, CleanupEvent, CustomVideoStream, ProvisioningStream, ThrottledStream, TimedClientStream, TransportStreamBuffer, }, }, model::AppConfig, - tools::atomic_once_flag::AtomicOnceFlag, }; use axum::response::IntoResponse; use log::trace; @@ -15,6 +15,7 @@ use reqwest::StatusCode; use serde::{Deserialize, Deserializer, Serialize, Serializer}; use shared::model::PlaylistItemType; use std::{fmt, net::SocketAddr, str::FromStr, sync::Arc}; +use tokio_util::sync::CancellationToken; #[derive(Debug, Copy, Clone)] pub enum CustomVideoStreamType { @@ -71,7 +72,37 @@ impl<'de> Deserialize<'de> for CustomVideoStreamType { } } +fn prepare_video_headers(headers: &[(String, String)]) -> Vec<(String, String)> { + let mut h: Vec<(String, String)> = headers + .iter() + .filter(|(key, _)| { + !(key.eq_ignore_ascii_case("content-type") + || key.eq_ignore_ascii_case("content-length") + || key.eq_ignore_ascii_case("range") + || key.eq_ignore_ascii_case("content-range") + || key.eq_ignore_ascii_case("accept-ranges")) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + h.push(("content-type".to_string(), "video/mp2t".to_string())); + h +} + +fn get_custom_stream_response_timeout_secs(cfg: &AppConfig) -> u32 { + cfg.config.load().custom_stream_response_timeout_secs +} + +fn apply_custom_stream_timeout(cfg: &AppConfig, stream: BoxedProviderStream) -> BoxedProviderStream { + let timeout_secs = get_custom_stream_response_timeout_secs(cfg); + if timeout_secs == 0 { + stream + } else { + Box::pin(TimedClientStream::new_without_kick(stream, timeout_secs)) + } +} + fn create_video_stream( + cfg: &AppConfig, stream_type: CustomVideoStreamType, video_buffer: Option<&TransportStreamBuffer>, headers: &[(String, String)], @@ -79,15 +110,11 @@ fn create_video_stream( ) -> ProviderStreamResponse { if let Some(video) = video_buffer { trace!("{log_message}"); - let mut response_headers: Vec<(String, String)> = headers - .iter() - .filter(|(key, _)| !(key.eq("content-type") || key.eq("content-length") || key.contains("range"))) - .map(|(key, value)| (key.clone(), value.clone())) - .collect(); - response_headers.push(("content-type".to_string(), "video/mp2t".to_string())); + let stream = + apply_custom_stream_timeout(cfg, Box::pin(ThrottledStream::new(CustomVideoStream::new(video.clone()), 8000))); ( - Some(Box::pin(ThrottledStream::new(CustomVideoStream::new(video.clone()), 8000))), - Some((response_headers, StatusCode::OK, None, Some(stream_type))), + Some(stream), + Some((prepare_video_headers(headers), StatusCode::OK, None, Some(stream_type))), ) } else { (None, None) @@ -102,6 +129,7 @@ pub fn create_channel_unavailable_stream( let custom_stream_response = cfg.custom_stream_response.load(); let video = custom_stream_response.as_ref().and_then(|c| c.channel_unavailable.as_ref()); create_video_stream( + cfg, CustomVideoStreamType::ChannelUnavailable, video, headers, @@ -116,6 +144,7 @@ pub fn create_user_connections_exhausted_stream( let custom_stream_response = cfg.custom_stream_response.load(); let video = custom_stream_response.as_ref().and_then(|c| c.user_connections_exhausted.as_ref()); create_video_stream( + cfg, CustomVideoStreamType::UserConnectionsExhausted, video, headers, @@ -130,6 +159,7 @@ pub fn create_provider_connections_exhausted_stream( let custom_stream_response = cfg.custom_stream_response.load(); let video = custom_stream_response.as_ref().and_then(|c| c.provider_connections_exhausted.as_ref()); create_video_stream( + cfg, CustomVideoStreamType::ProviderConnectionsExhausted, video, headers, @@ -141,6 +171,7 @@ pub fn create_user_account_expired_stream(cfg: &AppConfig, headers: &[(String, S let custom_stream_response = cfg.custom_stream_response.load(); let video = custom_stream_response.as_ref().and_then(|c| c.user_account_expired.as_ref()); create_video_stream( + cfg, CustomVideoStreamType::UserAccountExpired, video, headers, @@ -152,6 +183,7 @@ pub fn create_panel_api_provisioning_stream(cfg: &AppConfig, headers: &[(String, let custom_stream_response = cfg.custom_stream_response.load(); let video = custom_stream_response.as_ref().and_then(|c| c.panel_api_provisioning.as_ref()); create_video_stream( + cfg, CustomVideoStreamType::Provisioning, video, headers, @@ -162,29 +194,24 @@ pub fn create_panel_api_provisioning_stream(cfg: &AppConfig, headers: &[(String, pub fn create_panel_api_provisioning_stream_with_stop( cfg: &AppConfig, headers: &[(String, String)], - stop_signal: Arc, + stop_signal: CancellationToken, ) -> ProviderStreamResponse { let custom_stream_response = cfg.custom_stream_response.load(); let video = custom_stream_response.as_ref().and_then(|c| c.panel_api_provisioning.as_ref()); if let Some(video) = video { trace!("Streaming response panel api provisioning"); - let mut response_headers: Vec<(String, String)> = headers - .iter() - .filter(|(key, _)| !(key.eq("content-type") || key.eq("content-length") || key.contains("range"))) - .map(|(key, value)| (key.clone(), value.clone())) - .collect(); - response_headers.push(("content-type".to_string(), "video/mp2t".to_string())); let stream = ProvisioningStream::new(video.clone(), stop_signal); + let stream = apply_custom_stream_timeout(cfg, Box::pin(ThrottledStream::new(stream, 8000))); ( - Some(Box::pin(ThrottledStream::new(stream, 8000))), - Some((response_headers, StatusCode::OK, None, Some(CustomVideoStreamType::Provisioning))), + Some(stream), + Some((prepare_video_headers(headers), StatusCode::OK, None, Some(CustomVideoStreamType::Provisioning))), ) } else { (None, None) } } -pub async fn create_custom_video_stream_response( +pub fn create_custom_video_stream_response( app_state: &Arc, addr: &SocketAddr, video_response: CustomVideoStreamType, @@ -201,8 +228,10 @@ pub async fn create_custom_video_stream_response( CustomVideoStreamType::UserAccountExpired => create_user_account_expired_stream(config, &[]), CustomVideoStreamType::Provisioning => create_panel_api_provisioning_stream(config, &[]), } { - app_state.connection_manager.update_stream_detail(addr, video_response).await; - app_state.connection_manager.release_provider_connection(addr).await; + app_state.connection_manager.send_cleanup(CleanupEvent::UpdateDetailAndReleaseProviderConnection { + addr: *addr, + video_type: video_response, + }); let mut builder = axum::response::Response::builder().status(status_code); for (key, value) in headers { builder = builder.header(key, value); diff --git a/backend/src/api/model/streams/provider_stream_factory.rs b/backend/src/api/model/streams/provider_stream_factory.rs index d3bbf8af0..cab5a5491 100644 --- a/backend/src/api/model/streams/provider_stream_factory.rs +++ b/backend/src/api/model/streams/provider_stream_factory.rs @@ -8,7 +8,6 @@ use crate::{ }, }, model::{ConfigProvider, ReverseProxyDisabledHeaderConfig}, - tools::atomic_once_flag::AtomicOnceFlag, utils::{ debug_if_enabled, request::{ @@ -39,6 +38,7 @@ use std::{ }, time::{Duration, Instant}, }; +use tokio_util::sync::CancellationToken; use url::Url; use shared::model::PlaylistItemType; use shared::utils::DEFAULT_USER_AGENT; @@ -65,8 +65,8 @@ pub struct ProviderStreamFactoryOptions { url: Url, headers: HeaderMap, default_user_agent: Option, - range_bytes: Arc>, - reconnect_flag: Arc, + range_bytes: Option>, + reconnect_flag: CancellationToken, provider: Option>, } @@ -101,9 +101,9 @@ impl ProviderStreamFactoryOptions { let url = stream_url.clone(); let range_bytes = if matches!(item_type, PlaylistItemType::Live | PlaylistItemType::LiveUnknown) { - Arc::new(requested_range.map(AtomicUsize::new)) + requested_range.map(|v| Arc::new(AtomicUsize::new(v))) } else { - Arc::new(Some(AtomicUsize::new(requested_range.unwrap_or(0)))) + Some(Arc::new(AtomicUsize::new(requested_range.unwrap_or(0)))) }; let mut flags = ProviderStreamFactoryFlagsSet::new(); if stream_options.stream_retry { @@ -127,7 +127,7 @@ impl ProviderStreamFactoryOptions { addr, flags, buffer_size, - reconnect_flag: Arc::new(AtomicOnceFlag::new()), + reconnect_flag: CancellationToken::new(), url, headers, default_user_agent, @@ -146,17 +146,14 @@ impl ProviderStreamFactoryOptions { #[inline] fn is_buffer_enabled(&self) -> bool { self.flags.contains(ProviderStreamFactoryFlags::BufferEnabled) } - #[inline] - fn is_shared_stream(&self) -> bool { self.flags.contains(ProviderStreamFactoryFlags::ShareStream) } - #[inline] pub(crate) fn get_buffer_size(&self) -> usize { self.buffer_size } #[inline] - pub fn get_reconnect_flag_clone(&self) -> Arc { Arc::clone(&self.reconnect_flag) } + pub fn get_reconnect_flag_clone(&self) -> CancellationToken { self.reconnect_flag.clone() } #[inline] - pub fn cancel_reconnect(&self) { self.reconnect_flag.notify(); } + pub fn cancel_reconnect(&self) { self.reconnect_flag.cancel(); } #[inline] pub fn get_url(&self) -> &Url { &self.url } @@ -172,18 +169,14 @@ impl ProviderStreamFactoryOptions { #[inline] pub fn get_total_bytes_send(&self) -> Option { - self.range_bytes.as_ref().as_ref().map(|atomic| atomic.load(Ordering::Acquire)) + self.range_bytes.as_ref().map(|atomic| atomic.load(Ordering::Acquire)) } - // pub fn get_range_bytes(&self) -> &Arc> { - // &self.range_bytes - // } + #[inline] + pub fn get_range_bytes_clone(&self) -> Option> { self.range_bytes.clone() } #[inline] - pub fn get_range_bytes_clone(&self) -> Arc> { Arc::clone(&self.range_bytes) } - - #[inline] - pub fn should_continue(&self) -> bool { self.reconnect_flag.is_active() } + pub fn should_continue(&self) -> bool { !self.reconnect_flag.is_cancelled() } #[inline] pub fn was_range_requested(&self) -> bool { self.flags.contains(ProviderStreamFactoryFlags::RangeRequested) } @@ -582,18 +575,17 @@ pub async fn create_provider_stream( stream_options: ProviderStreamFactoryOptions, ) -> Option { let client_stream_factory = |stream, reconnect_flag, range_cnt| { - let stream = - if !stream_options.is_piped() && stream_options.is_buffer_enabled() && !stream_options.is_shared_stream() { - BufferedStream::new( - stream, - stream_options.get_buffer_size(), - stream_options.get_reconnect_flag_clone(), - stream_options.get_url_as_str(), - ) - .boxed() - } else { - stream - }; + let stream = if !stream_options.is_piped() && stream_options.is_buffer_enabled() { + BufferedStream::new( + stream, + stream_options.get_buffer_size(), + stream_options.get_reconnect_flag_clone(), + stream_options.get_url_as_str(), + ) + .boxed() + } else { + stream + }; ClientStream::new(stream, reconnect_flag, range_cnt, stream_options.get_url_as_str()).boxed() }; @@ -608,7 +600,7 @@ pub async fn create_provider_stream( let continue_signal = stream_options.get_reconnect_flag_clone(); if is_media_stream_or_not_piped && stream_options.should_reconnect() { - let continue_client_signal = Arc::clone(&continue_signal); + let continue_client_signal = continue_signal.clone(); let continue_streaming_signal = continue_client_signal.clone(); let stream_options_provider = stream_options.clone(); let app_state_clone = Arc::clone(app_state); @@ -619,15 +611,36 @@ pub async fn create_provider_stream( let continue_streaming = continue_streaming_signal.clone(); let app_state_clone = Arc::clone(&app_state_clone); async move { - if continue_streaming.is_active() { + if continue_streaming.is_cancelled() { + app_state_clone.connection_manager.release_provider_connection(&stream_opts.addr).await; + None + } else { match get_provider_stream(&app_state_clone, &client, &stream_opts).await { - Ok(Some((stream, _info))) => Some((stream, ())), + Ok(Some((stream, info))) => { + // If we reconnected with a byte offset and the provider responded + // 200 OK instead of 206 Partial Content, the stream would restart + // from byte 0, producing a corrupt video. Cancel the reconnect + // to avoid silently delivering corrupt data to the client. + if let Some((_, status, _, _)) = &info { + let current_pos = stream_opts.get_total_bytes_send().unwrap_or(0); + if current_pos > 0 && *status != StatusCode::PARTIAL_CONTENT { + warn!( + "Reconnect aborted: provider ignored Range request \ + (responded {status} instead of 206). URL: {}", + sanitize_sensitive_info(stream_opts.get_log_url().as_ref()) + ); + continue_streaming.cancel(); + return None; + } + } + Some((stream, ())) + } Ok(None) => { app_state_clone .connection_manager .release_provider_connection(&stream_opts.addr) .await; - continue_streaming.notify(); + continue_streaming.cancel(); if let (Some(boxed_provider_stream), _response_info) = create_channel_unavailable_stream( &app_state_clone.app_config, @@ -644,7 +657,7 @@ pub async fn create_provider_stream( .connection_manager .release_provider_connection(&stream_opts.addr) .await; - continue_streaming.notify(); + continue_streaming.cancel(); if let (Some(boxed_provider_stream), _response_info) = create_channel_unavailable_stream( &app_state_clone.app_config, @@ -657,9 +670,6 @@ pub async fn create_provider_stream( None } } - } else { - app_state_clone.connection_manager.release_provider_connection(&stream_opts.addr).await; - None } } }) @@ -668,7 +678,7 @@ pub async fn create_provider_stream( Some(( client_stream_factory( init_stream.chain(unfold).boxed(), - Arc::clone(&continue_client_signal), + continue_client_signal.clone(), stream_options.get_range_bytes_clone(), ) .boxed(), @@ -678,7 +688,7 @@ pub async fn create_provider_stream( Some(( client_stream_factory( init_stream.boxed(), - Arc::clone(&continue_signal), + continue_signal.clone(), stream_options.get_range_bytes_clone(), ) .boxed(), diff --git a/backend/src/api/model/streams/provisioning_stream.rs b/backend/src/api/model/streams/provisioning_stream.rs index ca9397168..1629cf6e5 100644 --- a/backend/src/api/model/streams/provisioning_stream.rs +++ b/backend/src/api/model/streams/provisioning_stream.rs @@ -1,36 +1,51 @@ -use crate::{ - api::model::{StreamError, TransportStreamBuffer}, - tools::atomic_once_flag::AtomicOnceFlag, -}; +use crate::api::model::{StreamError, TransportStreamBuffer, STREAM_IDLE_TIMEOUT}; use bytes::Bytes; use futures::Stream; +use log::debug; use std::{ + future::Future, pin::Pin, - sync::Arc, task::{Context, Poll}, + time::Duration, }; +use tokio::time::{sleep_until, Instant, Sleep}; +use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned}; pub struct ProvisioningStream { buffer: TransportStreamBuffer, - stop_signal: Arc, + stop_cancelled: Pin>, + idle_deadline: Pin>, } impl ProvisioningStream { - pub fn new(buffer: TransportStreamBuffer, stop_signal: Arc) -> Self { Self { buffer, stop_signal } } + pub fn new(buffer: TransportStreamBuffer, stop_signal: CancellationToken) -> Self { + let idle_deadline = Box::pin(sleep_until(Instant::now() + Duration::from_secs(STREAM_IDLE_TIMEOUT))); + Self { buffer, stop_cancelled: Box::pin(stop_signal.cancelled_owned()), idle_deadline } + } } impl Stream for ProvisioningStream { type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - if !self.stop_signal.is_active() { + if self.stop_cancelled.as_mut().poll(cx).is_ready() { return Poll::Ready(None); } - self.buffer.register_waker(cx.waker()); - match self.buffer.next_chunk() { - Some(chunk) => Poll::Ready(Some(Ok(chunk))), - None => Poll::Pending, + if self.idle_deadline.as_mut().poll(cx).is_ready() { + debug!("Provisioning stream idle timeout reached; terminating client stream"); + return Poll::Ready(None); + } + + if let Some(chunk) = self.buffer.next_chunk() { + // Reset idle deadline on each delivered chunk. + self.idle_deadline + .as_mut() + .reset(Instant::now() + Duration::from_secs(STREAM_IDLE_TIMEOUT)); + Poll::Ready(Some(Ok(chunk))) + } else { + self.buffer.register_waker(cx.waker()); + Poll::Pending } } } diff --git a/backend/src/api/model/streams/shared_stream_manager.rs b/backend/src/api/model/streams/shared_stream_manager.rs index 3b6de54a8..2cbb68f5b 100644 --- a/backend/src/api/model/streams/shared_stream_manager.rs +++ b/backend/src/api/model/streams/shared_stream_manager.rs @@ -109,16 +109,19 @@ async fn send_burst_buffer( start_buffer: &VecDeque>, client_tx: &Sender, cancellation_token: &CancellationToken, -) { +) -> usize { + let mut sent = 0_usize; for buf in start_buffer { if cancellation_token.is_cancelled() { - return; + return sent; } if let Err(err) = client_tx.send(buf.as_ref().clone()).await { debug!("Failed sending burst-buffer chunk to client: {err}"); - return; // stop on send error + return sent; // stop on send error } + sent = sent.saturating_add(1); } + sent } /// Represents the state of a shared provider URL. @@ -158,6 +161,7 @@ impl SharedStreamState { } } + #[allow(clippy::too_many_lines)] async fn subscribe( &self, addr: &SocketAddr, @@ -191,17 +195,31 @@ impl SharedStreamState { let timeout_duration = Duration::from_secs(300); // 5 minutes let mut last_active = Instant::now(); let mut last_lag_log = Instant::now().checked_sub(Duration::from_secs(10)).unwrap_or_else(Instant::now); + let mut consecutive_lag_count: u32 = 0; + let subscriber_buf_size = self.buf_size; let address = *addr; + let subscriber_started_at = Instant::now(); let handle = tokio::spawn(async move { // initial burst buffer let snapshot = { let buffer = burst_buffer.read().await; buffer.snapshot() }; - send_burst_buffer(&snapshot, &client_tx_clone, &cancel_token).await; + let sent_burst_chunks = send_burst_buffer(&snapshot, &client_tx_clone, &cancel_token).await; + if sent_burst_chunks > 0 { + debug_if_enabled!( + "Shared stream subscriber {} replayed {sent_burst_chunks} burst chunks after {} ms", + sanitize_sensitive_info(&address.to_string()), + subscriber_started_at.elapsed().as_millis() + ); + } let mut loop_cnt = 0; + let mut first_live_chunk_logged = false; + let mut startup_chunks_sent = 0usize; + let mut startup_bytes_sent = 0usize; + let mut startup_stats_logged = false; loop { tokio::select! { biased; @@ -225,15 +243,41 @@ impl SharedStreamState { result = broadcast_rx.recv() => { match result { Ok(data) => { + consecutive_lag_count = 0; // If the client press pause, skip if client_tx_clone.is_closed() { continue; } + let chunk_len = data.len(); if let Err(err) = client_tx.send(data).await { debug!("Shared stream client send error: {address} {err}"); break; } + if !first_live_chunk_logged { + debug_if_enabled!( + "Shared stream subscriber {} received first live chunk after {} ms", + sanitize_sensitive_info(&address.to_string()), + subscriber_started_at.elapsed().as_millis() + ); + first_live_chunk_logged = true; + } + if !startup_stats_logged { + startup_chunks_sent = startup_chunks_sent.saturating_add(1); + startup_bytes_sent = startup_bytes_sent.saturating_add(chunk_len); + if subscriber_started_at.elapsed() >= Duration::from_secs(5) { + debug_if_enabled!( + "Shared stream subscriber {} startup throughput: chunks={} bytes={} over {} ms (queue_used={}/{})", + sanitize_sensitive_info(&address.to_string()), + startup_chunks_sent, + startup_bytes_sent, + subscriber_started_at.elapsed().as_millis(), + subscriber_buf_size.saturating_sub(client_tx_clone.capacity()), + subscriber_buf_size + ); + startup_stats_logged = true; + } + } loop_cnt += 1; last_active = Instant::now(); @@ -243,17 +287,28 @@ impl SharedStreamState { } } Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => { + consecutive_lag_count = consecutive_lag_count.saturating_add(1); if last_lag_log.elapsed() > Duration::from_secs(5) { let buffered_bytes = { let buffer = burst_buffer_for_log.read().await; buffer.current_bytes }; - warn!("Shared stream client lagged behind {address}. Skipped {skipped} messages (buffered {buffered_bytes} bytes, yield counter {yield_counter})"); + warn!( + "Shared stream client lagged behind {address}. Skipped {skipped} messages \ + (buffered {buffered_bytes} bytes, yield counter {yield_counter}, \ + consecutive lags={consecutive_lag_count})" + ); last_lag_log = Instant::now(); } - // Sleep to prevent CPU spin when client is persistently lagging behind. - // This provides backpressure for slow clients that can't keep up. - sleep(Duration::from_millis(50)).await; + // Exponential backoff for persistently lagging clients: 50 ms base, + // doubling up to 1 600 ms, to provide increasing backpressure without + // permanently blocking the subscriber task. + let backoff_ms = 50_u64 + .saturating_mul( + 1_u64.checked_shl(consecutive_lag_count.min(5)).unwrap_or(u64::MAX), + ) + .min(1_600); + sleep(Duration::from_millis(backoff_ms)).await; } Err(_) => break, } @@ -270,6 +325,7 @@ impl SharedStreamState { (convert_stream(ReceiverStream::new(client_rx).boxed()), provider) } + #[allow(clippy::too_many_lines)] fn broadcast(&self, stream_url: &str, bytes_stream: S, shared_streams: Arc) where S: Stream> + Unpin + 'static + Send, @@ -280,12 +336,17 @@ impl SharedStreamState { let sender = self.broadcaster.clone(); let stop_token = self.stop_token.clone(); let burst_buffer = self.burst_buffer.clone(); + let broadcast_started_at = Instant::now(); tokio::spawn(async move { let mut counter = 0usize; let idle_timeout = Duration::from_secs(STREAM_IDLE_TIMEOUT); let idle = sleep(idle_timeout); tokio::pin!(idle); + let mut first_source_chunk_logged = false; + let mut startup_chunks_seen = 0usize; + let mut startup_bytes_seen = 0usize; + let mut startup_stats_logged = false; loop { tokio::select! { @@ -309,6 +370,28 @@ impl SharedStreamState { idle.as_mut().reset(Instant::now() + idle_timeout); match chunk { Some(Ok(data)) => { + if !first_source_chunk_logged { + debug_if_enabled!( + "Shared stream source produced first chunk for {} after {} ms", + sanitize_sensitive_info(&streaming_url), + broadcast_started_at.elapsed().as_millis() + ); + first_source_chunk_logged = true; + } + if !startup_stats_logged { + startup_chunks_seen = startup_chunks_seen.saturating_add(1); + startup_bytes_seen = startup_bytes_seen.saturating_add(data.len()); + if broadcast_started_at.elapsed() >= Duration::from_secs(5) { + debug_if_enabled!( + "Shared stream source startup throughput for {}: chunks={} bytes={} over {} ms", + sanitize_sensitive_info(&streaming_url), + startup_chunks_seen, + startup_bytes_seen, + broadcast_started_at.elapsed().as_millis() + ); + startup_stats_logged = true; + } + } let arc_data = Arc::new(data); { let mut buffer = burst_buffer.write().await; @@ -426,6 +509,46 @@ impl SharedStreamManager { } } + /// Tears down a shared stream that was preempted via priority eviction. + /// Stops the broadcast task and removes subscriber mappings, but does NOT + /// release the provider handle (the preemption path already handles that). + /// + /// # Precondition + /// The caller **must** have already released the provider allocation for this stream + /// before calling this method. Failing to do so will leak the provider slot because + /// `provider_guard` inside the shared state is intentionally not released here. + pub async fn teardown_preempted_stream(&self, stream_url: &str) { + let shared_state_opt = { + let mut shared_streams = self.shared_streams.write().await; + + let remove_keys: Vec = shared_streams + .key_by_addr + .iter() + .filter_map(|(addr, url)| if url == stream_url { Some(*addr) } else { None }) + .collect(); + for k in remove_keys { + shared_streams.key_by_addr.remove(&k); + } + + shared_streams.by_key.remove(stream_url) + }; + + if let Some(shared_state) = shared_state_opt { + debug_if_enabled!( + "Tearing down preempted shared stream {}", + sanitize_sensitive_info(stream_url) + ); + + for handle in shared_state.task_handles.write().await.drain(..) { + handle.abort(); + } + + // Cancel stop_token to terminate the broadcast task. + // Do NOT release provider_guard — the preemption caller already released the allocation. + shared_state.stop_token.cancel(); + } + } + pub async fn release_connection(&self, addr: &SocketAddr, send_stop_signal: bool) { let (stream_url, shared_state) = { let shared_streams = self.shared_streams.read().await; @@ -444,6 +567,21 @@ impl SharedStreamManager { (tx, is_empty, subs.len()) }; + let Some(client_stop_signal) = tx else { + trace_if_enabled!( + "Ignoring duplicate shared stream release for {} (already removed)", + sanitize_sensitive_info(&addr.to_string()) + ); + return; + }; + + // Remove the addr → url mapping eagerly so key_by_addr does not leak when + // other subscribers are still active (unregister only runs when is_empty). + { + let mut shared_streams = self.shared_streams.write().await; + shared_streams.key_by_addr.remove(addr); + } + debug_if_enabled!( "Shared stream subscriber removed {}; remaining subscribers={remaining}", sanitize_sensitive_info(&addr.to_string()) @@ -460,9 +598,7 @@ impl SharedStreamManager { } } - if let Some(client_stop_signal) = tx { - client_stop_signal.cancel(); - } + client_stop_signal.cancel(); } } @@ -518,6 +654,7 @@ impl SharedStreamManager { S: Stream> + Unpin + 'static + Send, E: std::fmt::Debug + Send, { + let registration_started_at = Instant::now(); let buf_size = CHANNEL_SIZE.max(buffer_size); let config = app_state.app_config.config.load(); let min_buffer_bytes = resolve_min_burst_buffer_bytes(&config); @@ -525,6 +662,11 @@ impl SharedStreamManager { app_state.shared_stream_manager.register(addr, stream_url, Arc::clone(&shared_state)).await; app_state.active_provider.make_shared_connection(addr, stream_url).await; let subscribed_stream = Self::subscribe_shared_stream(app_state, stream_url, addr).await; + debug_if_enabled!( + "Shared stream startup register+subscribe completed for {} in {} ms", + sanitize_sensitive_info(stream_url), + registration_started_at.elapsed().as_millis() + ); shared_state.broadcast(stream_url, bytes_stream, Arc::clone(&app_state.shared_stream_manager)); debug_if_enabled!( "Created shared provider stream {} (channel_capacity={buf_size}, burst_buffer_min={min_buffer_bytes} bytes)", @@ -548,3 +690,155 @@ impl SharedStreamManager { } } } + +#[cfg(test)] +mod tests { + use super::{SharedStreamManager, SharedStreamState, CHANNEL_SIZE}; + use crate::{ + api::model::{ActiveProviderManager, EventManager}, + model::{AppConfig, Config, ConfigInput, SourcesConfig}, + utils::FileLockManager, + }; + use arc_swap::{ArcSwap, ArcSwapOption}; + use shared::{ + model::{ConfigPaths, InputFetchMethod, InputType}, + utils::Internable, + }; + use std::{collections::HashMap, net::SocketAddr, sync::Arc}; + use tokio_util::sync::CancellationToken; + + fn create_test_app_config() -> AppConfig { + let input = Arc::new(ConfigInput { + id: 1, + name: "provider_1".intern(), + input_type: InputType::Xtream, + headers: HashMap::default(), + url: "http://provider-1.example".to_string(), + username: Some("user1".to_string()), + password: Some("pass1".to_string()), + enabled: true, + priority: 0, + max_connections: 1, + method: InputFetchMethod::default(), + aliases: None, + ..ConfigInput::default() + }); + + let sources = SourcesConfig { inputs: vec![input], ..SourcesConfig::default() }; + + AppConfig { + config: Arc::new(ArcSwap::from_pointee(Config::default())), + sources: Arc::new(ArcSwap::from_pointee(sources)), + hdhomerun: Arc::new(ArcSwapOption::default()), + api_proxy: Arc::new(ArcSwapOption::default()), + file_locks: Arc::new(FileLockManager::default()), + paths: Arc::new(ArcSwap::from_pointee(ConfigPaths { + home_path: String::new(), + config_path: String::new(), + storage_path: String::new(), + config_file_path: String::new(), + sources_file_path: String::new(), + mapping_file_path: None, + mapping_files_used: None, + template_file_path: None, + template_files_used: None, + api_proxy_file_path: String::new(), + custom_stream_response_path: None, + })), + custom_stream_response: Arc::new(ArcSwapOption::default()), + access_token_secret: [0; 32], + encrypt_secret: [0; 16], + ffprobe_available: Arc::default(), + } + } + + #[tokio::test] + async fn test_duplicate_release_connection_is_idempotent_with_remaining_subscribers() { + let app_cfg = create_test_app_config(); + let event_manager = Arc::new(EventManager::new()); + let provider_manager = Arc::new(ActiveProviderManager::new(&app_cfg, &event_manager)); + let shared_manager = Arc::new(SharedStreamManager::new(provider_manager)); + + let stream_url = "https://example.invalid/live/stream.ts"; + let addr_1: SocketAddr = "127.0.0.1:41001".parse().unwrap_or_else(|_| unreachable!()); + let addr_2: SocketAddr = "127.0.0.1:41002".parse().unwrap_or_else(|_| unreachable!()); + + let state = Arc::new(SharedStreamState::new(Vec::new(), CHANNEL_SIZE.max(8), None, 1024)); + + { + let mut reg = shared_manager.shared_streams.write().await; + reg.by_key.insert(stream_url.to_string(), Arc::clone(&state)); + reg.key_by_addr.insert(addr_1, stream_url.to_string()); + reg.key_by_addr.insert(addr_2, stream_url.to_string()); + } + + { + let mut subs = state.subscribers.write().await; + subs.insert(addr_1, CancellationToken::new()); + subs.insert(addr_2, CancellationToken::new()); + } + + shared_manager.release_connection(&addr_1, false).await; + { + let subs = state.subscribers.read().await; + assert_eq!(subs.len(), 1, "first release should remove exactly one subscriber"); + assert!(subs.contains_key(&addr_2), "remaining subscriber must stay registered"); + } + + // Duplicate release for the same address should be ignored. + shared_manager.release_connection(&addr_1, false).await; + { + let subs = state.subscribers.read().await; + assert_eq!(subs.len(), 1, "duplicate release must not modify subscribers"); + assert!(subs.contains_key(&addr_2), "remaining subscriber must stay registered"); + } + + let reg = shared_manager.shared_streams.read().await; + assert!(reg.by_key.contains_key(stream_url), "shared stream must stay registered while one subscriber remains"); + } + + #[tokio::test] + async fn test_duplicate_release_connection_is_idempotent_with_single_subscriber() { + let app_cfg = create_test_app_config(); + let event_manager = Arc::new(EventManager::new()); + let provider_manager = Arc::new(ActiveProviderManager::new(&app_cfg, &event_manager)); + let shared_manager = Arc::new(SharedStreamManager::new(provider_manager)); + + let stream_url = "https://example.invalid/live/single.ts"; + let addr_1: SocketAddr = "127.0.0.1:42001".parse().unwrap_or_else(|_| unreachable!()); + let state = Arc::new(SharedStreamState::new(Vec::new(), CHANNEL_SIZE.max(8), None, 1024)); + + { + let mut reg = shared_manager.shared_streams.write().await; + reg.by_key.insert(stream_url.to_string(), Arc::clone(&state)); + reg.key_by_addr.insert(addr_1, stream_url.to_string()); + } + + { + let mut subs = state.subscribers.write().await; + subs.insert(addr_1, CancellationToken::new()); + } + + shared_manager.release_connection(&addr_1, false).await; + { + let reg = shared_manager.shared_streams.read().await; + assert!(!reg.by_key.contains_key(stream_url), "stream should be unregistered after last subscriber leaves"); + assert!(!reg.key_by_addr.contains_key(&addr_1), "address mapping should be removed"); + } + { + let subs = state.subscribers.read().await; + assert!(subs.is_empty(), "subscriber list should be empty after release"); + } + + shared_manager.release_connection(&addr_1, false).await; + { + let reg = shared_manager.shared_streams.read().await; + assert!(!reg.by_key.contains_key(stream_url), "duplicate release must keep stream unregistered"); + assert!(!reg.key_by_addr.contains_key(&addr_1), "duplicate release must keep address mapping absent"); + } + { + let subs = state.subscribers.read().await; + assert!(subs.is_empty(), "duplicate release must keep subscribers empty"); + } + } +} diff --git a/backend/src/api/model/streams/timed_client_stream.rs b/backend/src/api/model/streams/timed_client_stream.rs index c57f70e72..3d1c68a67 100644 --- a/backend/src/api/model/streams/timed_client_stream.rs +++ b/backend/src/api/model/streams/timed_client_stream.rs @@ -9,19 +9,32 @@ use shared::{ utils::{default_kick_secs, sanitize_sensitive_info}, }; use std::{ + future::Future, net::SocketAddr, pin::Pin, sync::Arc, task::Poll, - time::{Duration, Instant}, + time::Duration, }; +use tokio::time::{sleep_until, Instant, Sleep}; + +enum TimeoutAction { + Kick { + app_state: Arc, + addr: SocketAddr, + virtual_id: VirtualId, + }, + Stop, +} pub struct TimedClientStream { inner: BoxedProviderStream, - deadline: Instant, - app_state: Arc, - addr: SocketAddr, - virtual_id: VirtualId, + /// `Sleep` future pinned in-place. Polling it registers a waker with the + /// Tokio runtime so this task is woken at the deadline regardless of whether + /// the inner stream is producing data — unlike a plain `Instant` comparison, + /// which is only evaluated when `poll_next` happens to be called. + deadline: Pin>, + timeout_action: TimeoutAction, } impl TimedClientStream { @@ -32,33 +45,47 @@ impl TimedClientStream { addr: SocketAddr, virtual_id: VirtualId, ) -> Self { - let deadline = Instant::now() + Duration::from_secs(u64::from(duration)); - Self { inner, deadline, app_state: Arc::clone(app_state), addr, virtual_id } + let deadline = Box::pin(sleep_until(Instant::now() + Duration::from_secs(u64::from(duration)))); + Self { + inner, + deadline, + timeout_action: TimeoutAction::Kick { app_state: Arc::clone(app_state), addr, virtual_id }, + } + } + + pub(crate) fn new_without_kick(inner: BoxedProviderStream, duration: u32) -> Self { + let deadline = Box::pin(sleep_until(Instant::now() + Duration::from_secs(u64::from(duration)))); + Self { inner, deadline, timeout_action: TimeoutAction::Stop } } } + impl Stream for TimedClientStream { type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll> { - if Instant::now() >= self.deadline { - let kick_secs = self - .app_state - .app_config - .config - .load() - .web_ui - .as_ref() - .map_or_else(default_kick_secs, |wc| wc.kick_secs); - let connection_manager = Arc::clone(&self.app_state.connection_manager); - let addr = self.addr; - let virtual_id = self.virtual_id; - debug_if_enabled!( - "TimedClient stream exceeds time limit. Closing stream with virtual_id {virtual_id} for addr: {}", - sanitize_sensitive_info(&addr.to_string()) - ); - tokio::spawn(async move { - connection_manager.kick_connection(&addr, virtual_id, kick_secs).await; - }); + // Poll the sleep future first. This registers a waker so the executor + // wakes this task exactly when the deadline fires — even if the upstream + // provider is stalled and emitting no data. + if self.deadline.as_mut().poll(cx).is_ready() { + if let TimeoutAction::Kick { app_state, addr, virtual_id } = &self.timeout_action { + let kick_secs = app_state + .app_config + .config + .load() + .web_ui + .as_ref() + .map_or_else(default_kick_secs, |wc| wc.kick_secs); + let connection_manager = Arc::clone(&app_state.connection_manager); + let addr = *addr; + let virtual_id = *virtual_id; + debug_if_enabled!( + "TimedClient stream exceeds time limit. Closing stream with virtual_id {virtual_id} for addr: {}", + sanitize_sensitive_info(&addr.to_string()) + ); + tokio::spawn(async move { + connection_manager.kick_connection(&addr, virtual_id, kick_secs).await; + }); + } return Poll::Ready(None); } Pin::as_mut(&mut self.inner).poll_next(cx) diff --git a/backend/src/api/panel_api.rs b/backend/src/api/panel_api.rs index 7014ba7eb..188b126c5 100644 --- a/backend/src/api/panel_api.rs +++ b/backend/src/api/panel_api.rs @@ -14,7 +14,6 @@ use crate::{ csv_patch_batch_append, csv_patch_batch_remove_expired, csv_patch_batch_sort_by_exp_date, csv_patch_batch_update_credentials, csv_patch_batch_update_exp_date, get_csv_file_path, }, - tools::atomic_once_flag::AtomicOnceFlag, utils::{debug_if_enabled, format_http_status, persist_source_config, read_sources_file_from_path}, }; use axum::http::{Method, StatusCode}; @@ -45,6 +44,7 @@ use std::{ sync::Arc, time::{Duration, Instant}, }; +use tokio_util::sync::CancellationToken; use url::Url; #[derive(Debug, Clone)] @@ -3479,16 +3479,20 @@ async fn wait_for_panel_api_account_ready( pub(crate) async fn run_panel_api_provisioning_probe( app_state: Arc, input: ConfigInput, - stop_signal: Arc, + stop_signal: CancellationToken, addr: SocketAddr, virtual_id: VirtualId, ) -> Result<(), TuliproxError> { + if stop_signal.is_cancelled() { + return Ok(()); + } + let Some(panel_cfg) = input.panel_api.as_ref() else { debug_if_enabled!( "panel_api provisioning probe skipped (missing config) for input {}", sanitize_sensitive_info(&input.name) ); - stop_signal.notify(); + stop_signal.cancel(); let _ = app_state.connection_manager.kick_connection(&addr, virtual_id, 0).await; return Ok(()); }; @@ -3497,7 +3501,7 @@ pub(crate) async fn run_panel_api_provisioning_probe( "panel_api provisioning probe skipped (panel_api.enabled false) for input {}", sanitize_sensitive_info(&input.name) ); - stop_signal.notify(); + stop_signal.cancel(); let _ = app_state.connection_manager.kick_connection(&addr, virtual_id, 0).await; return Ok(()); } @@ -3506,7 +3510,7 @@ pub(crate) async fn run_panel_api_provisioning_probe( "panel_api provisioning probe skipped (panel_api.url empty) for input {}", sanitize_sensitive_info(&input.name) ); - stop_signal.notify(); + stop_signal.cancel(); let _ = app_state.connection_manager.kick_connection(&addr, virtual_id, 0).await; return Ok(()); } @@ -3524,7 +3528,10 @@ pub(crate) async fn run_panel_api_provisioning_probe( ); let deadline = Instant::now() + Duration::from_secs(max_wait_secs); - let outcome = try_provision_account_on_exhausted(&app_state, &input).await; + let outcome = tokio::select! { + () = stop_signal.cancelled() => return Ok(()), + outcome = try_provision_account_on_exhausted(&app_state, &input) => outcome, + }; let credentials = outcome.as_ref().map(PanelApiProvisionOutcome::credentials); if let Some(outcome) = outcome.as_ref() { @@ -3542,13 +3549,16 @@ pub(crate) async fn run_panel_api_provisioning_probe( let Some((username, password)) = credentials else { if max_wait_secs > 0 { - tokio::time::sleep(Duration::from_secs(max_wait_secs)).await; + tokio::select! { + () = stop_signal.cancelled() => return Ok(()), + () = tokio::time::sleep(Duration::from_secs(max_wait_secs)) => {} + } } debug_if_enabled!( "panel_api provisioning probe timeout reached for input {} (no credentials)", sanitize_sensitive_info(&input.name) ); - stop_signal.notify(); + stop_signal.cancel(); debug_if_enabled!( "panel_api provisioning closing client connection for input {} addr={}", sanitize_sensitive_info(&input.name), @@ -3562,13 +3572,16 @@ pub(crate) async fn run_panel_api_provisioning_probe( let Some(test_url) = build_panel_api_test_url(&resolved_url, username, password) else { if max_wait_secs > 0 { - tokio::time::sleep(Duration::from_secs(max_wait_secs)).await; + tokio::select! { + () = stop_signal.cancelled() => return Ok(()), + () = tokio::time::sleep(Duration::from_secs(max_wait_secs)) => {} + } } debug_if_enabled!( "panel_api provisioning probe failed to build test url for input {}", sanitize_sensitive_info(&input.name) ); - stop_signal.notify(); + stop_signal.cancel(); let _ = app_state.connection_manager.kick_connection(&addr, virtual_id, 0).await; return Ok(()); }; @@ -3579,7 +3592,11 @@ pub(crate) async fn run_panel_api_provisioning_probe( while Instant::now() < deadline { attempt += 1; debug_if_enabled!("panel_api provisioning probe attempt {}", attempt); - match probe_panel_api_test_url(&app_state, &test_url, probe_method).await { + let probe_result = tokio::select! { + () = stop_signal.cancelled() => return Ok(()), + probe_result = probe_panel_api_test_url(&app_state, &test_url, probe_method) => probe_result, + }; + match probe_result { Ok(status) => { debug_if_enabled!( "panel_api provisioning probe status: '{}' url: {}", @@ -3612,7 +3629,10 @@ pub(crate) async fn run_panel_api_provisioning_probe( } let remaining = deadline.checked_duration_since(now).unwrap_or_default(); let sleep_for = if remaining < probe_delay { remaining } else { probe_delay }; - tokio::time::sleep(sleep_for).await; + tokio::select! { + () = stop_signal.cancelled() => return Ok(()), + () = tokio::time::sleep(sleep_for) => {} + } } if ready { @@ -3628,7 +3648,10 @@ pub(crate) async fn run_panel_api_provisioning_probe( attempt ); } - stop_signal.notify(); + if stop_signal.is_cancelled() { + return Ok(()); + } + stop_signal.cancel(); debug_if_enabled!( "panel_api provisioning closing client connection for input {} addr={}", sanitize_sensitive_info(&input.name), @@ -3646,10 +3669,13 @@ pub fn create_panel_api_provisioning_stream_details( addr: SocketAddr, virtual_id: VirtualId, ) -> StreamDetails { - let stop_signal = Arc::new(AtomicOnceFlag::new()); + let stop_signal = CancellationToken::new(); let headers = [("connection".to_string(), "close".to_string())]; - let (stream, stream_info) = - create_panel_api_provisioning_stream_with_stop(&app_state.app_config, &headers, Arc::clone(&stop_signal)); + let (stream, stream_info) = create_panel_api_provisioning_stream_with_stop( + &app_state.app_config, + &headers, + stop_signal.clone(), + ); if stream.is_none() { debug_if_enabled!( @@ -3663,6 +3689,7 @@ pub fn create_panel_api_provisioning_stream_details( provider_name, request_url: None, grace_period: *grace_period_options, + provider_grace_active: false, disable_provider_grace: true, reconnect_flag: None, provider_handle: None, @@ -3671,7 +3698,7 @@ pub fn create_panel_api_provisioning_stream_details( let app_state_clone = Arc::clone(app_state); let input_clone = input.clone(); - let stop_clone = Arc::clone(&stop_signal); + let stop_clone = stop_signal.clone(); tokio::spawn(async move { if let Err(err) = run_panel_api_provisioning_probe(app_state_clone, input_clone, stop_clone, addr, virtual_id).await @@ -3686,6 +3713,7 @@ pub fn create_panel_api_provisioning_stream_details( provider_name, request_url: None, grace_period: *grace_period_options, + provider_grace_active: false, disable_provider_grace: true, reconnect_flag: None, provider_handle: None, diff --git a/backend/src/api/scheduler.rs b/backend/src/api/scheduler.rs index b2e600016..4bdd53014 100644 --- a/backend/src/api/scheduler.rs +++ b/backend/src/api/scheduler.rs @@ -10,12 +10,13 @@ use crate::{ }; use chrono::{DateTime, FixedOffset, Local}; use cron::Schedule; -use shared::{model::ScheduleTaskType, utils::interner_gc}; +use shared::{model::ScheduleTaskType, utils::{interner_gc, interner_len}}; use std::{ str::FromStr, sync::Arc, time::{Duration, Instant, SystemTime}, }; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; pub fn datetime_to_instant(datetime: DateTime) -> Instant { @@ -67,19 +68,123 @@ pub fn exec_scheduler( let app_state_clone = Arc::clone(app_state); let http_client = client.clone(); let cancel_token = cancel.clone(); - tokio::spawn(async move { - start_scheduler(http_client, expression.as_str(), task_type, app_state_clone, exec_targets, cancel_token) - .await; - }); + + match task_type { + ScheduleTaskType::PlaylistUpdate => { + // Bounded channel with capacity 1: if an update is already pending or + // running and the scheduler fires again, the extra signal is dropped + // (deduplicated). This prevents redundant runs from piling up when + // updates are slow or blocked waiting for the playlist lock. + let (tx, rx) = mpsc::channel::<()>(1); + + // Cron trigger: fires at scheduled times and notifies the worker. + let trigger_cancel = cancel_token.clone(); + tokio::spawn(async move { + start_playlist_trigger(expression.as_str(), trigger_cancel, tx).await; + }); + + // Worker: processes triggers one at a time, blocking on the playlist + // lock when another update is active. Cancels cleanly on shutdown. + let worker_client = http_client; + let worker_state = app_state_clone; + let worker_targets = exec_targets; + tokio::spawn(async move { + run_playlist_update_worker(worker_client, worker_state, worker_targets, rx, cancel_token).await; + }); + } + ScheduleTaskType::LibraryScan | ScheduleTaskType::GeoIpUpdate => { + tokio::spawn(async move { + start_scheduler( + http_client, + expression.as_str(), + task_type, + app_state_clone, + exec_targets, + cancel_token, + ) + .await; + }); + } + } } } +/// Cron trigger for playlist updates. Fires at each scheduled time and sends a +/// unit signal to the worker via a bounded channel. If the channel is already +/// full (one update is pending), `try_send` silently drops the signal — +/// deduplication at zero cost. +async fn start_playlist_trigger(expression: &str, cancel: CancellationToken, tx: mpsc::Sender<()>) { + match Schedule::from_str(expression) { + Ok(schedule) => { + let offset = *Local::now().offset(); + loop { + let mut upcoming = schedule.upcoming(offset).take(1); + if let Some(datetime) = upcoming.next() { + tokio::select! { + biased; + () = cancel.cancelled() => break, + () = tokio::time::sleep_until(tokio::time::Instant::from(datetime_to_instant(datetime))) => { + // If the channel is full, there is already one pending run queued. + // This fire is covered by that pending run; drop it. + let _ = tx.try_send(()); + } + } + } + } + } + Err(err) => exit!("Failed to start scheduler: {err}"), + } +} + +/// Worker for playlist updates. Waits for trigger signals from the cron trigger +/// and runs updates one at a time. Blocks on `acquire_playlist_lock` while +/// another update source (manual trigger, metadata update, another schedule) holds +/// the lock — the worker resumes automatically once the lock is released. +async fn run_playlist_update_worker( + client: reqwest::Client, + app_state: Arc, + targets: Arc, + mut rx: mpsc::Receiver<()>, + cancel: CancellationToken, +) { + loop { + tokio::select! { + biased; + () = cancel.cancelled() => break, + msg = rx.recv() => { + if msg.is_none() { + break; // Sender dropped (scheduler task exited) + } + run_playlist_update_inner(&client, &app_state, &targets).await; + } + } + } +} + +async fn run_playlist_update_inner(client: &reqwest::Client, app_state: &Arc, targets: &Arc) { + exec_processing( + client, + Arc::clone(&app_state.app_config), + Arc::clone(targets), + Some(Arc::clone(&app_state.event_manager)), + Some(Arc::clone(app_state)), + Some(app_state.playlists.clone()), + Some(app_state.update_guard.clone()), + app_state.get_disabled_headers(), + Some(Arc::clone(&app_state.active_provider)), + Some(Arc::clone(&app_state.metadata_manager)), + None, + None, + ) + .await; +} + async fn start_scheduler( client: reqwest::Client, expression: &str, task_type: ScheduleTaskType, app_state: Arc, - targets: Arc, + _targets: Arc, cancel: CancellationToken, ) { match Schedule::from_str(expression) { @@ -91,9 +196,7 @@ async fn start_scheduler( tokio::select! { () = tokio::time::sleep_until(tokio::time::Instant::from(datetime_to_instant(datetime))) => { match task_type { - ScheduleTaskType::PlaylistUpdate => { - run_playlist_update(&client, &app_state, &targets); - } + ScheduleTaskType::PlaylistUpdate => unreachable!("handled by channel-based path"), ScheduleTaskType::LibraryScan => { run_library_scan(&client, &app_state); } @@ -113,35 +216,6 @@ async fn start_scheduler( } } -fn run_playlist_update(client: &reqwest::Client, app_state: &Arc, targets: &Arc) { - let client = client.clone(); - let app_state = Arc::clone(app_state); - let targets = Arc::clone(targets); - tokio::spawn(async move { - let app_config = Arc::clone(&app_state.app_config); - let event_manager = Arc::clone(&app_state.event_manager); - let playlist_state = app_state.playlists.clone(); - let provider_manager = Arc::clone(&app_state.active_provider); - let disabled_headers = app_state.get_disabled_headers(); - let metadata_manager = Arc::clone(&app_state.metadata_manager); - exec_processing( - &client, - app_config, - targets, - Some(event_manager), - Some(app_state.clone()), - Some(playlist_state), - Some(app_state.update_guard.clone()), - disabled_headers, - Some(provider_manager), - Some(metadata_manager), - None, - None, - ) - .await; - }); -} - fn run_library_scan(client: &reqwest::Client, app_state: &Arc) { let config = app_state.app_config.config.load(); if let Some(lib_config) = config.library.as_ref() { @@ -191,7 +265,7 @@ pub fn get_process_targets( let inputs: Vec = user_targets.inputs.iter().filter(|&id| process_targets.inputs.contains(id)).copied().collect(); let targets: Vec = - user_targets.targets.iter().filter(|&id| process_targets.inputs.contains(id)).copied().collect(); + user_targets.targets.iter().filter(|&id| process_targets.targets.contains(id)).copied().collect(); let target_names: Vec = user_targets .target_names .iter() @@ -206,12 +280,23 @@ pub fn get_process_targets( // TODO Consider making the GC interval configurable. // The 180-second interval is hardcoded. For deployments with different memory/performance characteristics, a configurable interval might be useful. + +/// Minimum number of interned strings required before GC runs. +/// Below this threshold the write-lock overhead outweighs the benefit. +const INTERNER_GC_MIN_POOL_SIZE: usize = 100; + pub fn exec_interner_prune(app_state: &Arc) { let app_state = Arc::clone(app_state); tokio::spawn({ async move { loop { tokio::time::sleep(Duration::from_secs(180)).await; + // Skip GC entirely when the pool is too small — acquiring a write + // lock on the global interner briefly blocks all concurrent interns, + // so the cost only pays off when there are enough strings to clean. + if interner_len() < INTERNER_GC_MIN_POOL_SIZE { + continue; + } if let Some(permit) = app_state.update_guard.try_playlist() { // Gate check: ensure updates aren't in progress; permit dropped to allow concurrent updates during GC drop(permit); diff --git a/backend/src/main.rs b/backend/src/main.rs index 87f4553be..ae0901cb8 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -227,6 +227,10 @@ fn print_info(app_config: &AppConfig) { info!("Api Proxy File: {:?}", &paths.api_proxy_file_path); info!("Mapping path: {:?}", &paths.mapping_file_path.as_ref().map_or_else(|| "not used", |v| v.as_str())); info!("Template path: {:?}", &paths.template_file_path.as_ref().map_or_else(|| "not used", |v| v.as_str())); + info!( + "Resources path: {:?}", + &paths.custom_stream_response_path.as_ref().map_or_else(|| "not used", |v| v.as_str()) + ); if let Some(mapping_paths) = paths.mapping_files_used.as_ref() { for mapping_path in mapping_paths { @@ -247,9 +251,6 @@ fn print_info(app_config: &AppConfig) { if let Some(metadata_update) = config.metadata_update.as_ref() { info!("Metadata path: {}", metadata_update.cache_path); } - if let Some(resource_path) = paths.custom_stream_response_path.as_ref() { - info!("Resource path: {resource_path}"); - } } fn get_file_paths(args: &Args) -> ConfigPaths { diff --git a/backend/src/model/config/app.rs b/backend/src/model/config/app.rs index 89745aebe..4cf3b3f5e 100644 --- a/backend/src/model/config/app.rs +++ b/backend/src/model/config/app.rs @@ -408,7 +408,8 @@ impl AppConfig { } let path = PathBuf::from(custom_stream_response_path); - let path = utils::make_path_absolute(&path, &config.storage_dir); + let home_path = self.paths.load().home_path.clone(); + let path = utils::make_path_absolute(&path, home_path.as_str()); let paths = self.paths.load_full(); let mut new_paths = paths.as_ref().clone(); diff --git a/backend/src/model/config/base.rs b/backend/src/model/config/base.rs index 3e425f2dd..2cc3bc337 100644 --- a/backend/src/model/config/base.rs +++ b/backend/src/model/config/base.rs @@ -72,6 +72,7 @@ pub struct Config { pub mapping_path: Option, pub template_path: Option, pub custom_stream_response_path: Option, + pub custom_stream_response_timeout_secs: u32, pub video: Option, pub metadata_update: Option, pub schedules: Option>, @@ -295,6 +296,7 @@ impl From<&ConfigDto> for Config { mapping_path: dto.mapping_path.clone(), template_path: dto.template_path.clone(), custom_stream_response_path: dto.custom_stream_response_path.clone(), + custom_stream_response_timeout_secs: dto.custom_stream_response_timeout_secs, video: dto.video.as_ref().map(Into::into), metadata_update: dto.metadata_update.as_ref().map(Into::into), schedules: dto.schedules.as_ref().map(|s| s.iter().map(Into::into).collect()), diff --git a/backend/src/model/config/library.rs b/backend/src/model/config/library.rs index fe9cdffc8..750f8b310 100644 --- a/backend/src/model/config/library.rs +++ b/backend/src/model/config/library.rs @@ -203,7 +203,7 @@ mod tests { library: Some(LibraryConfigDto { enabled: true, scan_directories: vec![LibraryScanDirectoryDto { - path: std::env::temp_dir().to_string_lossy().to_string(), + path: tempfile::env::temp_dir().to_string_lossy().to_string(), ..Default::default() }], ..Default::default() diff --git a/backend/src/repository/bplustree.rs b/backend/src/repository/bplustree.rs index 76d8cbef3..9e1694785 100644 --- a/backend/src/repository/bplustree.rs +++ b/backend/src/repository/bplustree.rs @@ -2242,7 +2242,17 @@ where /// Internal store without locking, used for compaction or initial save. fn store_internal(&mut self, filepath: &Path) -> io::Result { - let tempfile = NamedTempFile::new()?; + let tempfile = if let Some(parent_dir) = filepath.parent() { + if let Ok(file) = NamedTempFile::new_in(parent_dir) { + file + } else { + let temp_dir = tempfile::env::temp_dir(); + NamedTempFile::new_in(&temp_dir)? + } + } else { + let temp_dir = tempfile::env::temp_dir(); + NamedTempFile::new_in(&temp_dir)? + }; let mut file = utils::file_writer(&tempfile); let mut buffer = vec![0u8; PAGE_SIZE_USIZE]; diff --git a/backend/src/repository/bplustree_migration.rs b/backend/src/repository/bplustree_migration.rs index eaa50c614..34d022f05 100644 --- a/backend/src/repository/bplustree_migration.rs +++ b/backend/src/repository/bplustree_migration.rs @@ -524,12 +524,18 @@ fn migrate_user_db_schema(db_path: &Path, merge_guard_path: &Path) -> io::Result return Ok(false); } - // Already V3? (15-field sequence — rmp_serde fails on shorter sequences) - if BPlusTree::::load(db_path).is_ok() { - return Ok(false); + // Try legacy schemas first to preserve explicit upgrade behavior: + // V1 -> V3, then V2 -> V3. + if let Ok(tree) = BPlusTree::::load(db_path) { + let mut v3_tree: BPlusTree = BPlusTree::new(); + for (key, v1) in &tree { + v3_tree.insert(key.clone(), StoredApiUserV3::from_v1(v1)); + } + v3_tree.store(db_path)?; + create_user_db_merge_guard(merge_guard_path)?; + return Ok(true); } - // Try V2 → V3 if let Ok(tree) = BPlusTree::::load(db_path) { let mut v3_tree: BPlusTree = BPlusTree::new(); for (key, v2) in &tree { @@ -540,15 +546,10 @@ fn migrate_user_db_schema(db_path: &Path, merge_guard_path: &Path) -> io::Result return Ok(true); } - // Try V1 → V3 - if let Ok(tree) = BPlusTree::::load(db_path) { - let mut v3_tree: BPlusTree = BPlusTree::new(); - for (key, v1) in &tree { - v3_tree.insert(key.clone(), StoredApiUserV3::from_v1(v1)); - } - v3_tree.store(db_path)?; - create_user_db_merge_guard(merge_guard_path)?; - return Ok(true); + // If legacy decoding failed, accept that DB may already be V3. + // Do not rewrite or persist anything in that case. + if BPlusTree::::load(db_path).is_ok() { + return Ok(false); } Err(io::Error::new( @@ -809,4 +810,85 @@ mod tests { Ok(()) } + + #[test] + fn user_db_schema_migration_v2_to_v3_creates_merge_guard() -> io::Result<()> { + let temp = tempdir()?; + let db_path = temp.path().join(storage_const::API_USER_DB_FILE); + let merge_guard_path = user_db_merge_guard_path(temp.path()); + + let mut v2_tree: BPlusTree = BPlusTree::new(); + v2_tree.insert( + "alice".to_string(), + StoredApiUserV2 { + target: "channels".to_string(), + username: "alice".to_string(), + password: "secret".to_string(), + token: Some("token".to_string()), + proxy: ProxyType::Reverse(None), + server: Some("srv".to_string()), + epg_timeshift: Some("1".to_string()), + epg_request_timeshift: Some("2".to_string()), + created_at: Some(1), + exp_date: Some(2), + max_connections: Some(3), + status: Some(ProxyUserStatus::Active), + ui_enabled: true, + comment: Some("note".to_string()), + }, + ); + let _ = v2_tree.store(&db_path)?; + assert!(!merge_guard_path.exists()); + + let migrated = migrate_user_db_schema(&db_path, &merge_guard_path)?; + assert!(migrated); + assert!(merge_guard_path.exists()); + + let v3_tree = BPlusTree::::load(&db_path)?; + let user = v3_tree + .query(&"alice".to_string()) + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "alice missing after migration"))?; + assert_eq!(user.username, "alice"); + assert_eq!(user.epg_request_timeshift.as_deref(), Some("2")); + assert_eq!(user.priority, None); + + Ok(()) + } + + #[test] + fn user_db_schema_v3_is_detected_without_writing_merge_guard() -> io::Result<()> { + let temp = tempdir()?; + let db_path = temp.path().join(storage_const::API_USER_DB_FILE); + let merge_guard_path = user_db_merge_guard_path(temp.path()); + + let mut v3_tree: BPlusTree = BPlusTree::new(); + v3_tree.insert( + "bob".to_string(), + StoredApiUserV3 { + target: "channels".to_string(), + username: "bob".to_string(), + password: "secret".to_string(), + token: None, + proxy: ProxyType::Reverse(None), + server: None, + epg_timeshift: None, + epg_request_timeshift: None, + created_at: None, + exp_date: None, + max_connections: Some(1), + status: Some(ProxyUserStatus::Active), + ui_enabled: true, + comment: None, + priority: Some(5), + }, + ); + let _ = v3_tree.store(&db_path)?; + assert!(!merge_guard_path.exists()); + + let migrated = migrate_user_db_schema(&db_path, &merge_guard_path)?; + assert!(!migrated); + assert!(!merge_guard_path.exists()); + + Ok(()) + } } diff --git a/backend/src/tools/atomic_once_flag.rs b/backend/src/tools/atomic_once_flag.rs deleted file mode 100644 index 5228de997..000000000 --- a/backend/src/tools/atomic_once_flag.rs +++ /dev/null @@ -1,48 +0,0 @@ -use std::sync::atomic::{AtomicBool, Ordering}; - -/// A flag that is initially active (`true`) and can only be disabled once. -/// Once the flag is disabled by calling [`notify()`], it remains inactive (`false`) forever. -/// -/// ## Use Case -/// This type is useful when you need a one-way toggle to mark a resource or state as "finalized", -/// "shut down", or "disabled". -/// -/// ## Example -/// ```rust -/// let flag = AtomicOnceFlag::new(); -/// assert!(flag.is_active()); -/// flag.notify(); -/// assert!(!flag.is_active()); -#[derive(Debug)] -pub struct AtomicOnceFlag { - enabled: AtomicBool, -} - -impl Default for AtomicOnceFlag { - fn default() -> Self { - Self::new() - } -} - -impl AtomicOnceFlag { - /// Creates a new `AtomicOnceFlag`. - pub fn new() -> Self { - Self { - enabled: AtomicBool::new(true), - } - } - - /// Disables the flag. After calling this method, `is_active()` will always return `false`. - /// - /// This operation is atomic and uses the specified memory ordering. - pub fn notify(&self) { - self.enabled.store(false, Ordering::Release); - } - - /// Checks if the flag is still active. - /// - /// Returns `true` if the flag is active (initial state). Returns `false` if the flag has been disabled. - pub fn is_active(&self) -> bool { - self.enabled.load(Ordering::Acquire) - } -} \ No newline at end of file diff --git a/backend/src/tools/mod.rs b/backend/src/tools/mod.rs index e7dcbf55d..6c350a285 100644 --- a/backend/src/tools/mod.rs +++ b/backend/src/tools/mod.rs @@ -1,2 +1 @@ pub mod lru_cache; -pub mod atomic_once_flag; \ No newline at end of file diff --git a/backend/src/utils/file/config_reader.rs b/backend/src/utils/file/config_reader.rs index 85c79c64f..93d5459a0 100644 --- a/backend/src/utils/file/config_reader.rs +++ b/backend/src/utils/file/config_reader.rs @@ -601,8 +601,11 @@ pub async fn read_initial_app_config( } } - paths.template_files_used = template_files_used; - app_config.paths.store(Arc::new(paths.clone())); + let current_paths = app_config.paths.load_full(); + let mut merged_paths = current_paths.as_ref().clone(); + merged_paths.mapping_files_used = paths.mapping_files_used.clone(); + merged_paths.template_files_used = template_files_used; + app_config.paths.store(Arc::new(merged_paths)); if server_mode { match read_api_proxy_config(&app_config, resolve_env).await { diff --git a/frontend/public/assets/i18n/en.json b/frontend/public/assets/i18n/en.json index 01a8a684e..054f60759 100644 --- a/frontend/public/assets/i18n/en.json +++ b/frontend/public/assets/i18n/en.json @@ -420,6 +420,7 @@ "CONFIG_HOT_RELOAD": "If true, config files are reloaded automatically on change.", "CONNECT_TIMEOUT_SECS": "Maximum connect timeout in seconds.", "CUSTOM_STREAM_RESPONSE_PATH": "Path to a directory containing custom error response files for stream failures.", + "CUSTOM_STREAM_RESPONSE_TIMEOUT_SECS": "Maximum duration in seconds for serving custom stream response video. `0` disables timeout.", "DEFAULT_USER_AGENT": "The User-Agent string used for outgoing requests if not specified elsewhere.", "DISK_BASED_PROCESSING": "If enabled, input playlists are processed from disk to save RAM.\nIf you have enough RAM, you can let this option disabled to improve performance.", "MAPPING_PATH": "Path or directory containing mapping configuration files.", @@ -756,6 +757,7 @@ "CUSTOM_ATTRIBUTES": "Custom Attributes", "CUSTOM_HEADERS": "Custom headers", "CUSTOM_STREAM_RESPONSE_PATH": "Custom Stream Response dir", + "CUSTOM_STREAM_RESPONSE_TIMEOUT_SECS": "Custom Stream Response Timeout (secs)", "DASHBOARD": "Dashboard", "DEFAULT_EPG": "Normal Epg", "DEFAULT_USER_AGENT": "Default User-Agent", diff --git a/frontend/src/app/components/config/main_config_view.rs b/frontend/src/app/components/config/main_config_view.rs index faeb97f2a..ee6376bfb 100644 --- a/frontend/src/app/components/config/main_config_view.rs +++ b/frontend/src/app/components/config/main_config_view.rs @@ -27,6 +27,7 @@ const LABEL_USER_CONFIG_DIR: &str = "LABEL.USER_CONFIG_DIR"; const LABEL_SLEEP_TIMER_MINS: &str = "LABEL.SLEEP_TIMER_MINS"; const LABEL_CONNECT_TIMEOUT_SECS: &str = "LABEL.CONNECT_TIMEOUT_SECS"; const LABEL_CUSTOM_STREAM_RESPONSE_PATH: &str = "LABEL.CUSTOM_STREAM_RESPONSE_PATH"; +const LABEL_CUSTOM_STREAM_RESPONSE_TIMEOUT_SECS: &str = "LABEL.CUSTOM_STREAM_RESPONSE_TIMEOUT_SECS"; const LABEL_ACCEPT_INSECURE_SSL_CERTIFICATES: &str = "LABEL.ACCEPT_INSECURE_SSL_CERTIFICATES"; generate_form_reducer!( @@ -48,6 +49,7 @@ generate_form_reducer!( SleepTimerMins => sleep_timer_mins: Option, ConnectTimeoutSecs => connect_timeout_secs: u32, CustomStreamResponsePath => custom_stream_response_path: Option, + CustomStreamResponseTimeoutSecs => custom_stream_response_timeout_secs: u32, } ); @@ -108,6 +110,7 @@ pub fn MainConfigView() -> Html { { config_field_optional!(form_state.form, translate.t(LABEL_SLEEP_TIMER_MINS), sleep_timer_mins) } { config_field!(form_state.form, translate.t(LABEL_CONNECT_TIMEOUT_SECS), connect_timeout_secs) } { config_field_optional!(form_state.form, translate.t(LABEL_CUSTOM_STREAM_RESPONSE_PATH), custom_stream_response_path) } + { config_field!(form_state.form, translate.t(LABEL_CUSTOM_STREAM_RESPONSE_TIMEOUT_SECS), custom_stream_response_timeout_secs) } } }; @@ -130,6 +133,7 @@ pub fn MainConfigView() -> Html { { edit_field_number_option_u32!(form_state, translate.t(LABEL_SLEEP_TIMER_MINS), sleep_timer_mins, MainConfigFormAction::SleepTimerMins) } { edit_field_number!(form_state, translate.t(LABEL_CONNECT_TIMEOUT_SECS), connect_timeout_secs, MainConfigFormAction::ConnectTimeoutSecs) } { edit_field_text_option!(form_state, translate.t(LABEL_CUSTOM_STREAM_RESPONSE_PATH), custom_stream_response_path, MainConfigFormAction::CustomStreamResponsePath) } + { edit_field_number!(form_state, translate.t(LABEL_CUSTOM_STREAM_RESPONSE_TIMEOUT_SECS), custom_stream_response_timeout_secs, MainConfigFormAction::CustomStreamResponseTimeoutSecs) } } }; diff --git a/shared/src/model/config/base.rs b/shared/src/model/config/base.rs index 2faa42d3c..bd54ec731 100644 --- a/shared/src/model/config/base.rs +++ b/shared/src/model/config/base.rs @@ -12,7 +12,7 @@ use crate::{ is_blank_or_default_backup_dir, is_blank_or_default_custom_stream_response_path, is_blank_or_default_mapping_path, is_blank_or_default_storage_dir, is_blank_or_default_template_path, is_blank_or_default_user_config_dir, is_default_connect_timeout_secs, is_false, - is_none_or_empty_metadata_update, is_none_or_empty_video, normalize_optional_config_file_path, + is_none_or_empty_metadata_update, is_none_or_empty_video, is_zero_u32, normalize_optional_config_file_path, normalize_optional_dir, DEFAULT_BACKUP_DIR, DEFAULT_CUSTOM_STREAM_RESPONSE_PATH, DEFAULT_STORAGE_DIR, DEFAULT_USER_CONFIG_DIR, MAPPING_FILE, TEMPLATE_FILE, }, @@ -42,6 +42,8 @@ pub struct ConfigDto { skip_serializing_if = "is_blank_or_default_custom_stream_response_path" )] pub custom_stream_response_path: Option, + #[serde(default, skip_serializing_if = "is_zero_u32")] + pub custom_stream_response_timeout_secs: u32, #[serde(default, skip_serializing_if = "is_none_or_empty_video")] pub video: Option, #[serde(default, skip_serializing_if = "is_none_or_empty_metadata_update")] @@ -92,6 +94,7 @@ impl Default for ConfigDto { mapping_path: None, template_path: None, custom_stream_response_path: None, + custom_stream_response_timeout_secs: 0, video: None, metadata_update: None, schedules: None, @@ -134,6 +137,8 @@ pub struct MainConfigDto { pub template_path: Option, #[serde(default, skip_serializing_if = "is_blank_optional_string")] pub custom_stream_response_path: Option, + #[serde(default, skip_serializing_if = "is_zero_u32")] + pub custom_stream_response_timeout_secs: u32, #[serde(default, skip_serializing_if = "is_false")] pub user_access_control: bool, #[serde(default, skip_serializing_if = "is_false")] @@ -162,6 +167,7 @@ impl Default for MainConfigDto { mapping_path: default_main_mapping_path(), template_path: default_main_template_path(), custom_stream_response_path: None, + custom_stream_response_timeout_secs: 0, user_access_control: false, connect_timeout_secs: default_connect_timeout_secs(), sleep_timer_mins: None, @@ -184,6 +190,7 @@ impl From<&ConfigDto> for MainConfigDto { mapping_path: config.mapping_path.clone(), template_path: config.template_path.clone(), custom_stream_response_path: config.custom_stream_response_path.clone(), + custom_stream_response_timeout_secs: config.custom_stream_response_timeout_secs, user_access_control: config.user_access_control, connect_timeout_secs: config.connect_timeout_secs, sleep_timer_mins: config.sleep_timer_mins, @@ -357,6 +364,7 @@ impl ConfigDto { self.template_path = normalize_optional_config_file_path(&main_config.template_path, TEMPLATE_FILE); self.custom_stream_response_path = normalize_optional_dir(&main_config.custom_stream_response_path, DEFAULT_CUSTOM_STREAM_RESPONSE_PATH); + self.custom_stream_response_timeout_secs = main_config.custom_stream_response_timeout_secs; self.user_access_control = main_config.user_access_control; self.connect_timeout_secs = main_config.connect_timeout_secs; self.sleep_timer_mins = main_config.sleep_timer_mins; diff --git a/shared/src/utils/default_utils.rs b/shared/src/utils/default_utils.rs index db038a3e5..2bc929cae 100644 --- a/shared/src/utils/default_utils.rs +++ b/shared/src/utils/default_utils.rs @@ -10,6 +10,7 @@ use std::sync::Arc; pub const fn is_zero_u16(v: &u16) -> bool { *v == 0 } pub const fn is_zero_i16(v: &i16) -> bool { *v == 0 } +pub const fn is_zero_u32(v: &u32) -> bool { *v == 0 } pub const fn is_true(v: &bool) -> bool { *v } pub const fn is_false(v: &bool) -> bool { !*v } pub const fn default_as_true() -> bool { true } diff --git a/shared/src/utils/string_interner.rs b/shared/src/utils/string_interner.rs index 6e5308df1..0e1de7934 100644 --- a/shared/src/utils/string_interner.rs +++ b/shared/src/utils/string_interner.rs @@ -116,6 +116,10 @@ fn intern_string(s: String) -> Arc { Arc::from(s) } +/// Returns the current number of strings held in the interning pool. +/// Uses a read lock and is safe to call on hot paths for threshold checks. +pub fn interner_len() -> usize { INTERNER.read().map(|g| g.len()).unwrap_or(0) } + /// Garbage collection: removes strings that are only referenced by the cache. pub fn interner_gc() -> usize { if let Ok(mut guard) = INTERNER.write() {