From 2f52c14e48ef41786fe2f2caf056783b77c8c5b3 Mon Sep 17 00:00:00 2001 From: euzu Date: Fri, 6 Dec 2024 17:05:53 +0100 Subject: [PATCH] shared stream --- Cargo.lock | 14 +++++ Cargo.toml | 2 + src/api/api_utils.rs | 99 ++++++++++++++++++++++------------ src/api/m3u_api.rs | 3 +- src/api/model/shared_stream.rs | 64 +++++----------------- src/api/xtream_api.rs | 3 +- src/model/config.rs | 2 + 7 files changed, 99 insertions(+), 88 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bc7f3a456..591dfb2a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1794,6 +1794,8 @@ dependencies = [ "serde_json", "serde_yaml", "time", + "tokio", + "tokio-stream", "unidecode", "url", ] @@ -2931,6 +2933,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f4e6ce100d0eb49a2734f8c0812bcd324cf357d21810932c5df6b96ef2b86f1" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "tokio-util" version = "0.7.12" diff --git a/Cargo.toml b/Cargo.toml index 5b0b0f15f..1c3f1e802 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,3 +51,5 @@ time = "0.3" blake3 = "1.5" bytes = "1.8.0" async-std = "1.13" +tokio-stream = { version = "0.1", features = ["sync"] } +tokio = "1.41" diff --git a/src/api/api_utils.rs b/src/api/api_utils.rs index 5e2ccbc88..3258babb0 100644 --- a/src/api/api_utils.rs +++ b/src/api/api_utils.rs @@ -7,10 +7,15 @@ use crate::utils::request_utils; use crate::utils::request_utils::mask_sensitive_info; use actix_web::http::header::{HeaderValue, CACHE_CONTROL}; use actix_web::{HttpRequest, HttpResponse}; +use async_std::prelude::Stream; +use async_std::stream::StreamExt; +use bytes::Bytes; use log::{debug, error, log_enabled, Level}; use std::collections::HashMap; use std::path::Path; use std::sync::atomic::AtomicU32; +use std::sync::Arc; +use tokio_stream::wrappers::BroadcastStream; use url::Url; pub async fn serve_file(file_path: &Path, req: &HttpRequest, mime_type: mime::Mime) -> HttpResponse { @@ -52,60 +57,86 @@ pub fn get_user_server_info(cfg: &Config, user: &ProxyUserCredentials) -> ApiPro server_info_list.iter().find(|c| c.name.eq(server_info_name)).map_or_else(|| server_info_list.first().unwrap().clone(), std::clone::Clone::clone) } +async fn create_notify_stream(app_state: &AppState, stream_url: &str) -> Option>> { + let notify_stream_url = stream_url.to_string(); + let shared_streams_map = app_state.shared_streams.clone(); + let shared_streams = shared_streams_map.lock().await; + if let Some(shared_stream) = shared_streams.get(¬ify_stream_url) { + let rx = shared_stream.data_stream.subscribe(); + drop(shared_streams); + let (stream, notify) = NotifyStream::new(tokio_stream::wrappers::BroadcastStream::new(rx)); + actix_rt::spawn(async move { + let _ = notify.await; + let mut shared_streams = shared_streams_map.lock().await; + if let Some(shared_stream) = shared_streams.get(¬ify_stream_url) { + let cur_count = shared_stream.client_count.fetch_sub(1u32, std::sync::atomic::Ordering::SeqCst); + if cur_count == 1 { + shared_streams.remove(¬ify_stream_url); + } + } + }); + return Some(stream); + } + drop(shared_streams); + None +} + +async fn create_shared_stream(app_state: &AppState, bytes_stream: S, stream_url: &str) +where + S: Stream> + Unpin + 'static, +{ + let (tx, _) = tokio::sync::broadcast::channel(1); + let sender = Arc::new(tx); + let _ = app_state.shared_streams.lock().await.insert( + stream_url.to_string(), + SharedStream { + data_stream: sender.clone(), + client_count: AtomicU32::new(1), + }, + ); + let mut source_stream = Box::pin(bytes_stream); + actix_rt::spawn(async move { + while let Some(Ok(item)) = source_stream.next().await { + if sender.send(item).is_err() { + // ignore + } + } + }); +} + pub async fn stream_response(app_state: &AppState, stream_url: &str, req: &HttpRequest, input: Option<&ConfigInput>, share_stream: bool) -> HttpResponse { let req_headers: HashMap<&str, &[u8]> = req.headers().iter().map(|(k, v)| (k.as_str(), v.as_bytes())).collect(); if log_enabled!(Level::Debug) { debug!("Try to open stream {}", mask_sensitive_info(stream_url)); } if share_stream { - if let shared_streams = app_state.shared_streams.lock().await { - if let Some(shared_stream) = shared_streams.get(stream_url) { - shared_stream.client_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - } + let shared_streams = app_state.shared_streams.lock().await; + if let Some(shared_stream) = shared_streams.get(stream_url) { + shared_stream.client_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst); } } - if let Ok(url) = Url::parse(stream_url) { let client = request_utils::get_client_request(input, &url, Some(&req_headers)); match client.send().await { Ok(response) => { - if response.status().is_success() { + let status = response.status(); + if status.is_success() { let mut response_builder = HttpResponse::Ok(); response.headers().iter().for_each(|(k, v)| { response_builder.insert_header((k.as_str(), v.as_ref())); }); - let stream = if share_stream { - let byte_stream = Box::pin(response.bytes_stream()); - let cloned_byte_stream = Box::pin(byte_stream.clone()); - let _ = app_state.shared_streams.lock().await.insert(stream_url.to_string(), - SharedStream { - data_stream: byte_stream, - client_count: AtomicU32::new(0), - }, - ); - let (stream, notify) = NotifyStream::new(cloned_byte_stream); - let notify_stream_url = stream_url.to_string(); - actix_rt::spawn(async move { - let _ = notify.await; - debug!("connection closed {notify_stream_url}"); - if let mut shared_streams = app_state.shared_streams.lock().await { - if let Some(shared_stream) = shared_streams.get(stream_url) { - let cur_count = shared_stream.client_count.fetch_sub(1u32, std::sync::atomic::Ordering::SeqCst); - if cur_count == 1 { - shared_streams.remove(stream_url); - } - } - } - }); - stream + if share_stream { + create_shared_stream(app_state, response.bytes_stream(), stream_url).await; + if let Some(stream) = create_notify_stream(app_state, stream_url).await { + return response_builder.body(actix_web::body::BodyStream::new(stream)); + } } else { - response.bytes_stream() - }; - return response_builder.body(actix_web::body::BodyStream::new(stream)); + return response_builder.body(actix_web::body::BodyStream::new(response.bytes_stream())); + } } if log_enabled!(Level::Debug) { - debug!("Failed to open stream got status {} for {}", response.status(), mask_sensitive_info(stream_url)); + debug!("Failed to open stream got status {} for {}", status, mask_sensitive_info(stream_url)); } } Err(err) => { diff --git a/src/api/m3u_api.rs b/src/api/m3u_api.rs index 6f8c29b89..d67ce9ca1 100644 --- a/src/api/m3u_api.rs +++ b/src/api/m3u_api.rs @@ -68,7 +68,8 @@ async fn m3u_api_stream( debug!("Redirecting stream request to {}", mask_sensitive_info(&stream_url)); return HttpResponse::Found().insert_header(("Location", stream_url.to_string())).finish(); } - return stream_response(&app_state, m3u_item.url.as_str(), &req, None).await; + let share_live_streams = target.options.as_ref().map_or(false, |opt| opt.share_live_stream); + return stream_response(&app_state, m3u_item.url.as_str(), &req, None, share_live_streams).await; } Err(err) => { error!("Failed to get m3u url: {}", mask_sensitive_info(err.to_string().as_str())); diff --git a/src/api/model/shared_stream.rs b/src/api/model/shared_stream.rs index e2005dbc5..5bda0c58a 100644 --- a/src/api/model/shared_stream.rs +++ b/src/api/model/shared_stream.rs @@ -1,22 +1,23 @@ -use futures::stream::{Stream}; +use bytes::Bytes; +use futures::channel::oneshot::{channel, Receiver, Sender}; +use futures::stream::Stream; +use futures::StreamExt; +use std::sync::atomic::AtomicU32; +use std::sync::Arc; use std::{ error::Error as StdError, pin::Pin, task::{Context, Poll}, }; -use std::sync::atomic::AtomicU32; -use bytes::Bytes; -use futures::channel::oneshot::{channel, Receiver, Sender}; -use futures::StreamExt; pub struct NotifyStream { stream: S, tx: Option>, } -impl NotifyStream +impl NotifyStream where - S: Stream> + S: Stream>, { pub fn new(stream: S) -> (Self, Receiver<()>) { let (send, recv) = channel(); @@ -27,8 +28,9 @@ where } } -impl Stream for NotifyStream where - S: Stream> + Unpin, +impl Stream for NotifyStream +where + S: Stream> + Unpin, E: Into> + 'static, { type Item = S::Item; @@ -61,49 +63,7 @@ impl Stream for NotifyStream where } pub struct SharedStream { - pub data_stream: Pin>>>, + pub data_stream: Arc>, pub client_count: AtomicU32, } -// -// async fn stream_handler( -// state: web::Data, -// _req: HttpRequest, -// ) -> impl Responder { -// let mut shared_state = state.lock().unwrap(); -// -// // Check if a shared resource already exists -// if let Some(resource) = shared_state.as_mut() { -// // Increment client count -// resource.client_count += 1; -// -// // Create a stream for this client -// let client_stream = resource.data_stream.clone(); -// return HttpResponse::Ok().streaming(client_stream); -// } -// -// // If no shared resource, create one -// let data_stream = crate::api::model::shared_stream::create_data_stream(); // Replace with your actual stream logic -// let boxed_stream = Box::pin(data_stream); -// -// *shared_state = Some(SharedStream { -// data_stream: boxed_stream.clone(), -// client_count: 1, -// }); -// -// // Respond with the stream for this client -// HttpResponse::Ok().streaming(boxed_stream) -// } -// -// -// async fn cleanup_on_disconnect(state: web::Data) { -// let mut shared_state = state.lock().unwrap(); -// if let Some(resource) = shared_state.as_mut() { -// resource.client_count -= 1; -// if resource.client_count == 0 { -// // Drop the shared resource if no clients are left -// *shared_state = None; -// } -// } -// } - diff --git a/src/api/xtream_api.rs b/src/api/xtream_api.rs index 118f35c99..d42a09fad 100644 --- a/src/api/xtream_api.rs +++ b/src/api/xtream_api.rs @@ -235,7 +235,8 @@ async fn xtream_player_api_stream( if log_enabled!(Level::Debug) { debug!("Streaming stream request from {}", mask_sensitive_info(&stream_url)); } - stream_response(app_state, &stream_url, req, Some(input)).await + let share_live_streams = target.options.as_ref().map_or(false, |opt| opt.share_live_stream); + stream_response(app_state, &stream_url, req, Some(input), share_live_streams).await } diff --git a/src/model/config.rs b/src/model/config.rs index 150f65f0c..ff7e266de 100644 --- a/src/model/config.rs +++ b/src/model/config.rs @@ -307,6 +307,8 @@ pub struct ConfigTargetOptions { pub m3u_include_type_in_url: bool, #[serde(default)] pub m3u_mask_redirect_url: bool, + #[serde(default)] + pub share_live_stream: bool, } #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]