From dab98b8d392553b876f4338eeff7a67f6018a1ac Mon Sep 17 00:00:00 2001 From: euzu <33094714+euzu@users.noreply.github.com> Date: Thu, 19 Mar 2026 16:51:18 +0100 Subject: [PATCH] fixed connection exhausted because of user priority handling (#655) * fixed connection exhausted because of user priority handling --- .gitignore | 17 +- backend/src/api/api_utils.rs | 20 +- .../src/api/model/active_provider_manager.rs | 217 +++-- backend/src/api/model/stream.rs | 10 + .../api/model/streams/active_client_stream.rs | 865 ++++++++++++++---- 5 files changed, 872 insertions(+), 257 deletions(-) diff --git a/.gitignore b/.gitignore index a78ce6b00..d60b130e5 100644 --- a/.gitignore +++ b/.gitignore @@ -27,12 +27,6 @@ docs/**/* docker/binaries .DS_Store -# Ignore documentation files -AGENTS.md -PLANS.md -STEPS.md -CLAUDE.md - *.m3u .idea/ @@ -51,6 +45,15 @@ artifacts # Ignore VSCode settings and files /.vscode -.claude .envrc backend/Cargo.lock + +# Ignore ai documentation files +AGENTS.md +PLANS.md +STEPS.md +CLAUDE.md +CODEX.md +.claude +.dual-graph/ +.dual-graph-context/ diff --git a/backend/src/api/api_utils.rs b/backend/src/api/api_utils.rs index d428afbc0..eb12644bf 100644 --- a/backend/src/api/api_utils.rs +++ b/backend/src/api/api_utils.rs @@ -332,7 +332,7 @@ pub struct ForceStreamRequestContext<'a> { /// This means direct provider piping is enabled only when retry is disabled and buffering is disabled. /// /// Returns a `StreamOptions` instance with the resolved configuration. -fn get_stream_options(app_state: &Arc) -> StreamOptions { +pub(in crate::api) fn get_stream_options(app_state: &Arc) -> StreamOptions { let (stream_retry, buffer_enabled, buffer_size) = app_state .app_config .config @@ -680,8 +680,8 @@ async fn create_stream_response_details( } } - // if we have no stream, we should release the provider unless we intentionally defer - // opening it for provider-grace hold_stream behavior. + // If no upstream stream is ready, release the provider unless provider grace + // intentionally deferred the open until the grace check resolves. let provider_handle = if stream.is_none() && !defer_provider_stream_until_grace_check { let provider_handle = streaming_strategy.provider_handle.take(); app_state.connection_manager.release_provider_handle(provider_handle).await; @@ -896,8 +896,7 @@ pub async fn force_provider_stream_response( } }; - let deferred_grace_hold_stream = - !stream_details.has_stream() && stream_details.provider_grace_active && stream_details.grace_period.hold_stream; + let deferred_grace_hold_stream = stream_details.has_deferred_provider_open(); if stream_details.has_stream() || deferred_grace_hold_stream { let provider_response = @@ -1064,9 +1063,8 @@ pub async fn stream_response( // needs to resolve (provider-grace with hold_stream, or user-grace). The grace task will // determine the correct mode (UserExhausted / ProviderExhausted / Inner) and serve the // appropriate custom video or terminate cleanly. - let deferred_grace_hold_stream = !stream_details.has_stream() - && (stream_details.provider_grace_active - || connection_permission == UserConnectionPermission::GracePeriod); + let deferred_grace_hold_stream = + stream_details.has_deferred_provider_open() || connection_permission == UserConnectionPermission::GracePeriod; if stream_details.has_stream() || deferred_grace_hold_stream { // let content_length = get_stream_content_length(provider_response.as_ref()); @@ -1096,7 +1094,11 @@ pub async fn stream_response( is_stream_shared = false; } } - let provider_handle = if is_stream_shared { stream_details.provider_handle.take() } else { None }; + let provider_handle = if is_stream_shared && !stream_details.has_deferred_provider_open() { + stream_details.provider_handle.take() + } else { + None + }; stream_channel.shared = is_stream_shared; let stream = create_active_client_stream( diff --git a/backend/src/api/model/active_provider_manager.rs b/backend/src/api/model/active_provider_manager.rs index 3ae198b70..de3885198 100644 --- a/backend/src/api/model/active_provider_manager.rs +++ b/backend/src/api/model/active_provider_manager.rs @@ -27,6 +27,7 @@ static DUMMY_ADDR: LazyLock = LazyLock::new(|| SocketAddr::from(([12 pub type ClientConnectionId = SocketAddr; type AllocationId = u64; type SharedConnectionId = AllocationId; +type PreemptionCandidate = (PriorityOwner, AllocationId, i8, Instant); // Key for BTreeMap priority index: (priority, Reverse, AllocationId) // Semantics: lower numeric priority value = higher importance (0 = highest, 127 = lowest). // `.last()` on the BTreeMap returns the entry with the highest priority value, which is @@ -35,6 +36,21 @@ type SharedConnectionId = AllocationId; // (smallest `created_at`) sorts last and is evicted first. type PriorityKey = (i8, Reverse, AllocationId); +fn is_better_preemption_candidate( + current: Option, + candidate: PreemptionCandidate, +) -> bool { + match current { + None => true, + Some((_, current_allocation_id, current_priority, current_created_at)) => { + candidate.2 > current_priority + || (candidate.2 == current_priority + && (candidate.3 < current_created_at + || (candidate.3 == current_created_at && candidate.1 < current_allocation_id))) + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum PriorityOwner { Single(ClientConnectionId), @@ -292,25 +308,38 @@ impl ActiveProviderManager { } } - if matches!(&allocation, ProviderAllocation::GracePeriod(_)) - && self.has_non_preemptable_connection_on_input(input_name, priority).await - { - allocation.release().await; - return None; - } - self.register_allocation(allocation, addr, priority).await } - async fn has_non_preemptable_connection_on_input(&self, input_name: &Arc, new_priority: i8) -> bool { - let connections = self.connections.read().await; - connections.priority_index.iter().any(|(prov_name, tree)| { - self.providers.is_provider_for_input(prov_name, input_name) - && tree - .iter() - .next() - .is_some_and(|((prio, _, _), _)| *prio <= new_priority) - }) + fn select_preemption_candidate( + &self, + connections: &Connections, + input_name: &Arc, + new_priority: i8, + reserved_providers: &HashSet>, + ) -> Option { + let mut victim = None; + + for (prov_name, tree) in &connections.priority_index { + if !self.providers.is_provider_for_input(prov_name, input_name) || reserved_providers.contains(prov_name) { + continue; + } + + let Some(((victim_priority, Reverse(created_at), allocation_id), owner)) = tree.iter().next_back() else { + continue; + }; + + if *victim_priority <= new_priority { + continue; + } + + let candidate = (*owner, *allocation_id, *victim_priority, *created_at); + if is_better_preemption_candidate(victim, candidate) { + victim = Some(candidate); + } + } + + victim } async fn acquire_connection_inner( @@ -424,33 +453,9 @@ impl ActiveProviderManager { ) -> bool { let reserved_providers = self.reserved_provider_names_for_other(input_name, session_owner).await; - // Find best victim across all aliases under read lock - let victim: Option<(PriorityOwner, AllocationId, i8, Instant)> = { + let victim = { let connections = self.connections.read().await; - let mut found: Option<(PriorityOwner, AllocationId, i8, Instant)> = None; - for (prov_name, tree) in &connections.priority_index { - if !self.providers.is_provider_for_input(prov_name, input_name) { - continue; - } - if reserved_providers.contains(prov_name) { - continue; - } - if let Some(((prio, Reverse(created_at), alloc_id), owner)) = tree.iter().next_back() { - if *prio <= new_priority { - continue; - } - let is_better = match found { - None => true, - Some((_, _, v_prio, v_created)) => { - *prio > v_prio || (*prio == v_prio && *created_at < v_created) - } - }; - if is_better { - found = Some((*owner, *alloc_id, *prio, *created_at)); - } - } - } - found + self.select_preemption_candidate(&connections, input_name, new_priority, &reserved_providers) }; let Some((owner, alloc_id, v_prio, victim_created_at)) = victim else { @@ -569,37 +574,10 @@ impl ActiveProviderManager { session_owner: Option<&str>, ) -> Option { let reserved_providers = self.reserved_provider_names_for_other(input_name, session_owner).await; - // Victim: (owner, alloc_id, priority, created_at) - let mut victim: Option<(PriorityOwner, AllocationId, i8, Instant)> = None; - - { + let victim = { let connections = self.connections.read().await; - - // Search across ALL provider aliases of this input using the priority index - for (prov_name, tree) in &connections.priority_index { - if !self.providers.is_provider_for_input(prov_name, input_name) { - continue; - } - if reserved_providers.contains(prov_name) { - continue; - } - // Highest priority value (lowest importance) is the best victim candidate. - if let Some(((prio, Reverse(created_at), alloc_id), owner)) = tree.iter().next_back() { - if *prio <= new_priority { - continue; - } - let is_better = match victim { - None => true, - Some((_, _, v_prio, v_created)) => { - *prio > v_prio || (*prio == v_prio && *created_at < v_created) - } - }; - if is_better { - victim = Some((*owner, *alloc_id, *prio, *created_at)); - } - } - } - } + self.select_preemption_candidate(&connections, input_name, new_priority, &reserved_providers) + }; if let Some((owner, alloc_id, v_prio, victim_created_at)) = victim { match owner { @@ -1126,7 +1104,7 @@ impl ActiveProviderManager { mod tests { use super::ActiveProviderManager; use crate::{ - api::model::EventManager, + api::model::{EventManager, ProviderAllocation}, model::{AppConfig, Config, ConfigInput, ConfigInputAlias, SourcesConfig}, utils::FileLockManager, }; @@ -1526,7 +1504,7 @@ mod tests { } #[tokio::test] - async fn test_equal_priority_user_does_not_preempt_and_gets_rejected() { + async fn test_equal_priority_user_gets_grace_without_preempting_existing_stream() { let app_cfg = create_test_app_config_single_provider_pool(); let event_manager = Arc::new(EventManager::new()); let manager = ActiveProviderManager::new(&app_cfg, &event_manager); @@ -1546,20 +1524,101 @@ mod tests { // Provider is now exhausted assert!(manager.is_exhausted(&input_name).await); - // User 2 arrives with same priority. - // Equal priority must not preempt and must not receive a grace allocation. + // User 2 arrives with the same priority and should be granted grace instead of being rejected. let alloc2 = manager .acquire_connection(&input_name, &user_2_addr, default_user_priority()) - .await; + .await + .expect("same-priority user should get grace allocation"); + assert!(matches!(alloc2.allocation, ProviderAllocation::GracePeriod(_))); + assert_eq!(alloc2.allocation.get_provider_name().as_deref(), Some(input_name.as_ref())); - assert!(alloc2.is_none(), "same-priority user must be rejected when provider capacity is exhausted"); - // User 1 should NOT be cancelled + // User 1 should NOT be cancelled, and the provider should be temporarily over limit. assert!(!token1.is_cancelled(), "same-prio user should not be evicted"); + assert!(manager.is_over_limit(&input_name).await, "provider should be temporarily over limit during grace"); manager.release_connection(&user_1_addr).await; manager.release_connection(&user_2_addr).await; } + #[tokio::test] + async fn test_higher_priority_user_preempts_first_inserted_low_priority_victim_on_exact_tie() { + let app_cfg = create_test_app_config_with_dual_provider_pool(); + let event_manager = Arc::new(EventManager::new()); + let manager = ActiveProviderManager::new(&app_cfg, &event_manager); + + let input_name = "provider_1".intern(); + let old_low_addr: SocketAddr = "127.0.0.1:48101".parse().unwrap(); + let new_low_addr: SocketAddr = "127.0.0.1:48102".parse().unwrap(); + let high_prio_addr: SocketAddr = "127.0.0.1:48103".parse().unwrap(); + + // Place two equal-priority low-priority victims on different provider aliases. + // The test then normalizes their created_at timestamps to an exact tie so the + // selector must fall back to the final stable tie-break instead of clock order. + let old_alloc = manager + .acquire_exact_connection_with_grace(&"provider_2".intern(), &old_low_addr, false, 20) + .await + .expect("old low-priority allocation should succeed"); + let old_token = old_alloc.cancel_token.clone().expect("old allocation should have cancel token"); + assert_eq!(old_alloc.allocation.get_provider_name().as_deref(), Some("provider_2")); + + let new_alloc = manager + .acquire_exact_connection_with_grace(&"provider_1".intern(), &new_low_addr, false, 20) + .await + .expect("new low-priority allocation should succeed"); + let new_token = new_alloc.cancel_token.clone().expect("new allocation should have cancel token"); + assert_eq!(new_alloc.allocation.get_provider_name().as_deref(), Some("provider_1")); + + { + let mut connections = manager.connections.write().await; + let old_created_at = connections + .single + .get(&old_low_addr) + .and_then(|per_addr| per_addr.get(&old_alloc.allocation_id)) + .map(|info| info.created_at) + .expect("old allocation should still be registered"); + + let (new_created_at, new_priority) = { + let per_addr = connections + .single + .get_mut(&new_low_addr) + .expect("new allocation address should still be registered"); + let info = per_addr + .get_mut(&new_alloc.allocation_id) + .expect("new allocation should still be registered"); + let original_created_at = info.created_at; + info.created_at = old_created_at; + (original_created_at, info.priority) + }; + + let provider_name = "provider_1".intern(); + let tree = connections + .priority_index + .get_mut(&provider_name) + .expect("priority index for provider_1 should exist"); + let owner = tree + .remove(&(new_priority, std::cmp::Reverse(new_created_at), new_alloc.allocation_id)) + .expect("new allocation should still be indexed"); + tree.insert((new_priority, std::cmp::Reverse(old_created_at), new_alloc.allocation_id), owner); + } + + assert!(manager.is_exhausted(&input_name).await); + + // Higher-priority request should now select the first-inserted victim because + // priority and created_at are exactly tied across provider aliases. + let high_alloc = manager + .acquire_connection(&input_name, &high_prio_addr, 0) + .await + .expect("higher-priority request should preempt the first-inserted low-priority victim on exact tie"); + assert_eq!(high_alloc.allocation.get_provider_name().as_deref(), Some("provider_2")); + + assert!(old_token.is_cancelled(), "first-inserted low-priority victim should be canceled first on exact tie"); + assert!(!new_token.is_cancelled(), "later allocation should remain active after the stable tie-break"); + + manager.release_connection(&high_prio_addr).await; + manager.release_connection(&old_low_addr).await; + manager.release_connection(&new_low_addr).await; + } + #[tokio::test] async fn test_shared_priority_downgrades_after_high_priority_user_leaves() { let app_cfg = create_test_app_config_single_provider_pool(); diff --git a/backend/src/api/model/stream.rs b/backend/src/api/model/stream.rs index 0d8b3e870..0daf3d1af 100644 --- a/backend/src/api/model/stream.rs +++ b/backend/src/api/model/stream.rs @@ -57,6 +57,16 @@ impl StreamDetails { #[inline] pub fn has_grace_period(&self) -> bool { self.grace_period.period_millis > 0 } + + #[inline] + pub fn has_deferred_provider_open(&self) -> bool { + self.stream.is_none() + && self.provider_grace_active + && self.grace_period.hold_stream + && self.provider_handle.is_some() + && self.provider_name.is_some() + && self.request_url.is_some() + } } pub struct StreamingStrategy { diff --git a/backend/src/api/model/streams/active_client_stream.rs b/backend/src/api/model/streams/active_client_stream.rs index 825a1f801..6b9807454 100644 --- a/backend/src/api/model/streams/active_client_stream.rs +++ b/backend/src/api/model/streams/active_client_stream.rs @@ -1,8 +1,10 @@ use crate::{ api::{ + api_utils::get_stream_options, model::{ - AppState, BoxedProviderStream, CleanupEvent, ConnectionManager, CustomVideoStreamType, ProviderHandle, - StreamDetails, StreamError, TimedClientStream, TransportStreamBuffer, + create_provider_stream, AppState, BoxedProviderStream, CleanupEvent, ConnectionManager, + CustomVideoStreamType, ProviderHandle, ProviderStreamFactoryOptions, StreamDetails, StreamError, + TimedClientStream, TransportStreamBuffer, }, panel_api::{can_provision_on_exhausted, find_input_by_provider_name, run_panel_api_provisioning_probe}, }, @@ -20,6 +22,7 @@ use shared::{ }; use std::{ pin::Pin, + net::SocketAddr, sync::{ atomic::{AtomicU8, Ordering}, Arc, @@ -81,11 +84,36 @@ struct GraceProvisioningInfo { stop_signal: CancellationToken, } +#[derive(Clone)] +struct TimedStreamContext { + app_state: Arc, + duration_secs: u32, + virtual_id: VirtualId, +} + +struct DeferredProviderOpenContext { + app_state: Arc, + provider_stream_factory_options: ProviderStreamFactoryOptions, +} + +enum DeferredProviderOpenOutcome { + Stream(BoxedProviderStream), + Mode(StreamMode), + Failed, +} + +enum DeferredProviderOpenState { + Pending(Box), + Opening(Pin + Send>>), +} + #[allow(clippy::struct_excessive_bools)] struct ActiveClientStreamState { inner: Option, send_custom_stream_flag: Option>, provider_handle: Option, + deferred_provider_open: Option, + timed_stream_context: Option, preempt_cancelled: Option>>, grace_task_handle: Option>, provisionable: bool, @@ -104,6 +132,32 @@ struct ActiveClientStreamState { } impl ActiveClientStreamState { + fn mode_for_custom_video_type(video_type: CustomVideoStreamType) -> Option { + match video_type { + CustomVideoStreamType::ChannelUnavailable => Some(StreamMode::ChannelUnavailable), + CustomVideoStreamType::UserConnectionsExhausted => Some(StreamMode::UserExhausted), + CustomVideoStreamType::ProviderConnectionsExhausted => Some(StreamMode::ProviderExhausted), + CustomVideoStreamType::LowPriorityPreempted => Some(StreamMode::LowPriorityPreempted), + CustomVideoStreamType::Provisioning => Some(StreamMode::Provisioning), + CustomVideoStreamType::UserAccountExpired => None, + } + } + + fn wrap_provider_stream(&self, stream: BoxedProviderStream) -> BoxedProviderStream { + if let Some(ctx) = self.timed_stream_context.as_ref() { + TimedClientStream::new( + &ctx.app_state, + stream, + ctx.duration_secs, + self.fingerprint.addr, + ctx.virtual_id, + ) + .boxed() + } else { + stream + } + } + fn custom_video_type_for_mode(mode: StreamMode) -> CustomVideoStreamType { match mode { StreamMode::UserExhausted => CustomVideoStreamType::UserConnectionsExhausted, @@ -295,6 +349,74 @@ impl ActiveClientStreamState { } +fn wrap_timed_client_stream_if_needed( + app_state: &Arc, + stream: BoxedProviderStream, + addr: SocketAddr, + virtual_id: VirtualId, +) -> BoxedProviderStream { + let config = app_state.app_config.config.load(); + match config.sleep_timer_mins { + None => stream, + Some(mins) => { + let secs = u32::try_from((u64::from(mins) * 60).min(u64::from(u32::MAX))).unwrap_or(0); + if secs > 0 { + TimedClientStream::new(app_state, stream, secs, addr, virtual_id).boxed() + } else { + stream + } + } + } +} + +fn create_deferred_provider_open_future( + app_state: &Arc, + stream_details: &StreamDetails, + fingerprint: &Fingerprint, + stream_channel: &StreamChannel, + req_headers: &HeaderMap, +) -> Option { + if !stream_details.has_deferred_provider_open() { + return None; + } + + let provider_name = stream_details.provider_name.as_deref()?; + let request_url = stream_details.request_url.as_deref()?; + let input = find_input_by_provider_name(app_state.as_ref(), provider_name)?; + let stream_url = url::Url::parse(request_url).ok()?; + let stream_options = get_stream_options(app_state); + let default_user_agent = app_state.app_config.config.load().default_user_agent.clone(); + let disabled_headers = app_state.get_disabled_headers(); + let mut provider_stream_factory_options = ProviderStreamFactoryOptions::new( + fingerprint.addr, + stream_channel.item_type, + stream_channel.shared, + &stream_options, + &stream_url, + req_headers, + Some(&input.headers), + disabled_headers.as_ref(), + default_user_agent.as_deref(), + ); + provider_stream_factory_options.set_provider(input.get_resolve_provider(stream_url.as_ref())); + + Some(DeferredProviderOpenState::Pending(Box::new(DeferredProviderOpenContext { + app_state: Arc::clone(app_state), + provider_stream_factory_options, + }))) +} + +fn create_timed_stream_context(app_state: &Arc, virtual_id: VirtualId) -> Option { + let config = app_state.app_config.config.load(); + let mins = config.sleep_timer_mins?; + let duration_secs = u32::try_from((u64::from(mins) * 60).min(u64::from(u32::MAX))).unwrap_or(0); + (duration_secs > 0).then(|| TimedStreamContext { + app_state: Arc::clone(app_state), + duration_secs, + virtual_id, + }) +} + pub(in crate::api) struct ActiveClientStream { state: ActiveClientStreamState, } @@ -302,117 +424,170 @@ pub(in crate::api) struct ActiveClientStream { impl Stream for ActiveClientStream { type Item = Result; + #[allow(clippy::too_many_lines)] fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // 1. Preemption check (user priority feature) - if let Some(fut) = self.state.preempt_cancelled.as_mut() { - if fut.as_mut().poll(cx).is_ready() && !self.state.stop_provider_stream_preempted() { - return Poll::Ready(None); - } - } - - // 2. Grace task lifecycle management + waker registration - self.state.clear_finished_grace_task(); - if let Some(waker) = &self.state.waker { - waker.register(cx.waker()); - } - - // 3. Read atomic mode flag (set by grace task or stop_provider_stream) - let mode = match &self.state.send_custom_stream_flag { - Some(flag) => StreamMode::from_u8(flag.load(Ordering::Acquire)), - None => StreamMode::Inner, - }; - - // 4. Dispatch based on current streaming phase - match mode { - // Grace period: hold_stream=true, waiting for grace task to resolve - StreamMode::GracePending => { - self.state.reset_custom_video_timeout(); - Poll::Pending - } - - // Live streaming: forward bytes from upstream provider - StreamMode::Inner => { - self.state.reset_custom_video_timeout(); - // Poll the inner stream if present; a `None` inner (provider already - // stopped/dropped) is treated the same as a finished stream. - let inner_poll = match self.state.inner.as_mut() { - Some(inner) => Pin::new(inner).poll_next(cx), - None => Poll::Ready(None), - }; - match inner_poll { - Poll::Ready(Some(Ok(bytes))) => Poll::Ready(Some(Ok(bytes))), - Poll::Ready(Some(Err(e))) => { - error!("Inner stream error: {e:?}"); - if self.state.grace_task_handle.is_none() { - self.state.stop_provider_stream(StreamMode::ChannelUnavailable); - Poll::Ready(None) - } else { - // Grace task is still resolving; keep stream alive until it decides - Poll::Pending - } - } - Poll::Ready(None) => { - if self.state.grace_task_handle.is_none() { - self.state.stop_provider_stream(StreamMode::ChannelUnavailable); - Poll::Ready(None) - } else { - // Grace task is still resolving; keep stream alive until it decides - Poll::Pending - } - } - Poll::Pending => Poll::Pending, - } - } - - // Custom video modes: serve the appropriate buffer - video_mode => { - if self.state.custom_video_timeout_mode != Some(video_mode) { - self.state.enter_custom_mode(video_mode); - } - - if self.state.custom_video_timed_out(cx, video_mode) { - info!( - "Custom video {video_mode:?} timed out for {}, terminating stream", - sanitize_sensitive_info(&self.state.fingerprint.addr.to_string()) - ); + loop { + // 1. Preemption check (user priority feature) + if let Some(fut) = self.state.preempt_cancelled.as_mut() { + if fut.as_mut().poll(cx).is_ready() && !self.state.stop_provider_stream_preempted() { return Poll::Ready(None); } + } - let is_provisioning = video_mode == StreamMode::Provisioning && self.state.provisionable; + // 2. Grace task lifecycle management + waker registration + self.state.clear_finished_grace_task(); + if let Some(waker) = &self.state.waker { + waker.register(cx.waker()); + } - let buffer_opt = match video_mode { - StreamMode::UserExhausted => self.state.custom_video.user_exhausted.as_mut(), - StreamMode::ProviderExhausted => self.state.custom_video.provider_exhausted.as_mut(), - StreamMode::ChannelUnavailable => self.state.custom_video.unavailable.as_mut(), - StreamMode::Provisioning => self.state.custom_video.provisioning.as_mut(), - StreamMode::LowPriorityPreempted => self.state.custom_video.low_priority_preempted.as_mut(), - _ => None, - }; + // 3. Read atomic mode flag (set by grace task or stop_provider_stream) + let mode = match &self.state.send_custom_stream_flag { + Some(flag) => StreamMode::from_u8(flag.load(Ordering::Acquire)), + None => StreamMode::Inner, + }; - if let Some(buffer) = buffer_opt { - buffer.register_waker(cx.waker()); - match buffer.next_chunk() { - Some(chunk) => Poll::Ready(Some(Ok(chunk))), - None => { - // Provisioning loops until preemption fires; all others terminate. - if is_provisioning { - Poll::Pending - } else { - info!( - "Custom video {video_mode:?} buffer exhausted for {}, terminating stream", - sanitize_sensitive_info(&self.state.fingerprint.addr.to_string()) - ); - Poll::Ready(None) + // 4. Dispatch based on current streaming phase + match mode { + // Grace period: hold_stream=true, waiting for grace task to resolve + StreamMode::GracePending => { + self.state.reset_custom_video_timeout(); + return Poll::Pending; + } + + // Live streaming: forward bytes from upstream provider + StreamMode::Inner => { + self.state.reset_custom_video_timeout(); + + if self.state.inner.is_none() { + if let Some(deferred_provider_open) = self.state.deferred_provider_open.take() { + match deferred_provider_open { + DeferredProviderOpenState::Pending(context) => { + let app_state = Arc::clone(&context.app_state); + let client = { + let http_client = app_state.http_client.load(); + http_client.as_ref().clone() + }; + let future = Box::pin(async move { + match create_provider_stream( + &app_state, + &client, + context.provider_stream_factory_options, + ) + .await + { + Some((_stream, Some((_headers, _status, _response_url, Some(custom_video_type))))) => { + ActiveClientStreamState::mode_for_custom_video_type(custom_video_type) + .map_or(DeferredProviderOpenOutcome::Failed, DeferredProviderOpenOutcome::Mode) + } + Some((stream, _stream_info)) => DeferredProviderOpenOutcome::Stream(stream), + None => DeferredProviderOpenOutcome::Failed, + } + }); + self.state.deferred_provider_open = Some(DeferredProviderOpenState::Opening(future)); + continue; + } + DeferredProviderOpenState::Opening(mut future) => match future.as_mut().poll(cx) { + Poll::Pending => { + self.state.deferred_provider_open = + Some(DeferredProviderOpenState::Opening(future)); + return Poll::Pending; + } + Poll::Ready(DeferredProviderOpenOutcome::Stream(stream)) => { + self.state.inner = Some(self.state.wrap_provider_stream(stream)); + continue; + } + Poll::Ready(DeferredProviderOpenOutcome::Mode(mode)) => { + self.state.enter_custom_mode(mode); + continue; + } + Poll::Ready(DeferredProviderOpenOutcome::Failed) => { + self.state.enter_custom_mode(StreamMode::ChannelUnavailable); + continue; + } + }, } } + + if self.state.grace_task_handle.is_none() { + self.state.stop_provider_stream(StreamMode::ChannelUnavailable); + return Poll::Ready(None); + } + + return Poll::Pending; } - } else { + + match self.state.inner.as_mut().map(|inner| Pin::new(inner).poll_next(cx)) { + Some(Poll::Ready(Some(Ok(bytes)))) => return Poll::Ready(Some(Ok(bytes))), + Some(Poll::Ready(Some(Err(e)))) => { + error!("Inner stream error: {e:?}"); + if self.state.grace_task_handle.is_none() { + self.state.stop_provider_stream(StreamMode::ChannelUnavailable); + return Poll::Ready(None); + } + + return Poll::Pending; + } + Some(Poll::Ready(None)) | None => { + if self.state.grace_task_handle.is_none() { + self.state.stop_provider_stream(StreamMode::ChannelUnavailable); + return Poll::Ready(None); + } + + return Poll::Pending; + } + Some(Poll::Pending) => return Poll::Pending, + } + } + + // Custom video modes: serve the appropriate buffer + video_mode => { + if self.state.custom_video_timeout_mode != Some(video_mode) { + self.state.enter_custom_mode(video_mode); + } + + if self.state.custom_video_timed_out(cx, video_mode) { + info!( + "Custom video {video_mode:?} timed out for {}, terminating stream", + sanitize_sensitive_info(&self.state.fingerprint.addr.to_string()) + ); + return Poll::Ready(None); + } + + let is_provisioning = video_mode == StreamMode::Provisioning && self.state.provisionable; + + let buffer_opt = match video_mode { + StreamMode::UserExhausted => self.state.custom_video.user_exhausted.as_mut(), + StreamMode::ProviderExhausted => self.state.custom_video.provider_exhausted.as_mut(), + StreamMode::ChannelUnavailable => self.state.custom_video.unavailable.as_mut(), + StreamMode::Provisioning => self.state.custom_video.provisioning.as_mut(), + StreamMode::LowPriorityPreempted => self.state.custom_video.low_priority_preempted.as_mut(), + _ => None, + }; + + if let Some(buffer) = buffer_opt { + buffer.register_waker(cx.waker()); + if let Some(chunk) = buffer.next_chunk() { + return Poll::Ready(Some(Ok(chunk))); + } + + // Provisioning loops until preemption fires; all others terminate. + if is_provisioning { + return Poll::Pending; + } + + info!( + "Custom video {video_mode:?} buffer exhausted for {}, terminating stream", + sanitize_sensitive_info(&self.state.fingerprint.addr.to_string()) + ); + return Poll::Ready(None); + } + // No custom video configured for this mode -> terminate immediately. info!( "No custom video configured for {video_mode:?} mode for {}, terminating stream", sanitize_sensitive_info(&self.state.fingerprint.addr.to_string()) ); - Poll::Ready(None) + return Poll::Ready(None); } } } @@ -458,20 +633,20 @@ pub(crate) async fn create_active_client_stream( } let grant_user_grace_period = connection_permission == UserConnectionPermission::GracePeriod; let username = user.username.as_str(); - let provider_name = stream_details.provider_name.as_ref().map_or_else(String::new, ToString::to_string); + let provider_name = stream_details.provider_name.as_deref().unwrap_or(""); let user_agent = req_headers.get(USER_AGENT).map(|h| String::from_utf8_lossy(h.as_bytes())).unwrap_or_default(); let virtual_id = stream_channel.virtual_id; - let is_shared_source_stream = stream_channel.shared && stream_details.request_url.is_some(); + let is_shared_source_stream = stream_channel.shared && stream_details.stream.is_some(); app_state .connection_manager .update_connection( username, user.max_connections, fingerprint, - &provider_name, - stream_channel, + provider_name, + stream_channel.clone(), user_agent, session_token, ) @@ -483,28 +658,9 @@ pub(crate) async fn create_active_client_stream( // Shared broadcaster source (first subscriber path): feed provider bytes directly. // Grace/custom handling is not needed here because this stream is only the fan-out source. if is_shared_source_stream { - let stream = match stream_details.stream.take() { - None => { - let provider_handle = stream_details.provider_handle.take(); - app_state.connection_manager.release_provider_handle(provider_handle).await; - futures::stream::empty::>().boxed() - } - Some(stream) => { - let config = app_state.app_config.config.load(); - match config.sleep_timer_mins { - None => stream, - Some(mins) => { - let secs = u32::try_from((u64::from(mins) * 60).min(u64::from(u32::MAX))).unwrap_or(0); - if secs > 0 { - TimedClientStream::new(app_state, stream, secs, fingerprint.addr, virtual_id).boxed() - } else { - stream - } - } - } - } - }; - return stream; + if let Some(stream) = stream_details.stream.take() { + return wrap_timed_client_stream_if_needed(app_state, stream, fingerprint.addr, virtual_id); + } } let provisioning_info = resolve_grace_period_provisioning(app_state, &stream_details); @@ -561,25 +717,22 @@ pub(crate) async fn create_active_client_stream( }, ); + let deferred_provider_open = + create_deferred_provider_open_future(app_state, &stream_details, fingerprint, &stream_channel, req_headers); + let timed_stream_context = deferred_provider_open + .as_ref() + .and_then(|_| create_timed_stream_context(app_state, virtual_id)); + let stream: Option = match stream_details.stream.take() { None => { - let provider_handle = stream_details.provider_handle.take(); - app_state.connection_manager.release_provider_handle(provider_handle).await; + if !stream_details.has_deferred_provider_open() { + let provider_handle = stream_details.provider_handle.take(); + app_state.connection_manager.release_provider_handle(provider_handle).await; + } None } Some(stream) => { - let config = app_state.app_config.config.load(); - Some(match config.sleep_timer_mins { - None => stream, - Some(mins) => { - let secs = u32::try_from((u64::from(mins) * 60).min(u64::from(u32::MAX))).unwrap_or(0); - if secs > 0 { - TimedClientStream::new(app_state, stream, secs, fingerprint.addr, virtual_id).boxed() - } else { - stream - } - } - }) + Some(wrap_timed_client_stream_if_needed(app_state, stream, fingerprint.addr, virtual_id)) } }; @@ -598,6 +751,8 @@ pub(crate) async fn create_active_client_stream( let state = ActiveClientStreamState { inner: stream, + deferred_provider_open, + timed_stream_context, preempt_cancelled, grace_task_handle, provider_handle: stream_details.provider_handle, @@ -698,25 +853,28 @@ fn stream_grace_period( let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_millis(grace_period_millis); loop { + let capacity_wait = capacity_notify.notified(); + tokio::pin!(capacity_wait); + + let user_ok = match &user_grace_check { + Some((username, max_connections)) => { + user_manager.user_connections(username).await <= *max_connections + } + None => true, + }; + let provider_ok = match &provider_grace_check { + Some(provider_name) => { + !provider_manager.is_over_limit(provider_name).await + } + None => true, + }; + if user_ok && provider_ok { + break; + } + tokio::select! { () = tokio::time::sleep_until(deadline) => break, - () = capacity_notify.notified() => { - let user_ok = match &user_grace_check { - Some((username, max_connections)) => { - user_manager.user_connections(username).await <= *max_connections - } - None => true, - }; - let provider_ok = match &provider_grace_check { - Some(provider_name) => { - !provider_manager.is_over_limit(provider_name).await - } - None => true, - }; - if user_ok && provider_ok { - break; - } - } + () = &mut capacity_wait => {} } } @@ -813,29 +971,38 @@ fn stream_grace_period( #[cfg(test)] mod tests { - use super::{ActiveClientStream, ActiveClientStreamState, CustomVideoBuffers, StreamMode}; + use super::{ + create_active_client_stream, stream_grace_period, ActiveClientStream, ActiveClientStreamState, + CustomVideoBuffers, DeferredProviderOpenOutcome, DeferredProviderOpenState, StreamMode, TimedStreamContext, + }; use crate::{ api::model::{ - ActiveProviderManager, ActiveUserManager, ConnectionManager, CustomVideoStreamType, EventManager, - SharedStreamManager, + ActiveProviderManager, ActiveUserManager, AppState, CancelTokens, ConnectionManager, CustomVideoStreamType, + DownloadQueue, EventManager, MetadataUpdateManager, PlaylistStorageState, SharedStreamManager, StreamDetails, + StreamError, UpdateGuard, }, auth::Fingerprint, - model::{AppConfig, Config, ConfigInput, SourcesConfig}, + model::{AppConfig, Config, ConfigInput, GracePeriodOptions, ProcessTargets, ProxyUserCredentials, SourcesConfig}, utils::{FileLockManager, GeoIp}, }; use arc_swap::{ArcSwap, ArcSwapOption}; + use axum::http::HeaderMap; use futures::{pin_mut, StreamExt}; + use reqwest::Client; use shared::{ - model::{ConfigPaths, InputFetchMethod, InputType}, + model::{ConfigPaths, InputFetchMethod, InputType, PlaylistItemType, StreamChannel, UserConnectionPermission, XtreamCluster}, utils::Internable, }; use std::{ collections::HashMap, sync::{ - atomic::AtomicU8, + atomic::{AtomicU8, Ordering}, Arc, }, + time::Duration, }; + use tokio::sync::mpsc; + use bytes::Bytes; fn create_test_app_config() -> AppConfig { let input = Arc::new(ConfigInput { @@ -900,6 +1067,136 @@ mod tests { )) } + fn create_test_app_state() -> Arc { + let app_cfg = Arc::new(create_test_app_config()); + let event_manager = Arc::new(EventManager::new()); + let active_provider = Arc::new(ActiveProviderManager::new(&app_cfg, &event_manager)); + let shared_stream_manager = Arc::new(SharedStreamManager::new(Arc::clone(&active_provider))); + active_provider.set_shared_stream_manager(Arc::clone(&shared_stream_manager)); + + let geoip = Arc::new(ArcSwapOption::::default()); + let config = app_cfg.config.load(); + let active_users = Arc::new(ActiveUserManager::new(&config, &geoip, &event_manager)); + let connection_manager = + Arc::new(ConnectionManager::new(&active_users, &active_provider, &shared_stream_manager, &event_manager)); + + let tokens = CancelTokens::default(); + let metadata_manager = Arc::new(MetadataUpdateManager::new(tokens.metadata.clone())); + let (manual_update_sender, _) = mpsc::channel::>(1); + + Arc::new(AppState { + forced_targets: Arc::new(ArcSwap::from_pointee(ProcessTargets { + enabled: false, + inputs: Vec::new(), + targets: Vec::new(), + target_names: Vec::new(), + })), + app_config: app_cfg, + http_client: Arc::new(ArcSwap::from_pointee(Client::new())), + http_client_no_redirect: Arc::new(ArcSwap::from_pointee(Client::new())), + downloads: Arc::new(DownloadQueue::new()), + cache: Arc::new(ArcSwapOption::default()), + shared_stream_manager, + active_users, + active_provider, + connection_manager, + event_manager, + cancel_tokens: Arc::new(ArcSwap::from_pointee(tokens)), + playlists: Arc::new(PlaylistStorageState::new()), + geoip, + update_guard: UpdateGuard::new(), + metadata_manager, + manual_update_sender, + }) + } + + fn create_test_user(username: &str) -> ProxyUserCredentials { + let mut user = ProxyUserCredentials::default(); + user.username = username.to_string(); + user.max_connections = 1; + user + } + + fn create_test_fingerprint(addr: std::net::SocketAddr) -> Fingerprint { + Fingerprint::new(format!("fp-{addr}"), addr.ip().to_string(), addr) + } + + fn create_test_stream_channel(virtual_id: u32, url: &str) -> StreamChannel { + StreamChannel { + target_id: 1, + virtual_id, + provider_id: 1, + item_type: PlaylistItemType::Live, + cluster: XtreamCluster::Live, + group: "Live".intern(), + title: "Test Channel".intern(), + url: url.into(), + shared: false, + technical: None, + } + } + + fn create_test_shared_stream_channel(virtual_id: u32, url: &str) -> StreamChannel { + let mut channel = create_test_stream_channel(virtual_id, url); + channel.shared = true; + channel + } + + fn create_deferred_provider_grace_details( + provider_name: &Arc, + provider_handle: crate::api::model::ProviderHandle, + ) -> StreamDetails { + StreamDetails { + stream: None, + stream_info: None, + provider_name: Some(Arc::clone(provider_name)), + request_url: Some("http://provider-1.example/live/1".intern()), + grace_period: GracePeriodOptions { + period_millis: 100, + timeout_secs: 0, + hold_stream: true, + }, + provider_grace_active: true, + disable_provider_grace: false, + reconnect_flag: None, + provider_handle: Some(provider_handle), + } + } + + async fn start_deferred_provider_grace_resolution( + app_state: &Arc, + provider_name: &Arc, + deferred_addr: std::net::SocketAddr, + ) -> ( + Arc, + tokio::task::JoinHandle<()>, + crate::api::model::ProviderHandle, + ) { + let deferred_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(provider_name, &deferred_addr, true, 0) + .await + .expect("deferred client should receive provider grace allocation"); + let stream_details = create_deferred_provider_grace_details(provider_name, deferred_handle); + let (flag, task) = stream_grace_period( + app_state, + &stream_details, + false, + &create_test_user("grace-user"), + &create_test_fingerprint(deferred_addr), + 1, + None, + None, + true, + app_state.connection_manager.capacity_notified(), + ); + ( + flag.expect("provider grace should install a mode flag"), + task.expect("provider grace should spawn a grace-resolution task"), + stream_details.provider_handle.expect("deferred provider handle must be retained during grace"), + ) + } + async fn assert_missing_custom_video_terminates(mode: StreamMode, provisionable: bool) { let connection_manager = create_test_connection_manager(); let addr = "127.0.0.1:55001".parse().unwrap_or_else(|_| unreachable!()); @@ -908,6 +1205,8 @@ mod tests { inner: None, send_custom_stream_flag: Some(Arc::new(AtomicU8::new(mode as u8))), provider_handle: None, + deferred_provider_open: None, + timed_stream_context: None, preempt_cancelled: None, grace_task_handle: None, provisionable, @@ -988,4 +1287,246 @@ mod tests { async fn test_low_priority_preempted_without_custom_video_terminates_immediately() { assert_missing_custom_video_terminates(StreamMode::LowPriorityPreempted, false).await; } + + #[tokio::test(start_paused = true)] + async fn test_provider_grace_resolution_transitions_from_grace_pending_to_inner_when_capacity_notify_arrives() { + let app_state = create_test_app_state(); + let provider_name = "provider_1".intern(); + let holder_addr = "127.0.0.1:55010".parse().unwrap_or_else(|_| unreachable!()); + let deferred_addr = "127.0.0.1:55011".parse().unwrap_or_else(|_| unreachable!()); + + let holder_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(&provider_name, &holder_addr, false, 0) + .await + .expect("holder should consume the provider's live capacity"); + let (flag, grace_task, deferred_handle) = + start_deferred_provider_grace_resolution(&app_state, &provider_name, deferred_addr).await; + + assert_eq!( + StreamMode::from_u8(flag.load(Ordering::Acquire)), + StreamMode::GracePending, + "provider grace resolution must begin in GracePending while provider capacity is exhausted" + ); + + app_state.connection_manager.release_provider_handle(Some(holder_handle)).await; + let join_result = tokio::time::timeout(Duration::from_millis(1), grace_task).await; + + assert!( + join_result.is_ok(), + "provider grace resolution stayed pending after capacity_notify should have fired" + ); + assert_eq!( + StreamMode::from_u8(flag.load(Ordering::Acquire)), + StreamMode::Inner, + "capacity-notify should resolve provider grace from GracePending to Inner before the deadline" + ); + + app_state.connection_manager.release_provider_handle(Some(deferred_handle)).await; + } + + #[tokio::test(start_paused = true)] + async fn test_active_client_stream_deferred_provider_grace_retains_provider_handle_while_grace_pending() { + let app_state = create_test_app_state(); + let provider_name = "provider_1".intern(); + let holder_addr = "127.0.0.1:55012".parse().unwrap_or_else(|_| unreachable!()); + let deferred_addr = "127.0.0.1:55013".parse().unwrap_or_else(|_| unreachable!()); + let third_addr = "127.0.0.1:55014".parse().unwrap_or_else(|_| unreachable!()); + + let holder_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(&provider_name, &holder_addr, false, 0) + .await + .expect("holder should consume the provider's live capacity"); + let deferred_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(&provider_name, &deferred_addr, true, 0) + .await + .expect("deferred client should receive provider grace allocation"); + let stream_details = create_deferred_provider_grace_details(&provider_name, deferred_handle.clone()); + let stream = create_active_client_stream( + stream_details, + &app_state, + &create_test_user("grace-user"), + UserConnectionPermission::Allowed, + &create_test_fingerprint(deferred_addr), + create_test_stream_channel(1, "http://provider-1.example/live/1"), + None, + &HeaderMap::default(), + ) + .await; + pin_mut!(stream); + + assert!( + matches!(futures::poll!(stream.next()), std::task::Poll::Pending), + "deferred active-client-stream should park in GracePending while waiting for provider grace resolution" + ); + + let third_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(&provider_name, &third_addr, true, 0) + .await; + + assert!( + third_handle.is_none(), + "deferred active-client-stream should retain the deferred provider grace reservation while GracePending" + ); + + app_state.connection_manager.release_provider_handle(Some(holder_handle)).await; + app_state.connection_manager.release_provider_handle(Some(deferred_handle)).await; + } + + #[tokio::test(start_paused = true)] + async fn test_active_client_stream_shared_deferred_provider_grace_retains_provider_handle_while_grace_pending() { + let app_state = create_test_app_state(); + let provider_name = "provider_1".intern(); + let holder_addr = "127.0.0.1:55017".parse().unwrap_or_else(|_| unreachable!()); + let deferred_addr = "127.0.0.1:55018".parse().unwrap_or_else(|_| unreachable!()); + let third_addr = "127.0.0.1:55019".parse().unwrap_or_else(|_| unreachable!()); + + let holder_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(&provider_name, &holder_addr, false, 0) + .await + .expect("holder should consume the provider's live capacity"); + let deferred_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(&provider_name, &deferred_addr, true, 0) + .await + .expect("deferred shared client should receive provider grace allocation"); + let stream_details = create_deferred_provider_grace_details(&provider_name, deferred_handle.clone()); + let stream = create_active_client_stream( + stream_details, + &app_state, + &create_test_user("grace-user"), + UserConnectionPermission::Allowed, + &create_test_fingerprint(deferred_addr), + create_test_shared_stream_channel(1, "http://provider-1.example/live/1"), + None, + &HeaderMap::default(), + ) + .await; + pin_mut!(stream); + + assert!( + matches!(futures::poll!(stream.next()), std::task::Poll::Pending), + "shared deferred active-client-stream should stay pending instead of returning an empty stream" + ); + + let third_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(&provider_name, &third_addr, true, 0) + .await; + + assert!( + third_handle.is_none(), + "shared deferred active-client-stream should retain the deferred provider grace reservation while pending" + ); + + app_state.connection_manager.release_provider_handle(Some(holder_handle)).await; + app_state.connection_manager.release_provider_handle(Some(deferred_handle)).await; + } + + #[tokio::test(start_paused = true)] + async fn test_active_client_stream_deferred_provider_open_applies_sleep_timer_timeout() { + let app_state = create_test_app_state(); + let connection_manager = create_test_connection_manager(); + let addr = "127.0.0.1:55020".parse().unwrap_or_else(|_| unreachable!()); + let state = ActiveClientStreamState { + inner: None, + send_custom_stream_flag: Some(Arc::new(AtomicU8::new(StreamMode::Inner as u8))), + provider_handle: None, + deferred_provider_open: Some(DeferredProviderOpenState::Opening(Box::pin(async { + DeferredProviderOpenOutcome::Stream(futures::stream::pending::>().boxed()) + }))), + timed_stream_context: Some(TimedStreamContext { + app_state, + duration_secs: 1, + virtual_id: 1, + }), + preempt_cancelled: None, + grace_task_handle: None, + provisionable: false, + custom_video: CustomVideoBuffers { + user_exhausted: None, + provider_exhausted: None, + unavailable: None, + provisioning: None, + low_priority_preempted: None, + }, + waker: None, + connection_manager, + fingerprint: Arc::new(Fingerprint::new( + "fp-timeout".to_string(), + "127.0.0.1".to_string(), + addr, + )), + provider_stopped: false, + user_stream_released: true, + provider_handle_released: true, + custom_video_timeout_secs: 0, + custom_video_timeout_mode: None, + custom_video_timeout_sleep: None, + }; + let stream = ActiveClientStream { state }; + pin_mut!(stream); + + assert!( + matches!(futures::poll!(stream.next()), std::task::Poll::Pending), + "deferred-open success should first install the wrapped upstream stream and park pending" + ); + + tokio::time::advance(Duration::from_secs(2)).await; + + let result = tokio::time::timeout(Duration::from_millis(1), stream.next()).await; + assert!( + result.is_ok(), + "deferred-open stream should stop once the configured sleep timer expires" + ); + match result { + Ok(joined) => assert!( + joined.is_none(), + "sleep timer should terminate the deferred-open stream without yielding bytes" + ), + Err(_) => unreachable!("timeout already checked"), + } + } + + #[tokio::test(start_paused = true)] + async fn test_provider_grace_resolution_transitions_from_grace_pending_to_provider_exhausted_at_deadline() { + let app_state = create_test_app_state(); + let provider_name = "provider_1".intern(); + let holder_addr = "127.0.0.1:55015".parse().unwrap_or_else(|_| unreachable!()); + let deferred_addr = "127.0.0.1:55016".parse().unwrap_or_else(|_| unreachable!()); + + let holder_handle = app_state + .active_provider + .acquire_exact_connection_with_grace(&provider_name, &holder_addr, false, 0) + .await + .expect("holder should consume the provider's live capacity"); + let (flag, grace_task, deferred_handle) = + start_deferred_provider_grace_resolution(&app_state, &provider_name, deferred_addr).await; + + assert_eq!( + StreamMode::from_u8(flag.load(Ordering::Acquire)), + StreamMode::GracePending, + "provider grace resolution must begin in GracePending while provider capacity is exhausted" + ); + + tokio::time::advance(Duration::from_millis(101)).await; + let task_result = grace_task.await; + + assert!( + task_result.is_ok(), + "grace-resolution task should complete once the deadline expires without capacity becoming available" + ); + assert_eq!( + StreamMode::from_u8(flag.load(Ordering::Acquire)), + StreamMode::ProviderExhausted, + "provider grace resolution should transition from GracePending to ProviderExhausted when the deadline expires" + ); + + app_state.connection_manager.release_provider_handle(Some(holder_handle)).await; + app_state.connection_manager.release_provider_handle(Some(deferred_handle)).await; + } }