From 95dde2876dde19d1f2a35316a5ca187d3be2faa5 Mon Sep 17 00:00:00 2001 From: euzu Date: Tue, 4 Nov 2025 18:01:44 +0100 Subject: [PATCH] First step provider and user connections. -> Missing shared streams connections --- backend/src/api/api_utils.rs | 152 +-- backend/src/api/endpoints/hls_api.rs | 23 +- backend/src/api/main_api.rs | 30 +- .../src/api/model/active_provider_manager.rs | 1177 +---------------- backend/src/api/model/active_user_manager.rs | 222 ++-- backend/src/api/model/app_state.rs | 7 +- backend/src/api/model/mod.rs | 6 +- backend/src/api/model/provider_config.rs | 2 +- .../src/api/model/provider_lineup_manager.rs | 1066 +++++++++++++++ backend/src/api/model/release_task.rs | 9 - backend/src/api/model/stream.rs | 53 +- .../api/model/streams/active_client_stream.rs | 23 +- .../model/streams/shared_stream_manager.rs | 114 +- backend/src/api/serve.rs | 24 +- backend/src/auth/fingerprint.rs | 6 +- frontend/scss/_theme.scss | 5 + .../components/source_editor/editor_view.rs | 4 +- .../model/active_user_connection_change.rs | 3 +- shared/src/model/stream_info.rs | 7 +- 19 files changed, 1412 insertions(+), 1521 deletions(-) create mode 100644 backend/src/api/model/provider_lineup_manager.rs delete mode 100644 backend/src/api/model/release_task.rs diff --git a/backend/src/api/api_utils.rs b/backend/src/api/api_utils.rs index 728221129..d67dd21fb 100644 --- a/backend/src/api/api_utils.rs +++ b/backend/src/api/api_utils.rs @@ -1,17 +1,16 @@ use crate::api::endpoints::xtream_api::{get_xtream_player_api_stream_url, ApiStreamContext}; -use crate::api::model::ProviderConfig; +use crate::api::model::{ProviderAllocation, ProviderConfig, ProviderStreamState, StreamDetails, StreamingStrategy}; use crate::api::model::UserSession; use crate::api::model::{ create_channel_unavailable_stream, create_custom_video_stream_response, create_provider_connections_exhausted_stream, create_provider_stream, - get_stream_response_with_headers, ActiveClientStream, AppState, BoxedProviderStream, - CustomVideoStreamType, PersistPipeStream, ProviderAllocation, ProviderConnectionGuard, - ProviderStreamFactoryOptions, ProviderStreamInfo, ProviderStreamResponse, SharedStreamManager, + get_stream_response_with_headers, ActiveClientStream, AppState, + CustomVideoStreamType, PersistPipeStream, + ProviderStreamFactoryOptions, SharedStreamManager, StreamError, ThrottledStream, UserApiRequest, }; use crate::model::ConfigInput; use crate::model::{ConfigTarget, ProxyUserCredentials}; -use crate::tools::atomic_once_flag::AtomicOnceFlag; use crate::tools::lru_cache::LRUResourceCache; use crate::utils::create_new_file_for_write; use crate::utils::request; @@ -318,51 +317,6 @@ async fn get_redirect_alternative_url<'a>( Cow::Borrowed(redirect_url) } -type StreamUrl = String; -type ProviderName = String; - -enum ProviderStreamState { - Custom(ProviderStreamResponse), - Available(Option, StreamUrl), - GracePeriod(Option, StreamUrl), -} - -pub struct StreamDetails { - pub stream: Option, - stream_info: ProviderStreamInfo, - pub provider_name: Option, - pub grace_period_millis: u64, - pub reconnect_flag: Option>, - pub provider_connection_guard: Option>, -} - -impl StreamDetails { - pub fn from_stream(stream: BoxedProviderStream) -> Self { - Self { - stream: Some(stream), - stream_info: None, - provider_name: None, - grace_period_millis: default_grace_period_millis(), - reconnect_flag: None, - provider_connection_guard: None, - } - } - #[inline] - pub fn has_stream(&self) -> bool { - self.stream.is_some() - } - - #[inline] - pub fn has_grace_period(&self) -> bool { - self.grace_period_millis > 0 - } -} - -struct StreamingStrategy { - provider_connection_guard: Option>, - provider_stream_state: ProviderStreamState, - input_headers: Option>, -} /// Determines the appropriate streaming strategy for the given input and stream URL. /// @@ -388,52 +342,54 @@ async fn resolve_streaming_strategy( force_provider: Option<&str>, ) -> StreamingStrategy { // allocate a provider connection - let provider_connection_guard = match force_provider { - Some(provider) => { - app_state - .active_provider - .force_exact_acquire_connection(provider, &fingerprint.addr, app_state.get_release_sender()) - .await - } - None => { - app_state - .active_provider - .acquire_connection(&input.name, &fingerprint.addr, app_state.get_release_sender()) - .await - } + let provider_connection_handle = match force_provider { + Some(provider) => app_state.active_provider.force_exact_acquire_connection(provider, &fingerprint.addr).await, + None => app_state.active_provider.acquire_connection(&input.name, &fingerprint.addr).await, }; - let stream_response_params = match &**provider_connection_guard { - ProviderAllocation::Exhausted => { - debug!("Input {} is exhausted. No connections allowed.", input.name); - let stream = create_provider_connections_exhausted_stream(&app_state.app_config, &[]); - ProviderStreamState::Custom(stream) - } - ProviderAllocation::Available(_, ref provider) - | ProviderAllocation::GracePeriod(_, ref provider) => { - // force_stream_provider means we keep the url and the provider. - // If force_stream_provider or the input is the same as the config we don't need to get new url - let (provider, url) = if force_provider.is_some() || provider.id == input.id { - (input.name.clone(), stream_url.to_string()) - } else { - ( - provider.name.clone(), - get_stream_alternative_url(stream_url, input, provider), - ) - }; + let stream_response_params = + if let Some(allocation) = provider_connection_handle.as_ref().map(|ph| &ph.allocation) { + match allocation { + ProviderAllocation::Exhausted => { + debug!("Input {} is exhausted. No connections allowed.", input.name); + let stream = create_provider_connections_exhausted_stream(&app_state.app_config, &[]); + ProviderStreamState::Custom(stream) + } + ProviderAllocation::Available(ref provider) + | ProviderAllocation::GracePeriod(ref provider) => { + // force_stream_provider means we keep the url and the provider. + // If force_stream_provider or the input is the same as the config we don't need to get new url + let (provider, url) = if force_provider.is_some() || provider.id == input.id { + (input.name.clone(), stream_url.to_string()) + } else { + ( + provider.name.clone(), + get_stream_alternative_url(stream_url, input, provider), + ) + }; - if matches!( - &**provider_connection_guard, - ProviderAllocation::Available(_, _) - ) { - ProviderStreamState::Available(Some(provider), url) - } else { - ProviderStreamState::GracePeriod(Some(provider), url) + let available = match allocation { + ProviderAllocation::Exhausted + | ProviderAllocation::Available(_) => true, + ProviderAllocation::GracePeriod(_) => false, + }; + + if available { + ProviderStreamState::Available(Some(provider), url) + } else { + ProviderStreamState::GracePeriod(Some(provider), url) + } + } } - } + } else { + debug!("Input {} is exhausted. No connections allowed.", input.name); + let stream = create_provider_connections_exhausted_stream(&app_state.app_config, &[]); + ProviderStreamState::Custom(stream) }; + + StreamingStrategy { - provider_connection_guard: Some(provider_connection_guard), + provider_handle: provider_connection_handle, provider_stream_state: stream_response_params, input_headers: Some(input.headers.clone()), } @@ -485,9 +441,9 @@ async fn create_stream_response_details( config_grace_period_millis, ); let guard_provider_name = streaming_strategy - .provider_connection_guard + .provider_handle .as_ref() - .and_then(|guard| guard.get_provider_name()); + .and_then(|guard| guard.allocation.get_provider_name()); match streaming_strategy.provider_stream_state { // custom stream means we display our own stream like connection exhausted, channel-unavailable... @@ -499,7 +455,7 @@ async fn create_stream_response_details( provider_name: guard_provider_name.clone(), grace_period_millis, reconnect_flag: None, - provider_connection_guard: streaming_strategy.provider_connection_guard.clone(), + provider_handle: streaming_strategy.provider_handle.clone(), } } ProviderStreamState::Available(_provider_name, request_url) @@ -553,7 +509,7 @@ async fn create_stream_response_details( // if we have no stream, we should release the provider if stream.is_none() { - if let Some(guard) = streaming_strategy.provider_connection_guard.take() { + if let Some(guard) = streaming_strategy.provider_handle.take() { drop(guard); } error!("Cant open stream {}", sanitize_sensitive_info(&request_url)); @@ -565,7 +521,7 @@ async fn create_stream_response_details( provider_name: guard_provider_name.clone(), grace_period_millis, reconnect_flag, - provider_connection_guard: streaming_strategy.provider_connection_guard.take(), + provider_handle: streaming_strategy.provider_handle.take(), } } } @@ -800,7 +756,7 @@ pub async fn force_provider_stream_response( ); return try_unwrap_body!(response.body(body_stream)); } - drop(stream_details.provider_connection_guard.take()); + drop(stream_details.provider_handle.take()); if let (Some(stream), _stream_info) = create_channel_unavailable_stream( &app_state.app_config, &[], @@ -876,7 +832,7 @@ pub async fn stream_response( let provider_name = stream_details.provider_name.clone(); let provider_guard = if share_stream { - stream_details.provider_connection_guard.take() + stream_details.provider_handle.take() } else { None }; @@ -974,7 +930,7 @@ pub async fn stream_response( return stream_resp.into_response(); } - drop(stream_details.provider_connection_guard.take()); + drop(stream_details.provider_handle.take()); axum::http::StatusCode::BAD_REQUEST.into_response() } diff --git a/backend/src/api/endpoints/hls_api.rs b/backend/src/api/endpoints/hls_api.rs index 21d9758dd..b5d964a68 100644 --- a/backend/src/api/endpoints/hls_api.rs +++ b/backend/src/api/endpoints/hls_api.rs @@ -3,7 +3,7 @@ use crate::api::api_utils::{ force_provider_stream_response, get_stream_alternative_url, is_seek_request, }; use crate::api::api_utils::{get_headers_from_request, try_option_bad_request, HeaderFilter}; -use crate::api::model::UserSession; +use crate::api::model::{ProviderAllocation, UserSession}; use crate::api::model::{create_custom_video_stream_response, CustomVideoStreamType}; use crate::api::model::AppState; use crate::auth::Fingerprint; @@ -56,15 +56,20 @@ pub(in crate::api) async fn handle_hls_stream_request( let (request_url, session_token) = match user_session { Some(session) => { - match app_state + let handle = app_state .active_provider - .force_exact_acquire_connection(&session.provider, &fingerprint.addr, app_state.get_release_sender()) - .await - .get_provider_config() - { - Some(provider_cfg) => { - let stream_url = get_stream_alternative_url(&url, input, &provider_cfg); - (stream_url, Some(session.token.clone())) + .force_exact_acquire_connection(&session.provider, &fingerprint.addr) + .await; + match handle { + Some(provider_handle) => { + match provider_handle.allocation { + ProviderAllocation::Exhausted => (url, None), + ProviderAllocation::Available(cfg) + | ProviderAllocation::GracePeriod(cfg) => { + let stream_url = get_stream_alternative_url(&url, input, &cfg); + (stream_url, Some(session.token.clone())) + } + } } None => (url, None), } diff --git a/backend/src/api/main_api.rs b/backend/src/api/main_api.rs index 1a4cc1015..f5f6df317 100644 --- a/backend/src/api/main_api.rs +++ b/backend/src/api/main_api.rs @@ -11,7 +11,7 @@ use crate::api::endpoints::xmltv_api::xmltv_api_register; use crate::api::endpoints::xtream_api::xtream_api_register; use crate::api::hdhomerun_proprietary::spawn_proprietary_tasks; use crate::api::hdhomerun_ssdp::spawn_ssdp_discover_task; -use crate::api::model::{create_cache, create_http_client, ActiveProviderManager, ActiveUserManager, AppState, CancelTokens, DownloadQueue, EventManager, HdHomerunAppState, PlaylistStorageState, ProviderConnectionGuard, ReleaseTask, SharedStreamManager}; +use crate::api::model::{create_cache, create_http_client, ActiveProviderManager, ActiveUserManager, AppState, CancelTokens, DownloadQueue, EventManager, HdHomerunAppState, PlaylistStorageState, SharedStreamManager}; use crate::api::scheduler::exec_scheduler; use crate::api::serve::serve; use crate::model::{AppConfig, Config, Healthcheck, ProcessTargets, RateLimitConfig}; @@ -26,7 +26,6 @@ use std::io::ErrorKind; use std::path::PathBuf; use std::sync::atomic::AtomicI8; use std::sync::Arc; -use tokio::sync::mpsc::{unbounded_channel, UnboundedSender}; use tokio_util::sync::CancellationToken; use tower_governor::key_extractor::SmartIpKeyExtractor; use crate::repository::storage::get_geoip_path; @@ -85,9 +84,9 @@ fn create_shared_data( }; let cache = create_cache(&config); - let shared_stream_manager = Arc::new(SharedStreamManager::new()); let (provider_change_tx, provider_change_rx) = tokio::sync::mpsc::unbounded_channel(); let active_provider = Arc::new(ActiveProviderManager::new(app_config, provider_change_tx)); + let shared_stream_manager = Arc::new(SharedStreamManager::new(Arc::clone(&active_provider))); let (active_user_change_tx, active_user_change_rx) = tokio::sync::mpsc::unbounded_channel(); let active_users = ActiveUserManager::new( &config, @@ -99,7 +98,6 @@ fn create_shared_data( let event_manager = Arc::new(EventManager::new(active_user_change_rx, provider_change_rx, )); let client = create_http_client(app_config); - let release_tx = create_release_task(); AppState { forced_targets: Arc::new(ArcSwap::new(Arc::clone(forced_targets))), @@ -113,26 +111,10 @@ fn create_shared_data( event_manager, cancel_tokens: Arc::new(ArcSwap::from_pointee(CancelTokens::default())), playlists: Arc::new(PlaylistStorageState::new()), - geoip, - release_tx + geoip } } -fn create_release_task() -> UnboundedSender { - let (cleanup_tx, mut cleanup_rx) = unbounded_channel::(); - // Spawn the async cleanup worker - tokio::spawn(async move { - while let Some(release) = cleanup_rx.recv().await { - match release { - ReleaseTask::ForceProvider(guard) => guard.force_release().await, - ReleaseTask::ProviderAllocation(allocation) => ProviderConnectionGuard::release_allocation(allocation).await, - ReleaseTask::UserConnection(manager, addr) => manager.remove_connection(&addr).await, - } - } - }); - cleanup_tx -} - fn exec_update_on_boot( client: Arc, app_state: &Arc, @@ -228,6 +210,7 @@ pub(in crate::api) fn start_hdhomerun( )); let c_token = cancel_token.clone(); let active_user_manager = Arc::clone(&app_data.active_users); + let active_provider_manager = Arc::clone(&app_data.active_provider); tokio::spawn(async move { let router = axum::Router::>::new() .layer(create_cors_layer()) @@ -245,7 +228,7 @@ pub(in crate::api) fn start_hdhomerun( .await { Ok(listener) => { - serve(listener, router, Some(c_token), active_user_manager).await; + serve(listener, router, Some(c_token), &active_user_manager, &active_provider_manager).await; } Err(err) => error!("{err}"), } @@ -381,8 +364,7 @@ pub async fn start_server( let router: axum::Router<()> = router.with_state(shared_data.clone()); let listener = tokio::net::TcpListener::bind(format!("{host}:{port}")).await?; - let active_user_manager = Arc::clone(&shared_data.active_users); - serve(listener, router, None, active_user_manager).await; + serve(listener, router, None, &shared_data.active_users, &shared_data.active_provider).await; Ok(()) } diff --git a/backend/src/api/model/active_provider_manager.rs b/backend/src/api/model/active_provider_manager.rs index 85ed43de9..b3e1cc08c 100644 --- a/backend/src/api/model/active_provider_manager.rs +++ b/backend/src/api/model/active_provider_manager.rs @@ -1,799 +1,30 @@ -use crate::api::model::{ProviderConfig, ProviderConfigConnection, ProviderConfigWrapper, ProviderConnectionChangeCallback, ProviderConnectionChangeSender, ReleaseTask}; +use crate::api::model::provider_lineup_manager::{ProviderAllocation, ProviderLineupManager}; +use crate::api::model::{ProviderConfig, ProviderConnectionChangeSender}; use crate::model::{AppConfig, ConfigInput}; -use crate::utils::debug_if_enabled; -use arc_swap::ArcSwap; -use log::{debug, error, log_enabled, trace}; -use shared::utils::{default_grace_period_millis, default_grace_period_timeout_secs, display_vec, sanitize_sensitive_info}; +use log::{debug}; +use shared::utils::{default_grace_period_millis, default_grace_period_timeout_secs}; use std::collections::HashMap; -use std::fmt; -use std::ops::Deref; -use std::sync::atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering}; +use std::net::SocketAddr; use std::sync::Arc; -use tokio::runtime::Handle; -use tokio::sync::mpsc::UnboundedSender; use tokio::sync::RwLock; -const CONNECTION_STATE_ACTIVE: u8 = 0; -const CONNECTION_STATE_SHARED: u8 = 1; -const CONNECTION_STATE_RELEASED: u8 = 2; - -#[derive(Debug)] -pub struct ProviderConnectionGuard { - allocation: ProviderAllocation, - release_tx: UnboundedSender, -} - -impl ProviderConnectionGuard { - pub fn new(allocation: ProviderAllocation, release_tx: UnboundedSender) -> Self { - Self { allocation, release_tx } - } - - // for shared streams, we need to disable release - // The connection should be released when all shared streams close! - pub(crate) fn disable_release(&self) { - match &self.allocation { - ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(state, _) | - ProviderAllocation::GracePeriod(state, _) => { - let _ = state.compare_exchange(CONNECTION_STATE_ACTIVE, CONNECTION_STATE_SHARED, Ordering::SeqCst, Ordering::SeqCst); - } - } - } - - pub(crate) async fn release(&self) { - Self::release_allocation(self.allocation.clone()).await; - } - - pub async fn release_allocation(allocation: ProviderAllocation) { - match &allocation { - ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(state, config) | - ProviderAllocation::GracePeriod(state, config) => { - // we can't release shared state - if state.compare_exchange(CONNECTION_STATE_ACTIVE, CONNECTION_STATE_RELEASED, Ordering::SeqCst, Ordering::SeqCst).is_ok() { - config.release().await; - } - } - } - } - - // we need to ensure the connections is released - pub(crate) async fn force_release(&self) { - Self::force_release_allocation(self.allocation.clone()).await; - } - - async fn force_release_allocation(allocation: ProviderAllocation) { - match &allocation { - ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(state, config) - | ProviderAllocation::GracePeriod(state, config) => { - if state.load(Ordering::SeqCst) < CONNECTION_STATE_RELEASED { - state.store(CONNECTION_STATE_RELEASED, Ordering::SeqCst); - trace!("Forced releasing provider connection {:?}", config.name); - config.release().await; - } - } - } - } - - pub fn get_provider_name(&self) -> Option { - match self.allocation { - ProviderAllocation::Exhausted => None, - ProviderAllocation::Available(_, ref cfg) | - ProviderAllocation::GracePeriod(_, ref cfg) => { - Some(cfg.name.clone()) - } - } - } - pub fn get_provider_config(&self) -> Option> { - match self.allocation { - ProviderAllocation::Exhausted => None, - ProviderAllocation::Available(_, ref cfg) | - ProviderAllocation::GracePeriod(_, ref cfg) => { - Some(Arc::clone(cfg)) - } - } - } -} - -impl Deref for ProviderConnectionGuard { - type Target = ProviderAllocation; - fn deref(&self) -> &Self::Target { - &self.allocation - } -} - -impl Drop for ProviderConnectionGuard { - fn drop(&mut self) { - let allocation = self.allocation.clone(); - if let Err(_err) = self.release_tx.send(ReleaseTask::ProviderAllocation(allocation)) { - // fallback - if let Ok(handle) = Handle::try_current() { - let allocation = self.allocation.clone(); - handle.spawn(async move { - Self::release_allocation(allocation).await; - }); - } else { - // No runtime - error!("💥💥💥 Dropping ProviderConnectionGuard without async runtime — Provider Connection not freed"); - } - } - } -} +pub type ProviderConnectionId = SocketAddr; #[derive(Debug, Clone)] -pub enum ProviderAllocation { - Exhausted, - Available(Arc, Arc), - GracePeriod(Arc, Arc), +pub struct ProviderHandle { + pub id: ProviderConnectionId, + pub allocation: ProviderAllocation, } -impl ProviderAllocation { - pub fn new_available(config: Arc) -> Self { - ProviderAllocation::Available(Arc::new(AtomicU8::new(CONNECTION_STATE_ACTIVE)), config) - } - - pub fn new_grace_period(config: Arc) -> Self { - ProviderAllocation::GracePeriod(Arc::new(AtomicU8::new(CONNECTION_STATE_ACTIVE)), config) - } -} - -impl PartialEq for ProviderAllocation { - fn eq(&self, other: &Self) -> bool { - // Note: released flag ignored - match (self, other) { - (ProviderAllocation::Exhausted, ProviderAllocation::Exhausted) => true, - (ProviderAllocation::Available(_, cfg1), ProviderAllocation::Available(_, cfg2)) - | (ProviderAllocation::GracePeriod(_, cfg1), ProviderAllocation::GracePeriod(_, cfg2)) => cfg1 == cfg2, - _ => false, - } - } -} - -/// This manages different types of provider lineups: -/// -/// `Single(SingleProviderLineup)`: A single provider. -/// `Multi(MultiProviderLineup)`: A set of providers grouped by priority. -#[derive(Debug)] -enum ProviderLineup { - Single(SingleProviderLineup), - Multi(MultiProviderLineup), -} - -impl fmt::Display for ProviderLineup { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - ProviderLineup::Single(lineup) => { - write!(f, "SingleProviderLineup: {{ {} }}", lineup.provider) - } - ProviderLineup::Multi(lineup) => { - write!(f, "MultiProviderLineup: {{")?; - for (i, group) in lineup.providers.iter().enumerate() { - write!(f, " Group {}: {}", i + 1, group)?; - } - write!(f, " }}")?; - Ok(()) - } - } - } -} - -impl ProviderLineup { - async fn get_next(&self, grace_period_timeout_secs: u64) -> Option> { - match self { - ProviderLineup::Single(lineup) => lineup.get_next(grace_period_timeout_secs).await, - ProviderLineup::Multi(lineup) => lineup.get_next(grace_period_timeout_secs).await, - } - } - - async fn acquire(&self, with_grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { - match self { - ProviderLineup::Single(lineup) => lineup.acquire(with_grace, grace_period_timeout_secs).await, - ProviderLineup::Multi(lineup) => lineup.acquire(with_grace, grace_period_timeout_secs).await, - } - } - - // async fn release(&self, provider_name: &str) { - // match self { - // ProviderLineup::Single(lineup) => lineup.release(provider_name).await, - // ProviderLineup::Multi(lineup) => lineup.release(provider_name).await, - // } - // } -} - -/// Handles a single provider and ensures safe allocation/release of connections. -#[derive(Debug)] -struct SingleProviderLineup { - provider: ProviderConfigWrapper, -} - -impl SingleProviderLineup { - fn new<'a, F>(cfg: &ConfigInput, get_connection: Option, connection_change: &ProviderConnectionChangeCallback) -> Self - where - F: Fn(&str) -> Option<&'a ProviderConfigConnection>, - { - Self { - provider: ProviderConfigWrapper::new(ProviderConfig::new(cfg, get_connection, Arc::clone(connection_change))), - } - } - - async fn get_next(&self, grace_period_timeout_secs: u64) -> Option> { - self.provider.get_next(false, grace_period_timeout_secs).await - } - - async fn acquire(&self, with_grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { - self.provider.try_allocate(with_grace, grace_period_timeout_secs).await - } - - #[cfg(test)] - async fn release(&self, provider_name: &str) { - if self.provider.name == provider_name { - self.provider.release().await; - } - } -} - - -/// Manages provider groups based on priority: -/// -/// `SingleProviderGroup(ProviderConfig)`: A single provider. -/// `MultiProviderGroup(AtomicUsize, Vec)`: A list of providers with a priority index. -#[derive(Debug)] -enum ProviderPriorityGroup { - SingleProviderGroup(ProviderConfigWrapper), - MultiProviderGroup(AtomicUsize, Vec), -} - -impl fmt::Display for ProviderPriorityGroup { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - ProviderPriorityGroup::SingleProviderGroup(p) => { - write!(f, "Single({p})") - } - ProviderPriorityGroup::MultiProviderGroup(_, providers) => { - write!(f, "Multi({})", display_vec(providers)) - } - } - } -} - -impl ProviderPriorityGroup { - async fn is_exhausted(&self) -> bool { - match self { - ProviderPriorityGroup::SingleProviderGroup(g) => g.is_exhausted().await, - ProviderPriorityGroup::MultiProviderGroup(_, groups) => { - for g in groups { - if !g.is_exhausted().await { - return false; - } - } - true - } - } - } -} - - -/// Manages multiple providers, ensuring that connections are allocated in a round-robin manner based on priority. -#[repr(align(64))] -#[derive(Debug)] -struct MultiProviderLineup { - name: String, - providers: Vec, - index: AtomicUsize, -} - -impl MultiProviderLineup { - pub fn new<'a, F>(cfg_input: &ConfigInput, get_connection: Option, connection_change: &ProviderConnectionChangeCallback) -> Self - where - F: Fn(&str) -> Option<&'a ProviderConfigConnection> + Copy, - { - let mut inputs = vec![ProviderConfigWrapper::new(ProviderConfig::new(cfg_input, get_connection, Arc::clone(connection_change)))]; - if let Some(aliases) = &cfg_input.aliases { - for alias in aliases { - inputs.push(ProviderConfigWrapper::new(ProviderConfig::new_alias(cfg_input, alias, get_connection, Arc::clone(connection_change)))); - } - } - let mut providers = HashMap::new(); - for provider in inputs { - let priority = provider.get_priority(); - providers.entry(priority) - .or_insert_with(Vec::new) - .push(provider); - } - let mut values: Vec<(i16, Vec)> = providers.into_iter().collect(); - values.sort_by(|(p1, _), (p2, _)| p1.cmp(p2)); - let providers: Vec = values.into_iter().map(|(_, mut group)| { - if group.len() > 1 { - ProviderPriorityGroup::MultiProviderGroup(AtomicUsize::new(0), group) - } else { - ProviderPriorityGroup::SingleProviderGroup(group.remove(0)) - } - }).collect(); - - Self { - name: cfg_input.name.clone(), - providers, - index: AtomicUsize::new(0), - } - } - - /// Attempts to acquire the next available provider from a specific priority group. - /// - /// # Parameters - /// - `priority_group`: Thep rovider group to search within. - /// - /// # Returns - /// - `ProviderAllocation`: A reference to the next available provider in the specified group. - /// - /// # Behavior - /// - Iterates through the providers in the given group in a round-robin manner. - /// - Checks if a provider has available capacity before selecting it. - /// - Uses atomic operations to maintain fair provider selection. - /// - /// # Thread Safety - /// - Uses `RwLock` for safe concurrent access. - /// - Ensures fair provider allocation across multiple threads. - /// - /// # Example Usage - /// ```rust - /// let lineup = MultiProviderLineup::new(&config); - /// match lineup.acquire_next_provider_from_group(priority_group) { - /// ProviderAllocation::Exhausted => println!("All providers exhausted"), - /// ProviderAllocation::Available(provider) => println!("Provider available {}", provider.name), - /// ProviderAllocation::GracePeriodprovider) => println!("Provider with grace period {}", provider.name), - /// } - /// } - /// ``` - async fn acquire_next_provider_from_group(priority_group: &ProviderPriorityGroup, grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { - match priority_group { - ProviderPriorityGroup::SingleProviderGroup(p) => { - let result = p.try_allocate(grace, grace_period_timeout_secs).await; - if !matches!(result, ProviderAllocation::Exhausted) { - return result; - } - } - ProviderPriorityGroup::MultiProviderGroup(index, pg) => { - let provider_count = pg.len(); - let mut idx = index.load(Ordering::Relaxed) % provider_count; - let start = idx; - - for _ in start..provider_count { - let p = &pg[idx]; - let result = p.try_allocate(grace, grace_period_timeout_secs).await; - if !matches!(result, ProviderAllocation::Exhausted) { - index.store((idx + 1) % provider_count, Ordering::Relaxed); - return result; - } - idx = (idx + 1) % provider_count; - } - index.store(idx, Ordering::Relaxed); - } - } - ProviderAllocation::Exhausted - } - - // Used for redirect to cylce through provider - async fn get_next_provider_from_group(priority_group: &ProviderPriorityGroup, grace: bool, grace_period_timeout_secs: u64) -> Option> { - match priority_group { - ProviderPriorityGroup::SingleProviderGroup(p) => { - return p.get_next(grace, grace_period_timeout_secs).await; - } - ProviderPriorityGroup::MultiProviderGroup(index, pg) => { - let provider_count = pg.len(); - let mut idx = index.load(Ordering::Relaxed) % provider_count; - let start = idx; - for _ in start..provider_count { - if let Some(p) = pg.get(idx) { - let result = p.get_next(grace, grace_period_timeout_secs).await; - if result.is_some() { - index.store((idx + 1) % provider_count, Ordering::Relaxed); - return result; - } - } - idx = (idx + 1) % provider_count; - } - index.store(idx, Ordering::Relaxed); - } - } - None - } - - /// Attempts to acquire a provider from the lineup based on priority and availability. - /// - /// # Returns - /// - `ProviderAllocation`: A reference to the acquired provider if allocation was successful. - /// - /// # Behavior - /// - The method iterates through provider priority groups in a round-robin fashion. - /// - It attempts to allocate a provider from the highest priority group first. - /// - If a provider has available capacity, it is returned. - /// - If all providers in a group are exhausted, it moves to the next group. - /// - Updates the internal index to ensure fair distribution of requests. - /// - /// # Thread Safety - /// - Uses atomic operations (`AtomicUsize`) for thread-safe indexing. - /// - Uses `RwLock` for thread-safe provider allocation. - /// - /// # Example Usage - /// ```rust - /// let lineup = MultiProviderLineup::new(&config); - /// match lineup.acquire() { - /// ProviderAllocation::Exhausted => println!("All providers exhausted"), - /// ProviderAllocation::Available(provider) => println!("Provider available {}", provider.name), - /// ProviderAllocation::GracePeriodprovider) => println!("Provider with grace period {}", provider.name), - /// } - /// ``` - async fn acquire(&self, with_grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { - let main_idx = self.index.load(Ordering::SeqCst); - let provider_count = self.providers.len(); - - for index in main_idx..provider_count { - let priority_group = &self.providers[index]; - let allocation = { - let without_grace_allocation = Self::acquire_next_provider_from_group(priority_group, false, grace_period_timeout_secs).await; - if with_grace && matches!(without_grace_allocation, ProviderAllocation::Exhausted) { - Self::acquire_next_provider_from_group(priority_group, true, grace_period_timeout_secs).await - } else { - without_grace_allocation - } - }; - if !matches!(allocation, ProviderAllocation::Exhausted) { - if priority_group.is_exhausted().await { - self.index.store((index + 1) % provider_count, Ordering::SeqCst); - } - return allocation; - } - } - - ProviderAllocation::Exhausted - } - - // it intended to use with redirects to cycle through provider - async fn get_next(&self, grace_period_timeout_secs: u64) -> Option> { - let main_idx = self.index.load(Ordering::SeqCst); - let provider_count = self.providers.len(); - - for index in main_idx..provider_count { - let priority_group = &self.providers[index]; - let allocation = { - let config = Self::get_next_provider_from_group(priority_group, false, grace_period_timeout_secs).await; - if config.is_none() { - Self::get_next_provider_from_group(priority_group, true, grace_period_timeout_secs).await - } else { - config - } - }; - match allocation { - None => {} - Some(config) => { - if priority_group.is_exhausted().await { - self.index.store((index + 1) % provider_count, Ordering::SeqCst); - } - return Some(config); - } - } - } - - None - } - - - #[cfg(test)] - async fn release(&self, provider_name: &str) { - for g in &self.providers { - match g { - ProviderPriorityGroup::SingleProviderGroup(pc) => { - if pc.name == provider_name { - pc.release().await; - break; - } - } - ProviderPriorityGroup::MultiProviderGroup(_, group) => { - for pc in group { - if pc.name == provider_name { - pc.release().await; - return; - } - } - } - } - } - } - - pub async fn get_total_connections(&self) -> usize { - let mut total_connections = 0; - - for group in &self.providers { - match group { - ProviderPriorityGroup::SingleProviderGroup(provider) => { - total_connections += provider.get_current_connections().await; - } - ProviderPriorityGroup::MultiProviderGroup(_, providers) => { - for provider in providers { - total_connections += provider.get_current_connections().await; - } - } - } - } - total_connections - } -} - - -struct ProviderLineupManager { - grace_period_millis: AtomicU64, - grace_period_timeout_secs: AtomicU64, - inputs: Arc>>>, - providers: Arc>>, - connection_change_tx: ProviderConnectionChangeSender, -} - -impl ProviderLineupManager { - pub fn new(inputs: Vec>, grace_period_millis: u64, grace_period_timeout_secs: u64, connection_change_tx: ProviderConnectionChangeSender) -> Self { - let lineups = inputs.iter().map(|i| Self::create_lineup(i, None, connection_change_tx.clone())).collect(); - Self { - grace_period_millis: AtomicU64::new(grace_period_millis), - grace_period_timeout_secs: AtomicU64::new(grace_period_timeout_secs), - inputs: Arc::new(ArcSwap::from_pointee(inputs)), - providers: Arc::new(ArcSwap::from_pointee(lineups)), - connection_change_tx, - } - } - - fn create_lineup(cfg_input: &ConfigInput, provider_connections: Option<&HashMap<&str, ProviderConfigConnection>>, connection_change_sender: ProviderConnectionChangeSender) -> ProviderLineup { - let get_connections = provider_connections.map(|c| |name: &str| c.get(name)); - - let cfg_name = cfg_input.name.clone(); - let on_connection_change: ProviderConnectionChangeCallback = Arc::new(move |_name: &str, connections: usize| { - let connection_change_sender = connection_change_sender.clone(); - let provider_cfg_name = cfg_name.clone(); - if let Err(err) = connection_change_sender.send((provider_cfg_name.clone(), connections)) { - error!("Failed to send connection change: {provider_cfg_name}: {connections}, {err}"); - } - }); - - let on_connection_change = Arc::new(on_connection_change); - - if cfg_input.aliases.as_ref().is_some_and(|a| !a.is_empty()) { - ProviderLineup::Multi(MultiProviderLineup::new(cfg_input, get_connections, &on_connection_change)) - } else { - ProviderLineup::Single(SingleProviderLineup::new(cfg_input, get_connections, &on_connection_change)) - } - } - - fn inputs_differ(a: &ConfigInput, b: &ConfigInput) -> bool { - if a.enabled != b.enabled - || a.max_connections != b.max_connections - || a.priority != b.priority - || a.input_type != b.input_type - || a.username != b.username - || a.password != b.password - || a.url != b.url - { - return true; - } - - match (&a.aliases, &b.aliases) { - (None, None) => {} - (Some(_), None) | (None, Some(_)) => return true, - (Some(a_aliases), Some(b_aliases)) => { - if a_aliases.len() != b_aliases.len() { - return true; - } - - for b_alias in b_aliases { - let Some(a_alias) = a_aliases.iter().find(|a| a.name == b_alias.name) else { - return true; - }; - - if a_alias.max_connections != b_alias.max_connections - || a_alias.priority != b_alias.priority - || a_alias.username != b_alias.username - || a_alias.password != b_alias.password - || a_alias.url != b_alias.url - { - return true; - } - } - } - } - - false - } - - fn has_changed(&self, new_inputs: &[Arc]) -> bool { - let old_inputs = self.inputs.load(); - if old_inputs.len() != new_inputs.len() { - return true; - } - for new_input in new_inputs { - let Some(old_input) = old_inputs.iter().find(|i| i.name == new_input.name) else { - return true; - }; - - if Self::inputs_differ(old_input.as_ref(), new_input.as_ref()) { - return true; - } - } - - false - } - - pub async fn update_config(&self, new_inputs: Vec>, grace_period_millis: u64, grace_period_timeout_secs: u64) { - self.grace_period_millis.store(grace_period_millis, Ordering::Relaxed); - self.grace_period_timeout_secs.store(grace_period_timeout_secs, Ordering::Relaxed); - - if !self.has_changed(&new_inputs) { - return; - } - - let old_lineups = self.providers.load(); - let mut provider_connections = HashMap::new(); - for lineup in old_lineups.iter() { - match lineup { - ProviderLineup::Single(single) => { - provider_connections.insert(single.provider.name.as_str(), single.provider.get_connection_info().await); - } - ProviderLineup::Multi(multi) => { - for group in &multi.providers { - match group { - ProviderPriorityGroup::SingleProviderGroup(cfg) => { - provider_connections.insert(cfg.name.as_str(), cfg.get_connection_info().await); - } - ProviderPriorityGroup::MultiProviderGroup(_, cfgs) => { - for cfg in cfgs { - provider_connections.insert(cfg.name.as_str(), cfg.get_connection_info().await); - } - } - } - } - } - } - } - - let mut new_lineups: Vec = Vec::with_capacity(new_inputs.len()); - let connections = Some(provider_connections); - for input in &new_inputs { - new_lineups.push(Self::create_lineup(input, connections.as_ref(), self.connection_change_tx.clone())); - } - - debug_if_enabled!("inputs {}", sanitize_sensitive_info(&display_vec(&new_inputs))); - debug_if_enabled!("lineup {}", sanitize_sensitive_info(&display_vec(&new_lineups))); - - self.inputs.store(Arc::new(new_inputs)); - self.providers.store(Arc::new(new_lineups)); - } - - fn get_provider_config<'a>(name: &str, providers: &'a Vec) -> Option<(&'a ProviderLineup, &'a ProviderConfigWrapper)> { - for lineup in providers { - match lineup { - ProviderLineup::Single(single) => { - if single.provider.name == name { - return Some((lineup, &single.provider)); - } - } - ProviderLineup::Multi(multi) => { - for group in &multi.providers { - match group { - ProviderPriorityGroup::SingleProviderGroup(single) => { - if single.name == name { - return Some((lineup, single)); - } - } - ProviderPriorityGroup::MultiProviderGroup(_, configs) => { - for config in configs { - if config.name == name { - return Some((lineup, config)); - } - } - } - } - } - } - } - } - None - } - - async fn force_exact_acquire_connection(&self, provider_name: &str, release_tx: UnboundedSender) -> Arc { - let providers = self.providers.load(); - let allocation = match Self::get_provider_config(provider_name, &providers) { - None => ProviderAllocation::Exhausted, // No Name matched, we don't have this provider - Some((_lineup, config)) => config.force_allocate().await, - }; - - Arc::new(ProviderConnectionGuard::new(allocation, release_tx)) - } - - // Returns the next available provider connection - async fn acquire_connection(&self, input_name: &str, release_tx: UnboundedSender) -> Arc { - let providers = self.providers.load(); - let allocation = match Self::get_provider_config(input_name, &providers) { - None => ProviderAllocation::Exhausted, // No Name matched, we don't have this provider - Some((lineup, _config)) => lineup.acquire(self.grace_period_millis.load(Ordering::Relaxed) > 0, - self.grace_period_timeout_secs.load(Ordering::Relaxed)).await - }; - - if log_enabled!(log::Level::Debug) { - match allocation { - ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(_, ref cfg) | - ProviderAllocation::GracePeriod(_, ref cfg) => { - debug!("Using provider {}", cfg.name); - } - } - } - - Arc::new(ProviderConnectionGuard::new(allocation, release_tx)) - } - - // This method is used for redirects to cycle through provider - // - pub async fn get_next_provider(&self, input_name: &str) -> Option> { - let providers = self.providers.load(); - match Self::get_provider_config(input_name, &providers) { - None => None, - Some((lineup, _config)) => { - let cfg = lineup.get_next(self.grace_period_timeout_secs.load(Ordering::Relaxed)).await; - if log_enabled!(log::Level::Debug) { - if let Some(ref c) = cfg { - debug!("Using provider {}", c.name); - } - } - cfg - } - } - } - - // we need the provider_name to exactly release this provider - // pub async fn release_connection(&self, provider_name: &str) { - // let providers = self.providers.load(); - // if let Some((lineup, _config)) = Self::get_provider_config(provider_name, &providers) { - // lineup.release(provider_name).await; - // } - // } - - pub async fn active_connections(&self) -> Option> { - let mut result = HashMap::::new(); - let mut add_provider = |name: String, count: usize| { - if count > 0 { - result.insert(name, count); - } - }; - let providers = self.providers.load(); - for lineup in providers.iter() { - match lineup { - ProviderLineup::Single(provider_lineup) => { - let connections = provider_lineup.provider.get_current_connections().await; - add_provider(provider_lineup.provider.name.clone(), connections); - } - ProviderLineup::Multi(provider_lineup) => { - let connections = provider_lineup.get_total_connections().await; - add_provider(provider_lineup.name.clone(), connections); - } - } - } - if result.is_empty() { - None - } else { - Some(result) - } - } - - pub async fn is_over_limit(&self, provider_name: &str) -> bool { - let providers = self.providers.load(); - if let Some((_, config)) = Self::get_provider_config(provider_name, &providers) { - config.is_over_limit(self.grace_period_timeout_secs.load(Ordering::Relaxed)).await - } else { - false - } +impl ProviderHandle { + pub fn new(id: ProviderConnectionId, allocation: ProviderAllocation) -> Self { + Self { id, allocation } } } pub struct ActiveProviderManager { providers: ProviderLineupManager, - connections: RwLock>>, + connections: RwLock>, } impl ActiveProviderManager { @@ -824,17 +55,45 @@ impl ActiveProviderManager { self.providers.update_config(inputs, grace_period_millis, grace_period_timeout_secs).await; } - pub async fn force_exact_acquire_connection(&self, provider_name: &str, addr: &str, release_tx: UnboundedSender) -> Arc { - let guard = self.providers.force_exact_acquire_connection(provider_name, release_tx).await; - self.register_connection(addr, &guard).await; - guard + async fn acquire_connection_inner(&self, provider_or_input_name: &str, addr: &SocketAddr, force: bool) -> Option { + // Lock connections + let mut connections = self.connections.write().await; + + // Call the specific acquisition function + let allocation = if force { + self.providers.force_exact_acquire_connection(provider_or_input_name).await + } else { + self.providers.acquire_connection(provider_or_input_name).await + }; + + match &allocation { + ProviderAllocation::Exhausted => {} + ProviderAllocation::Available(_) | ProviderAllocation::GracePeriod(_) => { + let provider_name = allocation.get_provider_name().unwrap_or_default(); + + if let Some(old) = connections.insert(addr.clone(), allocation.clone()) { + crate::utils::trace_if_enabled!( + "register_connection: address {addr} already had a allocation for provider {:?} — forcing release on the old allocation", + old.get_provider_name().unwrap_or_default() + ); + old.release().await; + } + + debug!("Added provider connection {provider_name:?} for {addr}"); + return Some(ProviderHandle::new(addr.clone(), allocation)); + } + } + + None + } + + pub async fn force_exact_acquire_connection(&self, provider_name: &str, addr: &SocketAddr) -> Option { + self.acquire_connection_inner(provider_name, addr, true).await } // Returns the next available provider connection - pub async fn acquire_connection(&self, input_name: &str, addr: &str, release_tx: UnboundedSender) -> Arc { - let guard = self.providers.acquire_connection(input_name, release_tx).await; - self.register_connection(addr, &guard).await; - guard + pub async fn acquire_connection(&self, input_name: &str, addr: &SocketAddr) -> Option { + self.acquire_connection_inner(input_name, addr, false).await } // This method is used for redirects to cycle through provider @@ -842,11 +101,6 @@ impl ActiveProviderManager { self.providers.get_next_provider(provider_name).await } - // we need the provider_name to exactly release this provider - // pub async fn release_connection(&self, provider_name: &str) { - // self.providers.release_connection(provider_name).await; - // } - pub async fn active_connections(&self) -> Option> { self.providers.active_connections().await } @@ -855,332 +109,15 @@ impl ActiveProviderManager { self.providers.is_over_limit(provider_name).await } - async fn register_connection(&self, addr: &str, guard: &Arc) { - if !matches!(guard.allocation, ProviderAllocation::Exhausted) { - debug!("Added provider connection {:?} for {addr}", guard.get_provider_name().unwrap_or_default()); - let mut connections = self.connections.write().await; - if let Some(old) = connections.insert(addr.to_string(), Arc::clone(guard)) { - crate::utils::trace_if_enabled!("register_connection: address {addr} already had a guard for provider {:?} — forcing release on the old guard", - old.get_provider_name().unwrap_or_default()); - // safe force_release or release depending on semantics — force_release ensures decrement - old.force_release().await; - } + pub async fn release_connection(&self, addr: &SocketAddr) { + let handle = self.connections.write().await.remove(addr); + if let Some(allocation) = handle { + debug!("Released provider connection {:?} for {addr}", allocation.get_provider_name().unwrap_or_default()); + allocation.release().await; } } - pub async fn release_connection(&self, addr: &str) { - let guard = self.connections.write().await.remove(addr); - if let Some(guard) = guard { - debug!("Released provider connection {:?} for {addr}", guard.get_provider_name().unwrap_or_default()); - guard.release().await; - } + pub async fn release_handle(&self, handle: &ProviderHandle) { + self.release_connection(&handle.id).await } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::model::ConfigInputAlias; - use crate::Arc; - use shared::model::{InputFetchMethod, InputType}; - use std::sync::atomic::AtomicU16; - use std::thread; - - macro_rules! should_available { - ($lineup:expr, $provider_id:expr, $grace_period_timeout_secs: expr) => { - thread::sleep(std::time::Duration::from_millis(200)); - match $lineup.acquire(true, $grace_period_timeout_secs).await { - ProviderAllocation::Exhausted => assert!(false, "Should available and not exhausted"), - ProviderAllocation::Available(_, provider) => assert_eq!(provider.id, $provider_id), - ProviderAllocation::GracePeriod(_, provider) => assert!(false, "Should available and not grace period: {}", provider.id), - } - }; - } - macro_rules! should_grace_period { - ($lineup:expr, $provider_id:expr, $grace_period_timeout_secs: expr) => { - thread::sleep(std::time::Duration::from_millis(200)); - match $lineup.acquire(true, $grace_period_timeout_secs).await { - ProviderAllocation::Exhausted => assert!(false, "Should grace period and not exhausted"), - ProviderAllocation::Available(_, provider) => assert!(false, "Should grace period and not available: {}", provider.id), - ProviderAllocation::GracePeriod(_, provider) => assert_eq!(provider.id, $provider_id), - } - }; - } - - macro_rules! should_exhausted { - ($lineup:expr, $grace_period_timeout_secs: expr) => { - thread::sleep(std::time::Duration::from_millis(200)); - match $lineup.acquire(true, $grace_period_timeout_secs).await { - ProviderAllocation::Exhausted => {}, - ProviderAllocation::Available(_, provider) => assert!(false, "Should exhausted and not available: {}", provider.id), - ProviderAllocation::GracePeriod(_, provider) => assert!(false, "Should exhausted and not grace period: {}", provider.id), - } - }; - } - - // Helper function to create a ConfigInput instance - fn create_config_input(id: u16, name: &str, priority: i16, max_connections: u16) -> ConfigInput { - ConfigInput { - id, - name: name.to_string(), - url: "http://example.com".to_string(), - epg: Option::default(), - username: None, - password: None, - persist: None, - enabled: true, - input_type: InputType::Xtream, // You can use a default value here - max_connections, - priority, - aliases: None, - headers: HashMap::default(), - options: None, - method: InputFetchMethod::default(), - staged: None, - t_batch_url: None, - } - } - - // Helper function to create a ConfigInputAlias instance - fn create_config_input_alias(id: u16, url: &str, priority: i16, max_connections: u16) -> ConfigInputAlias { - ConfigInputAlias { - id, - name: format!("alias_{id}"), - url: url.to_string(), - username: Some("alias_user".to_string()), - password: Some("alias_pass".to_string()), - priority, - max_connections, - } - } - - fn dummy_callback(_: &str, _: usize) {} - - - // Test acquiring with an alias - #[test] - fn test_provider_with_alias() { - let mut input = create_config_input(1, "provider1_1", 1, 1); - let alias = create_config_input_alias(2, "http://alias1", 2, 2); - - // Adding alias to the provider - input.aliases = Some(vec![alias]); - - let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); - let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; - // Create MultiProviderLineup with the provider and alias - let lineup = MultiProviderLineup::new(&input, Some(dummy_get_connection), &change_callback); - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - // Test that the alias provider is available - should_available!(lineup, 1, 5); - // Try acquiring again - should_available!(lineup, 2, 5); - should_available!(lineup, 2, 5); - should_grace_period!(lineup, 1, 5); - should_grace_period!(lineup, 2, 5); - should_exhausted!(lineup, 5); - should_exhausted!(lineup, 5); - }); - } - - // Test acquiring from a MultiProviderLineup where the alias has a different priority - #[test] - fn test_provider_with_priority_alias() { - let mut input = create_config_input(1, "provider2_1", 1, 2); - let alias = create_config_input_alias(2, "http://alias.com", 0, 2); - // Adding alias with different priority - input.aliases = Some(vec![alias]); - let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); - let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; - let lineup = MultiProviderLineup::new(&input, Some(dummy_get_connection), &change_callback); - // The alias has a higher priority, so the alias should be acquired first - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - for _ in 0..2 { - should_available!(lineup, 2, 5); - } - should_available!(lineup, 1, 5); - }); - } - - // Test provider when there are multiple aliases, all with distinct priorities - #[test] - fn test_provider_with_multiple_aliases() { - let mut input = create_config_input(1, "provider3_1", 1, 1); - let alias1 = create_config_input_alias(2, "http://alias1.com", 1, 2); - let alias2 = create_config_input_alias(3, "http://alias2.com", 0, 1); - - // Adding multiple aliases - input.aliases = Some(vec![alias1, alias2]); - let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); - let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; - let lineup = MultiProviderLineup::new(&input, Some(dummy_get_connection), &change_callback); - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - // The alias with priority 0 should be acquired first (higher priority) - should_available!(lineup, 3, 5); - // Acquire again, and provider should still be available (with remaining capacity) - should_available!(lineup, 1, 5); - // // Check that the second alias with priority 2 is considered next - should_available!(lineup, 2, 5); - should_available!(lineup, 2, 5); - - should_grace_period!(lineup, 3, 5); - should_grace_period!(lineup, 1, 5); - should_grace_period!(lineup, 2, 5); - - should_exhausted!(lineup, 5); - }); - } - - - // Test acquiring when all aliases are exhausted - #[test] - fn test_provider_with_exhausted_aliases() { - let mut input = create_config_input(1, "provider4_1", 1, 1); - let alias1 = create_config_input_alias(2, "http://alias.com", 2, 1); - let alias2 = create_config_input_alias(3, "http://alias.com", -2, 1); - - // Adding alias - input.aliases = Some(vec![alias1, alias2]); - let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); - let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; - let lineup = MultiProviderLineup::new(&input, Some(dummy_get_connection), &change_callback); - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - // Acquire connection from alias2 - should_available!(lineup, 3, 5); - // Acquire connection from provider1 - should_available!(lineup, 1, 5); - // Acquire connection from alias1 - should_available!(lineup, 2, 5); - - // Acquire connection from alias2 - should_grace_period!(lineup, 3, 5); - // Acquire connection from provider1 - should_grace_period!(lineup, 1, 5); - // Acquire connection from alias1 - should_grace_period!(lineup, 2, 5); - - // Now, all are exhausted - should_exhausted!(lineup, 5); - }); - } - - // Test acquiring a connection when there is available capacity - #[test] - fn test_acquire_when_capacity_available() { - let cfg = create_config_input(1, "provider5_1", 1, 2); - let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); - let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; - let lineup = SingleProviderLineup::new(&cfg, Some(dummy_get_connection), &change_callback); - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - // First acquire attempt should succeed - should_available!(lineup, 1, 5); - // Second acquire attempt should succeed as well - should_available!(lineup, 1, 5); - // Third with grace time - should_grace_period!(lineup, 1, 5); - // Fourth acquire attempt should fail as the provider is exhausted - should_exhausted!(lineup, 5); - }); - } - - - // Test releasing a connection - #[test] - fn test_release_connection() { - let cfg = create_config_input(1, "provider7_1", 1, 2); - let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); - let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; - - let lineup = SingleProviderLineup::new(&cfg, Some(dummy_get_connection), &change_callback); - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - // Acquire two connections - should_available!(lineup, 1, 5); - should_available!(lineup, 1, 5); - should_grace_period!(lineup, 1, 5); - should_exhausted!(lineup, 5); - lineup.release("provider7_1").await; - should_grace_period!(lineup, 1, 5); - lineup.release("provider7_1").await; - lineup.release("provider7_1").await; - should_available!(lineup, 1, 5); - should_grace_period!(lineup, 1, 5); - should_exhausted!(lineup, 5); - }); - } - - // Test acquiring with MultiProviderLineup and round-robin allocation - #[test] - fn test_multi_provider_acquire() { - let mut cfg1 = create_config_input(1, "provider8_1", 1, 2); - let alias = create_config_input_alias(2, "http://alias1", 1, 1); - - // Adding alias to the provider - cfg1.aliases = Some(vec![alias]); - - // Create MultiProviderLineup with the provider and alias - let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); - let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; - let lineup = MultiProviderLineup::new(&cfg1, Some(dummy_get_connection), &change_callback); - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - // Test acquiring the first provider - should_available!(lineup, 1, 5); - - // Test acquiring the second provider - should_available!(lineup, 2, 5); - - // Test acquiring the first provider - should_available!(lineup, 1, 5); - - should_grace_period!(lineup, 1, 5); - should_grace_period!(lineup, 2, 5); - - lineup.release("provider8_1").await; - lineup.release("alias_2").await; - lineup.release("provider8_1").await; - - should_available!(lineup, 1, 5); - should_grace_period!(lineup, 1, 5); - should_grace_period!(lineup, 2, 5); - - should_exhausted!(lineup, 5); - }); - } - - // Test concurrent access to `acquire` using multiple threads - #[test] - fn test_concurrent_acquire() { - let cfg = create_config_input(1, "provider9_1", 1, 2); - let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); - let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; - let lineup = Arc::new(SingleProviderLineup::new(&cfg, Some(dummy_get_connection), &change_callback)); - - let available_count = Arc::new(AtomicU16::new(2)); - let grace_period_count = Arc::new(AtomicU16::new(1)); - let exhausted_count = Arc::new(AtomicU16::new(2)); - - for _ in 0..5 { - let lineup_clone = Arc::clone(&lineup); - let available = Arc::clone(&available_count); - let grace_period = Arc::clone(&grace_period_count); - let exhausted = Arc::clone(&exhausted_count); - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - match lineup_clone.acquire(true, 5).await { - ProviderAllocation::Exhausted => exhausted.fetch_sub(1, Ordering::SeqCst), - ProviderAllocation::Available(_, _) => available.fetch_sub(1, Ordering::SeqCst), - ProviderAllocation::GracePeriod(_, _) => grace_period.fetch_sub(1, Ordering::SeqCst), - } - }); - } - assert_eq!(exhausted_count.load(Ordering::SeqCst), 0); - assert_eq!(available_count.load(Ordering::SeqCst), 0); - assert_eq!(grace_period_count.load(Ordering::SeqCst), 0); - } -} +} \ No newline at end of file diff --git a/backend/src/api/model/active_user_manager.rs b/backend/src/api/model/active_user_manager.rs index 07679a54a..6f07d723d 100644 --- a/backend/src/api/model/active_user_manager.rs +++ b/backend/src/api/model/active_user_manager.rs @@ -1,21 +1,21 @@ -use std::borrow::Cow; -use crate::api::model::{ActiveProviderManager, ReleaseTask}; use crate::api::model::SharedStreamManager; +use crate::api::model::ActiveProviderManager; +use crate::auth::Fingerprint; use crate::model::Config; use crate::model::ProxyUserCredentials; +use crate::utils::GeoIp; +use arc_swap::ArcSwapOption; use jsonwebtoken::get_current_timestamp; use log::{debug, error, info}; use shared::model::{ActiveUserConnectionChange, StreamChannel, StreamInfo, UserConnectionPermission}; use shared::utils::{current_time_secs, default_grace_period_millis, default_grace_period_timeout_secs, sanitize_sensitive_info, strip_port}; +use std::borrow::Cow; use std::collections::HashMap; +use std::net::SocketAddr; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; -use arc_swap::ArcSwapOption; -use tokio::runtime::Handle; -use tokio::sync::mpsc::{UnboundedSender, UnboundedReceiver}; +use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio::sync::RwLock; -use crate::auth::Fingerprint; -use crate::utils::GeoIp; const USER_GC_TTL: u64 = 900; // 15 Min const USER_CON_TTL: u64 = 10_800; // 3 hours @@ -24,66 +24,6 @@ const USER_SESSION_LIMIT: usize = 50; type ActiveUserConnectionChangeSender = UnboundedSender; pub type ActiveUserConnectionChangeReceiver = UnboundedReceiver; -macro_rules! active_user_manager_shared_impl { - () => { - #[inline] - async fn get_active_connections(user: &Arc>>) -> usize { - user.read().await.iter().filter(|(_, c)| c.connections > 0).map(|(_, c)| c.connections as usize).sum() - } - - async fn log_active_user(&self) { - let user = Arc::clone(&self.user); - let is_log_user_enabled = self.is_log_user_enabled(); - let user_connection_count = Self::get_active_connections(&user).await; - let user_count = user.read().await.iter().filter(|(_, c)| c.connections > 0).count(); - if let Err(err) = self.connection_change_tx.send(ActiveUserConnectionChange::Connections(user_count, user_connection_count)) { - error!("Failed to send active user connection change: user-count: {user_count}, user-connection-count: {user_connection_count} {err:?}"); - } - if is_log_user_enabled { - info!("Active Users: {user_count}, Active User Connections: {user_connection_count}"); - } - } - - pub async fn remove_connection(&self, addr: &str) { - let username_opt = { - self.user_by_addr.write().await.remove(addr) - }; - - if let Some(username) = username_opt.as_ref() { - { - let mut user = self.user.write().await; - if let Some(connection_data) = user.get_mut(username) { - 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; - } - connection_data.streams.retain(|c| c.addr != addr); - } - } - } - - self.shared_stream_manager.release_connection(addr, true).await; - self.provider_manager.release_connection(addr).await; - - if let Err(e) = self.close_socket_signal_tx.send(addr.to_string()) { - debug!("No active receivers for close signal ({addr}): {e:?}"); - } - - if let Err(err) = self.connection_change_tx.send(ActiveUserConnectionChange::Disconnected(addr.to_string())) { - error!("Failed to send active user connection change: {err:?}"); - } - - if username_opt.is_some() { - self.log_active_user().await; - } - } - }; -} - fn get_grace_options(config: &Config) -> (u64, u64) { let (grace_period_millis, grace_period_timeout_secs) = config.reverse_proxy.as_ref() .and_then(|r| r.stream.as_ref()) @@ -91,67 +31,13 @@ fn get_grace_options(config: &Config) -> (u64, u64) { (grace_period_millis, grace_period_timeout_secs) } -#[derive(Clone)] -pub struct ConnectionGuardUserManager { - log_active_user: bool, - user: Arc>>, - user_by_addr: Arc>>, - shared_stream_manager: Arc, - provider_manager: Arc, - connection_change_tx: ActiveUserConnectionChangeSender, - close_socket_signal_tx: tokio::sync::broadcast::Sender, -} - -impl ConnectionGuardUserManager { - active_user_manager_shared_impl!(); - fn is_log_user_enabled(&self) -> bool { - self.log_active_user - } -} - -pub struct UserConnectionGuard { - manager: Arc, - // username: String, - addr: String, - release_tx: UnboundedSender, -} - -impl UserConnectionGuard { - pub fn new(manager: Arc, addr: &str, release_tx: UnboundedSender) -> Self { - Self { - manager, - addr: addr.to_string(), - release_tx, - } - } -} - -impl Drop for UserConnectionGuard { - fn drop(&mut self) { - let manager = self.manager.clone(); - let addr = self.addr.clone(); - if let Err(_err) = self.release_tx.send(ReleaseTask::UserConnection(manager, addr)) { - if let Ok(handle) = Handle::try_current() { - let manager = self.manager.clone(); - let addr = self.addr.clone(); - handle.spawn(async move { - manager.remove_connection(&addr).await; - }); - } else { - // No runtime - error!("💥💥💥 Dropping UserConnectionGuard without async runtime — User Connection not freed"); - } - } - } -} - #[derive(Clone, Debug)] pub struct UserSession { pub token: String, pub virtual_id: u32, pub provider: String, pub stream_url: String, - pub addr: String, + pub addr: SocketAddr, pub ts: u64, pub permission: UserConnectionPermission, } @@ -197,12 +83,12 @@ pub struct ActiveUserManager { grace_period_timeout_secs: AtomicU64, log_active_user: AtomicBool, user: Arc>>, - user_by_addr: Arc>>, + user_by_addr: Arc>>, gc_ts: Option, shared_stream_manager: Arc, provider_manager: Arc, geo_ip: Arc>, - close_socket_signal_tx: tokio::sync::broadcast::Sender, + close_socket_signal_tx: tokio::sync::broadcast::Sender, connection_change_tx: ActiveUserConnectionChangeSender, } @@ -230,7 +116,61 @@ impl ActiveUserManager { }) } - active_user_manager_shared_impl!(); + #[inline] + async fn get_active_connections(user: &Arc>>) -> usize { + user.read().await.iter().filter(|(_, c)| c.connections > 0).map(|(_, c)| c.connections as usize).sum() + } + + async fn log_active_user(&self) { + let user = Arc::clone(&self.user); + let is_log_user_enabled = self.is_log_user_enabled(); + let user_connection_count = Self::get_active_connections(&user).await; + let user_count = user.read().await.iter().filter(|(_, c)| c.connections > 0).count(); + if let Err(err) = self.connection_change_tx.send(ActiveUserConnectionChange::Connections(user_count, user_connection_count)) { + error!("Failed to send active user connection change: user-count: {user_count}, user-connection-count: {user_connection_count} {err:?}"); + } + if is_log_user_enabled { + info!("Active Users: {user_count}, Active User Connections: {user_connection_count}"); + } + } + + pub async fn remove_connection(&self, addr: &SocketAddr) { + let username_opt = { + self.user_by_addr.write().await.remove(addr) + }; + + if let Some(username) = username_opt.as_ref() { + { + let mut user = self.user.write().await; + if let Some(connection_data) = user.get_mut(username) { + 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; + } + connection_data.streams.retain(|c| c.addr != *addr); + } + } + } + + self.shared_stream_manager.release_connection(addr, true).await; + self.provider_manager.release_connection(addr).await; + + if let Err(e) = self.close_socket_signal_tx.send(addr.clone()) { + debug!("No active receivers for close signal ({addr}): {e:?}"); + } + + if let Err(err) = self.connection_change_tx.send(ActiveUserConnectionChange::Disconnected(addr.clone())) { + error!("Failed to send active user connection change: {err:?}"); + } + + if username_opt.is_some() { + self.log_active_user().await; + } + } pub fn update_config(&self, config: &Config) { let log_active_user = config.log.as_ref().is_some_and(|l| l.log_active_user); @@ -240,18 +180,6 @@ impl ActiveUserManager { self.log_active_user.store(log_active_user, Ordering::Relaxed); } - fn clone_inner(&self) -> ConnectionGuardUserManager { - ConnectionGuardUserManager { - log_active_user: self.log_active_user.load(Ordering::Relaxed), - user: Arc::clone(&self.user), - user_by_addr: Arc::clone(&self.user_by_addr), - shared_stream_manager: Arc::clone(&self.shared_stream_manager), - provider_manager: Arc::clone(&self.provider_manager), - connection_change_tx: self.connection_change_tx.clone(), - close_socket_signal_tx: self.close_socket_signal_tx.clone(), - } - } - pub async fn user_connections(&self, username: &str) -> u32 { if let Some(connection_data) = self.user.read().await.get(username) { return connection_data.connections; @@ -318,8 +246,7 @@ impl ActiveUserManager { #[allow(clippy::too_many_arguments)] pub async fn add_connection(&self, username: &str, max_connections: u32, fingerprint: &Fingerprint, - provider: &str, stream_channel: StreamChannel, user_agent: Cow<'_, str>, - release_tx: UnboundedSender) -> UserConnectionGuard { + provider: &str, stream_channel: StreamChannel, user_agent: Cow<'_, str>) { let country = { let geoip = self.geo_ip.load(); if let Some(geoip_db) = (*geoip).as_ref() { @@ -361,21 +288,20 @@ impl ActiveUserManager { } self.log_active_user().await; - UserConnectionGuard::new(Arc::new(self.clone_inner()), &fingerprint.addr, release_tx) } fn is_log_user_enabled(&self) -> bool { self.log_active_user.load(Ordering::Relaxed) } - fn new_user_session(session_token: &str, virtual_id: u32, provider: &str, stream_url: &str, addr: &str, + fn new_user_session(session_token: &str, virtual_id: u32, provider: &str, stream_url: &str, addr: &SocketAddr, connection_permission: UserConnectionPermission) -> UserSession { UserSession { token: session_token.to_string(), virtual_id, provider: provider.to_string(), stream_url: stream_url.to_string(), - addr: addr.to_string(), + addr: addr.clone(), ts: current_time_secs(), permission: connection_permission, } @@ -383,7 +309,7 @@ impl ActiveUserManager { #[allow(clippy::too_many_arguments)] pub async fn create_user_session(&self, user: &ProxyUserCredentials, session_token: &str, virtual_id: u32, - provider: &str, stream_url: &str, addr: &str, + provider: &str, stream_url: &str, addr: &SocketAddr, connection_permission: UserConnectionPermission) -> String { self.gc(); @@ -421,13 +347,13 @@ impl ActiveUserManager { token } - pub async fn update_session_addr(&self, username: &str, token: &str, addr: &str) { + pub async fn update_session_addr(&self, username: &str, token: &str, addr: &SocketAddr) { let mut user_map = self.user.write().await; user_map.get_mut(username).and_then(|connection_data| { connection_data.sessions.iter_mut().find_map(|session| { if session.token == token { - let old_addr = session.addr.clone(); - addr.clone_into(&mut session.addr); + let old_addr = session.addr; + session.addr = addr.clone(); Some(old_addr) } else { None @@ -436,7 +362,7 @@ impl ActiveUserManager { }); } - pub fn get_close_connection_channel(&self) -> tokio::sync::broadcast::Receiver { + pub fn get_close_connection_channel(&self) -> tokio::sync::broadcast::Receiver { self.close_socket_signal_tx.subscribe() } @@ -459,7 +385,7 @@ impl ActiveUserManager { connection_data.sessions[index].ts = current_time_secs(); // Only re-evaluate permission for limited users during grace if connection_data.max_connections > 0 - && connection_data.sessions[index].permission == UserConnectionPermission::GracePeriod + && connection_data.sessions[index].permission == UserConnectionPermission::GracePeriod { let new_permission = self.check_connection_permission(username, connection_data); connection_data.sessions[index].permission = new_permission; diff --git a/backend/src/api/model/app_state.rs b/backend/src/api/model/app_state.rs index 2accaab1d..46693319d 100644 --- a/backend/src/api/model/app_state.rs +++ b/backend/src/api/model/app_state.rs @@ -1,5 +1,5 @@ use crate::api::config_watch::exec_config_watch; -use crate::api::model::{ActiveProviderManager, EventManager, PlaylistStorage, PlaylistStorageState, ReleaseTask, SharedStreamManager}; +use crate::api::model::{ActiveProviderManager, EventManager, PlaylistStorage, PlaylistStorageState, SharedStreamManager}; use crate::api::model::{ActiveUserManager, DownloadQueue}; use crate::api::scheduler::exec_scheduler; use crate::model::{ @@ -19,7 +19,6 @@ use std::collections::HashMap; use std::sync::atomic::AtomicI8; use std::sync::Arc; use std::time::Duration; -use tokio::sync::mpsc::UnboundedSender; use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use crate::repository::storage::get_geoip_path; @@ -261,13 +260,9 @@ pub struct AppState { pub cancel_tokens: Arc>, pub playlists: Arc, pub geoip: Arc>, - pub(crate) release_tx: UnboundedSender } impl AppState { - pub fn get_release_sender(&self) -> UnboundedSender { - self.release_tx.clone() - } pub(in crate::api::model) async fn set_config(&self,config: Config) -> Result { let changes = self.detect_changes_for_config(&config); diff --git a/backend/src/api/model/mod.rs b/backend/src/api/model/mod.rs index 3e1802edf..d9340f3ac 100644 --- a/backend/src/api/model/mod.rs +++ b/backend/src/api/model/mod.rs @@ -11,7 +11,7 @@ mod stream; mod provider_config; mod event_manager; mod playlist_mem_cache; -mod release_task; +mod provider_lineup_manager; pub use self::app_state::*; pub use self::playlist_mem_cache::*; @@ -23,7 +23,7 @@ pub(in crate::api) use self::stream_error::*; pub(crate) use self::streams::*; pub(in crate::api) use self::active_user_manager::*; pub(in crate::api) use self::active_provider_manager::*; -pub(in crate::api) use self::stream::*; pub(in crate::api) use self::provider_config::*; +pub use self::stream::*; pub use self::event_manager::*; -pub use self::release_task::*; \ No newline at end of file +pub use self::provider_lineup_manager::*; diff --git a/backend/src/api/model/provider_config.rs b/backend/src/api/model/provider_config.rs index 9b60e1826..f389e3d18 100644 --- a/backend/src/api/model/provider_config.rs +++ b/backend/src/api/model/provider_config.rs @@ -1,5 +1,4 @@ use std::fmt; -use crate::api::model::ProviderAllocation; use crate::model::{ConfigInput, ConfigInputAlias, InputUserInfo}; use jsonwebtoken::get_current_timestamp; use log::{debug}; @@ -8,6 +7,7 @@ use std::sync::{Arc}; use tokio::sync::RwLock; use shared::model::InputType; use shared::write_if_some; +use crate::api::model::ProviderAllocation; pub type ProviderConnectionChangeSender = tokio::sync::mpsc::UnboundedSender<(String, usize)>; pub type ProviderConnectionChangeReceiver = tokio::sync::mpsc::UnboundedReceiver<(String, usize)>; diff --git a/backend/src/api/model/provider_lineup_manager.rs b/backend/src/api/model/provider_lineup_manager.rs new file mode 100644 index 000000000..d26b24402 --- /dev/null +++ b/backend/src/api/model/provider_lineup_manager.rs @@ -0,0 +1,1066 @@ +use crate::api::model::provider_config::ProviderConfigWrapper; +use crate::api::model::{ProviderConfig, ProviderConfigConnection, ProviderConnectionChangeCallback, ProviderConnectionChangeSender}; +use crate::utils::debug_if_enabled; +use crate::model::ConfigInput; +use arc_swap::ArcSwap; +use log::{debug, error, log_enabled}; +use shared::utils::{display_vec, sanitize_sensitive_info}; +use std::collections::HashMap; +use std::fmt; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::Arc; + +macro_rules! gen_provider_search { + ($fn_name:ident, $field: ident, $crit_type:ty) => { + fn $fn_name<'a>(criteria: $crit_type, providers: &'a Vec) -> Option<(&'a ProviderLineup, &'a ProviderConfigWrapper)> { + for lineup in providers { + match lineup { + ProviderLineup::Single(single) => { + if single.provider.$field == criteria { + return Some((lineup, &single.provider)); + } + } + ProviderLineup::Multi(multi) => { + for group in &multi.providers { + match group { + ProviderPriorityGroup::SingleProviderGroup(single) => { + if single.$field == criteria { + return Some((lineup, single)); + } + } + ProviderPriorityGroup::MultiProviderGroup(_, configs) => { + for config in configs { + if config.$field == criteria { + return Some((lineup, config)); + } + } + } + } + } + } + } + } + None + } + } +} + +#[derive(Debug, Clone)] +pub enum ProviderAllocation { + Exhausted, + Available(Arc), + GracePeriod(Arc), +} + +impl ProviderAllocation { + pub fn new_available(config: Arc) -> Self { + ProviderAllocation::Available(config) + } + + pub fn new_grace_period(config: Arc) -> Self { + ProviderAllocation::GracePeriod(config) + } + + pub fn get_provider_name(&self) -> Option { + match self { + ProviderAllocation::Exhausted => None, + ProviderAllocation::Available(ref cfg) | + ProviderAllocation::GracePeriod(ref cfg) => { + Some(cfg.name.clone()) + } + } + } + + pub fn get_provider_id(&self) -> Option { + match self { + ProviderAllocation::Exhausted => None, + ProviderAllocation::Available(ref cfg) | + ProviderAllocation::GracePeriod(ref cfg) => { + Some(cfg.id) + } + } + } + + pub fn get_provider_config(&self) -> Option> { + match self { + ProviderAllocation::Exhausted => None, + ProviderAllocation::Available(ref cfg) | + ProviderAllocation::GracePeriod(ref cfg) => { + Some(Arc::clone(cfg)) + } + } + } + + pub async fn release(&self) { + match &self { + ProviderAllocation::Exhausted => {} + ProviderAllocation::Available(config) | + ProviderAllocation::GracePeriod(config) => { + config.release().await; + } + } + } +} + +impl PartialEq for ProviderAllocation { + fn eq(&self, other: &Self) -> bool { + // Note: released flag ignored + match (self, other) { + (ProviderAllocation::Exhausted, ProviderAllocation::Exhausted) => true, + (ProviderAllocation::Available(cfg1), ProviderAllocation::Available(cfg2)) + | (ProviderAllocation::GracePeriod(cfg1), ProviderAllocation::GracePeriod(cfg2)) => cfg1 == cfg2, + _ => false, + } + } +} + +/// This manages different types of provider lineups: +/// +/// `Single(SingleProviderLineup)`: A single provider. +/// `Multi(MultiProviderLineup)`: A set of providers grouped by priority. +#[derive(Debug)] +enum ProviderLineup { + Single(SingleProviderLineup), + Multi(MultiProviderLineup), +} + +impl fmt::Display for ProviderLineup { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ProviderLineup::Single(lineup) => { + write!(f, "SingleProviderLineup: {{ {} }}", lineup.provider) + } + ProviderLineup::Multi(lineup) => { + write!(f, "MultiProviderLineup: {{")?; + for (i, group) in lineup.providers.iter().enumerate() { + write!(f, " Group {}: {}", i + 1, group)?; + } + write!(f, " }}")?; + Ok(()) + } + } + } +} + +impl ProviderLineup { + async fn get_next(&self, grace_period_timeout_secs: u64) -> Option> { + match self { + ProviderLineup::Single(lineup) => lineup.get_next(grace_period_timeout_secs).await, + ProviderLineup::Multi(lineup) => lineup.get_next(grace_period_timeout_secs).await, + } + } + + async fn acquire(&self, with_grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { + match self { + ProviderLineup::Single(lineup) => lineup.acquire(with_grace, grace_period_timeout_secs).await, + ProviderLineup::Multi(lineup) => lineup.acquire(with_grace, grace_period_timeout_secs).await, + } + } +} + +/// Handles a single provider and ensures safe allocation/release of connections. +#[derive(Debug)] +struct SingleProviderLineup { + provider: ProviderConfigWrapper, +} + +impl SingleProviderLineup { + fn new<'a, F>(cfg: &ConfigInput, get_connection: Option, connection_change: &ProviderConnectionChangeCallback) -> Self + where + F: Fn(&str) -> Option<&'a ProviderConfigConnection>, + { + Self { + provider: ProviderConfigWrapper::new(ProviderConfig::new(cfg, get_connection, Arc::clone(connection_change))), + } + } + + async fn get_next(&self, grace_period_timeout_secs: u64) -> Option> { + self.provider.get_next(false, grace_period_timeout_secs).await + } + + async fn acquire(&self, with_grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { + self.provider.try_allocate(with_grace, grace_period_timeout_secs).await + } + + #[cfg(test)] + async fn release(&self, provider_name: &str) { + if self.provider.name == provider_name { + self.provider.release().await; + } + } +} + + +/// Manages provider groups based on priority: +/// +/// `SingleProviderGroup(ProviderConfig)`: A single provider. +/// `MultiProviderGroup(AtomicUsize, Vec)`: A list of providers with a priority index. +#[derive(Debug)] +enum ProviderPriorityGroup { + SingleProviderGroup(ProviderConfigWrapper), + MultiProviderGroup(AtomicUsize, Vec), +} + +impl fmt::Display for ProviderPriorityGroup { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ProviderPriorityGroup::SingleProviderGroup(p) => { + write!(f, "Single({p})") + } + ProviderPriorityGroup::MultiProviderGroup(_, providers) => { + write!(f, "Multi({})", display_vec(providers)) + } + } + } +} + +impl ProviderPriorityGroup { + async fn is_exhausted(&self) -> bool { + match self { + ProviderPriorityGroup::SingleProviderGroup(g) => g.is_exhausted().await, + ProviderPriorityGroup::MultiProviderGroup(_, groups) => { + for g in groups { + if !g.is_exhausted().await { + return false; + } + } + true + } + } + } +} + + +/// Manages multiple providers, ensuring that connections are allocated in a round-robin manner based on priority. +#[repr(align(64))] +#[derive(Debug)] +struct MultiProviderLineup { + name: String, + providers: Vec, + index: AtomicUsize, +} + +impl MultiProviderLineup { + pub fn new<'a, F>(cfg_input: &ConfigInput, get_connection: Option, connection_change: &ProviderConnectionChangeCallback) -> Self + where + F: Fn(&str) -> Option<&'a ProviderConfigConnection> + Copy, + { + let mut inputs = vec![ProviderConfigWrapper::new(ProviderConfig::new(cfg_input, get_connection, Arc::clone(connection_change)))]; + if let Some(aliases) = &cfg_input.aliases { + for alias in aliases { + inputs.push(ProviderConfigWrapper::new(ProviderConfig::new_alias(cfg_input, alias, get_connection, Arc::clone(connection_change)))); + } + } + let mut providers = HashMap::new(); + for provider in inputs { + let priority = provider.get_priority(); + providers.entry(priority) + .or_insert_with(Vec::new) + .push(provider); + } + let mut values: Vec<(i16, Vec)> = providers.into_iter().collect(); + values.sort_by(|(p1, _), (p2, _)| p1.cmp(p2)); + let providers: Vec = values.into_iter().map(|(_, mut group)| { + if group.len() > 1 { + ProviderPriorityGroup::MultiProviderGroup(AtomicUsize::new(0), group) + } else { + ProviderPriorityGroup::SingleProviderGroup(group.remove(0)) + } + }).collect(); + + Self { + name: cfg_input.name.clone(), + providers, + index: AtomicUsize::new(0), + } + } + + /// Attempts to acquire the next available provider from a specific priority group. + /// + /// # Parameters + /// - `priority_group`: Thep rovider group to search within. + /// + /// # Returns + /// - `ProviderAllocation`: A reference to the next available provider in the specified group. + /// + /// # Behavior + /// - Iterates through the providers in the given group in a round-robin manner. + /// - Checks if a provider has available capacity before selecting it. + /// - Uses atomic operations to maintain fair provider selection. + /// + /// # Thread Safety + /// - Uses `RwLock` for safe concurrent access. + /// - Ensures fair provider allocation across multiple threads. + /// + /// # Example Usage + /// ```rust + /// let lineup = MultiProviderLineup::new(&config); + /// match lineup.acquire_next_provider_from_group(priority_group) { + /// ProviderAllocation::Exhausted => println!("All providers exhausted"), + /// ProviderAllocation::Available(provider) => println!("Provider available {}", provider.name), + /// ProviderAllocation::GracePeriodprovider) => println!("Provider with grace period {}", provider.name), + /// } + /// } + /// ``` + async fn acquire_next_provider_from_group(priority_group: &ProviderPriorityGroup, grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { + match priority_group { + ProviderPriorityGroup::SingleProviderGroup(p) => { + let result = p.try_allocate(grace, grace_period_timeout_secs).await; + if !matches!(result, ProviderAllocation::Exhausted) { + return result; + } + } + ProviderPriorityGroup::MultiProviderGroup(index, pg) => { + let provider_count = pg.len(); + let mut idx = index.load(Ordering::Relaxed) % provider_count; + let start = idx; + + for _ in start..provider_count { + let p = &pg[idx]; + let result = p.try_allocate(grace, grace_period_timeout_secs).await; + if !matches!(result, ProviderAllocation::Exhausted) { + index.store((idx + 1) % provider_count, Ordering::Relaxed); + return result; + } + idx = (idx + 1) % provider_count; + } + index.store(idx, Ordering::Relaxed); + } + } + ProviderAllocation::Exhausted + } + + // Used for redirect to cylce through provider + async fn get_next_provider_from_group(priority_group: &ProviderPriorityGroup, grace: bool, grace_period_timeout_secs: u64) -> Option> { + match priority_group { + ProviderPriorityGroup::SingleProviderGroup(p) => { + return p.get_next(grace, grace_period_timeout_secs).await; + } + ProviderPriorityGroup::MultiProviderGroup(index, pg) => { + let provider_count = pg.len(); + let mut idx = index.load(Ordering::Relaxed) % provider_count; + let start = idx; + for _ in start..provider_count { + if let Some(p) = pg.get(idx) { + let result = p.get_next(grace, grace_period_timeout_secs).await; + if result.is_some() { + index.store((idx + 1) % provider_count, Ordering::Relaxed); + return result; + } + } + idx = (idx + 1) % provider_count; + } + index.store(idx, Ordering::Relaxed); + } + } + None + } + + /// Attempts to acquire a provider from the lineup based on priority and availability. + /// + /// # Returns + /// - `ProviderAllocation`: A reference to the acquired provider if allocation was successful. + /// + /// # Behavior + /// - The method iterates through provider priority groups in a round-robin fashion. + /// - It attempts to allocate a provider from the highest priority group first. + /// - If a provider has available capacity, it is returned. + /// - If all providers in a group are exhausted, it moves to the next group. + /// - Updates the internal index to ensure fair distribution of requests. + /// + /// # Thread Safety + /// - Uses atomic operations (`AtomicUsize`) for thread-safe indexing. + /// - Uses `RwLock` for thread-safe provider allocation. + /// + /// # Example Usage + /// ```rust + /// let lineup = MultiProviderLineup::new(&config); + /// match lineup.acquire() { + /// ProviderAllocation::Exhausted => println!("All providers exhausted"), + /// ProviderAllocation::Available(provider) => println!("Provider available {}", provider.name), + /// ProviderAllocation::GracePeriodprovider) => println!("Provider with grace period {}", provider.name), + /// } + /// ``` + async fn acquire(&self, with_grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { + let main_idx = self.index.load(Ordering::SeqCst); + let provider_count = self.providers.len(); + + for index in main_idx..provider_count { + let priority_group = &self.providers[index]; + let allocation = { + let without_grace_allocation = Self::acquire_next_provider_from_group(priority_group, false, grace_period_timeout_secs).await; + if with_grace && matches!(without_grace_allocation, ProviderAllocation::Exhausted) { + Self::acquire_next_provider_from_group(priority_group, true, grace_period_timeout_secs).await + } else { + without_grace_allocation + } + }; + if !matches!(allocation, ProviderAllocation::Exhausted) { + if priority_group.is_exhausted().await { + self.index.store((index + 1) % provider_count, Ordering::SeqCst); + } + return allocation; + } + } + + ProviderAllocation::Exhausted + } + + // it intended to use with redirects to cycle through provider + async fn get_next(&self, grace_period_timeout_secs: u64) -> Option> { + let main_idx = self.index.load(Ordering::SeqCst); + let provider_count = self.providers.len(); + + for index in main_idx..provider_count { + let priority_group = &self.providers[index]; + let allocation = { + let config = Self::get_next_provider_from_group(priority_group, false, grace_period_timeout_secs).await; + if config.is_none() { + Self::get_next_provider_from_group(priority_group, true, grace_period_timeout_secs).await + } else { + config + } + }; + match allocation { + None => {} + Some(config) => { + if priority_group.is_exhausted().await { + self.index.store((index + 1) % provider_count, Ordering::SeqCst); + } + return Some(config); + } + } + } + + None + } + + + #[cfg(test)] + async fn release(&self, provider_name: &str) { + for g in &self.providers { + match g { + ProviderPriorityGroup::SingleProviderGroup(pc) => { + if pc.name == provider_name { + pc.release().await; + break; + } + } + ProviderPriorityGroup::MultiProviderGroup(_, group) => { + for pc in group { + if pc.name == provider_name { + pc.release().await; + return; + } + } + } + } + } + } + + pub async fn get_total_connections(&self) -> usize { + let mut total_connections = 0; + + for group in &self.providers { + match group { + ProviderPriorityGroup::SingleProviderGroup(provider) => { + total_connections += provider.get_current_connections().await; + } + ProviderPriorityGroup::MultiProviderGroup(_, providers) => { + for provider in providers { + total_connections += provider.get_current_connections().await; + } + } + } + } + total_connections + } +} + + +pub(in crate::api::model) struct ProviderLineupManager { + grace_period_millis: AtomicU64, + grace_period_timeout_secs: AtomicU64, + inputs: Arc>>>, + providers: Arc>>, + connection_change_tx: ProviderConnectionChangeSender, +} + +impl ProviderLineupManager { + pub fn new(inputs: Vec>, grace_period_millis: u64, grace_period_timeout_secs: u64, connection_change_tx: ProviderConnectionChangeSender) -> Self { + let lineups = inputs.iter().map(|i| Self::create_lineup(i, None, connection_change_tx.clone())).collect(); + Self { + grace_period_millis: AtomicU64::new(grace_period_millis), + grace_period_timeout_secs: AtomicU64::new(grace_period_timeout_secs), + inputs: Arc::new(ArcSwap::from_pointee(inputs)), + providers: Arc::new(ArcSwap::from_pointee(lineups)), + connection_change_tx, + } + } + + fn create_lineup(cfg_input: &ConfigInput, provider_connections: Option<&HashMap<&str, ProviderConfigConnection>>, connection_change_sender: ProviderConnectionChangeSender) -> ProviderLineup { + let get_connections = provider_connections.map(|c| |name: &str| c.get(name)); + + let cfg_name = cfg_input.name.clone(); + let on_connection_change: ProviderConnectionChangeCallback = Arc::new(move |_name: &str, connections: usize| { + let connection_change_sender = connection_change_sender.clone(); + let provider_cfg_name = cfg_name.clone(); + if let Err(err) = connection_change_sender.send((provider_cfg_name.clone(), connections)) { + error!("Failed to send connection change: {provider_cfg_name}: {connections}, {err}"); + } + }); + + let on_connection_change = Arc::new(on_connection_change); + + if cfg_input.aliases.as_ref().is_some_and(|a| !a.is_empty()) { + ProviderLineup::Multi(MultiProviderLineup::new(cfg_input, get_connections, &on_connection_change)) + } else { + ProviderLineup::Single(SingleProviderLineup::new(cfg_input, get_connections, &on_connection_change)) + } + } + + fn inputs_differ(a: &ConfigInput, b: &ConfigInput) -> bool { + if a.enabled != b.enabled + || a.max_connections != b.max_connections + || a.priority != b.priority + || a.input_type != b.input_type + || a.username != b.username + || a.password != b.password + || a.url != b.url + { + return true; + } + + match (&a.aliases, &b.aliases) { + (None, None) => {} + (Some(_), None) | (None, Some(_)) => return true, + (Some(a_aliases), Some(b_aliases)) => { + if a_aliases.len() != b_aliases.len() { + return true; + } + + for b_alias in b_aliases { + let Some(a_alias) = a_aliases.iter().find(|a| a.name == b_alias.name) else { + return true; + }; + + if a_alias.max_connections != b_alias.max_connections + || a_alias.priority != b_alias.priority + || a_alias.username != b_alias.username + || a_alias.password != b_alias.password + || a_alias.url != b_alias.url + { + return true; + } + } + } + } + + false + } + + fn has_changed(&self, new_inputs: &[Arc]) -> bool { + let old_inputs = self.inputs.load(); + if old_inputs.len() != new_inputs.len() { + return true; + } + for new_input in new_inputs { + let Some(old_input) = old_inputs.iter().find(|i| i.name == new_input.name) else { + return true; + }; + + if Self::inputs_differ(old_input.as_ref(), new_input.as_ref()) { + return true; + } + } + + false + } + + pub async fn update_config(&self, new_inputs: Vec>, grace_period_millis: u64, grace_period_timeout_secs: u64) { + self.grace_period_millis.store(grace_period_millis, Ordering::Relaxed); + self.grace_period_timeout_secs.store(grace_period_timeout_secs, Ordering::Relaxed); + + if !self.has_changed(&new_inputs) { + return; + } + + let old_lineups = self.providers.load(); + let mut provider_connections = HashMap::new(); + for lineup in old_lineups.iter() { + match lineup { + ProviderLineup::Single(single) => { + provider_connections.insert(single.provider.name.as_str(), single.provider.get_connection_info().await); + } + ProviderLineup::Multi(multi) => { + for group in &multi.providers { + match group { + ProviderPriorityGroup::SingleProviderGroup(cfg) => { + provider_connections.insert(cfg.name.as_str(), cfg.get_connection_info().await); + } + ProviderPriorityGroup::MultiProviderGroup(_, cfgs) => { + for cfg in cfgs { + provider_connections.insert(cfg.name.as_str(), cfg.get_connection_info().await); + } + } + } + } + } + } + } + + let mut new_lineups: Vec = Vec::with_capacity(new_inputs.len()); + let connections = Some(provider_connections); + for input in &new_inputs { + new_lineups.push(Self::create_lineup(input, connections.as_ref(), self.connection_change_tx.clone())); + } + + debug_if_enabled!("inputs {}", sanitize_sensitive_info(&display_vec(&new_inputs))); + debug_if_enabled!("lineup {}", sanitize_sensitive_info(&display_vec(&new_lineups))); + + self.inputs.store(Arc::new(new_inputs)); + self.providers.store(Arc::new(new_lineups)); + } + + gen_provider_search!(get_provider_config_by_name, name, &str); + + // + // fn get_provider_config<'a>(name: &str, providers: &'a Vec) -> Option<(&'a ProviderLineup, &'a ProviderConfigWrapper)> { + // for lineup in providers { + // match lineup { + // ProviderLineup::Single(single) => { + // if single.provider.name == name { + // return Some((lineup, &single.provider)); + // } + // } + // ProviderLineup::Multi(multi) => { + // for group in &multi.providers { + // match group { + // ProviderPriorityGroup::SingleProviderGroup(single) => { + // if single.name == name { + // return Some((lineup, single)); + // } + // } + // ProviderPriorityGroup::MultiProviderGroup(_, configs) => { + // for config in configs { + // if config.name == name { + // return Some((lineup, config)); + // } + // } + // } + // } + // } + // } + // } + // } + // None + // } + + fn log_allocation(allocation: &ProviderAllocation) { + if log_enabled!(log::Level::Debug) { + match allocation { + ProviderAllocation::Exhausted => {} + ProviderAllocation::Available(ref cfg) | + ProviderAllocation::GracePeriod(ref cfg) => { + debug!("Using provider {}", cfg.name); + } + } + } + } + + pub async fn force_exact_acquire_connection(&self, provider_name: &str) -> ProviderAllocation { + let providers = self.providers.load(); + let allocation = match Self::get_provider_config_by_name(provider_name, &providers) { + None => ProviderAllocation::Exhausted, // No Name matched, we don't have this provider + Some((_lineup, config)) => config.force_allocate().await, + }; + Self::log_allocation(&allocation); + allocation + } + + // Returns the next available provider connection + pub(crate) async fn acquire_connection(&self, input_name: &str) -> ProviderAllocation { + let providers = self.providers.load(); + let allocation = match Self::get_provider_config_by_name(input_name, &providers) { + None => ProviderAllocation::Exhausted, // No Name matched, we don't have this provider + Some((lineup, _config)) => lineup.acquire(self.grace_period_millis.load(Ordering::Relaxed) > 0, + self.grace_period_timeout_secs.load(Ordering::Relaxed)).await + }; + Self::log_allocation(&allocation); + allocation + } + + // This method is used for redirects to cycle through provider + // + pub async fn get_next_provider(&self, input_name: &str) -> Option> { + let providers = self.providers.load(); + match Self::get_provider_config_by_name(input_name, &providers) { + None => None, + Some((lineup, _config)) => { + let cfg = lineup.get_next(self.grace_period_timeout_secs.load(Ordering::Relaxed)).await; + if log_enabled!(log::Level::Debug) { + if let Some(ref c) = cfg { + debug!("Using provider {}", c.name); + } + } + cfg + } + } + } + + pub async fn active_connections(&self) -> Option> { + let mut result = HashMap::::new(); + let mut add_provider = |name: String, count: usize| { + if count > 0 { + result.insert(name, count); + } + }; + let providers = self.providers.load(); + for lineup in providers.iter() { + match lineup { + ProviderLineup::Single(provider_lineup) => { + let connections = provider_lineup.provider.get_current_connections().await; + add_provider(provider_lineup.provider.name.clone(), connections); + } + ProviderLineup::Multi(provider_lineup) => { + let connections = provider_lineup.get_total_connections().await; + add_provider(provider_lineup.name.clone(), connections); + } + } + } + if result.is_empty() { + None + } else { + Some(result) + } + } + + pub async fn is_over_limit(&self, provider_name: &str) -> bool { + let providers = self.providers.load(); + if let Some((_, config)) = Self::get_provider_config_by_name(provider_name, &providers) { + config.is_over_limit(self.grace_period_timeout_secs.load(Ordering::Relaxed)).await + } else { + false + } + } + + // gen_provider_search!(get_provider_config_by_id, id, u16); + + // pub async fn get_provider_by_id(&self, provider_id: u16) -> Option> { + // let providers = self.providers.load(); + // if let Some((_, config_wrapper)) = Self::get_provider_config_by_id(provider_id, &providers) { + // Some(config_wrapper.get_provider_config()) + // } else { + // None + // } + // } +} + + +#[cfg(test)] +mod tests { + use super::*; + use crate::model::ConfigInputAlias; + use crate::Arc; + use shared::model::{InputFetchMethod, InputType}; + use std::sync::atomic::AtomicU16; + use std::thread; + + macro_rules! should_available { + ($lineup:expr, $provider_id:expr, $grace_period_timeout_secs: expr) => { + thread::sleep(std::time::Duration::from_millis(200)); + match $lineup.acquire(true, $grace_period_timeout_secs).await { + ProviderAllocation::Exhausted => assert!(false, "Should available and not exhausted"), + ProviderAllocation::Available(_, provider) => assert_eq!(provider.id, $provider_id), + ProviderAllocation::GracePeriod(_, provider) => assert!(false, "Should available and not grace period: {}", provider.id), + } + }; + } + macro_rules! should_grace_period { + ($lineup:expr, $provider_id:expr, $grace_period_timeout_secs: expr) => { + thread::sleep(std::time::Duration::from_millis(200)); + match $lineup.acquire(true, $grace_period_timeout_secs).await { + ProviderAllocation::Exhausted => assert!(false, "Should grace period and not exhausted"), + ProviderAllocation::Available(_, provider) => assert!(false, "Should grace period and not available: {}", provider.id), + ProviderAllocation::GracePeriod(_, provider) => assert_eq!(provider.id, $provider_id), + } + }; + } + + macro_rules! should_exhausted { + ($lineup:expr, $grace_period_timeout_secs: expr) => { + thread::sleep(std::time::Duration::from_millis(200)); + match $lineup.acquire(true, $grace_period_timeout_secs).await { + ProviderAllocation::Exhausted => {}, + ProviderAllocation::Available(_, provider) => assert!(false, "Should exhausted and not available: {}", provider.id), + ProviderAllocation::GracePeriod(_, provider) => assert!(false, "Should exhausted and not grace period: {}", provider.id), + } + }; + } + + // Helper function to create a ConfigInput instance + fn create_config_input(id: u16, name: &str, priority: i16, max_connections: u16) -> ConfigInput { + ConfigInput { + id, + name: name.to_string(), + url: "http://example.com".to_string(), + epg: Option::default(), + username: None, + password: None, + persist: None, + enabled: true, + input_type: InputType::Xtream, // You can use a default value here + max_connections, + priority, + aliases: None, + headers: HashMap::default(), + options: None, + method: InputFetchMethod::default(), + staged: None, + t_batch_url: None, + } + } + + // Helper function to create a ConfigInputAlias instance + fn create_config_input_alias(id: u16, url: &str, priority: i16, max_connections: u16) -> ConfigInputAlias { + ConfigInputAlias { + id, + name: format!("alias_{id}"), + url: url.to_string(), + username: Some("alias_user".to_string()), + password: Some("alias_pass".to_string()), + priority, + max_connections, + } + } + + fn dummy_callback(_: &str, _: usize) {} + + + // Test acquiring with an alias + #[test] + fn test_provider_with_alias() { + let mut input = create_config_input(1, "provider1_1", 1, 1); + let alias = create_config_input_alias(2, "http://alias1", 2, 2); + + // Adding alias to the provider + input.aliases = Some(vec![alias]); + + let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); + let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; + // Create MultiProviderLineup with the provider and alias + let lineup = MultiProviderLineup::new(&input, Some(dummy_get_connection), &change_callback); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async move { + // Test that the alias provider is available + should_available!(lineup, 1, 5); + // Try acquiring again + should_available!(lineup, 2, 5); + should_available!(lineup, 2, 5); + should_grace_period!(lineup, 1, 5); + should_grace_period!(lineup, 2, 5); + should_exhausted!(lineup, 5); + should_exhausted!(lineup, 5); + }); + } + + // Test acquiring from a MultiProviderLineup where the alias has a different priority + #[test] + fn test_provider_with_priority_alias() { + let mut input = create_config_input(1, "provider2_1", 1, 2); + let alias = create_config_input_alias(2, "http://alias.com", 0, 2); + // Adding alias with different priority + input.aliases = Some(vec![alias]); + let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); + let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; + let lineup = MultiProviderLineup::new(&input, Some(dummy_get_connection), &change_callback); + // The alias has a higher priority, so the alias should be acquired first + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async move { + for _ in 0..2 { + should_available!(lineup, 2, 5); + } + should_available!(lineup, 1, 5); + }); + } + + // Test provider when there are multiple aliases, all with distinct priorities + #[test] + fn test_provider_with_multiple_aliases() { + let mut input = create_config_input(1, "provider3_1", 1, 1); + let alias1 = create_config_input_alias(2, "http://alias1.com", 1, 2); + let alias2 = create_config_input_alias(3, "http://alias2.com", 0, 1); + + // Adding multiple aliases + input.aliases = Some(vec![alias1, alias2]); + let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); + let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; + let lineup = MultiProviderLineup::new(&input, Some(dummy_get_connection), &change_callback); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async move { + // The alias with priority 0 should be acquired first (higher priority) + should_available!(lineup, 3, 5); + // Acquire again, and provider should still be available (with remaining capacity) + should_available!(lineup, 1, 5); + // // Check that the second alias with priority 2 is considered next + should_available!(lineup, 2, 5); + should_available!(lineup, 2, 5); + + should_grace_period!(lineup, 3, 5); + should_grace_period!(lineup, 1, 5); + should_grace_period!(lineup, 2, 5); + + should_exhausted!(lineup, 5); + }); + } + + + // Test acquiring when all aliases are exhausted + #[test] + fn test_provider_with_exhausted_aliases() { + let mut input = create_config_input(1, "provider4_1", 1, 1); + let alias1 = create_config_input_alias(2, "http://alias.com", 2, 1); + let alias2 = create_config_input_alias(3, "http://alias.com", -2, 1); + + // Adding alias + input.aliases = Some(vec![alias1, alias2]); + let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); + let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; + let lineup = MultiProviderLineup::new(&input, Some(dummy_get_connection), &change_callback); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async move { + // Acquire connection from alias2 + should_available!(lineup, 3, 5); + // Acquire connection from provider1 + should_available!(lineup, 1, 5); + // Acquire connection from alias1 + should_available!(lineup, 2, 5); + + // Acquire connection from alias2 + should_grace_period!(lineup, 3, 5); + // Acquire connection from provider1 + should_grace_period!(lineup, 1, 5); + // Acquire connection from alias1 + should_grace_period!(lineup, 2, 5); + + // Now, all are exhausted + should_exhausted!(lineup, 5); + }); + } + + // Test acquiring a connection when there is available capacity + #[test] + fn test_acquire_when_capacity_available() { + let cfg = create_config_input(1, "provider5_1", 1, 2); + let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); + let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; + let lineup = SingleProviderLineup::new(&cfg, Some(dummy_get_connection), &change_callback); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async move { + // First acquire attempt should succeed + should_available!(lineup, 1, 5); + // Second acquire attempt should succeed as well + should_available!(lineup, 1, 5); + // Third with grace time + should_grace_period!(lineup, 1, 5); + // Fourth acquire attempt should fail as the provider is exhausted + should_exhausted!(lineup, 5); + }); + } + + + // Test releasing a connection + #[test] + fn test_release_connection() { + let cfg = create_config_input(1, "provider7_1", 1, 2); + let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); + let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; + + let lineup = SingleProviderLineup::new(&cfg, Some(dummy_get_connection), &change_callback); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async move { + // Acquire two connections + should_available!(lineup, 1, 5); + should_available!(lineup, 1, 5); + should_grace_period!(lineup, 1, 5); + should_exhausted!(lineup, 5); + lineup.release("provider7_1").await; + should_grace_period!(lineup, 1, 5); + lineup.release("provider7_1").await; + lineup.release("provider7_1").await; + should_available!(lineup, 1, 5); + should_grace_period!(lineup, 1, 5); + should_exhausted!(lineup, 5); + }); + } + + // Test acquiring with MultiProviderLineup and round-robin allocation + #[test] + fn test_multi_provider_acquire() { + let mut cfg1 = create_config_input(1, "provider8_1", 1, 2); + let alias = create_config_input_alias(2, "http://alias1", 1, 1); + + // Adding alias to the provider + cfg1.aliases = Some(vec![alias]); + + // Create MultiProviderLineup with the provider and alias + let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); + let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; + let lineup = MultiProviderLineup::new(&cfg1, Some(dummy_get_connection), &change_callback); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async move { + // Test acquiring the first provider + should_available!(lineup, 1, 5); + + // Test acquiring the second provider + should_available!(lineup, 2, 5); + + // Test acquiring the first provider + should_available!(lineup, 1, 5); + + should_grace_period!(lineup, 1, 5); + should_grace_period!(lineup, 2, 5); + + lineup.release("provider8_1").await; + lineup.release("alias_2").await; + lineup.release("provider8_1").await; + + should_available!(lineup, 1, 5); + should_grace_period!(lineup, 1, 5); + should_grace_period!(lineup, 2, 5); + + should_exhausted!(lineup, 5); + }); + } + + // Test concurrent access to `acquire` using multiple threads + #[test] + fn test_concurrent_acquire() { + let cfg = create_config_input(1, "provider9_1", 1, 2); + let change_callback: ProviderConnectionChangeCallback = Arc::new(dummy_callback); + let dummy_get_connection = |_s: &str| -> Option<&ProviderConfigConnection> { None }; + let lineup = Arc::new(SingleProviderLineup::new(&cfg, Some(dummy_get_connection), &change_callback)); + + let available_count = Arc::new(AtomicU16::new(2)); + let grace_period_count = Arc::new(AtomicU16::new(1)); + let exhausted_count = Arc::new(AtomicU16::new(2)); + + for _ in 0..5 { + let lineup_clone = Arc::clone(&lineup); + let available = Arc::clone(&available_count); + let grace_period = Arc::clone(&grace_period_count); + let exhausted = Arc::clone(&exhausted_count); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async move { + match lineup_clone.acquire(true, 5).await { + ProviderAllocation::Exhausted => exhausted.fetch_sub(1, Ordering::SeqCst), + ProviderAllocation::Available(_) => available.fetch_sub(1, Ordering::SeqCst), + ProviderAllocation::GracePeriod(_) => grace_period.fetch_sub(1, Ordering::SeqCst), + } + }); + } + assert_eq!(exhausted_count.load(Ordering::SeqCst), 0); + assert_eq!(available_count.load(Ordering::SeqCst), 0); + assert_eq!(grace_period_count.load(Ordering::SeqCst), 0); + } +} diff --git a/backend/src/api/model/release_task.rs b/backend/src/api/model/release_task.rs deleted file mode 100644 index 6f6e036d1..000000000 --- a/backend/src/api/model/release_task.rs +++ /dev/null @@ -1,9 +0,0 @@ -use std::sync::Arc; -use crate::api::model::{ConnectionGuardUserManager, ProviderAllocation, ProviderConnectionGuard}; - -#[derive(Clone)] -pub enum ReleaseTask { - ForceProvider(Arc), - ProviderAllocation(ProviderAllocation), - UserConnection(Arc, String) -} \ No newline at end of file diff --git a/backend/src/api/model/stream.rs b/backend/src/api/model/stream.rs index 2c8d693fb..53f2c82a3 100644 --- a/backend/src/api/model/stream.rs +++ b/backend/src/api/model/stream.rs @@ -1,8 +1,12 @@ -use crate::api::model::StreamError; +use std::collections::HashMap; +use std::sync::Arc; +use crate::api::model::{ProviderHandle, StreamError}; use axum::http::StatusCode; use bytes::Bytes; use futures::stream::BoxStream; use url::Url; +use shared::utils::default_grace_period_millis; +use crate::tools::atomic_once_flag::AtomicOnceFlag; pub type BoxedProviderStream = BoxStream<'static, Result>; pub type ProviderStreamHeader = Vec<(String, String)>; @@ -11,3 +15,50 @@ pub type ProviderStreamInfo = Option<(ProviderStreamHeader, StatusCode, Option, ProviderStreamInfo); pub type ProviderStreamFactoryResponse = (BoxedProviderStream, ProviderStreamInfo); + + +type StreamUrl = String; +type ProviderName = String; + +pub enum ProviderStreamState { + Custom(ProviderStreamResponse), + Available(Option, StreamUrl), + GracePeriod(Option, StreamUrl), +} + +pub struct StreamDetails { + pub stream: Option, + pub(crate) stream_info: ProviderStreamInfo, + pub provider_name: Option, + pub grace_period_millis: u64, + pub reconnect_flag: Option>, + pub provider_handle: Option, +} + +impl StreamDetails { + pub fn from_stream(stream: BoxedProviderStream) -> Self { + Self { + stream: Some(stream), + stream_info: None, + provider_name: None, + grace_period_millis: default_grace_period_millis(), + reconnect_flag: None, + provider_handle: None, + } + } + #[inline] + pub fn has_stream(&self) -> bool { + self.stream.is_some() + } + + #[inline] + pub fn has_grace_period(&self) -> bool { + self.grace_period_millis > 0 + } +} + +pub struct StreamingStrategy { + pub provider_handle: Option, + pub provider_stream_state: ProviderStreamState, + pub input_headers: Option>, +} \ No newline at end of file diff --git a/backend/src/api/model/streams/active_client_stream.rs b/backend/src/api/model/streams/active_client_stream.rs index 786668bda..93800de66 100644 --- a/backend/src/api/model/streams/active_client_stream.rs +++ b/backend/src/api/model/streams/active_client_stream.rs @@ -1,10 +1,8 @@ -use crate::api::api_utils::StreamDetails; -use crate::api::model::AppState; +use crate::api::model::{AppState, ProviderHandle, StreamDetails}; use crate::api::model::BoxedProviderStream; use crate::api::model::StreamError; use crate::api::model::TimedClientStream; use crate::api::model::TransportStreamBuffer; -use crate::api::model::{ProviderConnectionGuard, UserConnectionGuard}; use crate::model::ProxyUserCredentials; use bytes::Bytes; use futures::Stream; @@ -28,10 +26,8 @@ const PROVIDER_EXHAUSTED_STREAM: u8 = 3_u8; pub(in crate::api) struct ActiveClientStream { inner: BoxedProviderStream, send_custom_stream_flag: Option>, - #[allow(unused)] - user_connection_guard: Option, #[allow(dead_code)] - provider_connection_guard: Option>, + provider_handle: Option, custom_video: (Option, Option), waker: Option>, } @@ -52,11 +48,11 @@ impl ActiveClientStream { let provider_name = stream_details.provider_name.as_ref().map_or_else(String::new, ToString::to_string); let user_agent = req_headers.get(USER_AGENT).map(|h| String::from_utf8_lossy(h.as_bytes())).unwrap_or_default(); - let user_connection_guard = Some(app_state.active_users.add_connection(username, user.max_connections, fingerprint, &provider_name, stream_channel, user_agent, app_state.get_release_sender()).await); + app_state.active_users.add_connection(username, user.max_connections, fingerprint, &provider_name, stream_channel, user_agent).await; let cfg = &app_state.app_config; let waker = Some(Arc::new(AtomicWaker::new())); let waker_clone = waker.clone(); - let grace_stop_flag = Self::stream_grace_period(app_state, &stream_details, grant_user_grace_period, user, &fingerprint.addr, waker_clone.clone()); + let grace_stop_flag = Self::stream_grace_period(app_state, &stream_details, grant_user_grace_period, user, &fingerprint, waker_clone.clone()); let custom_response = cfg.custom_stream_response.load(); let custom_video = custom_response.as_ref() .map_or((None, None), |c| @@ -67,8 +63,8 @@ impl ActiveClientStream { let stream = match stream_details.stream.take() { None => { - if let Some(guard) = stream_details.provider_connection_guard.as_ref() { - guard.release().await; + if let Some(guard) = stream_details.provider_handle.as_ref() { + app_state.active_provider.release_handle(guard).await; } futures::stream::empty::>().boxed() } @@ -89,8 +85,7 @@ impl ActiveClientStream { Self { inner: stream, - user_connection_guard, - provider_connection_guard: stream_details.provider_connection_guard, + provider_handle: stream_details.provider_handle, send_custom_stream_flag: grace_stop_flag, custom_video, waker, @@ -101,7 +96,7 @@ impl ActiveClientStream { stream_details: &StreamDetails, user_grace_period: bool, user: &ProxyUserCredentials, - addr: &str, + fingerprint: &Fingerprint, waker: Option>) -> Option> { let active_users = Arc::clone(&app_state.active_users); let active_provider = Arc::clone(&app_state.active_provider); @@ -127,7 +122,7 @@ impl ActiveClientStream { let stream_strategy_flag_copy = Arc::clone(&stream_strategy_flag); let grace_period_millis = stream_details.grace_period_millis; - let address = addr.to_string(); + let address = fingerprint.addr.clone(); let user_manager = Arc::clone(&active_users); let provider_manager = Arc::clone(&active_provider); let share_manager = Arc::clone(&shared_stream_manager); diff --git a/backend/src/api/model/streams/shared_stream_manager.rs b/backend/src/api/model/streams/shared_stream_manager.rs index 1a9ce2d69..4f1e7d1bf 100644 --- a/backend/src/api/model/streams/shared_stream_manager.rs +++ b/backend/src/api/model/streams/shared_stream_manager.rs @@ -1,4 +1,4 @@ -use crate::api::model::{ReleaseTask, StreamError}; +use crate::api::model::{ActiveProviderManager, ProviderHandle, StreamError}; use crate::api::model::AppState; use crate::utils::debug_if_enabled; use bytes::Bytes; @@ -7,18 +7,17 @@ use futures::{Stream, StreamExt}; use std::collections::{HashMap, VecDeque}; use std::fmt; use std::fmt::{Debug, Formatter}; +use std::net::SocketAddr; use std::sync::Arc; use crate::api::model::streams::buffered_stream::CHANNEL_SIZE; use crate::api::model::BoxedProviderStream; -use crate::api::model::ProviderConnectionGuard; use crate::utils::trace_if_enabled; -use log::{debug, error, trace}; +use log::{debug, trace}; use shared::utils::sanitize_sensitive_info; use std::pin::Pin; use std::task::{Context, Poll}; -use tokio::runtime::Handle; -use tokio::sync::mpsc::{Sender, UnboundedSender}; +use tokio::sync::mpsc::{Sender}; use tokio::sync::{mpsc, Mutex, RwLock}; use tokio_stream::wrappers::ReceiverStream; use tokio_util::sync::CancellationToken; @@ -54,7 +53,7 @@ fn convert_stream(stream: BoxStream) -> BoxStream, + client_tx: &Sender, + cancellation_token: &CancellationToken) { + for buf in start_buffer { + if cancellation_token.is_cancelled() { return; } + if let Err(err) = client_tx.send(buf.clone()).await { + debug!("Error sending current chunk: {err}"); + return; // stop on send error + } + } +} + /// Represents the state of a shared provider URL. /// /// - `headers`: The initial connection headers used during the setup of the shared stream. @@ -107,40 +121,16 @@ impl BurstBuffer { pub struct SharedStreamState { headers: Vec<(String, String)>, buf_size: usize, - provider_guard: Option>, + provider_guard: Option, subscribers: RwLock>, broadcaster: tokio::sync::broadcast::Sender, stop_token: CancellationToken, burst_buffer: Arc>, - release_tx: UnboundedSender, -} - -impl Drop for SharedStreamState { - fn drop(&mut self) { - if let Some(guard) = self.provider_guard.as_ref() { - if let Err(_err) = self.release_tx.send(ReleaseTask::ForceProvider(Arc::clone(guard))) { - // fallback - if let Ok(handle) = Handle::try_current() { - let guard = guard.clone(); - handle.spawn(async move { - guard.force_release().await; - }); - } else { - // No runtime - error!("💥💥💥 Dropping Shared Stream ProviderConnectionGuard without async runtime — Provider Connection not freed"); - } - } - } - } } impl SharedStreamState { fn new(headers: Vec<(String, String)>, buf_size: usize, - provider_guard: Option>, - release_tx: UnboundedSender) -> Self { - if let Some(guard) = &provider_guard { - guard.disable_release(); - } + provider_guard: Option) -> Self { let (broadcaster, _) = tokio::sync::broadcast::channel(buf_size); // TODO channel size versus byte size, channels are chunk sized, burst_buffer byte sized let burst_buffer_size_in_bytes = MIN_SHARED_BUFFER_SIZE.max(buf_size * 1024 * 12); @@ -152,20 +142,19 @@ impl SharedStreamState { broadcaster, stop_token: CancellationToken::new(), burst_buffer: Arc::new(Mutex::new(BurstBuffer::new(burst_buffer_size_in_bytes))), - release_tx } } - async fn subscribe(&self, addr: &str, manager: Arc) -> (BoxedProviderStream, Option) { + async fn subscribe(&self, addr: &SocketAddr, manager: Arc) -> (BoxedProviderStream, Option) { let (client_tx, client_rx) = mpsc::channel(self.buf_size); let mut broadcast_rx = self.broadcaster.subscribe(); let cancel_token = CancellationToken::new(); - self.subscribers.write().await.insert(addr.to_string(), cancel_token.clone()); + self.subscribers.write().await.insert(addr.clone(), cancel_token.clone()); - let address = addr.to_string(); let client_tx_clone = client_tx.clone(); let burst_buffer = self.burst_buffer.clone(); + let address = addr.clone(); tokio::spawn(async move { let snapshot = { let buffer = burst_buffer.lock().await; @@ -179,14 +168,14 @@ impl SharedStreamState { biased; () = cancel_token.cancelled() => { - debug!("Client disconnected from shared stream: {address}"); + debug!("Client disconnected from shared stream: {}", address.to_string()); break; } result = broadcast_rx.recv() => { match result { Ok(data) => { if let Err(err) = client_tx.send(data).await { - debug!("Shared stream client send error: {address} {err}"); + debug!("Shared stream client send error: {} {err}", address.to_string()); break; } loop_cnt += 1; @@ -196,7 +185,7 @@ impl SharedStreamState { } } Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => { - trace!("Client lagged behind. Skipped {skipped} messages. {address}"); + trace!("Client lagged behind. Skipped {skipped} messages. {}", address.to_string()); loop_cnt += 1; if loop_cnt >= YIELD_COUNTER { tokio::task::yield_now().await; @@ -211,10 +200,7 @@ impl SharedStreamState { manager.release_connection(&address, false).await; }); - let provider = match &self.provider_guard { - None => None, - Some(connection_guard) => connection_guard.get_provider_name() - }; + let provider = self.provider_guard.as_ref().and_then(|h| h.allocation.get_provider_name()); (convert_stream(ReceiverStream::new(client_rx).boxed()), provider) } @@ -292,17 +278,19 @@ impl SharedStreamState { } type SharedStreamRegister = HashMap>; -type SharedStreamAddrRegister = HashMap; +type SharedStreamAddrRegister = HashMap; pub struct SharedStreamManager { + provider_manager: Arc, shared_streams: RwLock, shared_streams_by_addr: RwLock, } impl SharedStreamManager { - pub(crate) fn new() -> Self { + pub(crate) fn new(provider_manager: Arc) -> Self { Self { + provider_manager, shared_streams: RwLock::new(SharedStreamRegister::new()), shared_streams_by_addr: RwLock::new(SharedStreamAddrRegister::new()), } @@ -321,7 +309,7 @@ impl SharedStreamManager { let shared_state = self.shared_streams.write().await.remove(stream_url); { let mut by_addr = self.shared_streams_by_addr.write().await; - let remove_keys: Vec = by_addr.iter() + let remove_keys: Vec = by_addr.iter() .filter_map(|(addr, url)| if url == stream_url { Some(addr.clone()) } else { None }) .collect(); for k in remove_keys { @@ -332,8 +320,8 @@ impl SharedStreamManager { if let Some(shared_state) = shared_state { debug_if_enabled!("Unregistering shared stream {}", sanitize_sensitive_info(stream_url)); - if let Some(guard) = &shared_state.provider_guard { - guard.force_release().await; + if let Some(provider_handle) = &shared_state.provider_guard { + self.provider_manager.release_handle(provider_handle).await; } if send_stop_signal { @@ -343,7 +331,7 @@ impl SharedStreamManager { } } - pub async fn release_connection(&self, addr: &str, send_stop_signal: bool) { + pub async fn release_connection(&self, addr: &SocketAddr, send_stop_signal: bool) { let stream_url = { self.shared_streams_by_addr.write().await.remove(addr) }; @@ -379,7 +367,7 @@ impl SharedStreamManager { } } - async fn subscribe_stream(&self, stream_url: &str, addr: &str, manager: Arc) -> Option<(BoxedProviderStream, Option)> { + async fn subscribe_stream(&self, stream_url: &str, addr: &SocketAddr, manager: Arc) -> Option<(BoxedProviderStream, Option)> { let shared_stream_state = { self.shared_streams.read().await.get(stream_url).map(Arc::clone) }; @@ -387,7 +375,7 @@ impl SharedStreamManager { None => None, Some(stream_state) => { debug_if_enabled!("Responding to existing shared client stream {addr} {}", sanitize_sensitive_info(stream_url)); - self.shared_streams_by_addr.write().await.insert(addr.to_string(), stream_url.to_owned()); + self.shared_streams_by_addr.write().await.insert(addr.clone(), stream_url.to_owned()); Some(stream_state.subscribe(addr, manager).await) } } @@ -401,20 +389,20 @@ impl SharedStreamManager { app_state: &AppState, stream_url: &str, bytes_stream: S, - addr: &str, + addr: &SocketAddr, headers: Vec<(String, String)>, buffer_size: usize, - provider_guard: Option>) -> Option<(BoxedProviderStream, Option)> + provider_guard: Option) -> Option<(BoxedProviderStream, Option)> where S: Stream> + Unpin + 'static + Send, E: std::fmt::Debug + Send, { let buf_size = CHANNEL_SIZE.max(buffer_size); - let shared_state = Arc::new(SharedStreamState::new(headers, buf_size, provider_guard, app_state.get_release_sender())); + let shared_state = Arc::new(SharedStreamState::new(headers, buf_size, provider_guard)); app_state.shared_stream_manager.register(stream_url, Arc::clone(&shared_state)).await; - debug_if_enabled!("Created shared provider stream {}", sanitize_sensitive_info(stream_url)); let subscribed_stream = Self::subscribe_shared_stream(app_state, stream_url, addr).await; shared_state.broadcast(stream_url, bytes_stream, Arc::clone(&app_state.shared_stream_manager)); + debug_if_enabled!("Created shared provider stream {}", sanitize_sensitive_info(stream_url)); subscribed_stream } @@ -422,23 +410,9 @@ impl SharedStreamManager { pub async fn subscribe_shared_stream( app_state: &AppState, stream_url: &str, - addr: &str, + addr: &SocketAddr, ) -> Option<(BoxedProviderStream, Option)> { let manager = Arc::clone(&app_state.shared_stream_manager); app_state.shared_stream_manager.subscribe_stream(stream_url, addr, manager).await } } - - -async fn send_burst_buffer( - start_buffer: &VecDeque, - client_tx: &Sender, - cancellation_token: &CancellationToken) { - for buf in start_buffer { - if cancellation_token.is_cancelled() { return; } - if let Err(err) = client_tx.send(buf.clone()).await { - debug!("Error sending current chunk: {err}"); - return; // stop on send error - } - } -} \ No newline at end of file diff --git a/backend/src/api/serve.rs b/backend/src/api/serve.rs index aefb5b730..7d8b279ac 100644 --- a/backend/src/api/serve.rs +++ b/backend/src/api/serve.rs @@ -17,7 +17,7 @@ use std::time::Duration; use tokio::sync::watch; use tokio_util::sync::CancellationToken; use tower::{Service, ServiceExt}; -use crate::api::model::ActiveUserManager; +use crate::api::model::{ActiveProviderManager, ActiveUserManager}; #[derive(Debug)] struct IncomingStream @@ -41,7 +41,8 @@ impl axum::extract::connect_info::Connected for SocketAddr { pub async fn serve(listener: tokio::net::TcpListener, router: axum::Router<()>, cancel_token: Option, - user_manager: Arc) { + user_manager: &Arc, + provider_manager: &Arc) { let (signal_tx, _signal_rx) = watch::channel(()); let mut make_service = router.into_make_service_with_connect_info::(); @@ -54,7 +55,7 @@ pub async fn serve(listener: tokio::net::TcpListener, } accept_result = listener.accept() => { let Ok((socket, remote_addr)) = accept_result else { continue }; - handle_connection(&mut make_service, &signal_tx, socket, remote_addr, Arc::clone(&user_manager)).await; + handle_connection(&mut make_service, &signal_tx, socket, remote_addr, Arc::clone(user_manager), Arc::clone(provider_manager)).await; } } } @@ -62,7 +63,7 @@ pub async fn serve(listener: tokio::net::TcpListener, None => { loop { let Ok((socket, remote_addr)) = listener.accept().await else { continue }; - handle_connection(&mut make_service, &signal_tx, socket, remote_addr, Arc::clone(&user_manager)).await; + handle_connection(&mut make_service, &signal_tx, socket, remote_addr, Arc::clone(user_manager), Arc::clone(provider_manager)).await; } } } @@ -74,6 +75,7 @@ async fn handle_connection( socket: tokio::net::TcpStream, remote_addr: SocketAddr, user_manager: Arc, + provider_manager: Arc, ) where M: for<'a> Service + Send + 'static, @@ -124,7 +126,7 @@ where let hyper_service = TowerToHyperService::new(tower_service); let signal_tx = signal_tx.clone(); - let addr_str = remote_addr.to_string(); + let addr = remote_addr; tokio::spawn(async move { #[allow(unused_mut)] @@ -133,9 +135,10 @@ where let mut signal_closed = pin!(signal_tx.closed().fuse()); let user_manager_clone = Arc::clone(&user_manager); + let provider_manager_clone = Arc::clone(&provider_manager); let mut addr_close_rx = user_manager_clone.get_close_connection_channel(); - trace!("Connection opened: {addr_str}"); + trace!("Connection opened: {}", addr.to_string()); loop { tokio::select! { @@ -143,18 +146,21 @@ where if let Err(err) = result { trace!("failed to serve connection: {err:#}"); } - user_manager_clone.remove_connection(&addr_str).await; + user_manager_clone.remove_connection(&addr).await; + provider_manager_clone.release_connection(&addr).await; break; } () = &mut signal_closed => { - user_manager_clone.remove_connection(&addr_str).await; + user_manager_clone.remove_connection(&addr).await; + provider_manager_clone.release_connection(&addr).await; debug!("Connection gracefully closed: {remote_addr}"); conn.as_mut().graceful_shutdown(); } Ok(msg) = addr_close_rx.recv() => { // this comes from user manager itself when a user connection is closed // no need to call `user_manager.remove_connection()` - if msg == addr_str { + if msg == addr { + provider_manager_clone.release_connection(&addr).await; debug!("Forced client disconnect {msg}"); conn.as_mut().graceful_shutdown(); break; diff --git a/backend/src/auth/fingerprint.rs b/backend/src/auth/fingerprint.rs index 02af019e6..18c5f6172 100644 --- a/backend/src/auth/fingerprint.rs +++ b/backend/src/auth/fingerprint.rs @@ -19,11 +19,11 @@ fn validate_header(value: &str) -> Option { pub struct Fingerprint { pub key: String, pub client_ip: String, - pub addr: String, + pub addr: SocketAddr, } impl Fingerprint { - pub fn new(key: String, client_ip: String, addr: String) -> Self { + pub fn new(key: String, client_ip: String, addr: SocketAddr) -> Self { Self { key, client_ip, @@ -82,6 +82,6 @@ impl Fingerprint { // debug!("{key}, {client_ip}, {addr}"); - Ok(Fingerprint::new(key, client_ip, addr.to_string())) + Ok(Fingerprint::new(key, client_ip, addr)) } } \ No newline at end of file diff --git a/frontend/scss/_theme.scss b/frontend/scss/_theme.scss index da3d5fd29..13fbbce26 100644 --- a/frontend/scss/_theme.scss +++ b/frontend/scss/_theme.scss @@ -361,6 +361,8 @@ --brick-port-invalid-background-color: #e74c3c; --source-editor-delete-color: #FF0000FF; + --source-editor-line-color: #ffffff; + --source-editor-pending-line-color: yellow; } @@ -628,4 +630,7 @@ body[data-theme="bright"] { --floating-particle-gradient-color-1: rgba(0,0,0,0.3); --floating-particle-gradient-color-2: rgba(0,0,0,0); + --source-editor-line-color: #000000; + --source-editor-pending-line-color: yellow; + } \ No newline at end of file diff --git a/frontend/src/app/components/source_editor/editor_view.rs b/frontend/src/app/components/source_editor/editor_view.rs index f4a3ddc0d..441d5b35b 100644 --- a/frontend/src/app/components/source_editor/editor_view.rs +++ b/frontend/src/app/components/source_editor/editor_view.rs @@ -427,7 +427,7 @@ pub fn SourceEditor() -> Html { Some(html! { - + { if *delete_mode { let mid_x = (from_x + to_x) / 2.0; let mid_y = (from_y + to_y) / 2.0; @@ -454,7 +454,7 @@ pub fn SourceEditor() -> Html { } diff --git a/shared/src/model/active_user_connection_change.rs b/shared/src/model/active_user_connection_change.rs index c3dff197d..a1845e8dc 100644 --- a/shared/src/model/active_user_connection_change.rs +++ b/shared/src/model/active_user_connection_change.rs @@ -1,3 +1,4 @@ +use std::net::SocketAddr; use crate::model::StreamInfo; #[allow(clippy::large_enum_variant)] @@ -5,6 +6,6 @@ use crate::model::StreamInfo; #[serde(tag = "type", content = "payload", rename_all = "camelCase")] pub enum ActiveUserConnectionChange { Connected(StreamInfo), - Disconnected(String), // addr + Disconnected(SocketAddr), // addr Connections(usize, usize) // user_count, connection_count } diff --git a/shared/src/model/stream_info.rs b/shared/src/model/stream_info.rs index 4579f8c6b..3c0acf3a7 100644 --- a/shared/src/model/stream_info.rs +++ b/shared/src/model/stream_info.rs @@ -1,3 +1,4 @@ +use std::net::SocketAddr; use serde::{Deserialize, Serialize}; use crate::model::{M3uPlaylistItem, PlaylistEntry, PlaylistItemType, XtreamCluster, XtreamPlaylistItem}; use crate::utils::{current_time_secs, longest}; @@ -55,7 +56,7 @@ pub struct StreamInfo { pub username: String, pub channel: StreamChannel, pub provider: String, - pub addr: String, + pub addr: SocketAddr, pub client_ip: String, #[serde(default)] pub user_agent: String, @@ -66,12 +67,12 @@ pub struct StreamInfo { } impl StreamInfo { - pub fn new(username: &str, addr: &str, client_ip: &str, provider: &str, stream_channel: StreamChannel, user_agent: String, country: Option) -> Self { + pub fn new(username: &str, addr: &SocketAddr, client_ip: &str, provider: &str, stream_channel: StreamChannel, user_agent: String, country: Option) -> Self { Self { username: username.to_string(), channel: stream_channel, provider: provider.to_string(), - addr: addr.to_string(), + addr: addr.clone(), client_ip: client_ip.to_string(), user_agent, ts: current_time_secs(),