From 69caca628c31d9528a2ed53cb9ddd14c9955be0f Mon Sep 17 00:00:00 2001 From: ForgQi <34411314+ForgQi@users.noreply.github.com> Date: Fri, 26 Sep 2025 01:42:03 +0800 Subject: [PATCH] Refactor segment event handling and add download task module Introduces a new `download` module to encapsulate download task logic, including segment event processing and file validation. Refactors segment event handling to use a new `SegmentInfo` struct and updates related interfaces and modules for improved clarity and maintainability. Adjusts uploader and downloader logic to work with the new event structure, and enhances FFmpeg and stream gears downloaders to emit and process segment events accordingly. --- Cargo.lock | 2 +- Cargo.toml | 2 +- .../biliup-cli/src/server/common/download.rs | 338 ++++++++++++++++++ crates/biliup-cli/src/server/common/mod.rs | 1 + crates/biliup-cli/src/server/common/upload.rs | 22 +- .../src/server/core/download_manager.rs | 108 +----- .../biliup-cli/src/server/core/downloader.rs | 42 ++- .../core/downloader/ffmpeg_downloader.rs | 104 ++++-- .../server/core/downloader/stream_gears.rs | 31 +- crates/biliup/src/downloader/util.rs | 1 + 10 files changed, 489 insertions(+), 162 deletions(-) create mode 100644 crates/biliup-cli/src/server/common/download.rs diff --git a/Cargo.lock b/Cargo.lock index e4975cd..8d8b1fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3626,7 +3626,7 @@ checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" [[package]] name = "stream-gears" -version = "1.1.5" +version = "1.1.6" dependencies = [ "async-channel", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index a546b8b..c32abac 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,7 +7,7 @@ members = [ resolver = "3" [workspace.package] -version = "1.1.5" +version = "1.1.6" [workspace.dependencies] serde = { version = "1.0", features = ["derive"] } diff --git a/crates/biliup-cli/src/server/common/download.rs b/crates/biliup-cli/src/server/common/download.rs new file mode 100644 index 0000000..c0d5147 --- /dev/null +++ b/crates/biliup-cli/src/server/common/download.rs @@ -0,0 +1,338 @@ +use crate::server::common::util::Recorder; +use crate::server::core::download_manager::{DownloaderMessage, UploaderMessage}; +use crate::server::core::downloader::{DownloadStatus, Downloader, SegmentEvent, SegmentInfo}; +use crate::server::core::monitor::RoomsHandle; +use crate::server::core::plugin::{DownloadPlugin, StreamInfo}; +use crate::server::errors::{AppError, AppResult}; +use crate::server::infrastructure::context::{Context, Stage, Worker, WorkerStatus}; +use async_channel::{Receiver, Sender}; +use error_stack::{FutureExt, ResultExt, bail}; +use std::fs; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; +use tracing::{error, info, warn}; + +// Configuration and retry policy +#[derive(Debug, Clone)] +pub struct RetryPolicy { + pub max_attempts: u32, + pub base_delay: Duration, + pub max_delay: Duration, +} + +impl RetryPolicy { + pub fn exponential(max_attempts: u32) -> Self { + Self { + max_attempts, + base_delay: Duration::from_millis(100), + max_delay: Duration::from_secs(30), + } + } +} + +/// 下载任务守卫，确保资源清理 +pub struct DownloadGuard { + worker: Arc, + danmaku_client: Option>, + rooms_handle: Arc, +} + +impl DownloadGuard { + fn new( + worker: Arc, + danmaku_client: Option>, + rooms_handle: Arc, + downloader: Arc, + ) -> Self { + // 更新工作器状态为工作中 + worker.change_status(Stage::Download, WorkerStatus::Working(downloader.clone())); + Self { + worker, + danmaku_client, + rooms_handle, + } + } +} + +impl Drop for DownloadGuard { + fn drop(&mut self) { + // 确保状态更新和资源清理 + self.worker + .change_status(Stage::Download, WorkerStatus::Idle); + + // 异步清理任务 + let danmaku = self.danmaku_client.clone(); + let rooms_handle = self.rooms_handle.clone(); + let worker = self.worker.clone(); + + tokio::spawn(async move { + if let Some(client) = danmaku { + if let Err(e) = client.stop().await { + error!("Error stopping danmaku client: {}", e); + } + } + rooms_handle.toggle(worker).await + }); + } +} + +/// 文件验证配置 +#[derive(Clone)] +pub struct FileValidator { + min_size: u64, + check_format: bool, +} + +impl Default for FileValidator { + fn default() -> Self { + Self { + min_size: 1024 * 1024 * 100, // 100MB minimum + check_format: true, + } + } +} + +impl FileValidator { + /// 验证文件有效性 + pub fn validate(&self, path: &Path) -> AppResult<()> { + let metadata = fs::metadata(path).change_context(AppError::Unknown)?; + + let size = metadata.len(); + + if size < self.min_size { + bail!(AppError::Custom(format!( + "File too small: {size} bytes, minimum: {} bytes", + self.min_size + ))); + } + + // 可选：检查文件格式 + if self.check_format { + self.validate_format(path)?; + } + + Ok(()) + } + + fn validate_format(&self, path: &Path) -> AppResult<()> { + // 简单的格式验证 - 检查扩展名 + if let Some(extension) = path.extension() { + let ext = extension.to_string_lossy().to_lowercase(); + match ext.as_str() { + "mp4" | "flv" | "ts" | "m3u8" => Ok(()), + _ => bail!(AppError::Custom(format!("Unsupported format: {}", ext))), + } + } else { + bail!(AppError::Custom("No file extension found".to_string())) + } + } +} + +/// 分段事件处理器 +#[derive(Clone)] +pub struct SegmentEventProcessor { + tx: Sender, + rx: Receiver, + uploader: Sender, + worker: Arc, + file_validator: FileValidator, +} + +impl SegmentEventProcessor { + /// 创建处理器 + pub fn new(uploader: Sender, worker: Arc) -> Self { + let (tx, rx) = async_channel::bounded(32); // Use tokio channel for async + + Self { + tx, + rx, + uploader, + worker, + file_validator: FileValidator::default(), + } + } + + /// 处理分段事件 + pub fn process(&self, event: SegmentInfo) -> AppResult<()> { + // 验证文件有效性 + self.file_validator.validate(&event.prev_file_path)?; + if event.segment_index == 0 { + // 发送到上传器 + let res = self + .uploader + .force_send(UploaderMessage::SegmentEvent( + self.rx.clone(), + self.worker.clone(), + )) + .change_context(AppError::Custom("Failed to send to uploader".to_string()))?; + if let Some(prev) = res { + warn!(SegmentEvent = ?prev, "replace an existing message in the channel"); + } + } + // 发送到缓冲区 + let res = self + .tx + .force_send(event) + .change_context(AppError::Custom("Failed to send to buffer".to_string()))?; + if let Some(prev) = res { + warn!(SegmentEvent = ?prev, "replace an existing message in the channel"); + } + Ok(()) + } + + /// 创建事件钩子 + pub fn create_hook( + &self, + danmaku: Option>, + ) -> impl Fn(SegmentEvent) + Clone + use<> { + let processor = self.clone(); + + move |event| { + match event { + SegmentEvent::Start { next_file_path } => { + // 触发弹幕滚动保存 + if let Some(ref client) = danmaku { + if let Err(e) = client + .rolling(&next_file_path.with_extension("xml").display().to_string()) + { + error!("Danmaku rolling error: {}", e); + } + } + } + SegmentEvent::Segment(event) => { + // 触发弹幕滚动保存 + if let Some(ref client) = danmaku { + if let Err(e) = client.rolling( + &event + .prev_file_path + .with_extension("xml") + .display() + .to_string(), + ) { + error!("Danmaku rolling error: {}", e); + } + } + // 异步处理事件 + let processor = processor.clone(); + if let Err(e) = processor.process(event) { + error!("Failed to process segment event: {}", e); + } + } + } + } + } +} + +/// 下载任务 +pub struct DownloadTask { + plugin: Arc, + stream_info: StreamInfo, + ctx: Context, + rooms_handle: Arc, +} + +struct DownloadComponents { + downloader: Arc, + danmaku_client: Option>, +} + +impl DownloadTask { + pub fn new( + plugin: Arc, + stream_info: StreamInfo, + ctx: Context, + rooms_handle: Arc, + ) -> Self { + Self { + plugin, + stream_info, + ctx, + rooms_handle, + } + } + + pub async fn execute(self, uploader: &Sender) -> AppResult<()> { + // 初始化组件 + let components = self.initialize_components().await?; + + // 创建守卫确保清理 + let _guard = DownloadGuard::new( + self.ctx.worker.clone(), + components.danmaku_client.clone(), + self.rooms_handle.clone(), + components.downloader.clone(), + ); + + // 创建事件处理器 + let processor = SegmentEventProcessor::new(uploader.clone(), self.ctx.worker.clone()); + + // 启动弹幕客户端 + if let Some(ref client) = components.danmaku_client { + self.start_danmaku(client).await?; + } + + // 执行下载 + let hook = processor.create_hook(components.danmaku_client.clone()); + let result = components + .downloader + .download(Box::new(hook)) + .await + .change_context(AppError::Custom("Failed to download segment".into()))?; + + // 处理结果 + match result { + DownloadStatus::Downloading => {} + DownloadStatus::SegmentCompleted => {} + DownloadStatus::StreamEnded => {} + DownloadStatus::Error(_) => {} + } + info!(result=?result, "finished downloading"); + Ok(()) + } + + async fn initialize_components(&self) -> AppResult { + // 获取配置和主播信息 + let config = self.ctx.worker.get_config(); + let streamer = self.ctx.worker.get_streamer(); + // 确定文件格式后缀 + let suffix = streamer + .format + .unwrap_or_else(|| self.stream_info.suffix.clone()); + // 创建录制器 + let recorder = Recorder::new( + streamer.filename_prefix, + &streamer.remark, + &self.stream_info.title, + &suffix, + ); + // 可选的弹幕客户端 + // 初始化弹幕客户端（如果存在） + let danmaku_client = self + .ctx + .extension + .get::>() + .map(Arc::clone); + + // 创建下载器实例 + let downloader = self + .plugin + .create_downloader(&self.stream_info, config, recorder) + .await; + + Ok(DownloadComponents { + downloader, + danmaku_client, + }) + } + + async fn start_danmaku(&self, client: &Arc) -> AppResult<()> { + // 启动弹幕下载逻辑 + println!( + "Starting danmaku client for stream: {}", + self.stream_info.url + ); + client.download(Box::new(|_| {})).await?; + Ok(()) + } +} diff --git a/crates/biliup-cli/src/server/common/mod.rs b/crates/biliup-cli/src/server/common/mod.rs index 4ef47b6..1ecf78d 100644 --- a/crates/biliup-cli/src/server/common/mod.rs +++ b/crates/biliup-cli/src/server/common/mod.rs @@ -1,3 +1,4 @@ +pub mod download; pub mod upload; /// 通用工具函数 pub mod util; diff --git a/crates/biliup-cli/src/server/common/upload.rs b/crates/biliup-cli/src/server/common/upload.rs index 5c52749..fa86ae0 100644 --- a/crates/biliup-cli/src/server/common/upload.rs +++ b/crates/biliup-cli/src/server/common/upload.rs @@ -1,5 +1,5 @@ use crate::server::core::download_manager::UActor; -use crate::server::core::downloader::SegmentEvent; +use crate::server::core::downloader::{SegmentEvent, SegmentInfo}; use crate::server::errors::{AppError, AppResult}; use crate::server::infrastructure::context::Worker; use crate::server::infrastructure::models::UploadStreamer; @@ -34,16 +34,16 @@ struct UploadedVideos { } pub async fn process_with_upload( - first_event: SegmentEvent, - rx: Receiver, + rx: Receiver, worker: &Worker, upload_config: UploadStreamer, ) -> AppResult<()> { + info!(upload_config=?upload_config, "Starting process with upload"); // 1. 初始化上传环境 let upload_context = initialize_upload_context(&worker, upload_config).await?; // 2. 流水线处理视频上传 - let uploaded_videos = pipeline_upload_videos(first_event, rx, &upload_context).await?; + let uploaded_videos = pipeline_upload_videos(rx, &upload_context).await?; // 3. 提交到B站 if !uploaded_videos.videos.is_empty() { @@ -98,8 +98,7 @@ async fn get_upload_line(client: &StatelessClient, line: &str) -> AppResult, + rx: Receiver, context: &UploadContext, ) -> AppResult { // let mut desc_v2 = Vec::new(); @@ -113,16 +112,11 @@ async fn pipeline_upload_videos( let mut uploaded = UploadedVideos::default(); - // 处理第一个事件 - let video = upload_single_file(&first_event.file_path, context).await?; - uploaded.videos.push(video); - uploaded.paths.push(first_event.file_path); - // 流式处理后续事件 while let Ok(event) = rx.recv().await { - let video = upload_single_file(&event.file_path, context).await?; + let video = upload_single_file(&event.prev_file_path, context).await?; uploaded.videos.push(video); - uploaded.paths.push(event.file_path); + uploaded.paths.push(event.prev_file_path); // 失败的文件不加入路径列表，避免后处理出错 } @@ -136,7 +130,7 @@ async fn upload_single_file(file_path: &Path, context: &UploadContext) -> AppRes let line = &context.line; let video_path = file_path; - println!( + info!( "{:?}", video_path .canonicalize() diff --git a/crates/biliup-cli/src/server/core/download_manager.rs b/crates/biliup-cli/src/server/core/download_manager.rs index 5f56269..e90c240 100644 --- a/crates/biliup-cli/src/server/core/download_manager.rs +++ b/crates/biliup-cli/src/server/core/download_manager.rs @@ -1,6 +1,7 @@ +use crate::server::common::download::DownloadTask; use crate::server::common::upload::{execute_postprocessor, process_with_upload}; use crate::server::common::util::Recorder; -use crate::server::core::downloader::{Downloader, SegmentEvent}; +use crate::server::core::downloader::{Downloader, SegmentEvent, SegmentInfo}; use crate::server::core::monitor::{Monitor, RoomsHandle}; use crate::server::core::plugin::{DownloadPlugin, StreamInfo}; use crate::server::errors::{AppError, AppResult}; @@ -114,95 +115,12 @@ impl DActor { async fn handle_message(&mut self, msg: DownloaderMessage) -> AppResult<()> { match msg { DownloaderMessage::Start(plugin, stream_info, ctx, rooms_handle) => { - // 获取配置和主播信息 - let config = ctx.worker.get_config(); - let streamer = ctx.worker.get_streamer(); - // 确定文件格式后缀 - let suffix = streamer - .format - .unwrap_or_else(|| stream_info.suffix.clone()); - // 创建录制器 - let recorder = Recorder::new( - streamer.filename_prefix, - &streamer.remark, - &stream_info.title, - &suffix, - ); - // 创建下载器实例 - let downloader = plugin - .create_downloader(&stream_info, config, recorder) - .await; + // 创建下载任务 + let task = DownloadTask::new(plugin, stream_info, ctx, rooms_handle); - // 更新工作器状态为工作中 - ctx.worker - .change_status(Stage::Download, WorkerStatus::Working(downloader.clone())); - // 初始化弹幕客户端（如果存在） - let mut danmaku_client = None; - if let Some(danmaku) = ctx.extension.get::>() { - let _ = danmaku - .download(Box::new(|_| {})) - .await - .inspect_err(|e| error!(e=?e)); - danmaku_client = Some(danmaku.clone()) - } + // 执行下载（使用 Result 链式处理） + task.execute(&self.sender).await?; - // 创建分段事件处理通道 - let (tx, rx) = bounded(16); - let hook = { - let room = ctx.worker.clone(); - let sender = self.sender.clone(); - let danmaku_client = danmaku_client.clone(); - move |event: SegmentEvent| { - // 如果有弹幕客户端，触发滚动保存 - if let Some(danmaku) = &danmaku_client { - let _ = danmaku - .rolling(&event.file_path.display().to_string()) - .inspect_err(|e| error!(e)); - } - // 处理分段事件 - if event.segment_index == 0 { - // 第一个分段，发送到上传器 - match sender.force_send(UploaderMessage::SegmentEvent( - event, - rx.clone(), - room.clone(), - )) { - Ok(Some(ret)) => { - warn!(SegmentEvent = ?ret, "replace an existing message in the channel"); - } - Err(_) => {} - Ok(None) => {} - }; - } else { - // 后续分段，发送到内部通道 - match tx.clone().force_send(event) { - Ok(Some(ret)) => { - warn!(SegmentEvent = ?ret, "replace an existing message in the channel"); - } - Err(_) => {} - Ok(None) => {} - } - } - } - }; - // 开始下载 - match downloader.download(Box::new(hook)).await { - Ok(status) => { - println!("Download completed with status: {:?}", status); - } - Err(err) => { - error!("download error: {:?}", err); - } - }; - // 停止弹幕客户端 - if let Some(danmaku) = &danmaku_client { - let _ = danmaku.stop().await.inspect_err(|e| error!(e)); - } - // 更新工作器状态为空闲 - ctx.worker - .change_status(Stage::Download, WorkerStatus::Idle); - // 切换房间状态 - rooms_handle.toggle(ctx.worker).await; Ok(()) } } @@ -234,12 +152,17 @@ impl UActor { /// * `msg` - 要处理的上传消息 async fn handle_message(&mut self, msg: UploaderMessage) { match msg { - UploaderMessage::SegmentEvent(event, rx, worker) => { + UploaderMessage::SegmentEvent(rx, worker) => { + worker.change_status(Stage::Upload, WorkerStatus::Pending); let result = match worker.get_upload_config() { - Some(config) => process_with_upload(event, rx, &worker, config).await, + Some(config) => process_with_upload(rx, &worker, config).await, None => { + let mut paths = Vec::new(); + while let Ok(event) = rx.recv().await { + paths.push(event.prev_file_path); + } // 无上传配置时，直接执行后处理 - execute_postprocessor(vec![event.file_path], &worker).await + execute_postprocessor(paths, &worker).await } }; @@ -247,6 +170,7 @@ impl UActor { error!("Process segment event failed: {}", e); // 可以添加错误通知机制 } + worker.change_status(Stage::Upload, WorkerStatus::Idle); } } } @@ -310,7 +234,7 @@ impl ActorHandle { #[derive(Debug)] pub enum UploaderMessage { /// 分段事件消息，包含事件、接收器和工作器 - SegmentEvent(SegmentEvent, Receiver, Arc), + SegmentEvent(Receiver, Arc), } /// 下载消息枚举 diff --git a/crates/biliup-cli/src/server/core/downloader.rs b/crates/biliup-cli/src/server/core/downloader.rs index d23dd47..5f7d480 100644 --- a/crates/biliup-cli/src/server/core/downloader.rs +++ b/crates/biliup-cli/src/server/core/downloader.rs @@ -65,18 +65,42 @@ pub enum DownloaderType { YtDlp, } +#[derive(Debug, Clone)] +pub struct SegmentInfo { + /// 分段文件路径 + pub prev_file_path: PathBuf, + pub next_file_path: Option, + /// 分段序号 + pub segment_index: usize, + // /// 分段开始时间戳 + // start_time: std::time::SystemTime, + // /// 分段结束时间戳 + // end_time: std::time::SystemTime, +} + +impl SegmentInfo { + pub fn new( + prev_file_path: PathBuf, + next_file_path: Option, + segment_index: usize, + ) -> Self { + Self { + prev_file_path, + next_file_path, + segment_index, + } + } +} + /// 分段事件 /// 当下载器完成一个分段时触发的事件 #[derive(Debug, Clone)] -pub struct SegmentEvent { - /// 分段文件路径 - pub file_path: PathBuf, - /// 分段序号 - pub segment_index: usize, - /// 分段开始时间戳 - pub start_time: std::time::SystemTime, - /// 分段结束时间戳 - pub end_time: std::time::SystemTime, +pub enum SegmentEvent { + Start { + /// 分段文件路径 + next_file_path: PathBuf, + }, + Segment(SegmentInfo), } /// 下载状态 diff --git a/crates/biliup-cli/src/server/core/downloader/ffmpeg_downloader.rs b/crates/biliup-cli/src/server/core/downloader/ffmpeg_downloader.rs index d134ffe..153350c 100644 --- a/crates/biliup-cli/src/server/core/downloader/ffmpeg_downloader.rs +++ b/crates/biliup-cli/src/server/core/downloader/ffmpeg_downloader.rs @@ -1,6 +1,6 @@ use crate::server::core::downloader; use crate::server::core::downloader::{ - DownloadConfig, DownloadStatus, Downloader, DownloaderType, SegmentEvent, + DownloadConfig, DownloadStatus, Downloader, DownloaderType, SegmentEvent, SegmentInfo, }; use crate::server::errors::{AppError, AppResult}; use async_trait::async_trait; @@ -70,10 +70,14 @@ impl FfmpegDownloader { // 内部分段特定的输出参数 // -f segment: 使用segment muxer进行自动分段 args.extend(["-f".to_string(), "segment".to_string()]); - // args.extend(["-segment_format".to_string(), "flv".to_string()]); + args.extend([ + "-segment_format".to_string(), + self.config.recorder.suffix.to_string(), + ]); // -segment_list pipe:1: 将分段文件名输出到stdout // 这样我们可以实时获取新生成的分段文件 args.extend(["-segment_list".to_string(), "pipe:1".to_string()]); + args.extend(["-map".to_string(), "0".to_string()]); // -segment_list_type flat: 输出格式为纯文件名列表 args.extend(["-segment_list_type".to_string(), "flat".to_string()]); @@ -81,6 +85,10 @@ impl FfmpegDownloader { // -reset_timestamps 1: 每个分段重置时间戳从0开始 // 确保每个分段文件可以独立播放 args.extend(["-reset_timestamps".to_string(), "1".to_string()]); + // %Y-%m-%dT%H_%M_%S 是 strftime 的时间占位符（需要配合 -strftime 1） + // %d 是序号占位符（printf 风格，默认模式） + // segment 复用器不能同时用这两种 + args.extend(["-strftime".to_string(), "1".to_string()]); // -segment_time: 分段时长（秒） if let Some(segment_time) = &self.config.segment_time { @@ -227,21 +235,6 @@ impl FfmpegDownloader { }; // let (tx, rx) = bounded(16); // 分段回调 - let segment_callback = |event: SegmentEvent| { - println!("New segment: {:?}", event.file_path); - // if i == 0 { - // match sender.force_send(UploaderMessage::SegmentEvent(event)) { - // Ok(Some(ret)) => { - // warn!(SegmentEvent = ?ret, "replace an existing message in the channel"); - // } - // Err(_) => {} - // Ok(None) => {} - // }; - // } else { - // i += 1; - // } - // 这里可以触发上传等后续处理 - }; // 根据退出码判断状态 match status.code() { @@ -254,12 +247,12 @@ impl FfmpegDownloader { // 触发分段回调 - segment_callback(SegmentEvent { - file_path: output_file, - segment_index: 0, - start_time: std::time::SystemTime::now(), - end_time: std::time::SystemTime::now(), - }); + // segment_callback(SegmentEvent { + // file_path: output_file, + // segment_index: 0, + // start_time: std::time::SystemTime::now(), + // end_time: std::time::SystemTime::now(), + // }); Ok(DownloadStatus::SegmentCompleted) } @@ -279,8 +272,8 @@ impl FfmpegDownloader { let mut cmd = Command::new("ffmpeg"); cmd.args(&args) .arg(format!( - "{}_%d.{}", - self.config.recorder.generate_filename(), + "{}.{}.part", + self.config.recorder.filename_template(), self.config.recorder.suffix )) .stdin(Stdio::null()) @@ -318,6 +311,7 @@ impl FfmpegDownloader { // 异步读取stdout let mut reader = BufReader::new(stdout).lines(); let mut segment_index = 0; + let mut prev_file_path: Option = None; while let Some(line) = reader.next_line().await.change_context(AppError::Unknown)? { info!("Received line: {line}"); @@ -329,19 +323,61 @@ impl FfmpegDownloader { // 触发分段回调 - callback(SegmentEvent { - file_path: file_path.clone(), + // 重命名文件 + let no_ext = file_path.with_extension(""); + tokio::fs::rename(&file_path, &no_ext) + .await + .change_context(AppError::Unknown)?; + info!("renamed file: from {file_path:?} to {no_ext:?}"); + + callback(SegmentEvent::Segment(SegmentInfo { + prev_file_path: no_ext, + next_file_path: None, segment_index, - start_time: std::time::SystemTime::now(), - end_time: std::time::SystemTime::now(), - }); + // start_time: std::time::SystemTime::now(), + // end_time: std::time::SystemTime::now(), + })); segment_index += 1; + prev_file_path = Some(file_path); } // 确保读任务结束（忽略它们的返回错误以避免因提前关闭管道导致的 join 错） let _ = stderr_task.await; - Ok(DownloadStatus::StreamEnded) + // 等待进程结束 + let status = { + let mut handle = self.process_handle.write().await; + if let Some(mut child) = handle.take() { + child.wait().await.change_context(AppError::Unknown)? + } else { + bail!(AppError::Custom("Process handle not found".to_string())); + } + }; + + if let Some(file_path) = prev_file_path { + // 重命名文件 + let no_ext = file_path.with_extension(""); + tokio::fs::rename(&file_path, &no_ext) + .await + .change_context(AppError::Unknown)?; + callback(SegmentEvent::Segment(SegmentInfo { + prev_file_path: no_ext, + next_file_path: None, + segment_index, + // start_time: std::time::SystemTime::now(), + // end_time: std::time::SystemTime::now(), + })); + } + + // 根据退出码判断状态 + match status.code() { + Some(0) => { + // 正常退出 + Ok(DownloadStatus::SegmentCompleted) + } + Some(255) => Ok(DownloadStatus::StreamEnded), + err => Ok(DownloadStatus::Error(format!("FFmpeg error: {err:?}"))), + } } } @@ -366,10 +402,12 @@ impl Downloader for FfmpegDownloader { async fn stop(&self) -> Result<(), Box> { let mut handle = self.process_handle.write().await; - if let Some(mut child) = handle.take() { + if let Some(child) = &mut *handle { child.kill().await?; + Ok(()) + } else { + Err(AppError::Custom("Process handle not found".to_string()).into()) } - Ok(()) } // async fn get_status(&self) -> DownloadStatus { diff --git a/crates/biliup-cli/src/server/core/downloader/stream_gears.rs b/crates/biliup-cli/src/server/core/downloader/stream_gears.rs index 327c5ae..b362ed3 100644 --- a/crates/biliup-cli/src/server/core/downloader/stream_gears.rs +++ b/crates/biliup-cli/src/server/core/downloader/stream_gears.rs @@ -1,4 +1,4 @@ -use crate::server::core::downloader::{DownloadStatus, Downloader, SegmentEvent}; +use crate::server::core::downloader::{DownloadStatus, Downloader, SegmentEvent, SegmentInfo}; use crate::server::errors::{AppError, AppResult}; use async_trait::async_trait; use axum::http::HeaderMap; @@ -85,19 +85,26 @@ impl StreamGears { .read_frame(9) .await .change_context(AppError::Unknown)?; - let mut i = 0; + // let mut i = 0; + // let mut prev_file_path = None; // 创建分段回调钩子 - let hook = move |s: &str| { - i += 1; - let event = SegmentEvent { - file_path: PathBuf::from(s), - segment_index: i, - start_time: SystemTime::now(), - end_time: SystemTime::now(), - }; - callback(event); - }; + let hook = { + let mut i = 0; + let mut prev_file_path = None::; + move |s: &str| { + let file_path = PathBuf::from(s); + let event = SegmentInfo { + prev_file_path: file_path.clone(), + next_file_path: None, + segment_index: i, + }; + callback(SegmentEvent::Segment(event)); + + i += 1; + prev_file_path = Some(file_path); + } + }; // 解析流头部，判断流类型 match header(&bytes) { Ok((_i, header)) => { diff --git a/crates/biliup/src/downloader/util.rs b/crates/biliup/src/downloader/util.rs index 90afe06..0d60950 100644 --- a/crates/biliup/src/downloader/util.rs +++ b/crates/biliup/src/downloader/util.rs @@ -270,6 +270,7 @@ impl LifecycleFile { } pub fn rename(&mut self) { + // 去掉 .part 后缀 match fs::rename(&self.path, &self.file_name) { Ok(_) => (self.hook)(&self.file_name), Err(e) => {