Refactor downloader architecture and config handling
Refactored the downloader system to pass DownloadConfig explicitly to all downloaders, removing internal state from FfmpegDownloader and StreamGears. Simplified DownloadPlugin::create_downloader to only require the downloader type, and updated DownloadTask to manage cancellation and retries more robustly. Adjusted related code to support the new DownloadConfig flow, improved resource cleanup, and updated trait signatures and usages throughout the codebase.
This commit is contained in:
1 parent
b497749744
commit
6051dc8fe9
9 files changed
+235
-219
No files matched your search
@@ -1,15 +1,19 @@
|
||||
use crate::server::common::upload::UploaderMessage;
|
||||
use crate::server::core::downloader::{DownloadStatus, Downloader, SegmentEvent, SegmentInfo};
|
||||
use crate::server::config::default_segment_time;
|
||||
use crate::server::core::downloader::{
|
||||
DownloadConfig, DownloadStatus, Downloader, SegmentEvent, SegmentInfo,
|
||||
};
|
||||
use crate::server::core::monitor::RoomsHandle;
|
||||
use crate::server::core::plugin::{DownloadPlugin, StreamStatus};
|
||||
use crate::server::errors::{AppError, AppResult};
|
||||
use crate::server::infrastructure::context::{Context, Worker, WorkerStatus};
|
||||
use async_channel::{Receiver, Sender};
|
||||
use error_stack::{ResultExt, bail, Report};
|
||||
use error_stack::{Report, ResultExt, bail};
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
// Configuration and retry policy
|
||||
@@ -50,29 +54,6 @@ impl DownloadGuard {
|
||||
rooms_handle,
|
||||
}
|
||||
}
|
||||
|
||||
async fn cleanup(&mut self) {
|
||||
// 异步清理任务
|
||||
let danmaku = self.danmaku_client.clone();
|
||||
let rooms_handle = self.rooms_handle.clone();
|
||||
let worker = self.worker.clone();
|
||||
|
||||
if let Some(client) = danmaku
|
||||
&& let Err(e) = client.stop().await
|
||||
{
|
||||
error!("Error stopping danmaku client: {}", e);
|
||||
}
|
||||
|
||||
// 确保状态更新和资源清理
|
||||
rooms_handle
|
||||
.toggle(worker.clone(), WorkerStatus::Idle)
|
||||
.await;
|
||||
info!(
|
||||
"{} => {:?}",
|
||||
worker.live_streamer.url,
|
||||
worker.downloader_status.read().await
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DownloadGuard {
|
||||
@@ -254,85 +235,168 @@ impl SegmentEventProcessor {
|
||||
/// 下载任务
|
||||
pub struct DownloadTask {
|
||||
plugin: Arc<dyn DownloadPlugin + Send + Sync>,
|
||||
ctx: Context,
|
||||
rooms_handle: Arc<RoomsHandle>,
|
||||
}
|
||||
|
||||
struct DownloadComponents {
|
||||
downloader: Arc<dyn Downloader + Send + Sync>,
|
||||
danmaku_client: Option<Arc<dyn Downloader>>,
|
||||
uploader: Sender<UploaderMessage>,
|
||||
token: CancellationToken,
|
||||
downloader: Arc<dyn Downloader>,
|
||||
}
|
||||
|
||||
impl DownloadTask {
|
||||
pub fn new(
|
||||
plugin: Arc<dyn DownloadPlugin + Send + Sync>,
|
||||
ctx: Context,
|
||||
rooms_handle: Arc<RoomsHandle>,
|
||||
downloader: Arc<dyn Downloader>,
|
||||
) -> Self {
|
||||
Self {
|
||||
plugin,
|
||||
ctx,
|
||||
rooms_handle,
|
||||
token: CancellationToken::new(),
|
||||
downloader,
|
||||
}
|
||||
}
|
||||
|
||||
pub(self) async fn execute(self, components: DownloadComponents) -> AppResult<DownloadStatus> {
|
||||
// 创建事件处理器
|
||||
let processor = SegmentEventProcessor::new(components.uploader.clone(), self.ctx.clone());
|
||||
pub(self) async fn execute(
|
||||
&self,
|
||||
mut ctx: Context,
|
||||
sender: Sender<UploaderMessage>,
|
||||
rooms_handle: &RoomsHandle,
|
||||
) -> AppResult<DownloadStatus> {
|
||||
let worker = ctx.worker.clone();
|
||||
// 重试配置
|
||||
let mut retry_count = 0;
|
||||
let max_retries = 5; // 最大重试次数
|
||||
let base_delay = Duration::from_secs(0); // 基础延迟时间(2秒)
|
||||
let max_delay = Duration::from_secs(worker.get_config().delay); // 最大延迟时间(60秒)
|
||||
let result = loop {
|
||||
// 创建守卫确保清理
|
||||
// 执行下载
|
||||
// 初始化组件
|
||||
let components = self
|
||||
.initialize_components(sender.clone(), ctx.clone())
|
||||
.await;
|
||||
info!("Download task completed: {:?}", components);
|
||||
ctx = ctx.clone();
|
||||
// 检查流状态
|
||||
match self.plugin.check_status(&mut ctx).await {
|
||||
Ok(StreamStatus::Live { stream_info }) => {
|
||||
ctx.stream_info.raw_stream_url = stream_info.raw_stream_url;
|
||||
info!(
|
||||
"Stream is still live, preparing to retry. Attempt: {}/{}",
|
||||
retry_count + 1,
|
||||
max_retries
|
||||
);
|
||||
|
||||
// 成功下载后重置计数
|
||||
retry_count = 0;
|
||||
}
|
||||
Ok(StreamStatus::Offline) => {
|
||||
retry_count += 1;
|
||||
// 继续循环,重新执行下载
|
||||
info!("Stream went offline, stopping download");
|
||||
}
|
||||
Ok(StreamStatus::Unknown) => {
|
||||
retry_count += 1;
|
||||
// 继续循环,重新执行下载
|
||||
info!("Stream status unknown, stopping download");
|
||||
}
|
||||
Err(e) => {
|
||||
retry_count += 1;
|
||||
// 继续循环,重新执行下载
|
||||
warn!("Failed to check stream status: {:?}, stopping download", e);
|
||||
}
|
||||
}
|
||||
|
||||
if self.token.is_cancelled() {
|
||||
info!("用户手动停止");
|
||||
break components;
|
||||
}
|
||||
|
||||
if retry_count >= max_retries {
|
||||
warn!("Maximum retry attempts ({}) reached, stopping", max_retries);
|
||||
break components;
|
||||
}
|
||||
|
||||
// 计算指数退避延迟: delay = base_delay * 2^retry_count
|
||||
let delay = if retry_count != 0 {
|
||||
base_delay * 2_u32.pow(retry_count)
|
||||
} else {
|
||||
Duration::ZERO
|
||||
};
|
||||
let delay = delay.min(max_delay); // 限制最大延迟时间
|
||||
|
||||
info!("Retrying download in {:?}...", delay);
|
||||
tokio::time::sleep(delay).await;
|
||||
};
|
||||
|
||||
info!("Download task completed: {:?}", result);
|
||||
|
||||
// 清理资源
|
||||
// 确保状态更新和资源清理
|
||||
rooms_handle
|
||||
.toggle(worker.clone(), WorkerStatus::Idle)
|
||||
.await;
|
||||
result
|
||||
}
|
||||
|
||||
async fn initialize_components(
|
||||
&self,
|
||||
uploader: Sender<UploaderMessage>,
|
||||
ctx: Context,
|
||||
) -> AppResult<DownloadStatus> {
|
||||
// 获取配置和主播信息
|
||||
let config = ctx.worker.get_config();
|
||||
let streamer = ctx.worker.get_streamer();
|
||||
let stream_info = &ctx.stream_info;
|
||||
let raw_stream_url = &stream_info.raw_stream_url;
|
||||
// 可选的弹幕客户端
|
||||
// 初始化弹幕客户端(如果存在)
|
||||
let danmaku_client = ctx.extension.get::<Arc<dyn Downloader>>().map(Arc::clone);
|
||||
|
||||
// 创建事件处理器
|
||||
let processor = SegmentEventProcessor::new(uploader.clone(), ctx.clone());
|
||||
let download_config = DownloadConfig {
|
||||
/// 流URL
|
||||
url: raw_stream_url.to_string(),
|
||||
segment_time: config.segment_time.or_else(default_segment_time),
|
||||
file_size: Some(config.file_size), // 2GB
|
||||
headers: stream_info.stream_headers.clone(),
|
||||
recorder: ctx.recorder.clone(),
|
||||
// output_dir: PathBuf::from("./downloads")
|
||||
output_dir: PathBuf::from("."),
|
||||
};
|
||||
// 启动弹幕客户端
|
||||
if let Some(ref client) = components.danmaku_client {
|
||||
self.start_danmaku(client).await?;
|
||||
if let Some(ref client) = danmaku_client {
|
||||
// 启动弹幕下载逻辑
|
||||
info!(
|
||||
"Starting danmaku client for stream: {}",
|
||||
ctx.stream_info.streamer_info.url
|
||||
);
|
||||
client
|
||||
.download(Box::new(|_| {}), Default::default())
|
||||
.await?;
|
||||
}
|
||||
|
||||
// 执行下载
|
||||
let hook = processor.create_hook(components.danmaku_client.clone());
|
||||
let result = components
|
||||
let hook = processor.create_hook(danmaku_client.clone());
|
||||
let result = self
|
||||
.downloader
|
||||
.download(Box::new(hook))
|
||||
.download(Box::new(hook), download_config)
|
||||
.await
|
||||
.change_context(AppError::Custom("Failed to download segment".into()))?;
|
||||
|
||||
// 异步清理任务
|
||||
if let Some(client) = danmaku_client
|
||||
&& let Err(e) = client.stop().await
|
||||
{
|
||||
error!("Error stopping danmaku client: {}", e);
|
||||
}
|
||||
// 处理结果
|
||||
info!(result=?result, "finished downloading");
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn initialize_components(&self, uploader: Sender<UploaderMessage>) -> DownloadComponents {
|
||||
// 获取配置和主播信息
|
||||
let config = self.ctx.worker.get_config();
|
||||
let streamer = self.ctx.worker.get_streamer();
|
||||
let stream_info = &self.ctx.stream_info;
|
||||
|
||||
// 可选的弹幕客户端
|
||||
// 初始化弹幕客户端(如果存在)
|
||||
let danmaku_client = self
|
||||
.ctx
|
||||
.extension
|
||||
.get::<Arc<dyn Downloader>>()
|
||||
.map(Arc::clone);
|
||||
|
||||
// 创建下载器实例
|
||||
let downloader = self
|
||||
.plugin
|
||||
.create_downloader(stream_info, config, self.ctx.recorder.clone())
|
||||
.await;
|
||||
|
||||
DownloadComponents {
|
||||
downloader,
|
||||
danmaku_client,
|
||||
uploader,
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_danmaku(&self, client: &Arc<dyn Downloader>) -> AppResult<()> {
|
||||
// 启动弹幕下载逻辑
|
||||
info!(
|
||||
"Starting danmaku client for stream: {}",
|
||||
self.ctx.stream_info.streamer_info.url
|
||||
);
|
||||
client.download(Box::new(|_| {})).await?;
|
||||
pub(crate) async fn stop(&self) -> AppResult<()> {
|
||||
// 仅发出取消信号并更新状态
|
||||
// 如果底层下载函数不支持取消,这里不能真正中断正在进行的下载
|
||||
self.token.cancel();
|
||||
// self.token.run_until_cancelled();
|
||||
self.downloader.stop().await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -365,41 +429,28 @@ impl DActor {
|
||||
/// * `msg` - 要处理的下载消息
|
||||
async fn handle_message(&mut self, msg: DownloaderMessage) {
|
||||
match msg {
|
||||
DownloaderMessage::Start(plugin, mut ctx, rooms_handle) => {
|
||||
DownloaderMessage::Start(plugin, ctx, rooms_handle) => {
|
||||
let worker = ctx.worker.clone();
|
||||
// 创建下载任务
|
||||
let task = DownloadTask::new(plugin.clone(), ctx.clone(), rooms_handle.clone());
|
||||
|
||||
// 初始化组件
|
||||
let components = task.initialize_components(self.sender.clone()).await;
|
||||
// 创建守卫确保清理
|
||||
let mut guard = DownloadGuard::new(
|
||||
worker.clone(),
|
||||
components.danmaku_client.clone(),
|
||||
rooms_handle.clone(),
|
||||
components.downloader.clone(),
|
||||
);
|
||||
// 创建下载器实例
|
||||
let downloader = plugin
|
||||
.create_downloader(ctx.worker.get_config().downloader)
|
||||
.await;
|
||||
// 创建下载任务
|
||||
let task = Arc::new(DownloadTask::new(plugin.clone(), downloader));
|
||||
|
||||
// 更新工作器状态为工作中
|
||||
rooms_handle
|
||||
.toggle(worker, WorkerStatus::Working(components.downloader.clone()))
|
||||
.toggle(worker.clone(), WorkerStatus::Working(task.clone()))
|
||||
.await;
|
||||
|
||||
// 执行下载
|
||||
let result = task.execute(components).await;
|
||||
let result = task.execute(ctx, self.sender.clone(), &rooms_handle).await;
|
||||
|
||||
match plugin.check_status(&mut ctx).await {
|
||||
Ok(StreamStatus::Live {stream_info}) => {
|
||||
ctx.stream_info.raw_stream_url = stream_info.raw_stream_url;
|
||||
println!("stream_info=?stream_info");
|
||||
}
|
||||
Ok(StreamStatus::Offline) => {}
|
||||
Ok(StreamStatus::Unknown) => {}
|
||||
Err(_) => {}
|
||||
}
|
||||
|
||||
guard.cleanup().await;
|
||||
info!("Handling message: {:?} done", result);
|
||||
info!(
|
||||
"Download workflow completed {} => {:?}",
|
||||
worker.live_streamer.url,
|
||||
worker.downloader_status.read().await
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,8 +12,10 @@ use std::path::PathBuf;
|
||||
|
||||
/// 下载器配置
|
||||
/// 包含下载过程中需要的各种参数和设置
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, Default)]
|
||||
pub struct DownloadConfig {
|
||||
/// 流 URL
|
||||
pub(crate) url: String,
|
||||
/// 分段时长 (格式: "HH:MM:SS")
|
||||
pub segment_time: Option<String>,
|
||||
|
||||
@@ -42,7 +44,7 @@ impl DownloadConfig {
|
||||
|
||||
/// 下载器类型枚举
|
||||
/// 定义支持的各种下载器类型
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[derive(Debug, Clone, Copy, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum DownloaderType {
|
||||
/// Ytarchive下载器
|
||||
@@ -131,6 +133,7 @@ pub trait Downloader: Send + Sync {
|
||||
async fn download(
|
||||
&self,
|
||||
callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
|
||||
download_config: DownloadConfig,
|
||||
) -> AppResult<DownloadStatus>;
|
||||
|
||||
/// 停止下载
|
||||
|
||||
@@ -16,10 +16,6 @@ use tracing::info;
|
||||
/// FFmpeg下载器实现
|
||||
/// 使用FFmpeg进行直播流下载,支持内部和外部分段
|
||||
pub struct FfmpegDownloader {
|
||||
/// 下载配置
|
||||
config: DownloadConfig,
|
||||
/// 流URL
|
||||
url: String,
|
||||
/// 下载状态
|
||||
status: Arc<RwLock<DownloadStatus>>,
|
||||
/// 进程句柄
|
||||
@@ -40,15 +36,8 @@ impl FfmpegDownloader {
|
||||
/// * `config` - 下载配置
|
||||
/// * `extra_args` - 额外的FFmpeg参数
|
||||
/// * `downloader_type` - 下载器类型
|
||||
pub fn new(
|
||||
url: &str,
|
||||
config: DownloadConfig,
|
||||
extra_args: Vec<String>,
|
||||
downloader_type: DownloaderType,
|
||||
) -> Self {
|
||||
pub fn new(extra_args: Vec<String>, downloader_type: DownloaderType) -> Self {
|
||||
Self {
|
||||
config,
|
||||
url: url.to_string(),
|
||||
status: Arc::new(RwLock::new(DownloadStatus::Downloading)),
|
||||
process_handle: Arc::new(RwLock::new(None)),
|
||||
extra_args,
|
||||
@@ -58,21 +47,21 @@ impl FfmpegDownloader {
|
||||
|
||||
/// 构建内部分段模式的FFmpeg命令参数
|
||||
/// 使用FFmpeg的segment muxer进行自动分段
|
||||
fn build_ffmpeg_args_internal_segment(&self) -> Vec<String> {
|
||||
fn build_ffmpeg_args_internal_segment(&self, download_config: &DownloadConfig) -> Vec<String> {
|
||||
let mut args = Vec::new();
|
||||
|
||||
// 内部分段使用info级别日志以获取分段信息
|
||||
args.extend(["-loglevel".to_string(), "info".to_string()]);
|
||||
|
||||
// 添加通用输入参数
|
||||
self.append_common_input_args(&mut args);
|
||||
self.append_common_input_args(&mut args, download_config);
|
||||
|
||||
// 内部分段特定的输出参数
|
||||
// -f segment: 使用segment muxer进行自动分段
|
||||
args.extend(["-f".to_string(), "segment".to_string()]);
|
||||
args.extend([
|
||||
"-segment_format".to_string(),
|
||||
self.config.recorder.suffix.to_string(),
|
||||
download_config.recorder.suffix.to_string(),
|
||||
]);
|
||||
// -segment_list pipe:1: 将分段文件名输出到stdout
|
||||
// 这样我们可以实时获取新生成的分段文件
|
||||
@@ -91,7 +80,7 @@ impl FfmpegDownloader {
|
||||
args.extend(["-strftime".to_string(), "1".to_string()]);
|
||||
|
||||
// -segment_time: 分段时长(秒)
|
||||
if let Some(segment_time) = &self.config.segment_time {
|
||||
if let Some(segment_time) = &download_config.segment_time {
|
||||
let seconds = downloader::parse_duration(segment_time);
|
||||
args.extend(["-segment_time".to_string(), seconds.to_string()]);
|
||||
}
|
||||
@@ -104,43 +93,42 @@ impl FfmpegDownloader {
|
||||
|
||||
/// 构建外部分段模式的FFmpeg命令参数
|
||||
/// 通过外部控制进行分段,每次录制固定时长或大小
|
||||
fn build_ffmpeg_args_external_segment(&self) -> Vec<String> {
|
||||
fn build_ffmpeg_args_external_segment(&self, download_config: &DownloadConfig) -> Vec<String> {
|
||||
let mut args = Vec::new();
|
||||
|
||||
// 外部分段使用quiet减少日志
|
||||
args.extend(["-loglevel".to_string(), "quiet".to_string()]);
|
||||
|
||||
// 添加通用输入参数
|
||||
self.append_common_input_args(&mut args);
|
||||
self.append_common_input_args(&mut args, download_config);
|
||||
|
||||
// 外部分段特定的输出参数
|
||||
// -to: 限制录制时长
|
||||
if let Some(segment_time) = &self.config.segment_time {
|
||||
if let Some(segment_time) = &download_config.segment_time {
|
||||
args.extend(["-to".to_string(), segment_time.clone()]);
|
||||
}
|
||||
|
||||
// -fs: 限制文件大小(字节)
|
||||
if let Some(file_size) = self.config.file_size {
|
||||
if let Some(file_size) = download_config.file_size {
|
||||
args.extend(["-fs".to_string(), file_size.to_string()]);
|
||||
}
|
||||
|
||||
// 添加通用输出参数
|
||||
self.append_common_output_args(&mut args, &self.config.recorder.suffix);
|
||||
self.append_common_output_args(&mut args, &download_config.recorder.suffix);
|
||||
|
||||
args
|
||||
}
|
||||
|
||||
/// 添加通用的输入参数
|
||||
/// 包括覆盖文件、HTTP头、超时设置等
|
||||
fn append_common_input_args(&self, args: &mut Vec<String>) {
|
||||
fn append_common_input_args(&self, args: &mut Vec<String>, download_config: &DownloadConfig) {
|
||||
args.push("-y".to_string()); // 覆盖已存在文件
|
||||
|
||||
// HTTP headers
|
||||
// -headers: 设置HTTP请求头,格式为"Key: Value\r\n"
|
||||
// 用于传递User-Agent、Cookie等信息
|
||||
if !self.config.headers.is_empty() {
|
||||
let headers_str = self
|
||||
.config
|
||||
if !download_config.headers.is_empty() {
|
||||
let headers_str = download_config
|
||||
.headers
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{}: {}\r\n", k, v))
|
||||
@@ -153,14 +141,14 @@ impl FfmpegDownloader {
|
||||
args.extend(["-rw_timeout".to_string(), "20000000".to_string()]);
|
||||
|
||||
// 对于m3u8流的特殊处理
|
||||
if self.url.contains(".m3u8") {
|
||||
if download_config.url.contains(".m3u8") {
|
||||
// -max_reload: HLS播放列表最大重载次数
|
||||
// 对于直播流需要设置较大值以持续获取新片段
|
||||
args.extend(["-max_reload".to_string(), "1000".to_string()]);
|
||||
}
|
||||
|
||||
// 输入URL
|
||||
args.extend(["-i".to_string(), self.url.clone()]);
|
||||
args.extend(["-i".to_string(), download_config.url.clone()]);
|
||||
}
|
||||
|
||||
/// 添加通用的输出参数
|
||||
@@ -204,9 +192,10 @@ impl FfmpegDownloader {
|
||||
async fn download_external(
|
||||
&self,
|
||||
callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
|
||||
download_config: DownloadConfig,
|
||||
) -> AppResult<DownloadStatus> {
|
||||
let args = self.build_ffmpeg_args_external_segment();
|
||||
let output_file = self.config.generate_output_filename();
|
||||
let args = self.build_ffmpeg_args_external_segment(&download_config);
|
||||
let output_file = download_config.generate_output_filename();
|
||||
|
||||
let mut cmd = Command::new("ffmpeg");
|
||||
cmd.args(&args)
|
||||
@@ -242,18 +231,14 @@ impl FfmpegDownloader {
|
||||
// 分段回调
|
||||
// 触发分段回调
|
||||
|
||||
callback(SegmentEvent::Segment(
|
||||
SegmentInfo {
|
||||
prev_file_path: output_file,
|
||||
segment_index: 0,
|
||||
next_file_path: None,
|
||||
}
|
||||
));
|
||||
callback(SegmentEvent::Segment(SegmentInfo {
|
||||
prev_file_path: output_file,
|
||||
segment_index: 0,
|
||||
next_file_path: None,
|
||||
}));
|
||||
// 根据退出码判断状态
|
||||
match status.code() {
|
||||
Some(0) => {
|
||||
Ok(DownloadStatus::SegmentCompleted)
|
||||
}
|
||||
Some(0) => Ok(DownloadStatus::SegmentCompleted),
|
||||
Some(255) => Ok(DownloadStatus::StreamEnded),
|
||||
err => Ok(DownloadStatus::Error(format!("FFmpeg error: {err:?}"))),
|
||||
}
|
||||
@@ -264,15 +249,16 @@ impl FfmpegDownloader {
|
||||
async fn download_internal(
|
||||
&self,
|
||||
callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
|
||||
download_config: DownloadConfig,
|
||||
) -> AppResult<DownloadStatus> {
|
||||
let args = self.build_ffmpeg_args_internal_segment();
|
||||
let args = self.build_ffmpeg_args_internal_segment(&download_config);
|
||||
|
||||
let mut cmd = Command::new("ffmpeg");
|
||||
cmd.args(&args)
|
||||
.arg(format!(
|
||||
"{}.{}.part",
|
||||
self.config.recorder.filename_template(),
|
||||
self.config.recorder.suffix
|
||||
download_config.recorder.filename_template(),
|
||||
download_config.recorder.suffix
|
||||
))
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
@@ -383,14 +369,15 @@ impl Downloader for FfmpegDownloader {
|
||||
async fn download(
|
||||
&self,
|
||||
callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
|
||||
download_config: DownloadConfig,
|
||||
) -> AppResult<DownloadStatus> {
|
||||
match self.downloader_type {
|
||||
DownloaderType::FfmpegExternal => self
|
||||
.download_external(callback)
|
||||
.download_external(callback, download_config)
|
||||
.await
|
||||
.change_context(AppError::Unknown),
|
||||
DownloaderType::FfmpegInternal => self
|
||||
.download_internal(callback)
|
||||
.download_internal(callback, download_config)
|
||||
.await
|
||||
.change_context(AppError::Unknown),
|
||||
_ => bail!(AppError::Custom("Unsupported downloader type".to_string())),
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
use crate::server::core::downloader::{DownloadStatus, Downloader, SegmentEvent, SegmentInfo};
|
||||
use crate::server::common::construct_headers;
|
||||
use crate::server::common::util::parse_time;
|
||||
use crate::server::core::downloader::{
|
||||
DownloadConfig, DownloadStatus, Downloader, SegmentEvent, SegmentInfo,
|
||||
};
|
||||
use crate::server::errors::{AppError, AppResult};
|
||||
use async_trait::async_trait;
|
||||
use axum::http::HeaderMap;
|
||||
@@ -17,14 +21,6 @@ use tracing::{debug, error, info};
|
||||
/// Stream-gears下载器实现
|
||||
/// 使用stream-gears库进行直播流下载
|
||||
pub struct StreamGears {
|
||||
/// 流URL
|
||||
url: String,
|
||||
/// HTTP请求头
|
||||
header_map: HeaderMap,
|
||||
/// 文件名
|
||||
file_name: String,
|
||||
/// 分段配置
|
||||
segment: Segmentable,
|
||||
/// 代理设置(可选)
|
||||
proxy: Option<String>,
|
||||
/// 下载状态
|
||||
@@ -42,18 +38,8 @@ impl StreamGears {
|
||||
/// * `file_name` - 输出文件名
|
||||
/// * `segment` - 分段配置
|
||||
/// * `proxy` - 代理设置(可选)
|
||||
pub fn new(
|
||||
url: &str,
|
||||
header_map: HeaderMap,
|
||||
file_name: String,
|
||||
segment: Segmentable,
|
||||
proxy: Option<String>,
|
||||
) -> Self {
|
||||
pub fn new(proxy: Option<String>) -> Self {
|
||||
Self {
|
||||
url: url.into(),
|
||||
header_map,
|
||||
file_name,
|
||||
segment,
|
||||
proxy,
|
||||
status: Arc::new(RwLock::new(DownloadStatus::Downloading)),
|
||||
token: CancellationToken::new(),
|
||||
@@ -63,15 +49,20 @@ impl StreamGears {
|
||||
async fn start_download(
|
||||
&self,
|
||||
callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
|
||||
download_config: DownloadConfig,
|
||||
) -> AppResult<DownloadStatus> {
|
||||
let url = self.url.clone();
|
||||
let file_name = self.file_name.clone();
|
||||
let _headers_in = self.header_map.clone();
|
||||
let url = download_config.url.clone();
|
||||
let file_name = download_config.recorder.filename_template();
|
||||
let headers_in = construct_headers(&download_config.headers);
|
||||
let proxy = self.proxy.clone();
|
||||
let segment = Segmentable::new(
|
||||
download_config.segment_time.as_deref().map(parse_time),
|
||||
download_config.file_size,
|
||||
);
|
||||
let _status = Arc::clone(&self.status);
|
||||
|
||||
// 创建HTTP客户端
|
||||
let client = StatelessClient::new(self.header_map.clone(), proxy.as_deref());
|
||||
let client = StatelessClient::new(headers_in, proxy.as_deref());
|
||||
// 获取可重试的响应
|
||||
let response = client
|
||||
.retryable(&url)
|
||||
@@ -111,7 +102,7 @@ impl StreamGears {
|
||||
info!("Downloading {}...", url);
|
||||
// FLV流下载
|
||||
let file = LifecycleFile::with_hook(&file_name, "flv", hook);
|
||||
httpflv::download(connection, file, self.segment.clone()).await;
|
||||
httpflv::download(connection, file, segment.clone()).await;
|
||||
}
|
||||
Err(Err::Incomplete(needed)) => {
|
||||
error!("needed: {needed:?}")
|
||||
@@ -120,7 +111,7 @@ impl StreamGears {
|
||||
error!("{e}");
|
||||
// HLS流下载
|
||||
let file = LifecycleFile::with_hook(&file_name, "ts", hook);
|
||||
hls::download(&url, &client, file, self.segment.clone())
|
||||
hls::download(&url, &client, file, segment.clone())
|
||||
.await
|
||||
.change_context(AppError::Unknown)?;
|
||||
}
|
||||
@@ -138,12 +129,13 @@ impl Downloader for StreamGears {
|
||||
async fn download(
|
||||
&self,
|
||||
callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
|
||||
download_config: DownloadConfig,
|
||||
) -> AppResult<DownloadStatus> {
|
||||
tokio::select! {
|
||||
_ = self.token.cancelled() => {
|
||||
bail!(AppError::Custom("StreamGears token cancelled".into()))
|
||||
}
|
||||
res = self.start_download(callback) => {res}
|
||||
res = self.start_download(callback, download_config) => {res}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -69,44 +69,16 @@ pub trait DownloadPlugin {
|
||||
/// 检查流状态
|
||||
async fn check_status(&self, ctx: &mut Context) -> Result<StreamStatus, Report<AppError>>;
|
||||
/// 创建下载器实例
|
||||
async fn create_downloader(
|
||||
&self,
|
||||
stream_info: &StreamInfoExt,
|
||||
config: Config,
|
||||
recorder: Recorder,
|
||||
) -> Arc<dyn Downloader> {
|
||||
let raw_stream_url = &stream_info.raw_stream_url;
|
||||
match config.downloader {
|
||||
Some(DownloaderType::Ffmpeg) => {
|
||||
let config = DownloadConfig {
|
||||
segment_time: config.segment_time.or_else(default_segment_time),
|
||||
file_size: Some(config.file_size), // 2GB
|
||||
headers: stream_info.stream_headers.clone(),
|
||||
recorder,
|
||||
// output_dir: PathBuf::from("./downloads")
|
||||
output_dir: PathBuf::from("."),
|
||||
};
|
||||
|
||||
Arc::new(FfmpegDownloader::new(
|
||||
raw_stream_url,
|
||||
config,
|
||||
Vec::new(),
|
||||
DownloaderType::FfmpegExternal,
|
||||
))
|
||||
}
|
||||
async fn create_downloader(&self, config: Option<DownloaderType>) -> Arc<dyn Downloader> {
|
||||
match config {
|
||||
Some(DownloaderType::Ffmpeg) => Arc::new(FfmpegDownloader::new(
|
||||
Vec::new(),
|
||||
DownloaderType::FfmpegExternal,
|
||||
)),
|
||||
// Some(DownloaderType::StreamGears) => {
|
||||
//
|
||||
// },
|
||||
_ => Arc::new(StreamGears::new(
|
||||
raw_stream_url,
|
||||
construct_headers(&stream_info.stream_headers),
|
||||
recorder.filename_template(),
|
||||
Segmentable::new(
|
||||
config.segment_time.as_deref().map(parse_time),
|
||||
Some(config.file_size),
|
||||
),
|
||||
None,
|
||||
)),
|
||||
_ => Arc::new(StreamGears::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ use crate::server::infrastructure::context::Context;
|
||||
use crate::server::infrastructure::models::StreamerInfo;
|
||||
use async_trait::async_trait;
|
||||
use chrono::Utc;
|
||||
use error_stack::{bail, IntoReport, Report, ResultExt};
|
||||
use error_stack::{IntoReport, Report, ResultExt, bail};
|
||||
use regex::Regex;
|
||||
use reqwest::header::HeaderMap;
|
||||
use serde_json::json;
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use crate::server::common::download::DownloadTask;
|
||||
use crate::server::common::util::Recorder;
|
||||
use crate::server::config::Config;
|
||||
use crate::server::core::downloader::Downloader;
|
||||
@@ -149,7 +150,7 @@ pub enum Stage {
|
||||
#[derive(Default, Clone)]
|
||||
pub enum WorkerStatus {
|
||||
/// 正在工作
|
||||
Working(Arc<dyn Downloader>),
|
||||
Working(Arc<DownloadTask>),
|
||||
/// 等待中
|
||||
Pending,
|
||||
/// 空闲状态(默认)
|
||||
|
||||
@@ -1,7 +1,14 @@
|
||||
use crate::server::api::bilibili_endpoints::{
|
||||
archive_pre_endpoint, get_myinfo_endpoint, get_proxy_endpoint,
|
||||
};
|
||||
use crate::server::api::endpoints::{add_upload_streamer_endpoint, add_user_endpoint, delete_streamers_endpoint, delete_template_endpoint, delete_user_endpoint, get_configuration, get_qrcode, get_status, get_streamer_info, get_streamer_info_files, get_streamers_endpoint, get_upload_streamer_endpoint, get_upload_streamers_endpoint, get_users_endpoint, get_videos, login_by_qrcode, pause_streamers_endpoint, post_streamers_endpoint, post_uploads, put_configuration, put_streamers_endpoint};
|
||||
use crate::server::api::endpoints::{
|
||||
add_upload_streamer_endpoint, add_user_endpoint, delete_streamers_endpoint,
|
||||
delete_template_endpoint, delete_user_endpoint, get_configuration, get_qrcode, get_status,
|
||||
get_streamer_info, get_streamer_info_files, get_streamers_endpoint,
|
||||
get_upload_streamer_endpoint, get_upload_streamers_endpoint, get_users_endpoint, get_videos,
|
||||
login_by_qrcode, pause_streamers_endpoint, post_streamers_endpoint, post_uploads,
|
||||
put_configuration, put_streamers_endpoint,
|
||||
};
|
||||
use crate::server::infrastructure::service_register::ServiceRegister;
|
||||
use axum::Router;
|
||||
use axum::body::Body;
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use async_trait::async_trait;
|
||||
use biliup_cli::server::core::downloader::{DownloadStatus, Downloader, SegmentEvent};
|
||||
use biliup_cli::server::core::downloader::{
|
||||
DownloadConfig, DownloadStatus, Downloader, SegmentEvent,
|
||||
};
|
||||
use biliup_cli::server::errors::{AppError, AppResult};
|
||||
use error_stack::ResultExt;
|
||||
use pyo3::prelude::*;
|
||||
@@ -24,6 +26,7 @@ impl Downloader for DanmakuClient {
|
||||
async fn download(
|
||||
&self,
|
||||
_callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
|
||||
download_config: DownloadConfig,
|
||||
) -> AppResult<DownloadStatus> {
|
||||
let py_client = self.py_client.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
|
||||
Reference in new issue
Block a user