diff --git a/backend/src/api/api_utils.rs b/backend/src/api/api_utils.rs index 8f36e6679..2d9967632 100644 --- a/backend/src/api/api_utils.rs +++ b/backend/src/api/api_utils.rs @@ -287,12 +287,12 @@ struct StreamingStrategy { /// - and optional HTTP headers to include in the request. /// /// This logic helps abstract the decision-making behind provider selection and stream URL resolution. -async fn resolve_streaming_strategy(app_state: &AppState, stream_url: &str, input: &ConfigInput, force_provider: Option<&str>) +async fn resolve_streaming_strategy(app_state: &AppState, stream_url: &str, addr: &str, input: &ConfigInput, 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).await, - None => app_state.active_provider.acquire_connection(&input.name).await + Some(provider) => app_state.active_provider.force_exact_acquire_connection(provider, addr).await, + None => app_state.active_provider.acquire_connection(&input.name, addr).await }; let stream_response_params = match &*provider_connection_guard { ProviderAllocation::Exhausted => { @@ -300,8 +300,8 @@ async fn resolve_streaming_strategy(app_state: &AppState, stream_url: &str, inpu let stream = create_provider_connections_exhausted_stream(&app_state.app_config, &[]); ProviderStreamState::Custom(stream) } - ProviderAllocation::Available(ref provider) - | ProviderAllocation::GracePeriod(ref provider) => { + 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 { @@ -310,7 +310,7 @@ async fn resolve_streaming_strategy(app_state: &AppState, stream_url: &str, inpu (provider.name.to_string(), get_stream_alternative_url(stream_url, input, provider)) }; - if matches!(&*provider_connection_guard, ProviderAllocation::Available(_)) { + if matches!(&*provider_connection_guard, ProviderAllocation::Available(_, _)) { ProviderStreamState::Available(Some(provider), url) } else { ProviderStreamState::GracePeriod(Some(provider), url) @@ -336,6 +336,7 @@ fn get_grace_period_millis(connection_permission: UserConnectionPermission, stre async fn create_stream_response_details(app_state: &AppState, stream_options: &StreamOptions, stream_url: &str, + addr: &str, req_headers: &HeaderMap, input: &ConfigInput, item_type: PlaylistItemType, @@ -343,7 +344,7 @@ async fn create_stream_response_details(app_state: &AppState, connection_permission: UserConnectionPermission, force_provider: Option<&str>) -> StreamDetails { let mut streaming_strategy = - resolve_streaming_strategy(app_state, stream_url, input, force_provider).await; + resolve_streaming_strategy(app_state, stream_url, addr, input, force_provider).await; let config_grace_period_millis = app_state.app_config.config.load().reverse_proxy.as_ref() .and_then(|r| r.stream.as_ref()).map_or_else(default_grace_period_millis, |s| s.grace_period_millis); let grace_period_millis = get_grace_period_millis(connection_permission, &streaming_strategy.provider_stream_state, config_grace_period_millis); @@ -537,12 +538,12 @@ pub async fn force_provider_stream_response(addr: &str, let connection_permission = UserConnectionPermission::Allowed; let mut stream_details = - create_stream_response_details(app_state, &stream_options, &user_session.stream_url, req_headers, input, item_type, share_stream, connection_permission, Some(&user_session.provider)).await; + create_stream_response_details(app_state, &stream_options, &user_session.stream_url, addr, req_headers, input, item_type, share_stream, connection_permission, Some(&user_session.provider)).await; if stream_details.has_stream() { let provider_response = stream_details.stream_info.as_ref().map(|(h, sc, url)| (h.clone(), *sc, url.clone())); - app_state.active_users.update_session_addr(&user.username, &user_session.token, addr).await; - let stream = ActiveClientStream::new(stream_details, app_state, user, connection_permission, addr).await; + app_state.active_users.update_session_addr(&user.username, &user_session.token, addr); + let stream = ActiveClientStream::new(stream_details, app_state, user, connection_permission, addr); let (status_code, header_map) = get_stream_response_with_headers(provider_response.map(|(h, s, _)| (h, s))); let mut response = axum::response::Response::builder().status(status_code); @@ -580,26 +581,27 @@ pub async fn stream_response(addr: &str, let share_stream = is_stream_share_enabled(item_type, target); if share_stream { - if let Some(value) = shared_stream_response(app_state, stream_url, user, connection_permission, addr).await { + if let Some(value) = shared_stream_response(app_state, stream_url, addr, user, connection_permission) { return value.into_response(); } } let stream_options = get_stream_options(app_state); let mut stream_details = - create_stream_response_details(app_state, &stream_options, stream_url, req_headers, input, item_type, share_stream, connection_permission, None).await; + create_stream_response_details(app_state, &stream_options, stream_url, addr, req_headers, input, item_type, share_stream, connection_permission, None).await; if stream_details.has_stream() { // let content_length = get_stream_content_length(provider_response.as_ref()); let provider_response = stream_details.stream_info.as_ref().map(|(h, sc, response_url)| (h.clone(), *sc, response_url.clone())); let provider_name = stream_details.provider_connection_guard.as_ref().and_then(ProviderConnectionGuard::get_provider_name); - let stream = ActiveClientStream::new(stream_details, app_state, user, connection_permission, addr).await; + let provider_guard = if share_stream { stream_details.provider_connection_guard.take() } else { None }; + let stream = ActiveClientStream::new(stream_details, app_state, user, connection_permission, addr); let stream_resp = if share_stream { debug_if_enabled!("Streaming shared stream request from {}", sanitize_sensitive_info(stream_url)); // Shared Stream response let shared_headers = provider_response.as_ref().map_or_else(Vec::new, |(h, _, _)| h.clone()); - SharedStreamManager::subscribe(app_state, stream_url, stream, shared_headers, stream_options.buffer_size).await; - if let Some(broadcast_stream) = SharedStreamManager::subscribe_shared_stream(app_state, stream_url).await { + SharedStreamManager::subscribe(app_state, stream_url, addr, stream, shared_headers, stream_options.buffer_size, provider_guard); + if let Some(broadcast_stream) = SharedStreamManager::subscribe_shared_stream(app_state, stream_url, addr) { let (status_code, header_map) = get_stream_response_with_headers(provider_response.map(|(h, s, _)| (h, s))); let mut response = axum::response::Response::builder() .status(status_code); @@ -627,7 +629,7 @@ pub async fn stream_response(addr: &str, if let Some(provider) = provider_name { if matches!(item_type, PlaylistItemType::LiveHls | PlaylistItemType::LiveDash | PlaylistItemType::Video | PlaylistItemType::Series | PlaylistItemType::Catchup) { - let _ = app_state.active_users.create_user_session(user, session_token, virtual_id, &provider, &session_url, addr, connection_permission).await; + let _ = app_state.active_users.create_user_session(user, session_token, virtual_id, &provider, &session_url, addr, connection_permission); } } @@ -649,13 +651,13 @@ fn get_stream_throttle(app_state: &AppState) -> u64 { .map(|stream| stream.throttle_kbps).unwrap_or_default() } -async fn shared_stream_response(app_state: &AppState, stream_url: &str, user: &ProxyUserCredentials, connect_permission: UserConnectionPermission, addr: &str) -> Option { - if let Some(stream) = SharedStreamManager::subscribe_shared_stream(app_state, stream_url).await { +fn shared_stream_response(app_state: &AppState, stream_url: &str, addr: &str, user: &ProxyUserCredentials, connect_permission: UserConnectionPermission) -> Option { + if let Some(stream) = SharedStreamManager::subscribe_shared_stream(app_state, stream_url, addr) { debug_if_enabled!("Using shared stream {}", sanitize_sensitive_info(stream_url)); - if let Some(headers) = app_state.shared_stream_manager.get_shared_state_headers(stream_url).await { + if let Some(headers) = app_state.shared_stream_manager.get_shared_state_headers(stream_url) { let (status_code, header_map) = get_stream_response_with_headers(Some((headers.clone(), axum::http::StatusCode::OK))); let stream_details = StreamDetails::from_stream(stream); - let stream = ActiveClientStream::new(stream_details, app_state, user, connect_permission, addr).await.boxed(); + let stream = ActiveClientStream::new(stream_details, app_state, user, connect_permission, addr).boxed(); let mut response = axum::response::Response::builder() .status(status_code); for (key, value) in &header_map { diff --git a/backend/src/api/endpoints/hls_api.rs b/backend/src/api/endpoints/hls_api.rs index 70deb9c28..c1a4680be 100644 --- a/backend/src/api/endpoints/hls_api.rs +++ b/backend/src/api/endpoints/hls_api.rs @@ -48,7 +48,7 @@ pub(in crate::api) async fn handle_hls_stream_request( let (request_url, session_token) = match user_session { Some(session) => { - match app_state.active_provider.force_exact_acquire_connection(&session.provider).await.get_provider_config() { + match app_state.active_provider.force_exact_acquire_connection(&session.provider, addr).await.get_provider_config() { Some(provider_cfg) => { let stream_url = get_stream_alternative_url(&url, input, &provider_cfg); (stream_url, Some(session.token.to_string())) @@ -61,8 +61,8 @@ pub(in crate::api) async fn handle_hls_stream_request( Some(provider_cfg) => { let stream_url = get_stream_alternative_url(&url, input, &provider_cfg); let user_session_token = format!("{fingerprint}{virtual_id}"); - let session_token= app_state.active_users.create_user_session(user, &user_session_token, virtual_id, &provider_cfg.name, &stream_url, addr, connection_permission).await; - (stream_url, session_token) + let session_token= app_state.active_users.create_user_session(user, &user_session_token, virtual_id, &provider_cfg.name, &stream_url, addr, connection_permission); + (stream_url, Some(session_token)) }, None => (url, None), } @@ -108,7 +108,7 @@ async fn hls_api_stream( let input = try_option_bad_request!(app_state.app_config.get_input_by_id(params.input_id), true, format!("Cant find input for target {target_name}, context {}, stream_id {virtual_id}", XtreamCluster::Live)); let user_session_token = format!("{fingerprint}{virtual_id}"); - let mut user_session = app_state.active_users.get_user_session(&user.username, &user_session_token).await; + let mut user_session = app_state.active_users.get_user_session(&user.username, &user_session_token); if let Some(session) = &mut user_session { if session.permission == UserConnectionPermission::Exhausted { @@ -134,7 +134,7 @@ async fn hls_api_stream( return axum::http::StatusCode::BAD_REQUEST.into_response(); } - let connection_permission = user.connection_permission(&app_state).await; + let connection_permission = user.connection_permission(&app_state); if connection_permission == UserConnectionPermission::Exhausted { return create_custom_video_stream_response(&app_state.app_config, CustomVideoStreamType::UserConnectionsExhausted).into_response(); } diff --git a/backend/src/api/endpoints/m3u_api.rs b/backend/src/api/endpoints/m3u_api.rs index 756fdfc19..5778825be 100644 --- a/backend/src/api/endpoints/m3u_api.rs +++ b/backend/src/api/endpoints/m3u_api.rs @@ -86,7 +86,7 @@ async fn m3u_api_stream( let session_key = format!("{fingerprint}{virtual_id}"); - let user_session = app_state.active_users.get_user_session(&user.username, &session_key).await; + let user_session = app_state.active_users.get_user_session(&user.username, &session_key); let session_url = if let Some(session) = &user_session { if session.permission == UserConnectionPermission::Exhausted { @@ -105,7 +105,7 @@ async fn m3u_api_stream( pli.url.as_str() }; - let connection_permission = user.connection_permission(app_state).await; + let connection_permission = user.connection_permission(app_state); if connection_permission == UserConnectionPermission::Exhausted { return create_custom_video_stream_response(&app_state.app_config, CustomVideoStreamType::UserConnectionsExhausted).into_response(); } diff --git a/backend/src/api/endpoints/v1_api.rs b/backend/src/api/endpoints/v1_api.rs index c1f5a550f..ddd89ba5d 100644 --- a/backend/src/api/endpoints/v1_api.rs +++ b/backend/src/api/endpoints/v1_api.rs @@ -298,7 +298,7 @@ async fn create_status_check(app_state: &Arc) -> StatusCheck { }; let (active_users, active_user_connections) = { let active_user = &app_state.active_users; - (active_user.active_users().await, active_user.active_connections().await) + (active_user.active_users(), active_user.active_connections()) }; let active_provider_connections = app_state.active_provider.active_connections().await.map(|c| c.into_iter().collect::>()); diff --git a/backend/src/api/endpoints/xtream_api.rs b/backend/src/api/endpoints/xtream_api.rs index c644df7fd..21d26d258 100644 --- a/backend/src/api/endpoints/xtream_api.rs +++ b/backend/src/api/endpoints/xtream_api.rs @@ -167,9 +167,9 @@ pub(in crate::api) fn get_xtream_player_api_stream_url( } } -async fn get_user_info(user: &ProxyUserCredentials, app_state: &AppState) -> XtreamAuthorizationResponse { +fn get_user_info(user: &ProxyUserCredentials, app_state: &AppState) -> XtreamAuthorizationResponse { let server_info = app_state.app_config.get_user_server_info(user); - let active_connections = app_state.get_active_connections_for_user(&user.username).await; + let active_connections = app_state.get_active_connections_for_user(&user.username); XtreamAuthorizationResponse::new(&server_info, user, active_connections, app_state.app_config.config.load().user_access_control) } @@ -201,7 +201,7 @@ async fn xtream_player_api_stream( let item_type = if stream_req.context == ApiStreamContext::Timeshift { PlaylistItemType::Catchup } else { pli.item_type }; let session_key = format!("{fingerprint}{virtual_id}"); - let user_session = app_state.active_users.get_user_session(&user.username, &session_key).await; + let user_session = app_state.active_users.get_user_session(&user.username, &session_key); let session_url = if let Some(session) = &user_session { if session.permission == UserConnectionPermission::Exhausted { @@ -222,7 +222,7 @@ async fn xtream_player_api_stream( pli.url.as_str() }; - let connection_permission = user.connection_permission(app_state).await; + let connection_permission = user.connection_permission(app_state); if connection_permission == UserConnectionPermission::Exhausted { return create_custom_video_stream_response(&app_state.app_config, CustomVideoStreamType::UserConnectionsExhausted).into_response(); } @@ -790,12 +790,12 @@ async fn xtream_player_api( let user_target = get_user_target(&api_req, app_state); if let Some((user, target)) = user_target { if !target.has_output(&TargetType::Xtream) { - return axum::response::Json(get_user_info(&user, app_state).await).into_response(); + return axum::response::Json(get_user_info(&user, app_state)).into_response(); } let action = api_req.action.trim(); if action.is_empty() { - return axum::response::Json(get_user_info(&user, app_state).await).into_response(); + return axum::response::Json(get_user_info(&user, app_state)).into_response(); } if user.permission_denied(app_state) { @@ -814,7 +814,7 @@ async fn xtream_player_api( match action { crate::model::XC_ACTION_GET_ACCOUNT_INFO => { - return axum::response::Json(get_user_info(&user, app_state).await).into_response(); + return axum::response::Json(get_user_info(&user, app_state)).into_response(); } crate::model::XC_ACTION_GET_SERIES_INFO => { skip_json_response_if_flag_set!(skip_series, xtream_get_stream_info_response(app_state, &user, &target, api_req.series_id.trim(), XtreamCluster::Series).await); diff --git a/backend/src/api/main_api.rs b/backend/src/api/main_api.rs index f6df98992..bbda2db1a 100644 --- a/backend/src/api/main_api.rs +++ b/backend/src/api/main_api.rs @@ -54,8 +54,9 @@ async fn healthcheck() -> impl axum::response::IntoResponse { fn create_shared_data(app_config: &Arc, forced_targets: &Arc) -> AppState { let config = app_config.config.load(); let cache = create_cache(&config); - let active_users = Arc::new(ActiveUserManager::new(&config)); + let shared_stream_manager = Arc::new(SharedStreamManager::new()); let active_provider = Arc::new(ActiveProviderManager::new(app_config)); + let active_users = Arc::new(ActiveUserManager::new(&config, &shared_stream_manager, &active_provider)); let client = create_http_client(app_config); AppState { @@ -64,7 +65,7 @@ fn create_shared_data(app_config: &Arc, forced_targets: &Arc {} + ProviderAllocation::Available(state, _) | + ProviderAllocation::GracePeriod(state, _) => { + let _ = state.compare_exchange(CONNECTION_STATE_ACTIVE, CONNECTION_STATE_SHARED, Ordering::SeqCst, Ordering::SeqCst); + } + } + } + pub(crate) fn release(&self) { + match &self.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() { + let provider_config = Arc::clone(config); + trace!("Releasing provider connection {:?}", provider_config.name); + tokio::spawn(async move { + provider_config.release().await; + }); + } + } + } + } + + // wen need to ensure the connections is released + pub(crate) fn force_release(&self) { + match &self.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); + let provider_config = Arc::clone(config); + trace!("Forced releasing provider connection {:?}", provider_config.name); + tokio::spawn(async move { + provider_config.release().await; + }); + } + } + } + } +} + impl ProviderConnectionGuard { pub fn new(allocation: ProviderAllocation) -> Self { Self { @@ -23,8 +76,8 @@ impl ProviderConnectionGuard { pub fn get_provider_name(&self) -> Option { match self.allocation { ProviderAllocation::Exhausted => None, - ProviderAllocation::Available(ref cfg) | - ProviderAllocation::GracePeriod(ref cfg) => { + ProviderAllocation::Available(_, ref cfg) | + ProviderAllocation::GracePeriod(_, ref cfg) => { Some(cfg.name.clone()) } } @@ -32,8 +85,8 @@ impl ProviderConnectionGuard { pub fn get_provider_config(&self) -> Option> { match self.allocation { ProviderAllocation::Exhausted => None, - ProviderAllocation::Available(ref cfg) | - ProviderAllocation::GracePeriod(ref cfg) => { + ProviderAllocation::Available(_, ref cfg) | + ProviderAllocation::GracePeriod(_, ref cfg) => { Some(Arc::clone(cfg)) } } @@ -49,28 +102,38 @@ impl Deref for ProviderConnectionGuard { impl Drop for ProviderConnectionGuard { fn drop(&mut self) { - match &self.allocation { - ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(config) | - ProviderAllocation::GracePeriod(config) => { - // let manager = self.manager.clone(); - let provider_config = Arc::clone(config); - tokio::spawn(async move { - provider_config.release().await; - // manager.release_connection(&provider_config.name).await; - }); - } - } + self.release(); } } #[derive(Debug, Clone)] pub enum ProviderAllocation { Exhausted, - Available(Arc), - GracePeriod(Arc), + Available(Arc, Arc), + GracePeriod(Arc, Arc), } +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: /// @@ -97,12 +160,12 @@ impl ProviderLineup { } } - 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, - } - } + // 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. @@ -129,11 +192,11 @@ impl SingleProviderLineup { self.provider.try_allocate(with_grace, grace_period_timeout_secs).await } - async fn release(&self, provider_name: &str) { - if self.provider.name == provider_name { - self.provider.release().await; - } - } + // async fn release(&self, provider_name: &str) { + // if self.provider.name == provider_name { + // self.provider.release().await; + // } + // } } @@ -144,7 +207,7 @@ impl SingleProviderLineup { #[derive(Debug)] enum ProviderPriorityGroup { SingleProviderGroup(ProviderConfigWrapper), - MultiProviderGroup(Mutex, Vec), + MultiProviderGroup(AtomicUsize, Vec), } impl ProviderPriorityGroup { @@ -194,7 +257,7 @@ impl MultiProviderLineup { values.sort_by(|(p1, _), (p2, _)| p1.cmp(p2)); let providers: Vec = values.into_iter().map(|(_, mut group)| { if group.len() > 1 { - ProviderPriorityGroup::MultiProviderGroup(Mutex::new(0), group) + ProviderPriorityGroup::MultiProviderGroup(AtomicUsize::new(0), group) } else { ProviderPriorityGroup::SingleProviderGroup(group.remove(0)) } @@ -237,31 +300,25 @@ impl MultiProviderLineup { match priority_group { ProviderPriorityGroup::SingleProviderGroup(p) => { let result = p.try_allocate(grace, grace_period_timeout_secs).await; - match result { - ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(_) | ProviderAllocation::GracePeriod(_) => return result + if !matches!(result, ProviderAllocation::Exhausted) { + return result; } } ProviderPriorityGroup::MultiProviderGroup(index, pg) => { let provider_count = pg.len(); - let mut idx = { - *index.lock().await - }; + let mut idx = index.load(Ordering::Relaxed) % provider_count; let start = idx; for _ in start..provider_count { - let p = pg.get(idx).unwrap(); - idx = (idx + 1) % provider_count; + let p = &pg[idx]; let result = p.try_allocate(grace, grace_period_timeout_secs).await; - match result { - ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(_) | ProviderAllocation::GracePeriod(_) => { - *index.lock().await = idx; - return result; - } + if !matches!(result, ProviderAllocation::Exhausted) { + index.store((idx + 1) % provider_count, Ordering::Relaxed); + return result; } + idx = (idx + 1) % provider_count; } - *index.lock().await = idx; + index.store(idx, Ordering::Relaxed); } } ProviderAllocation::Exhausted @@ -274,20 +331,19 @@ impl MultiProviderLineup { return p.get_next(grace, grace_period_timeout_secs).await; } ProviderPriorityGroup::MultiProviderGroup(index, pg) => { - let mut idx_guard = index.lock().await; - let mut idx = *idx_guard; 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.get(idx).unwrap(); - idx = (idx + 1) % provider_count; let result = p.get_next(grace, grace_period_timeout_secs).await; if result.is_some() { - *idx_guard = idx; + index.store((idx + 1) % provider_count, Ordering::Relaxed); return result; } + idx = (idx + 1) % provider_count; } - *idx_guard = idx; + index.store(idx, Ordering::Relaxed); } } None @@ -332,15 +388,11 @@ impl MultiProviderLineup { without_grace_allocation } }; - match allocation { - ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(_) | - ProviderAllocation::GracePeriod(_) => { - if priority_group.is_exhausted().await { - self.index.store((index + 1) % provider_count, Ordering::SeqCst); - } - return allocation; + if !matches!(allocation, ProviderAllocation::Exhausted) { + if priority_group.is_exhausted().await { + self.index.store((index + 1) % provider_count, Ordering::SeqCst); } + return allocation; } } @@ -377,26 +429,26 @@ impl MultiProviderLineup { } - 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; - } - } - } - } - } - } + // 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; + // } + // } + // } + // } + // } + // } } @@ -527,7 +579,6 @@ impl ProviderLineupManager { 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)> { @@ -561,7 +612,7 @@ impl ProviderLineupManager { None } - pub async fn force_exact_acquire_connection(&self, provider_name: &str) -> ProviderConnectionGuard { + async fn force_exact_acquire_connection(&self, provider_name: &str) -> ProviderConnectionGuard { 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 @@ -572,7 +623,7 @@ impl ProviderLineupManager { } // Returns the next available provider connection - pub async fn acquire_connection(&self, input_name: &str) -> ProviderConnectionGuard { + async fn acquire_connection(&self, input_name: &str) -> ProviderConnectionGuard { 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 @@ -583,8 +634,8 @@ impl ProviderLineupManager { if log_enabled!(log::Level::Debug) { match allocation { ProviderAllocation::Exhausted => {} - ProviderAllocation::Available(ref cfg) | - ProviderAllocation::GracePeriod(ref cfg) => { + ProviderAllocation::Available(_, ref cfg) | + ProviderAllocation::GracePeriod(_, ref cfg) => { debug!("Using provider {}", cfg.name); } } @@ -612,12 +663,12 @@ impl ProviderLineupManager { } // 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 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(); @@ -668,6 +719,7 @@ impl ProviderLineupManager { pub struct ActiveProviderManager { providers: ProviderLineupManager, + connections: DashMap, } impl ActiveProviderManager { @@ -676,6 +728,7 @@ impl ActiveProviderManager { let inputs = Self::get_config_inputs(cfg); Self { providers: ProviderLineupManager::new(inputs, grace_period_millis, grace_period_timeout_secs), + connections: DashMap::new(), } } @@ -697,13 +750,17 @@ 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) -> ProviderConnectionGuard { - self.providers.force_exact_acquire_connection(provider_name).await + pub async fn force_exact_acquire_connection(&self, provider_name: &str, addr: &str) -> ProviderConnectionGuard { + let guard = self.providers.force_exact_acquire_connection(provider_name).await; + self.register_connection(addr, &guard); + guard } // Returns the next available provider connection - pub async fn acquire_connection(&self, input_name: &str) -> ProviderConnectionGuard { - self.providers.acquire_connection(input_name).await + pub async fn acquire_connection(&self, input_name: &str, addr: &str) -> ProviderConnectionGuard { + let guard = self.providers.acquire_connection(input_name).await; + self.register_connection(addr, &guard); + guard } // This method is used for redirects to cycle through provider @@ -712,9 +769,9 @@ impl ActiveProviderManager { } // 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 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 @@ -723,6 +780,19 @@ impl ActiveProviderManager { pub async fn is_over_limit(&self, provider_name: &str) -> bool { self.providers.is_over_limit(provider_name).await } + + fn register_connection(&self, addr: &str, guard: &ProviderConnectionGuard) { + if !matches!(guard.allocation, ProviderAllocation::Exhausted) { + trace!("Added provider connection {:?}", guard.get_provider_name().unwrap_or_default()); + self.connections.insert(addr.to_string(), guard.clone()); + } + } + + pub fn release_connection(&self, addr: &str) { + if let Some((_, guard)) = self.connections.remove(addr) { + guard.release(); + } + } } #[cfg(test)] @@ -924,64 +994,64 @@ mod tests { } - // Test releasing a connection - #[test] - fn test_release_connection() { - let cfg = create_config_input(1, "provider7_1", 1, 2); - let lineup = SingleProviderLineup::new(&cfg, None); - 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 lineup = MultiProviderLineup::new(&cfg1, None); - 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 releasing a connection + // #[test] + // fn test_release_connection() { + // let cfg = create_config_input(1, "provider7_1", 1, 2); + // let lineup = SingleProviderLineup::new(&cfg, None); + // 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 lineup = MultiProviderLineup::new(&cfg1, None); + // 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] @@ -1002,8 +1072,8 @@ mod tests { 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), + ProviderAllocation::Available(_, _) => available.fetch_sub(1, Ordering::SeqCst), + ProviderAllocation::GracePeriod(_, _) => grace_period.fetch_sub(1, Ordering::SeqCst), } }); } diff --git a/backend/src/api/model/active_user_manager.rs b/backend/src/api/model/active_user_manager.rs index fa6999774..32c82b6d0 100644 --- a/backend/src/api/model/active_user_manager.rs +++ b/backend/src/api/model/active_user_manager.rs @@ -1,55 +1,55 @@ +use crate::api::model::streams::shared_stream_manager::SharedStreamManager; use crate::model::Config; -use crate::model::{ProxyUserCredentials}; -use shared::utils::{current_time_secs, default_grace_period_millis, default_grace_period_timeout_secs, sanitize_sensitive_info}; +use crate::model::ProxyUserCredentials; +use dashmap::DashMap; use jsonwebtoken::get_current_timestamp; use log::{debug, info}; -use std::collections::HashMap; +use shared::model::UserConnectionPermission; +use shared::utils::{current_time_secs, default_grace_period_millis, default_grace_period_timeout_secs, sanitize_sensitive_info}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; -use tokio::sync::RwLock; -use shared::model::UserConnectionPermission; +use crate::api::model::active_provider_manager::ActiveProviderManager; macro_rules! active_user_manager_shared_impl { () => { #[inline] - async fn get_active_connections(user: &Arc>>) -> usize { - user.read().await.values().map(|c| c.connections as usize).sum() + fn get_active_connections(user: &Arc>) -> usize { + user.iter().map(|c| c.connections as usize).sum() } fn log_active_user(&self) { if self.is_log_user_enabled() { let user = Arc::clone(&self.user); - tokio::spawn(async move { - let user_count = user.read().await.len(); - let user_connection_count = Self::get_active_connections(&user).await; - info!("Active Users: {user_count}, Active User Connections: {user_connection_count}"); - }); + let user_count = user.len(); + let user_connection_count = Self::get_active_connections(&user); + info!("Active Users: {user_count}, Active User Connections: {user_connection_count}"); } } - pub async fn remove_connection(&self, addr: &str) { - let mut addr_lock = self.user_by_addr.write().await; - let username_opt = addr_lock.get(addr).cloned(); + pub fn remove_connection(&self, addr: &str) { + let username_opt = self.user_by_addr.get(addr).as_deref().cloned(); if let Some(username) = username_opt { - addr_lock.remove(addr); - drop(addr_lock); - let mut lock = self.user.write().await; - if let Some(connection_data) = lock.get_mut(&username) { + self.user_by_addr.remove(addr); + let mut remove_user = false; + if let Some(mut entry) = self.user.get_mut(&username) { + let connection_data = &mut *entry; if connection_data.connections > 0 { connection_data.connections -= 1; } + if connection_data.connections == 0 { - lock.remove(&username); + remove_user = true; } else if connection_data.connections < connection_data.max_connections { - // Grace timeout expired, reset grace counters connection_data.granted_grace = false; connection_data.grace_ts = 0; } - } - - drop(lock); + if remove_user { + self.user.remove(&username); + } } + self.shared_stream_manager.release_connection(addr); + self.provider_manager.release_connection(addr); self.log_active_user(); } }; @@ -67,13 +67,15 @@ fn get_grace_options(config: &Config) -> (u64, u64) { struct ConnectionGuardUserManager { log_active_user: bool, - user: Arc>>, - user_by_addr: Arc>>, + user: Arc>, + user_by_addr: Arc>, + shared_stream_manager: Arc, + provider_manager: Arc, } impl ConnectionGuardUserManager { active_user_manager_shared_impl!(); - fn is_log_user_enabled(&self)-> bool { + fn is_log_user_enabled(&self) -> bool { self.log_active_user } } @@ -87,9 +89,7 @@ impl Drop for UserConnectionGuard { fn drop(&mut self) { let manager = self.manager.clone(); let addr = self.addr.clone(); - tokio::spawn(async move { - manager.remove_connection(&addr).await; - }); + manager.remove_connection(&addr); } } @@ -139,28 +139,34 @@ pub struct ActiveUserManager { grace_period_millis: AtomicU64, grace_period_timeout_secs: AtomicU64, log_active_user: AtomicBool, - user: Arc>>, - user_by_addr: Arc>>, + user: Arc>, + user_by_addr: Arc>, gc_ts: Option, close_signal_tx: tokio::sync::broadcast::Sender, + shared_stream_manager: Arc, + provider_manager: Arc, } impl ActiveUserManager { - pub fn new(config: &Config) -> Self { + pub fn new(config: &Config, shared_stream_manager: &Arc, provider_manager: &Arc) -> Self { let log_active_user = config.log.as_ref().is_some_and(|l| l.log_active_user); let (grace_period_millis, grace_period_timeout_secs) = get_grace_options(config); let (close_signal_tx, _) = tokio::sync::broadcast::channel(10); Self { - grace_period_millis : AtomicU64::new(grace_period_millis), + grace_period_millis: AtomicU64::new(grace_period_millis), grace_period_timeout_secs: AtomicU64::new(grace_period_timeout_secs), log_active_user: AtomicBool::new(log_active_user), - user: Arc::new(RwLock::new(HashMap::new())), - user_by_addr: Arc::new(RwLock::new(HashMap::new())), + user: Arc::new(DashMap::new()), + user_by_addr: Arc::new(DashMap::new()), gc_ts: Some(AtomicU64::new(current_time_secs())), - close_signal_tx + close_signal_tx, + shared_stream_manager: Arc::clone(shared_stream_manager), + provider_manager: Arc::clone(provider_manager), } } + active_user_manager_shared_impl!(); + pub fn update_config(&self, config: &Config) { let log_active_user = config.log.as_ref().is_some_and(|l| l.log_active_user); let (grace_period_millis, grace_period_timeout_secs) = get_grace_options(config); @@ -174,11 +180,13 @@ impl ActiveUserManager { 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), } } - pub async fn user_connections(&self, username: &str) -> u32 { - if let Some(connection_data) = self.user.read().await.get(username) { + pub fn user_connections(&self, username: &str) -> u32 { + if let Some(connection_data) = self.user.get(username) { return connection_data.connections; } 0 @@ -220,40 +228,37 @@ impl ActiveUserManager { UserConnectionPermission::Exhausted } - pub async fn connection_permission( + pub fn connection_permission( &self, username: &str, max_connections: u32, ) -> UserConnectionPermission { if max_connections > 0 { - if let Some(connection_data) = self.user.write().await.get_mut(username) { - return self.check_connection_permission(username, connection_data); + if let Some(mut entry) = self.user.get_mut(username) { + return self.check_connection_permission(username, &mut entry); } } UserConnectionPermission::Allowed } - - pub async fn active_users(&self) -> usize { - self.user.read().await.len() + pub fn active_users(&self) -> usize { + self.user.len() } - pub async fn active_connections(&self) -> usize { - Self::get_active_connections(&self.user).await + pub fn active_connections(&self) -> usize { + Self::get_active_connections(&self.user) } - pub async fn add_connection(&self, username: &str, max_connections: u32, addr: &str) -> UserConnectionGuard { - let mut lock = self.user.write().await; - if let Some(connection_data) = lock.get_mut(username) { - connection_data.connections += 1; - connection_data.max_connections = max_connections; - } else { - lock.insert(username.to_string(), UserConnectionData::new(1, max_connections)); - } - drop(lock); - let mut lock = self.user_by_addr.write().await; - lock.insert(addr.to_string(), username.to_string()); - drop(lock); + pub fn add_connection(&self, username: &str, max_connections: u32, addr: &str) -> UserConnectionGuard { + self.user + .entry(username.to_string()) + .and_modify(|connection_data| { + connection_data.connections += 1; + connection_data.max_connections = max_connections; + }) + .or_insert_with(|| UserConnectionData::new(1, max_connections)); + + self.user_by_addr.insert(addr.to_string(), username.to_string()); self.log_active_user(); @@ -264,8 +269,7 @@ impl ActiveUserManager { } } - active_user_manager_shared_impl!(); - fn is_log_user_enabled(&self)-> bool { + fn is_log_user_enabled(&self) -> bool { self.log_active_user.load(Ordering::Relaxed) } @@ -287,54 +291,59 @@ 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, - connection_permission: UserConnectionPermission) -> Option { - self.gc().await; - let mut lock = self.user.write().await; - if let Some(connection_data) = lock.get_mut(&user.username) { - // check existing session - for session in &mut connection_data.sessions { - if session.token.eq(&session_token) { - session.ts = current_time_secs(); - if !session.stream_url.eq(&stream_url) { - session.stream_url = stream_url.to_string(); - } - if !provider.eq(&session.provider) { - session.provider = provider.to_string(); - } - session.permission = connection_permission; - debug!("Using session for user {} with token {session_token} {}", user.username, sanitize_sensitive_info(stream_url)); - return Some(session.token.to_string()); - } - } + pub fn create_user_session(&self, user: &ProxyUserCredentials, session_token: &str, virtual_id: u32, + provider: &str, stream_url: &str, addr: &str, + connection_permission: UserConnectionPermission) -> String { + self.gc(); + let entry = self.user.entry(user.username.to_string()); + let mut connection_data = entry.or_insert_with(|| { + debug!("Creating session for user {} with token {session_token} {}", user.username, sanitize_sensitive_info(stream_url)); + let mut data = UserConnectionData::new(0, user.max_connections); + let session = Self::new_user_session(session_token, virtual_id, provider, stream_url, addr, connection_permission); + data.add_session(session); + data + }); - // no session creates new one - debug!("Creating session for user {} with token {session_token} {}", user.username, sanitize_sensitive_info(stream_url)); - let session = Self::new_user_session(session_token, virtual_id, provider, stream_url, addr, connection_permission); - let token = session.token.clone(); - connection_data.add_session(session); - Some(token) - } else { - debug!("Creating session for user {} with token {session_token} {}", user.username, sanitize_sensitive_info(stream_url)); - let mut connection_data = UserConnectionData::new(0, user.max_connections); - let session = Self::new_user_session(session_token, virtual_id, provider, stream_url, addr, connection_permission); - let token = session.token.clone(); - connection_data.add_session(session); - lock.insert(user.username.to_string(), connection_data); - Some(token) + // If a session exists, update it + for session in &mut connection_data.sessions { + if session.token == session_token { + session.ts = current_time_secs(); + if session.stream_url != stream_url { + session.stream_url = stream_url.to_string(); + } + if session.provider != provider { + session.provider = provider.to_string(); + } + session.permission = connection_permission; + debug!("Using session for user {} with token {session_token} {}", user.username, sanitize_sensitive_info(stream_url)); + return session.token.clone(); + } } + + // If not session exists, create one + debug!("Creating session for user {} with token {session_token} {}", user.username, sanitize_sensitive_info(stream_url)); + let session = Self::new_user_session(session_token, virtual_id, provider, stream_url, addr, connection_permission); + let token = session.token.clone(); + connection_data.add_session(session); + token } - pub async fn update_session_addr(&self, username: &str, token: &str, addr: &str) { - let mut lock = self.user.write().await; - if let Some(connection_data) = lock.get_mut(username) { - for session in &mut connection_data.sessions { - if session.token.eq(token) { - self.drop_connection(&session.addr); + pub fn update_session_addr(&self, username: &str, token: &str, addr: &str) { + let drop_addr = self.user.get_mut(username).and_then(|mut entry| { + let connection_data = &mut *entry; + connection_data.sessions.iter_mut().find_map(|session| { + if session.token == token { + let old_addr = session.addr.clone(); session.addr = addr.to_string(); + Some(old_addr) + } else { + None } - } + }) + }); + + if let Some(session_addr) = drop_addr { + self.drop_connection(&session_addr); } } @@ -346,13 +355,14 @@ impl ActiveUserManager { self.close_signal_tx.subscribe() } - pub async fn get_user_session(&self, username: &str, token: &str) -> Option { - self.update_user_session(username, token).await + pub fn get_user_session(&self, username: &str, token: &str) -> Option { + self.update_user_session(username, token) } - async fn update_user_session(&self, username: &str, token: &str) -> Option { - let mut lock = self.user.write().await; - if let Some(connection_data) = lock.get_mut(username) { + fn update_user_session(&self, username: &str, token: &str) -> Option { + if let Some(mut entry) = self.user.get_mut(username) { + let connection_data = &mut *entry; + if connection_data.max_connections == 0 { return Self::find_user_session(token, &connection_data.sessions).cloned(); } @@ -378,13 +388,13 @@ impl ActiveUserManager { None } - async fn gc(&self) { + fn gc(&self) { if let Some(gc_ts) = &self.gc_ts { let ts = gc_ts.load(Ordering::Acquire); let now = current_time_secs(); if now - ts > USER_CON_TTL { - let mut lock = self.user.write().await; - for (_, connection_data) in lock.iter_mut() { + for mut entry in self.user.iter_mut() { + let connection_data = &mut *entry; connection_data.sessions.retain(|s| now - s.ts < USER_CON_TTL); } gc_ts.store(now, Ordering::Release); diff --git a/backend/src/api/model/app_state.rs b/backend/src/api/model/app_state.rs index 11b4718bd..77318aaeb 100644 --- a/backend/src/api/model/app_state.rs +++ b/backend/src/api/model/app_state.rs @@ -215,12 +215,12 @@ impl AppState { Ok(changes) } - pub async fn get_active_connections_for_user(&self, username: &str) -> u32 { - self.active_users.user_connections(username).await + pub fn get_active_connections_for_user(&self, username: &str) -> u32 { + self.active_users.user_connections(username) } - pub async fn get_connection_permission(&self, username: &str, max_connections: u32) -> UserConnectionPermission { - self.active_users.connection_permission(username, max_connections).await + pub fn get_connection_permission(&self, username: &str, max_connections: u32) -> UserConnectionPermission { + self.active_users.connection_permission(username, max_connections) } fn detect_changes_for_config(&self, config: &Config) -> UpdateChanges { diff --git a/backend/src/api/model/provider_config.rs b/backend/src/api/model/provider_config.rs index 98455b7fb..de5c9bf07 100644 --- a/backend/src/api/model/provider_config.rs +++ b/backend/src/api/model/provider_config.rs @@ -41,6 +41,20 @@ pub struct ProviderConfig { connection: RwLock, } +impl PartialEq for ProviderConfig { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + && self.name == other.name + && self.url == other.url + && self.username == other.username + && self.password == other.password + && self.input_type == other.input_type + && self.max_connections == other.max_connections + && self.priority == other.priority + // Note: self.connection is skipped + } +} + impl ProviderConfig { pub fn new<'a, F>(cfg: &ConfigInput, get_connection: Option) -> Self where @@ -216,7 +230,6 @@ pub(in crate::api::model) struct ProviderConfigWrapper { inner: Arc, } - impl ProviderConfigWrapper { pub fn new(cfg: ProviderConfig) -> Self { Self { @@ -226,13 +239,13 @@ impl ProviderConfigWrapper { pub async fn force_allocate(&self) -> ProviderAllocation { self.inner.force_allocate().await; - ProviderAllocation::Available(Arc::clone(&self.inner)) + ProviderAllocation::new_available(Arc::clone(&self.inner)) } pub async fn try_allocate(&self, grace: bool, grace_period_timeout_secs: u64) -> ProviderAllocation { match self.inner.try_allocate(grace, grace_period_timeout_secs).await { - ProviderConfigAllocation::Available => ProviderAllocation::Available(Arc::clone(&self.inner)), - ProviderConfigAllocation::GracePeriod => ProviderAllocation::GracePeriod(Arc::clone(&self.inner)), + ProviderConfigAllocation::Available => ProviderAllocation::new_available(Arc::clone(&self.inner)), + ProviderConfigAllocation::GracePeriod => ProviderAllocation::new_grace_period(Arc::clone(&self.inner)), ProviderConfigAllocation::Exhausted => ProviderAllocation::Exhausted, } } diff --git a/backend/src/api/model/streams/active_client_stream.rs b/backend/src/api/model/streams/active_client_stream.rs index 29d818ac5..ca6689c88 100644 --- a/backend/src/api/model/streams/active_client_stream.rs +++ b/backend/src/api/model/streams/active_client_stream.rs @@ -35,7 +35,7 @@ pub(in crate::api) struct ActiveClientStream { } impl ActiveClientStream { - pub(crate) async fn new(mut stream_details: StreamDetails, + pub(crate) fn new(mut stream_details: StreamDetails, app_state: &AppState, user: &ProxyUserCredentials, connection_permission: UserConnectionPermission, @@ -47,7 +47,7 @@ impl ActiveClientStream { } let grant_user_grace_period = connection_permission == UserConnectionPermission::GracePeriod; let username = user.username.as_str(); - let user_connection_guard = Some(active_user.add_connection(username, user.max_connections, addr).await); + let user_connection_guard = Some(active_user.add_connection(username, user.max_connections, addr)); let cfg = &app_state.app_config; let waker = Arc::new(Mutex::new(None)); let waker_clone = Arc::clone(&waker); @@ -119,12 +119,12 @@ impl ActiveClientStream { let mut updated = false; if let Some((username, user_manager, max_connections, reconnect_flag)) = user_grace_check { - let active_connections = user_manager.user_connections(&username).await; + let active_connections = user_manager.user_connections(&username); if active_connections > max_connections { info!("User connections exhausted for active clients: {username}"); stream_strategy_flag_copy.store(USER_EXHAUSTED_STREAM, std::sync::atomic::Ordering::SeqCst); if let Some(flag) = reconnect_flag { - info!("Stopped reconnecting, user connections exhausted"); + info!("Stopped reconnecting, user connections exhausted: {username}"); flag.notify(); } updated = true; @@ -137,7 +137,7 @@ impl ActiveClientStream { info!("Provider connections exhausted for active clients: {provider_name}"); stream_strategy_flag_copy.store(PROVIDER_EXHAUSTED_STREAM, std::sync::atomic::Ordering::SeqCst); if let Some(flag) = reconnect_flag { - info!("Stopped reconnecting, provider connections exhausted"); + info!("Stopped reconnecting, provider connections exhausted {provider_name}"); flag.notify(); } updated = true; diff --git a/backend/src/api/model/streams/shared_stream_manager.rs b/backend/src/api/model/streams/shared_stream_manager.rs index f61a2c1e7..b3169f5e5 100644 --- a/backend/src/api/model/streams/shared_stream_manager.rs +++ b/backend/src/api/model/streams/shared_stream_manager.rs @@ -5,21 +5,19 @@ use crate::utils::debug_if_enabled; use bytes::Bytes; use futures::stream::{BoxStream, FuturesUnordered}; use futures::{Stream, StreamExt}; -use std::collections::HashMap; use std::sync::Arc; use tokio::sync::mpsc::Sender; -use tokio::sync::RwLock; use crate::api::model::stream::BoxedProviderStream; use dashmap::DashMap; use log::trace; -use std::pin::Pin; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::task::{Context, Poll}; -use tokio::sync::mpsc::error::TrySendError; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; use shared::utils::sanitize_sensitive_info; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::sync::mpsc; +use tokio::sync::mpsc::error::TrySendError; +use tokio_stream::wrappers::ReceiverStream; +use crate::api::model::active_provider_manager::ProviderConnectionGuard; /// /// Wraps a `ReceiverStream` as Stream> @@ -47,7 +45,7 @@ fn convert_stream(stream: BoxStream) -> BoxStream, buf_size: usize, + provider_guard: Option, subscribers: Arc>>, //Arc>>>, - next_subscriber_id: AtomicUsize, } impl SharedStreamState { - fn new(headers: Vec<(String, String)>, - buf_size: usize) -> Self { + fn new(headers: Vec<(String, String)>, buf_size: usize, + provider_guard: Option) -> Self { + if let Some(guard) = &provider_guard { + guard.disable_release(); + } Self { headers, buf_size, + provider_guard, subscribers: Arc::new(DashMap::new()), //Arc::new(RwLock::new(Vec::new())), - next_subscriber_id: AtomicUsize::new(1), } } - fn subscribe(&self) -> BoxedProviderStream { + fn subscribe(&self, addr: &str) -> BoxedProviderStream { let (tx, rx) = mpsc::channel(self.buf_size); - let id = self.next_subscriber_id.fetch_add(1, Ordering::AcqRel); - self.subscribers.insert(id, tx); + self.subscribers.insert(addr.to_string(), tx); convert_stream(ReceiverStream::new(rx).boxed()) } @@ -100,7 +100,7 @@ impl SharedStreamState { if subscribers.is_empty() { debug_if_enabled!("No active subscribers. Closing shared provider stream {}", sanitize_sensitive_info(&streaming_url)); - shared_streams.unregister(&streaming_url).await; + shared_streams.unregister(&streaming_url); break; } @@ -143,7 +143,7 @@ impl SharedStreamState { if !has_fillable_subscriber { debug_if_enabled!("All subscribers closed. Shutting down shared provider stream {}",sanitize_sensitive_info(&streaming_url)); - shared_streams.unregister(&streaming_url).await; + shared_streams.unregister(&streaming_url); break; } @@ -159,65 +159,93 @@ impl SharedStreamState { } debug_if_enabled!("Shared stream exhausted. Closing shared provider stream {}", sanitize_sensitive_info(&streaming_url)); - shared_streams.unregister(&streaming_url).await; + shared_streams.unregister(&streaming_url); }); } } -type SharedStreamRegister = RwLock>; +type SharedStreamRegister = DashMap; pub struct SharedStreamManager { shared_streams: SharedStreamRegister, + shared_streams_by_addr: DashMap, } impl SharedStreamManager { pub(crate) fn new() -> Self { Self { - shared_streams: RwLock::new(HashMap::new()), + shared_streams: DashMap::new(), + shared_streams_by_addr: DashMap::new(), } } - pub async fn get_shared_state_headers(&self, stream_url: &str) -> Option> { - self.shared_streams.read().await.get(stream_url).map(|s| s.headers.clone()) + pub fn get_shared_state_headers(&self, stream_url: &str) -> Option> { + self.shared_streams.get(stream_url).map(|s| s.headers.clone()) } - async fn unregister(&self, stream_url: &str) { - let _ = self.shared_streams.write().await.remove(stream_url); + fn unregister(&self, stream_url: &str) { + if let Some((_, shared_state)) = self.shared_streams.remove(stream_url) { + if let Some(guard) = &shared_state.provider_guard { + guard.force_release(); + } + } } - async fn subscribe_stream(&self, stream_url: &str) -> Option { - let stream_data = self.shared_streams.read().await.get(stream_url)?.subscribe(); + pub fn release_connection(&self, addr: &str) { + let stream_url = self.shared_streams_by_addr.get(addr).map(|a| a.value().clone()); + if let Some(stream_url) = stream_url { + let mut drop_state = false; + if let Some(entry) = self.shared_streams.get_mut(&stream_url) { + let shared_state = &*entry; + if let Some((_, tx)) = shared_state.subscribers.remove(addr) { + drop(tx); + } + if shared_state.subscribers.is_empty() { + drop_state = true; + } + } + if drop_state { + self.unregister(&stream_url); + } + } + } + + fn subscribe_stream(&self, stream_url: &str, addr: &str) -> Option { + let stream_data = self.shared_streams.get(stream_url)?.subscribe(addr); Some(stream_data) } - async fn register(&self, stream_url: &str, shared_state: SharedStreamState) { - let _ = self.shared_streams.write().await.insert(stream_url.to_string(), shared_state); + fn register(&self, stream_url: &str, shared_state: SharedStreamState) { + let _ = self.shared_streams.insert(stream_url.to_string(), shared_state); } - pub(crate) async fn subscribe( + pub(crate) fn subscribe( app_state: &AppState, stream_url: &str, + addr: &str, bytes_stream: S, headers: Vec<(String, String)>, - buffer_size: usize, ) -> Option + buffer_size: usize, + provider_guard: Option) -> Option where - S: Stream> + Unpin + 'static + std::marker::Send, - E: std::fmt::Debug + std::marker::Send, + S: Stream> + Unpin + 'static + Send, + E: std::fmt::Debug + Send, { let buf_size = std::cmp::max(buffer_size, STREAM_QUEUE_SIZE); - let shared_state = SharedStreamState::new(headers, buf_size); + let shared_state = SharedStreamState::new(headers, buf_size, provider_guard); shared_state.broadcast(stream_url, bytes_stream, Arc::clone(&app_state.shared_stream_manager)); - app_state.shared_stream_manager.register(stream_url, shared_state).await; + app_state.shared_stream_manager.register(stream_url, shared_state); debug_if_enabled!("Created shared provider stream {}", sanitize_sensitive_info(stream_url)); - Self::subscribe_shared_stream(app_state, stream_url).await + Self::subscribe_shared_stream(app_state, stream_url, addr) } /// Creates a broadcast notify stream for the given URL if a shared stream exists. - pub async fn subscribe_shared_stream( + pub fn subscribe_shared_stream( app_state: &AppState, stream_url: &str, + addr: &str, ) -> Option { debug_if_enabled!("Responding existing shared client stream {}", sanitize_sensitive_info(stream_url)); - app_state.shared_stream_manager.subscribe_stream(stream_url).await + app_state.shared_stream_manager.subscribe_stream(stream_url, addr) } } \ No newline at end of file diff --git a/backend/src/api/serve.rs b/backend/src/api/serve.rs index 8bcafc6b7..61b44dbfa 100644 --- a/backend/src/api/serve.rs +++ b/backend/src/api/serve.rs @@ -135,23 +135,23 @@ where let user_manager_clone = Arc::clone(&user_manager); let mut addr_close_rx = user_manager_clone.get_close_connection_channel(); - let connection_closed = async move || { + let connection_closed = move || { debug!("Connection closed: {remote_addr}"); let addr = remote_addr.to_string(); - user_manager_clone.remove_connection(&addr).await; + user_manager_clone.remove_connection(&addr); }; loop { tokio::select! { result = conn.as_mut() => { if let Err(err) = result { - connection_closed().await; + connection_closed(); trace!("failed to serve connection: {err:#}"); } break; } () = &mut signal_closed => { - connection_closed().await; + connection_closed(); trace!("signal received in task, starting graceful shutdown"); conn.as_mut().graceful_shutdown(); } diff --git a/backend/src/model/config/api_user.rs b/backend/src/model/config/api_user.rs index d6ae68983..65fb79204 100644 --- a/backend/src/model/config/api_user.rs +++ b/backend/src/model/config/api_user.rs @@ -102,11 +102,11 @@ impl ProxyUserCredentials { !self.has_permissions(app_state) } - pub async fn connection_permission(&self, app_state: &AppState) -> UserConnectionPermission { + pub fn connection_permission(&self, app_state: &AppState) -> UserConnectionPermission { let config = > as Access>::load(&app_state.app_config.config); if self.max_connections > 0 && config.user_access_control { // we allow requests with max connection reached, but we should block streaming after grace period - return app_state.get_connection_permission(&self.username, self.max_connections).await; + return app_state.get_connection_permission(&self.username, self.max_connections); } UserConnectionPermission::Allowed } diff --git a/backend/src/processing/parser/xmltv.rs b/backend/src/processing/parser/xmltv.rs index c40443215..0ef275686 100644 --- a/backend/src/processing/parser/xmltv.rs +++ b/backend/src/processing/parser/xmltv.rs @@ -10,11 +10,11 @@ use shared::model::EpgNamePrefix; use shared::utils::CONSTANTS; use std::borrow::Cow; use std::cmp::min; -use std::collections::hash_map::Entry; use std::collections::HashMap; use std::mem; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex, RwLock}; +use std::sync::{Arc, Mutex}; +use dashmap::DashMap; /// Splits a string at the first delimiter if the prefix matches a known country code. /// @@ -423,7 +423,7 @@ pub fn flatten_tvguide(tv_guides: &[Epg]) -> Option { let epg_children = Mutex::new(Vec::new()); let epg_attributes = tv_guides.first().and_then(|t| t.attributes.clone()); let count = tv_guides.iter().map(|tvg| tvg.children.len()).sum(); - let channel_mapping: RwLock> = RwLock::new(HashMap::with_capacity(count)); + let channel_mapping: DashMap = DashMap::with_capacity(count); let mut sorted_guides = tv_guides.to_vec(); // sort by priority @@ -435,24 +435,20 @@ pub fn flatten_tvguide(tv_guides: &[Epg]) -> Option { if c.name.as_str() == EPG_TAG_CHANNEL { if let Some(chan_id) = c.get_attribute_value(EPG_ATTRIB_ID) { let should_add = { - let channel_map = channel_mapping.read().unwrap(); // if not stored - !channel_map.contains_key(chan_id) || + !channel_mapping.contains_key(chan_id) || // or if priority is higher (less means higher priority) - channel_map.get(chan_id).is_none_or(|&priority| guide.priority < priority) + channel_mapping.get(chan_id).as_deref().is_none_or(|&priority| guide.priority < priority) }; if should_add { - let mut channel_map = channel_mapping.write().unwrap(); - match channel_map.entry(chan_id.to_string()) { - Entry::Vacant(e) => { - e.insert(guide.priority); + if let Some(mut existing) = channel_mapping.get_mut(chan_id) { + if guide.priority < *existing { + *existing = guide.priority; children.push(c.clone()); } - Entry::Occupied(mut e) if guide.priority < *e.get() => { - e.insert(guide.priority); - children.push(c.clone()); - } - Entry::Occupied(_) => {} + } else { + channel_mapping.insert(chan_id.clone(), guide.priority); + children.push(c.clone()); } } } @@ -461,9 +457,8 @@ pub fn flatten_tvguide(tv_guides: &[Epg]) -> Option { guide.children.iter().for_each(|c| { if c.name.as_str() == EPG_TAG_PROGRAMME { if let Some(chan_id) = c.get_attribute_value(EPG_ATTRIB_CHANNEL) { - let channel_map = channel_mapping.read().unwrap(); - if let Some(&stored_priority) = channel_map.get(chan_id) { - if stored_priority == guide.priority { + if let Some(stored_priority) = channel_mapping.get(chan_id) { + if *stored_priority == guide.priority { children.push(c.clone()); } } @@ -483,7 +478,6 @@ pub fn flatten_tvguide(tv_guides: &[Epg]) -> Option { } } - #[cfg(test)] mod tests { use crate::model::EpgSmartMatchConfig; diff --git a/webui/src/app/components/userlist/max_connections.rs b/webui/src/app/components/userlist/max_connections.rs index e3d1cb4e3..9793dc5a8 100644 --- a/webui/src/app/components/userlist/max_connections.rs +++ b/webui/src/app/components/userlist/max_connections.rs @@ -9,7 +9,7 @@ pub struct MaxConnectionsProps { #[function_component] pub fn MaxConnections(props: &MaxConnectionsProps) -> Html { - if props.value <= 0 { + if props.value == 0 { html! { } } else { html! { { props.value } } diff --git a/webui/src/app/components/userlist/user_status.rs b/webui/src/app/components/userlist/user_status.rs index f59a497ea..b694d943d 100644 --- a/webui/src/app/components/userlist/user_status.rs +++ b/webui/src/app/components/userlist/user_status.rs @@ -1,6 +1,5 @@ use shared::model::{ProxyUserStatus}; use yew::prelude::*; -use yew_i18n::use_translation; use crate::app::components::Chip; fn convert_status_to_chip_style(status: &ProxyUserStatus) -> String { diff --git a/webui/src/app/context.rs b/webui/src/app/context.rs index 596ec4a25..a474cf73f 100644 --- a/webui/src/app/context.rs +++ b/webui/src/app/context.rs @@ -3,9 +3,11 @@ use yew::UseStateHandle; use shared::model::{AppConfigDto, ConfigTargetDto, ProxyUserCredentialsDto, StatusCheck}; use crate::app::components::{InputRow, PlaylistPage, UserlistPage}; +type SingleSource = (Vec>, Vec>); + #[derive(Clone, PartialEq)] pub struct PlaylistContext { - pub sources: Rc>, Vec>)>>>>, + pub sources: Rc>>>, pub active_page: UseStateHandle, }