From e905a49e2e2be34b3d5347403c438521b84140ed Mon Sep 17 00:00:00 2001 From: euzu Date: Tue, 14 Jan 2025 10:47:37 +0100 Subject: [PATCH] provider direct stream implemented. --- src/api/api_utils.rs | 60 ++++++++++++++------- src/api/model/buffered_stream.rs | 79 +++++++++++++++++++++------- src/api/model/persist_pipe_stream.rs | 1 - src/api/model/shared_stream.rs | 6 ++- src/utils/lru_cache.rs | 3 +- 5 files changed, 109 insertions(+), 40 deletions(-) diff --git a/src/api/api_utils.rs b/src/api/api_utils.rs index 0b33a2875..181c4f46b 100644 --- a/src/api/api_utils.rs +++ b/src/api/api_utils.rs @@ -1,6 +1,6 @@ use crate::api::model::app_state::AppState; use crate::api::model::buffered_stream; -use crate::api::model::buffered_stream::get_stream_response_with_headers; +use crate::api::model::buffered_stream::{get_provider_stream, get_stream_response_with_headers}; use crate::api::model::request::UserApiRequest; use crate::api::model::shared_stream::SharedStream; use crate::debug_if_enabled; @@ -60,13 +60,12 @@ pub fn get_user_target<'a>(api_req: &'a UserApiRequest, app_state: &'a AppState) get_user_target_by_credentials(username, password, api_req, app_state) } -/// Creates a notify stream for the given URL if a shared stream exists. -async fn create_notify_stream( +/// Creates a broadcast notify stream for the given URL if a shared stream exists. +async fn create_broadcast_stream( app_state: &AppState, stream_url: &str, ) -> Option> { let notify_stream_url = stream_url.to_string(); - // Acquire lock and check for existing stream let shared_streams = app_state.shared_streams.lock().await; if let Some(shared_stream) = shared_streams.get(¬ify_stream_url) { @@ -77,7 +76,9 @@ async fn create_notify_stream( } } -pub async fn stream_response(app_state: &AppState, stream_url: &str, req: &HttpRequest, input: Option<&ConfigInput>, item_type: PlaylistItemType, target: &ConfigTarget) -> HttpResponse { +pub async fn stream_response(app_state: &AppState, stream_url: &str, + req: &HttpRequest, input: Option<&ConfigInput>, + item_type: PlaylistItemType, target: &ConfigTarget) -> HttpResponse { if log_enabled!(log::Level::Trace) { trace!("Try to open stream {}", mask_sensitive_info(stream_url)); } let share_stream = is_stream_share_enabled(item_type, target); @@ -87,28 +88,51 @@ pub async fn stream_response(app_state: &AppState, stream_url: &str, req: &HttpR } } + let (stream_retry, buffer_enabled, buffer_size) = app_state + .config + .reverse_proxy + .as_ref() + .and_then(|reverse_proxy| reverse_proxy.stream.as_ref()) + .map_or((false, false, 0), |stream| { + let (buffer_enabled, buffer_size) = stream + .buffer + .as_ref() + .map_or((false, 0), |buffer| (buffer.enabled, buffer.size)); + (stream.retry, buffer_enabled, buffer_size) + }); + + if let Ok(url) = Url::parse(stream_url) { - let (stream, provider_response) = buffered_stream::get_buffered_stream(&app_state.http_client, &url, req, input, item_type).await; - return if share_stream { - SharedStream::register(app_state, stream_url, stream).await; - if let Some(broadcast_stream) = create_notify_stream(app_state, stream_url).await { - let body_stream = BodyStream::new(broadcast_stream); - let mut response_builder = get_stream_response_with_headers(provider_response, stream_url); - response_builder.body(body_stream) - } else { - HttpResponse::BadRequest().finish() - } + let direct_pipe_provider_stream = !stream_retry && !buffer_enabled; + let (stream_opt, provider_response) = if direct_pipe_provider_stream { + get_provider_stream(&app_state.http_client, &url, req, input).await } else { - let mut response_builder = get_stream_response_with_headers(provider_response, stream_url); - response_builder.streaming(stream) + let buffer_stream_options = (item_type, stream_retry, buffer_enabled, buffer_size); + buffered_stream::get_buffered_stream(&app_state.http_client, &url, req, input, buffer_stream_options).await }; + if let Some(stream) = stream_opt { + let use_buffer = !buffer_enabled || direct_pipe_provider_stream; + return if share_stream { + SharedStream::register(app_state, stream_url, stream, use_buffer).await; + if let Some(broadcast_stream) = create_broadcast_stream(app_state, stream_url).await { + let body_stream = BodyStream::new(broadcast_stream); + let mut response_builder = get_stream_response_with_headers(provider_response, stream_url); + response_builder.body(body_stream) + } else { + HttpResponse::BadRequest().finish() + } + } else { + let mut response_builder = get_stream_response_with_headers(provider_response, stream_url); + response_builder.streaming(stream) + }; + } } error!("Url is malformed {}", mask_sensitive_info(stream_url)); HttpResponse::BadRequest().finish() } async fn shared_stream_response(app_state: &AppState, stream_url: &str, headers: Option<(Vec<(String, String)>, reqwest::StatusCode)>) -> Option { - if let Some(stream) = create_notify_stream(app_state, stream_url).await { + if let Some(stream) = create_broadcast_stream(app_state, stream_url).await { debug_if_enabled!("Using shared channel {}", mask_sensitive_info(stream_url)); if app_state.shared_streams.lock().await.get(stream_url).is_some() { let mut response_builder = get_stream_response_with_headers(headers, stream_url); diff --git a/src/api/model/buffered_stream.rs b/src/api/model/buffered_stream.rs index 6933670ef..8f5f1c224 100644 --- a/src/api/model/buffered_stream.rs +++ b/src/api/model/buffered_stream.rs @@ -1,28 +1,28 @@ use crate::api::api_utils::get_headers_from_request; +use crate::debug_if_enabled; use crate::model::config::ConfigInput; +use crate::model::playlist::PlaylistItemType; use crate::utils::request_utils::{get_request_headers, mask_sensitive_info}; use actix_web::http::header::{HeaderName, HeaderValue}; use actix_web::{HttpRequest, HttpResponseBuilder}; use bytes::Bytes; use core::time::Duration; +use log::{debug, error}; use reqwest::header::RANGE; -use reqwest::{Error, StatusCode}; +use reqwest::{Error, Response, StatusCode}; use std::collections::{HashMap, HashSet}; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; use std::task::{Context, Poll}; -use log::debug; use tokio::sync::mpsc; use tokio::sync::mpsc::Receiver; use tokio_stream::{Stream, StreamExt}; use url::Url; -use crate::debug_if_enabled; -use crate::model::playlist::PlaylistItemType; // TODO make this configurable -const STREAM_QUEUE_SIZE: usize = 1024; // mpsc channel holding messages. with 8092byte chunks and 2Mbit/s approx 8MB +pub const STREAM_QUEUE_SIZE: usize = 1024; // mpsc channel holding messages. with 8092byte chunks and 2Mbit/s approx 8MB const MEDIA_STREAM_HEADERS: &[&str] = &["content-type", "content-length", "connection", "accept-ranges", "content-range", "vary"]; fn get_request_range_start_bytes(req_headers: &HashMap>) -> Option { @@ -72,14 +72,55 @@ impl Drop for BufferedReceiverStream { } } +type ProviderStreamResponse = (Option> + Unpin + 'static>>>, Option<(Vec<(String, String)>, StatusCode)>); + +pub async fn get_provider_stream(http_client: &Arc, + stream_url: &Url, + req: &HttpRequest, + input: Option<&ConfigInput>) -> ProviderStreamResponse { + let req_headers = get_headers_from_request(req, &None); + debug_if_enabled!("Stream requested with headers: {:?}", req_headers.iter().map(|header| (header.0, String::from_utf8_lossy(header.1))).collect::>()); + // These are the configured headers for this input. + let input_headers = input.map(|i| i.headers.clone()); + // The stream url, we need to clone it because of move to async block. + // We merge configured input headers with the headers from the request. + let headers = get_request_headers(input_headers.as_ref(), Some(&req_headers)); + let client = http_client.get(stream_url.clone()).headers(headers.clone()); + match client.send().await { + Ok(mut response) => { + let response_headers = get_response_headers(&mut response); + // debug!("First headers {headers:?} {} {}", mask_sensitive_info(url.as_str())); + let status = response.status(); + if status.is_success() { + (Some(Box::pin(response.bytes_stream())), Some((response_headers, status))) + } else { + (None, Some((response_headers, status))) + } + } + Err(err) => { + let masked_url = mask_sensitive_info(stream_url.as_str()); + error!("Failed to open stream {masked_url} {err}"); + (None, None) + } + } +} + +fn get_response_headers(response: &mut Response) -> Vec<(String, String)> { + let response_headers: Vec<(String, String)> = response.headers_mut().iter() + .filter(|(key, _)| MEDIA_STREAM_HEADERS.contains(&key.as_str())) + .map(|(key, value)| (key.to_string(), value.to_str().unwrap().to_string())).collect(); + response_headers +} + #[allow(clippy::too_many_lines)] pub async fn get_buffered_stream(http_client: &Arc, stream_url: &Url, req: &HttpRequest, input: Option<&ConfigInput>, - item_type: PlaylistItemType) -> - (impl Stream> + Unpin + 'static, Option<(Vec<(String, String)>, StatusCode)>) { - let (tx, rx) = mpsc::channel::>(STREAM_QUEUE_SIZE); + options: (PlaylistItemType, bool, bool, usize)) -> ProviderStreamResponse { + let (item_type, retry_enabled, buffer_enabled, buffer_size) = options; + let channel_size = if buffer_enabled { if buffer_size > 0 { buffer_size } else { STREAM_QUEUE_SIZE } } else { 1 }; + let (tx, rx) = mpsc::channel::>(channel_size); let mut req_headers = get_headers_from_request(req, &None); debug_if_enabled!("Stream requested with headers: {:?}", req_headers.iter().map(|header| (header.0, String::from_utf8_lossy(header.1))).collect::>()); // we need the range bytes from client request for seek ing to the right position @@ -129,7 +170,7 @@ pub async fn get_buffered_stream(http_client: &Arc, let status = response.status(); if !status.is_success() { debug!("Failed to connect to provider stream. Status:{status} {masked_url}"); - if status.is_client_error() || status.is_server_error() { + if !retry_enabled || status.is_client_error() || status.is_server_error() { // We stop reconnecting, it seems the stream is not available continue_retry_signal.store(false, Ordering::Relaxed); } @@ -140,15 +181,13 @@ pub async fn get_buffered_stream(http_client: &Arc, // We need some header information from the provider, we extract the neccessary headers and forard them to the client debug_if_enabled!("Provider response headers: {:?}", response.headers_mut()); first_run = false; - let headers: Vec<(String, String)> = response.headers_mut().iter() - .filter(|(key, _)| MEDIA_STREAM_HEADERS.contains(&key.as_str())) - .map(|(key, value)| (key.to_string(), value.to_str().unwrap().to_string())).collect(); + let response_headers: Vec<(String, String)> = get_response_headers(&mut response); // debug!("First headers {headers:?} {} {}", mask_sensitive_info(url.as_str())); let mut status = response.status(); if partial_content && status.is_success() { status = StatusCode::PARTIAL_CONTENT; } - let _ = first_run_response_sender.send(Some((headers, status))).await; + let _ = first_run_response_sender.send(Some((response_headers, status))).await; } let mut byte_stream = response.bytes_stream(); while continue_retry_signal.load(Ordering::Relaxed) { @@ -175,14 +214,14 @@ pub async fn get_buffered_stream(http_client: &Arc, } Some(Err(err)) => { debug!("Provider stream error {masked_url} {err:?}"); - if item_type != PlaylistItemType::Live { + if !retry_enabled || item_type != PlaylistItemType::Live { continue_retry_signal.store(false, Ordering::Relaxed); } break; } None => { debug!("Provider stream finished no data available {masked_url}"); - if item_type != PlaylistItemType::Live { + if !retry_enabled || item_type != PlaylistItemType::Live { continue_retry_signal.store(false, Ordering::Relaxed); } break; @@ -193,13 +232,17 @@ pub async fn get_buffered_stream(http_client: &Arc, } Err(err) => { debug!("Provider stream finished with error {masked_url} {err}"); - if item_type != PlaylistItemType::Live { + if !retry_enabled || item_type != PlaylistItemType::Live { continue_retry_signal.store(false, Ordering::Relaxed); } continue; } } - actix_web::rt::time::sleep(Duration::from_millis(100)).await; + if retry_enabled { + actix_web::rt::time::sleep(Duration::from_millis(100)).await; + } else { + continue_retry_signal.store(false, Ordering::Relaxed); + } } debug!("Streaming stopped and no reconnect for {masked_url}"); drop(tx); @@ -218,7 +261,7 @@ pub async fn get_buffered_stream(http_client: &Arc, let provider_response = provider_response_receiver.recv().await.and_then(|o| o); drop(provider_response_receiver); - (BufferedReceiverStream::new(rx, continue_retry_signal_on_client_disconnect), provider_response) + (Some(Box::pin(BufferedReceiverStream::new(rx, continue_retry_signal_on_client_disconnect))), provider_response) } pub fn get_stream_response_with_headers(custom: Option<(Vec<(String, String)>, StatusCode)>, stream_url: &str) -> HttpResponseBuilder { diff --git a/src/api/model/persist_pipe_stream.rs b/src/api/model/persist_pipe_stream.rs index d9fbb0f89..c0838e02c 100644 --- a/src/api/model/persist_pipe_stream.rs +++ b/src/api/model/persist_pipe_stream.rs @@ -20,7 +20,6 @@ use tokio_stream::Stream; /// - Updates the size tracker. /// - When the stream is exhausted: /// - Calls `on_complete()` to finalize the operation and trigger the callback. - pub struct PersistPipeStream { inner: S, completed: bool, diff --git a/src/api/model/shared_stream.rs b/src/api/model/shared_stream.rs index 8aee72317..ec9fef7e1 100644 --- a/src/api/model/shared_stream.rs +++ b/src/api/model/shared_stream.rs @@ -4,9 +4,12 @@ use std::time::Duration; use async_std::prelude::{Stream, StreamExt}; use tokio::sync::broadcast; use crate::api::model::app_state::AppState; +use crate::api::model::buffered_stream::STREAM_QUEUE_SIZE; use crate::debug_if_enabled; use crate::utils::request_utils::mask_sensitive_info; +const MIN_STREAM_QUEUE_SIZE: usize = 10; + pub struct SharedStream { pub data_stream: Arc>, } @@ -18,11 +21,12 @@ impl SharedStream { app_state: &AppState, stream_url: &str, bytes_stream: S, + use_buffer: bool, ) where S: Stream> + Unpin + 'static, { // Create a broadcast channel for the shared stream - let (tx, _) = broadcast::channel(10); + let (tx, _) = broadcast::channel(if use_buffer { STREAM_QUEUE_SIZE } else { MIN_STREAM_QUEUE_SIZE }); let sender = Arc::new(tx); // Insert the shared stream into the shared state diff --git a/src/utils/lru_cache.rs b/src/utils/lru_cache.rs index a34525ce3..a62d07ea0 100644 --- a/src/utils/lru_cache.rs +++ b/src/utils/lru_cache.rs @@ -21,7 +21,6 @@ use std::path::{Path, PathBuf}; /// - `cache`: A `HashMap` that maps a unique key to a tuple containing the file path and its size. /// - `usage_order`: A `VecDeque` that tracks the access order of keys, with the oldest at the front. /// - `lock`: An `RwLock` to ensure thread-safe access to the cache during read and write operations. - pub struct LRUResourceCache { capacity: usize, // Maximum size in bytes cache_dir: PathBuf, @@ -50,7 +49,7 @@ impl LRUResourceCache { /// - Scans the cache directory and populates the internal data structures with existing files and their sizes. /// - Updates the `current_size` and `usage_order` fields based on the scanned files. - /// The use/access order is not restored!!! + /// The use/access order is not restored!!! pub async fn scan(&mut self) -> std::io::Result<()> { let _write_lock = self.lock.write().await; let mut visit = |entry: &std::fs::DirEntry, metadata: &std::fs::Metadata| {