From aea8fb56a8a66f7c6902d29baf848316d471e77e Mon Sep 17 00:00:00 2001 From: ForgQi <34411314+ForgQi@users.noreply.github.com> Date: Tue, 23 Sep 2025 04:26:20 +0800 Subject: [PATCH] Refactor server context and downloader integration Refactored the server context to use a unified Context struct with extension support, updated Worker to encapsulate streamer, config, and client, and adjusted all usages accordingly. Moved util.rs to common/util.rs and updated imports. Enhanced downloader and plugin interfaces to support richer context and danmaku integration. Added DanmakuClient implementation in stream-gears and improved post-processing hooks to support multiple files. Various bug fixes and improvements for error handling, file operations, and configuration defaults. --- crates/biliup-cli/src/server.rs | 4 +- .../src/server/api/bilibili_endpoints.rs | 2 +- crates/biliup-cli/src/server/api/endpoints.rs | 53 ++++--- crates/biliup-cli/src/server/app.rs | 1 - crates/biliup-cli/src/server/common/mod.rs | 1 + .../src/server/{ => common}/util.rs | 6 +- crates/biliup-cli/src/server/config.rs | 7 +- .../src/server/core/download_manager.rs | 82 +++++++--- .../biliup-cli/src/server/core/downloader.rs | 30 ++-- .../core/downloader/ffmpeg_downloader.rs | 86 ++++++----- .../server/core/downloader/stream_gears.rs | 22 ++- crates/biliup-cli/src/server/core/monitor.rs | 17 ++- crates/biliup-cli/src/server/core/plugin.rs | 13 +- .../src/server/infrastructure/context.rs | 79 ++++++---- .../server/infrastructure/models/hook_step.rs | 47 +++--- .../src/server/infrastructure/repositories.rs | 16 ++ .../server/infrastructure/service_register.rs | 28 ++-- crates/stream-gears/src/danmaku.rs | 79 ++++++++++ crates/stream-gears/src/lib.rs | 44 +++++- crates/stream-gears/src/server.rs | 144 ++++++++---------- 20 files changed, 499 insertions(+), 262 deletions(-) create mode 100644 crates/biliup-cli/src/server/common/mod.rs rename crates/biliup-cli/src/server/{ => common}/util.rs (96%) create mode 100644 crates/stream-gears/src/danmaku.rs diff --git a/crates/biliup-cli/src/server.rs b/crates/biliup-cli/src/server.rs index 8be2d10..dc76b9a 100644 --- a/crates/biliup-cli/src/server.rs +++ b/crates/biliup-cli/src/server.rs @@ -1,11 +1,9 @@ pub mod api; pub mod app; +pub mod common; pub mod config; pub mod core; pub mod errors; pub mod infrastructure; mod router; -pub mod util; - - // use tokio::sync::mpsc::Receiver; diff --git a/crates/biliup-cli/src/server/api/bilibili_endpoints.rs b/crates/biliup-cli/src/server/api/bilibili_endpoints.rs index eba32db..db5bbfe 100644 --- a/crates/biliup-cli/src/server/api/bilibili_endpoints.rs +++ b/crates/biliup-cli/src/server/api/bilibili_endpoints.rs @@ -1,9 +1,9 @@ use crate::server::errors::{AppError, report_to_response}; use crate::server::infrastructure::connection_pool::ConnectionPool; use crate::server::infrastructure::models::Configuration; +use axum::Json; use axum::extract::{Query, State}; use axum::response::Response; -use axum::Json; use biliup::client::StatelessClient; use biliup::uploader::credential::login_by_cookies; use bytes::Bytes; diff --git a/crates/biliup-cli/src/server/api/endpoints.rs b/crates/biliup-cli/src/server/api/endpoints.rs index 0263172..81527f6 100644 --- a/crates/biliup-cli/src/server/api/endpoints.rs +++ b/crates/biliup-cli/src/server/api/endpoints.rs @@ -7,12 +7,14 @@ use crate::server::infrastructure::models::{ Configuration, FileItem, InsertConfiguration, InsertLiveStreamer, InsertUploadStreamer, LiveStreamer, StreamerInfo, UploadStreamer, }; -use crate::server::infrastructure::repositories::{del_streamer, get_all_streamer}; +use crate::server::infrastructure::repositories::{ + del_streamer, get_all_streamer, get_upload_config, +}; use crate::server::infrastructure::service_register::ServiceRegister; +use axum::Json; use axum::extract::{Path, State}; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; -use axum::Json; use biliup::credential::Credential; use error_stack::{Report, ResultExt}; use ormlite::{Insert, Model}; @@ -41,7 +43,7 @@ pub async fn get_streamers_endpoint( .unwrap() .iter() .find_map(|worker| { - if worker.id == x.id { + if worker.live_streamer.id == x.id { Some(*worker.downloader_status.read().unwrap()) } else { None @@ -71,8 +73,11 @@ pub async fn post_streamers_endpoint( .await .change_context(AppError::Unknown) .map_err(report_to_response)?; + let upload_config = get_upload_config(&pool, live_streamers.id) + .await + .map_err(report_to_response)?; service_register - .add_room(live_streamers.id, &live_streamers.url, monitor) + .add_room(monitor, live_streamers.clone(), upload_config) .await .map_err(report_to_response)?; @@ -94,15 +99,20 @@ pub async fn put_streamers_endpoint( info!("not supported url: {}", &streamer.url); return Err((StatusCode::BAD_REQUEST, "Not supported url").into_response()); }; - let monitor = manager.ensure_monitor(); service_register .del_room(streamer.id) .await .map_err(report_to_response)?; + let upload_config = get_upload_config(&pool, streamer.id) + .await + .map_err(report_to_response)?; + + let monitor = manager.ensure_monitor(); + service_register - .add_room(streamer.id, &streamer.url, monitor) + .add_room(monitor, streamer.clone(), upload_config) .await .map_err(report_to_response)?; @@ -417,22 +427,23 @@ pub async fn get_videos() -> Result>, Response> { && media_extensions .iter() .any(|allowed| ext == allowed.trim_start_matches('.')) - && let Ok(metadata) = entry.metadata().await { - let mtime = metadata - .modified() - .ok() - .and_then(|time| time.duration_since(UNIX_EPOCH).ok()) - .map(|d| d.as_secs()) - .unwrap_or(0); + && let Ok(metadata) = entry.metadata().await + { + let mtime = metadata + .modified() + .ok() + .and_then(|time| time.duration_since(UNIX_EPOCH).ok()) + .map(|d| d.as_secs()) + .unwrap_or(0); - file_list.push(serde_json::json!({ - "key": index, - "name": file_name, - "updateTime": mtime, - "size": metadata.len(), - })); - index += 1; - } + file_list.push(serde_json::json!({ + "key": index, + "name": file_name, + "updateTime": mtime, + "size": metadata.len(), + })); + index += 1; + } } } Ok(Json(file_list)) diff --git a/crates/biliup-cli/src/server/app.rs b/crates/biliup-cli/src/server/app.rs index ae80834..13cf0e5 100644 --- a/crates/biliup-cli/src/server/app.rs +++ b/crates/biliup-cli/src/server/app.rs @@ -1,4 +1,3 @@ - use axum::http; use crate::server; diff --git a/crates/biliup-cli/src/server/common/mod.rs b/crates/biliup-cli/src/server/common/mod.rs new file mode 100644 index 0000000..812d1ed --- /dev/null +++ b/crates/biliup-cli/src/server/common/mod.rs @@ -0,0 +1 @@ +pub mod util; diff --git a/crates/biliup-cli/src/server/util.rs b/crates/biliup-cli/src/server/common/util.rs similarity index 96% rename from crates/biliup-cli/src/server/util.rs rename to crates/biliup-cli/src/server/common/util.rs index 4fcf270..f5c4573 100644 --- a/crates/biliup-cli/src/server/util.rs +++ b/crates/biliup-cli/src/server/common/util.rs @@ -1,13 +1,15 @@ // main.rs use chrono::{Duration, Local}; +use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; use url::Url; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Deserialize, Serialize)] pub struct Recorder { pub filename_prefix: Option, pub fname: String, pub room_title: String, + /// 录制后保存文件格式 (mp4, ts, mkv, flv) pub suffix: String, // 不含点，如 "mp4"、"flv" } @@ -143,7 +145,7 @@ pub fn parse_time(segment_time: &str) -> std::time::Duration { #[cfg(test)] mod tests { - use crate::server::util::media_ext_from_url; + use crate::server::common::util::media_ext_from_url; #[test] fn it_works() { diff --git a/crates/biliup-cli/src/server/config.rs b/crates/biliup-cli/src/server/config.rs index 9277fa7..57ad36a 100644 --- a/crates/biliup-cli/src/server/config.rs +++ b/crates/biliup-cli/src/server/config.rs @@ -16,7 +16,7 @@ pub struct Config { pub file_size: u64, // Byte // 形如 "00:00:00"；保留为字符串以保持直观（你也可以替换为Duration并写自定义反序列化） - #[serde(default)] + #[serde(default = "default_segment_time")] pub segment_time: Option, #[builder(default = default_filtering_threshold())] @@ -288,6 +288,11 @@ pub struct UserConfig { fn default_file_size() -> u64 { 2_621_440_000 } + +pub fn default_segment_time() -> Option { + Some("02:00:00".to_string()) +} + fn default_filtering_threshold() -> u64 { 20 } diff --git a/crates/biliup-cli/src/server/core/download_manager.rs b/crates/biliup-cli/src/server/core/download_manager.rs index 9c9da62..42ed920 100644 --- a/crates/biliup-cli/src/server/core/download_manager.rs +++ b/crates/biliup-cli/src/server/core/download_manager.rs @@ -1,8 +1,9 @@ -use crate::server::core::downloader::SegmentEvent; +use crate::server::common::util::Recorder; +use crate::server::core::downloader::{Downloader, SegmentEvent}; use crate::server::core::monitor::{Monitor, RoomsHandle}; use crate::server::core::plugin::{DownloadPlugin, StreamInfo}; use crate::server::errors::{AppError, AppResult}; -use crate::server::infrastructure::context::{Worker, WorkerStatus}; +use crate::server::infrastructure::context::{Context, Stage, Worker, WorkerStatus}; use crate::server::infrastructure::models::hook_step::process_video; use async_channel::{Receiver, Sender, bounded}; use biliup::bilibili::{BiliBili, Studio, Video}; @@ -73,22 +74,54 @@ impl DActor { async fn run(&mut self) { while let Ok(msg) = self.receiver.recv().await { - self.handle_message(msg).await; + if let Err(e) = self.handle_message(msg).await { + error!("Error handling message: {}", e); + } } } - async fn handle_message(&mut self, msg: DownloaderMessage) { + async fn handle_message(&mut self, msg: DownloaderMessage) -> AppResult<()> { match msg { - DownloaderMessage::Start(plugin, stream_info, room, rooms_handle) => { - let downloader = plugin.create_downloader(&stream_info, &room).await.unwrap(); + DownloaderMessage::Start(plugin, stream_info, ctx, rooms_handle) => { + let config = ctx.worker.get_config(); + let streamer = ctx.worker.get_streamer(); + // info!(stream_info=?stream_info, "Create downloader"); + let suffix = streamer + .format + .unwrap_or_else(|| stream_info.suffix.clone()); + // println!("suffix: {suffix}"); + let recorder = Recorder::new( + streamer.filename_prefix, + &streamer.remark, + &stream_info.title, + &suffix, + ); + let downloader = plugin + .create_downloader(&stream_info, config, recorder) + .await; - *room.downloader_status.write().unwrap() = WorkerStatus::Working; + ctx.worker + .change_status(Stage::Download, WorkerStatus::Working); + let mut danmaku_client = None; + if let Some(danmaku) = ctx.extension.get::>() { + danmaku + .download(Box::new(|_| {})) + .await + .inspect_err(|e| error!(e=?e)); + danmaku_client = Some(danmaku.clone()) + } let (tx, rx) = bounded(16); let hook = { - let room = room.clone(); + 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 { + 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, @@ -120,8 +153,13 @@ impl DActor { error!("download error: {:?}", err); } }; - *room.downloader_status.write().unwrap() = WorkerStatus::Idle; - rooms_handle.toggle(room).await; + if let Some(danmaku) = &danmaku_client { + danmaku.stop().await.inspect_err(|e| error!(e)); + } + ctx.worker + .change_status(Stage::Download, WorkerStatus::Idle); + rooms_handle.toggle(ctx.worker).await; + Ok(()) } } } @@ -144,18 +182,18 @@ impl UActor { async fn handle_message(&mut self, msg: UploaderMessage) { match msg { UploaderMessage::SegmentEvent(event, rx, worker) => { - let Some(upload_config) = worker.get_upload_config().await.unwrap() else { + let Some(upload_config) = worker.get_upload_config() else { return; }; let cookie_file = upload_config .user_cookie .unwrap_or("cookies.json".to_string()); let bilibili = login_by_cookies(&cookie_file, None).await; - let Ok(bilibili) = bilibili.map_err(|e| error!(e=?e)) else { + let Ok(bilibili) = bilibili.inspect_err(|e| error!(e=?e)) else { return; }; - let config = worker.get_config().await.unwrap(); - let client = worker.context.client.clone(); + let config = worker.get_config(); + let client = worker.client.clone(); let line = config.lines; let line = match line.as_str() { "bda2" => line::bda2(), @@ -176,7 +214,7 @@ impl UActor { &event.file_path, ) .await - .map_err(|e| error!(e=?e)) + .inspect_err(|e| error!(e=?e)) { videos.push(video); } @@ -253,13 +291,13 @@ impl UActor { info!(submit_result=?submit_result); if submit_result.is_ok() { info!("Submit successful"); - let streamer = worker.get_streamer().await.unwrap(); + let streamer = worker.get_streamer(); if let Some(post_processor) = streamer.postprocessor { - for video_path in &video_paths { - let _ = process_video(video_path, &post_processor) - .await - .map_err(|e| error!(e=?e)); - } + + let _ = process_video(&(video_paths.iter().map(|p|p.as_path()).collect::>()), &post_processor) + .await + .map_err(|e| error!(e=?e)); + } } } @@ -357,7 +395,7 @@ pub enum DownloaderMessage { Start( Arc, StreamInfo, - Arc, + Context, Arc, ), } diff --git a/crates/biliup-cli/src/server/core/downloader.rs b/crates/biliup-cli/src/server/core/downloader.rs index 1aaea35..36506d0 100644 --- a/crates/biliup-cli/src/server/core/downloader.rs +++ b/crates/biliup-cli/src/server/core/downloader.rs @@ -1,6 +1,8 @@ pub mod ffmpeg_downloader; pub mod stream_gears; +use crate::server::common::util::Recorder; +use crate::server::errors::AppResult; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -9,9 +11,6 @@ use std::path::PathBuf; /// 下载器配置 #[derive(Debug, Clone, Deserialize, Serialize)] pub struct DownloadConfig { - /// 录制后保存文件格式 (mp4, ts, mkv, flv) - pub format: String, - /// 分段时长 (格式: "HH:MM:SS") pub segment_time: Option, @@ -21,14 +20,18 @@ pub struct DownloadConfig { /// HTTP请求头 pub headers: HashMap, - /// 额外的FFmpeg参数 - pub extra_args: Vec, - - /// 下载器类型 - pub downloader_type: DownloaderType, - /// 输出文件名前缀 - pub filename_prefix: String, + pub recorder: Recorder, + + // 输出文件名前缀 + // pub filename_prefix: String, + pub output_dir: PathBuf, +} + +impl DownloadConfig { + fn generate_output_filename(&self) -> PathBuf { + self.output_dir.join(self.recorder.generate_path()) + } } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -86,13 +89,16 @@ pub trait Downloader: Send + Sync { // sender: Sender, // worker: Arc, callback: Box, - ) -> Result>; + ) -> AppResult; /// 停止下载 async fn stop(&self) -> Result<(), Box>; /// 获取当前下载状态 - async fn get_status(&self) -> DownloadStatus; + /// async fn get_status(&self) -> DownloadStatus; + fn rolling(&self, file_name: &str) -> Result<(), Box> { + Ok(()) + } } /// 解析时长字符串 "HH:MM:SS" 为秒数 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 33ac1d9..7d5bd3f 100644 --- a/crates/biliup-cli/src/server/core/downloader/ffmpeg_downloader.rs +++ b/crates/biliup-cli/src/server/core/downloader/ffmpeg_downloader.rs @@ -2,7 +2,9 @@ use crate::server::core::downloader; use crate::server::core::downloader::{ DownloadConfig, DownloadStatus, Downloader, DownloaderType, SegmentEvent, }; +use crate::server::errors::{AppError, AppResult}; use async_trait::async_trait; +use error_stack::{ResultExt, bail}; use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; @@ -15,19 +17,30 @@ use tracing::info; pub struct FfmpegDownloader { config: DownloadConfig, url: String, - output_dir: PathBuf, status: Arc>, process_handle: Arc>>, + + /// 额外的FFmpeg参数 + pub extra_args: Vec, + + /// 下载器类型 + pub downloader_type: DownloaderType, } impl FfmpegDownloader { - pub fn new(url: &str, config: DownloadConfig, output_dir: PathBuf) -> Self { + pub fn new( + url: &str, + config: DownloadConfig, + extra_args: Vec, + downloader_type: DownloaderType, + ) -> Self { Self { config, url: url.to_string(), - output_dir, status: Arc::new(RwLock::new(DownloadStatus::Downloading)), process_handle: Arc::new(RwLock::new(None)), + extra_args, + downloader_type, } } @@ -90,7 +103,7 @@ impl FfmpegDownloader { } // 添加通用输出参数 - self.append_common_output_args(&mut args, &self.config.format); + self.append_common_output_args(&mut args, &self.config.recorder.suffix); args } @@ -159,16 +172,16 @@ impl FfmpegDownloader { } // 添加额外参数 - args.extend(self.config.extra_args.clone()); + args.extend(self.extra_args.clone()); } /// 执行外部分段下载 async fn download_external( &self, callback: Box, - ) -> Result> { + ) -> AppResult { let args = self.build_ffmpeg_args_external_segment(); - let output_file = self.generate_output_filename(); + let output_file = self.config.generate_output_filename(); let mut cmd = Command::new("ffmpeg"); cmd.args(&args) @@ -178,7 +191,7 @@ impl FfmpegDownloader { .stderr(Stdio::piped()) .kill_on_drop(true); - let child = cmd.spawn()?; + let child = cmd.spawn().change_context(AppError::Unknown)?; // 保存进程句柄 { @@ -190,9 +203,9 @@ impl FfmpegDownloader { let status = { let mut handle = self.process_handle.write().await; if let Some(mut child) = handle.take() { - child.wait().await? + child.wait().await.change_context(AppError::Unknown)? } else { - return Err("Process handle not found".into()); + bail!(AppError::Custom("Process handle not found".to_string())); } }; // let (tx, rx) = bounded(16); @@ -218,7 +231,9 @@ impl FfmpegDownloader { Some(0) => { // 正常退出，重命名文件 let part_file = format!("{}.part", output_file.display()); - tokio::fs::rename(&part_file, &output_file).await?; + tokio::fs::rename(&part_file, &output_file) + .await + .change_context(AppError::Unknown)?; // 触发分段回调 @@ -240,15 +255,15 @@ impl FfmpegDownloader { async fn download_internal( &self, callback: Box, - ) -> Result> { + ) -> AppResult { let args = self.build_ffmpeg_args_internal_segment(); - let _output_pattern = self.output_dir.join("%d.{}"); let mut cmd = Command::new("ffmpeg"); cmd.args(&args) .arg(format!( "{}_%d.{}", - self.config.filename_prefix, self.config.format + self.config.recorder.generate_filename(), + self.config.recorder.suffix )) .stdin(Stdio::null()) .stdout(Stdio::piped()) @@ -256,12 +271,17 @@ impl FfmpegDownloader { .kill_on_drop(true); info!("FFmpeg cmd: {:?}", cmd); - let mut child = cmd.spawn()?; + let mut child = cmd.spawn().change_context(AppError::Unknown)?; // 获取stdout用于读取分段文件名 - let stdout = child.stdout.take().ok_or("Failed to capture stdout")?; + let stdout = child + .stdout + .take() + .ok_or(AppError::Custom("Failed to capture stdout".to_string()))?; - let stderr = child.stderr.take().ok_or("failed to capture stderr pipe")?; + let stderr = child.stderr.take().ok_or(AppError::Custom( + "failed to capture stderr pipe".to_string(), + ))?; // 保存进程句柄 { @@ -281,7 +301,7 @@ impl FfmpegDownloader { let mut reader = BufReader::new(stdout).lines(); let mut segment_index = 0; - while let Some(line) = reader.next_line().await? { + while let Some(line) = reader.next_line().await.change_context(AppError::Unknown)? { info!("Received line: {line}"); // 解析文件名 let file_path = PathBuf::from(line.trim()); @@ -305,14 +325,6 @@ impl FfmpegDownloader { Ok(DownloadStatus::StreamEnded) } - - fn generate_output_filename(&self) -> PathBuf { - let timestamp = chrono::Local::now().format("%Y-%m-%dT%H_%M_%S"); - self.output_dir.join(format!( - "{}_{}.{}", - self.config.filename_prefix, timestamp, self.config.format - )) - } } #[async_trait] @@ -320,11 +332,17 @@ impl Downloader for FfmpegDownloader { async fn download( &self, callback: Box, - ) -> Result> { - match self.config.downloader_type { - DownloaderType::FfmpegExternal => self.download_external(callback).await, - DownloaderType::FfmpegInternal => self.download_internal(callback).await, - _ => Err("Unsupported downloader type".into()), + ) -> AppResult { + match self.downloader_type { + DownloaderType::FfmpegExternal => self + .download_external(callback) + .await + .change_context(AppError::Unknown), + DownloaderType::FfmpegInternal => self + .download_internal(callback) + .await + .change_context(AppError::Unknown), + _ => bail!(AppError::Custom("Unsupported downloader type".to_string())), } } @@ -336,7 +354,7 @@ impl Downloader for FfmpegDownloader { Ok(()) } - async fn get_status(&self) -> DownloadStatus { - self.status.read().await.clone() - } + // async fn get_status(&self) -> DownloadStatus { + // self.status.read().await.clone() + // } } 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 dc67320..76a81b2 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,5 @@ use crate::server::core::downloader::{DownloadStatus, Downloader, SegmentEvent}; +use crate::server::errors::{AppError, AppResult}; use async_trait::async_trait; use axum::http::HeaderMap; use biliup::client::StatelessClient; @@ -6,6 +7,7 @@ use biliup::downloader::flv_parser::header; use biliup::downloader::httpflv::Connection; use biliup::downloader::util::{LifecycleFile, Segmentable}; use biliup::downloader::{hls, httpflv}; +use error_stack::ResultExt; use nom::Err; use std::path::PathBuf; use std::sync::{Arc, RwLock}; @@ -46,7 +48,7 @@ impl Downloader for StreamGears { async fn download( &self, callback: Box, - ) -> Result> { + ) -> AppResult { let url = self.url.clone(); let file_name = self.file_name.clone(); let headers_in = self.header_map.clone(); @@ -59,10 +61,16 @@ impl Downloader for StreamGears { // response.copy_to(&mut writer)?; // io::copy(&mut resp, &mut out).expect("Unable to copy the content."); let client = StatelessClient::new(self.header_map.clone(), proxy.as_deref()); - let response = client.retryable(&url).await?; + let response = client + .retryable(&url) + .await + .change_context(AppError::Unknown)?; let mut connection = Connection::new(response); // let buf = &mut [0u8; 9]; - let bytes = connection.read_frame(9).await?; + let bytes = connection + .read_frame(9) + .await + .change_context(AppError::Unknown)?; let mut i = 0; let hook = move |s: &str| { i += 1; @@ -88,7 +96,9 @@ impl Downloader for StreamGears { Err(e) => { error!("{e}"); let file = LifecycleFile::with_hook(&file_name, "ts", hook); - hls::download(&url, &client, file, self.segment.clone()).await?; + hls::download(&url, &client, file, self.segment.clone()) + .await + .change_context(AppError::Unknown)?; } } Ok(DownloadStatus::StreamEnded) @@ -99,8 +109,4 @@ impl Downloader for StreamGears { // 如上所述，如果底层下载函数不支持取消，这里不能真正中断正在进行的下载。 todo!() } - - async fn get_status(&self) -> DownloadStatus { - todo!() - } } diff --git a/crates/biliup-cli/src/server/core/monitor.rs b/crates/biliup-cli/src/server/core/monitor.rs index 213fecf..b42173d 100644 --- a/crates/biliup-cli/src/server/core/monitor.rs +++ b/crates/biliup-cli/src/server/core/monitor.rs @@ -1,6 +1,6 @@ use crate::server::core::download_manager::{ActorHandle, DownloaderMessage}; use crate::server::core::plugin::{DownloadPlugin, StreamStatus}; -use crate::server::infrastructure::context::{Worker, WorkerStatus}; +use crate::server::infrastructure::context::{Context, Stage, Worker, WorkerStatus}; use async_channel::{Receiver, Sender, bounded}; use std::sync::Arc; use std::time::Duration; @@ -18,19 +18,20 @@ async fn start_client( info!("start -> [{platform_name}]"); loop { if let Some(room) = rooms_handle.next().await { - let ulr = room.get_streamer().await.unwrap().url; - interval = room.get_config().await.unwrap().event_loop_interval; + let ulr = room.get_streamer().url; + interval = room.get_config().event_loop_interval; info!("[{platform_name}] room: {ulr}"); - match plugin.check_status(&ulr).await.unwrap() { + let mut ctx = Context::new(room.clone()); + match plugin.check_status(&mut ctx).await.unwrap() { StreamStatus::Live { stream_info } => { info!("room: {ulr} is live -> {:?}", stream_info); - *room.downloader_status.write().unwrap() = WorkerStatus::Pending; + room.change_status(Stage::Download, WorkerStatus::Pending); if actor_handle .down_sender .send(DownloaderMessage::Start( plugin.clone(), stream_info, - room.clone(), + ctx, rooms_handle.clone(), )) .await @@ -218,9 +219,9 @@ impl RoomsActor { } fn del(&mut self, id: i64) -> usize { - if let Some(i) = self.rooms.iter().position(|x| x.id == id) { + if let Some(i) = self.rooms.iter().position(|x| x.live_streamer.id == id) { self.rooms.remove(i); // 保序，但 O(n) - } else if let Some(i) = self.waiting.iter().position(|x| x.id == id) { + } else if let Some(i) = self.waiting.iter().position(|x| x.live_streamer.id == id) { self.waiting.swap_remove(i); }; info!("Removed room [{:?}] {}", self.rooms, self.rooms.len()); diff --git a/crates/biliup-cli/src/server/core/plugin.rs b/crates/biliup-cli/src/server/core/plugin.rs index a693a7a..55232a9 100644 --- a/crates/biliup-cli/src/server/core/plugin.rs +++ b/crates/biliup-cli/src/server/core/plugin.rs @@ -1,6 +1,8 @@ +use crate::server::common::util::Recorder; +use crate::server::config::Config; use crate::server::core::downloader::Downloader; use crate::server::errors::{AppError, AppResult}; -use crate::server::infrastructure::context::Worker; +use crate::server::infrastructure::context::{Context, Worker}; use async_trait::async_trait; use error_stack::Report; use serde::{Deserialize, Serialize}; @@ -43,12 +45,15 @@ pub trait DownloadBase: Send + Sync { #[async_trait] pub trait DownloadPlugin { fn matches(&self, url: &str) -> bool; - async fn check_status(&self, url: &str) -> Result>; + async fn check_status(&self, ctx: &mut Context) -> Result>; async fn create_downloader( &self, stream_info: &StreamInfo, - worker: &Worker, - ) -> AppResult>; + config: Config, + recorder: Recorder, + ) -> Box; + + fn danmaku_init(&self) -> Option>; fn name(&self) -> &str; } diff --git a/crates/biliup-cli/src/server/infrastructure/context.rs b/crates/biliup-cli/src/server/infrastructure/context.rs index e8910c6..eb3f515 100644 --- a/crates/biliup-cli/src/server/infrastructure/context.rs +++ b/crates/biliup-cli/src/server/infrastructure/context.rs @@ -3,68 +3,81 @@ use crate::server::errors::{AppError, AppResult}; use crate::server::infrastructure::connection_pool::ConnectionPool; use crate::server::infrastructure::models::{LiveStreamer, UploadStreamer}; use crate::server::infrastructure::repositories::{get_config, get_streamer}; +use axum::http::Extensions; use biliup::client::StatelessClient; use error_stack::ResultExt; use ormlite::Model; use serde::{Deserialize, Serialize}; -use std::sync::RwLock; +use std::sync::{Arc, RwLock}; use tracing::info; -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct Context { - pool: ConnectionPool, - pub client: StatelessClient, + pub worker: Arc, + pub extension: Extensions, } impl Context { - pub fn new(pool: ConnectionPool, client: StatelessClient) -> Self { - Self { pool, client } + pub fn new(worker: Arc) -> Self { + Self { + worker, + extension: Default::default(), + } } - pub async fn get_config(&self) -> AppResult { - get_config(&self.pool).await - } + // pub async fn get_config(&self) -> AppResult { + // get_config(&self.pool).await + // } } #[derive(Debug)] pub struct Worker { - pub id: i64, - pub url: String, - pub context: Context, pub downloader_status: RwLock, pub uploader_status: RwLock, + pub live_streamer: LiveStreamer, + pub upload_streamer: Option, + pub config: Arc>, + pub client: StatelessClient, } impl Worker { - pub fn new(id: i64, url: &str, context: Context) -> Self { + pub fn new( + live_streamer: LiveStreamer, + upload_streamer: Option, + config: Arc>, + client: StatelessClient, + ) -> Self { Self { - id, - context, downloader_status: RwLock::new(Default::default()), uploader_status: Default::default(), - url: url.to_string(), + live_streamer, + upload_streamer, + config, + client, } } - pub async fn get_streamer(&self) -> AppResult { - get_streamer(&self.context.pool, self.id).await + pub fn get_streamer(&self) -> LiveStreamer { + // get_streamer(&self.context.pool, self.id).await + self.live_streamer.clone() } - pub async fn get_config(&self) -> AppResult { - self.context.get_config().await + pub fn get_upload_config(&self) -> Option { + // let Some(id) = self.get_streamer().await?.upload_streamers_id else { + // return Ok(None); + // }; + // + // UploadStreamer::select() + // .where_("id = ?") + // .bind(id) + // .fetch_optional(&self.context.pool) + // .await + // .change_context(AppError::Unknown) + self.upload_streamer.clone() } - pub async fn get_upload_config(&self) -> AppResult> { - let Some(id) = self.get_streamer().await?.upload_streamers_id else { - return Ok(None); - }; - - UploadStreamer::select() - .where_("id = ?") - .bind(id) - .fetch_optional(&self.context.pool) - .await - .change_context(AppError::Unknown) + pub fn get_config(&self) -> Config { + self.config.read().unwrap().clone() } pub fn change_status(&self, stage: Stage, status: WorkerStatus) { @@ -81,13 +94,13 @@ impl Worker { impl Drop for Worker { fn drop(&mut self) { - info!("Dropping worker {}", self.id); + info!("Dropping worker {}", self.live_streamer.id); } } impl PartialEq for Worker { fn eq(&self, other: &Self) -> bool { - self.id == other.id + self.live_streamer.id == other.live_streamer.id } } diff --git a/crates/biliup-cli/src/server/infrastructure/models/hook_step.rs b/crates/biliup-cli/src/server/infrastructure/models/hook_step.rs index 70d9197..1ae3358 100644 --- a/crates/biliup-cli/src/server/infrastructure/models/hook_step.rs +++ b/crates/biliup-cli/src/server/infrastructure/models/hook_step.rs @@ -1,8 +1,10 @@ +use std::borrow::Cow; use serde::{Deserialize, Serialize}; use std::path::Path; use tokio::fs; use tokio::io::AsyncWriteExt; use tokio::process::Command; +use tracing::info; // 钩子步骤：既支持 key-value 形式（如 {run: "..."}），也支持纯字符串（如 "rm"） #[derive(Debug, Clone, Serialize, Deserialize)] @@ -18,16 +20,16 @@ pub enum HookStep { impl HookStep { /// 执行后处理操作 - pub async fn execute(&self, video_path: &Path) -> Result<(), Box> { + pub async fn execute(&self, video_paths: &[&Path]) -> Result<(), Box> { match self { HookStep::Run { run } => { - self.execute_command(run, video_path).await?; + self.execute_command(run, video_paths).await?; } HookStep::Move { mv } => { - self.move_file(video_path, mv).await?; + self.move_file(video_paths, mv).await?; } HookStep::Remove(cmd) if cmd == "rm" => { - self.remove_file(video_path).await?; + self.remove_file(video_paths).await?; } HookStep::Remove(cmd) => { return Err(format!("Unknown command: {}", cmd).into()); @@ -40,7 +42,7 @@ impl HookStep { async fn execute_command( &self, cmd: &str, - video_path: &Path, + video_paths: &[&Path], ) -> Result<(), Box> { println!("Executing: {}", cmd); @@ -57,7 +59,7 @@ impl HookStep { // 将视频文件路径写入标准输入 if let Some(mut stdin) = process.stdin.take() { stdin - .write_all(video_path.to_string_lossy().as_bytes()) + .write_all(video_paths.iter().map(|p| p.to_string_lossy()).reduce(|acc, e| Cow::from(acc.to_string() + "\n" + &*e)).expect("REASON").as_bytes()) .await?; } @@ -72,7 +74,7 @@ impl HookStep { /// 移动文件到指定目录 async fn move_file( &self, - video_path: &Path, + video_paths: &[&Path], target_dir: &str, ) -> Result<(), Box> { let target_path = Path::new(target_dir); @@ -81,34 +83,37 @@ impl HookStep { if !target_path.exists() { fs::create_dir_all(target_path).await?; } + for video_path in video_paths { + let file_name = video_path.file_name().ok_or("Invalid file name")?; - let file_name = video_path.file_name().ok_or("Invalid file name")?; - - let destination = target_path.join(file_name); - - println!( - "Moving {} to {}", - video_path.display(), - destination.display() - ); - fs::rename(video_path, destination).await?; + let destination = target_path.join(file_name); + println!( + "Moving {} to {}", + video_path.display(), + destination.display() + ); + fs::rename(video_path, destination).await?; + } Ok(()) } /// 删除文件 - async fn remove_file(&self, video_path: &Path) -> Result<(), Box> { - println!("Removing: {}", video_path.display()); - fs::remove_file(video_path).await?; + async fn remove_file(&self, video_paths: &[&Path]) -> Result<(), Box> { + for video_path in video_paths { + println!("Removing: {}", video_path.display()); + fs::remove_file(video_path).await?; + } Ok(()) } } /// 处理所有后处理器 pub async fn process_video( - video_path: &Path, + video_path: &[&Path], processors: &[HookStep], ) -> Result<(), Box> { + info!("Starting processing..."); for processor in processors { processor.execute(video_path).await?; } diff --git a/crates/biliup-cli/src/server/infrastructure/repositories.rs b/crates/biliup-cli/src/server/infrastructure/repositories.rs index 832c9f9..a361841 100644 --- a/crates/biliup-cli/src/server/infrastructure/repositories.rs +++ b/crates/biliup-cli/src/server/infrastructure/repositories.rs @@ -16,6 +16,22 @@ pub async fn get_streamer(pool: &ConnectionPool, id: i64) -> AppResult AppResult> { + let Some(id) = get_streamer(pool, id).await?.upload_streamers_id else { + return Ok(None); + }; + + UploadStreamer::select() + .where_("id = ?") + .bind(id) + .fetch_optional(pool) + .await + .change_context(AppError::Unknown) +} + pub async fn del_streamer(pool: &ConnectionPool, id: i64) -> AppResult { let streamer = get_streamer(pool, id).await?; streamer diff --git a/crates/biliup-cli/src/server/infrastructure/service_register.rs b/crates/biliup-cli/src/server/infrastructure/service_register.rs index d2ca6d8..80dac67 100644 --- a/crates/biliup-cli/src/server/infrastructure/service_register.rs +++ b/crates/biliup-cli/src/server/infrastructure/service_register.rs @@ -4,6 +4,7 @@ use crate::server::core::monitor::Monitor; use crate::server::errors::{AppError, AppResult}; use crate::server::infrastructure::connection_pool::ConnectionPool; use crate::server::infrastructure::context::{Context, Worker}; +use crate::server::infrastructure::models::{LiveStreamer, UploadStreamer}; use axum::extract::FromRef; use biliup::client::StatelessClient; use error_stack::bail; @@ -56,35 +57,42 @@ impl ServiceRegister { pub async fn add_room( &self, - id: i64, - url: &str, monitor: Arc, + live_streamer: LiveStreamer, + upload_streamer: Option, ) -> AppResult> { let worker = Arc::new(Worker::new( - id, - url, - Context::new(self.pool.clone(), self.client.clone()), + live_streamer, + upload_streamer, + self.config.clone(), + self.client.clone(), )); monitor.rooms_handle.add(worker.clone()).await; - self.workers.write().unwrap().push(worker); - info!("add {url} success"); + self.workers.write().unwrap().push(worker.clone()); + info!("add {worker:?} success"); Ok(Some(())) } pub async fn del_room(&self, id: i64) -> AppResult<()> { - let Some(i) = self.workers.read().unwrap().iter().position(|x| x.id == id) else { + let Some(i) = self + .workers + .read() + .unwrap() + .iter() + .position(|x| x.live_streamer.id == id) + else { return Err(error_stack::Report::new(AppError::Unknown)); }; let removed = self.workers.write().unwrap().swap_remove(i); - let url = &removed.url; + let url = &removed.live_streamer.url; let Some(manager) = self.get_manager(url) else { info!("not found url: {url}"); bail!(AppError::Unknown) }; let monitor = manager.ensure_monitor(); let len = monitor.rooms_handle.del(id).await; - info!("{id} removed, remained len {len}"); + info!("id: {id} removed, remained len {len}"); if len == 0 { *manager.monitor.lock().unwrap() = None; } diff --git a/crates/stream-gears/src/danmaku.rs b/crates/stream-gears/src/danmaku.rs new file mode 100644 index 0000000..a1fb5b4 --- /dev/null +++ b/crates/stream-gears/src/danmaku.rs @@ -0,0 +1,79 @@ +use async_trait::async_trait; +use biliup_cli::server::core::downloader::{DownloadStatus, Downloader, SegmentEvent}; +use biliup_cli::server::errors::{AppError, AppResult}; +use error_stack::ResultExt; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use std::collections::HashMap; +use std::sync::Arc; + +/// DanmakuClient provides Rust bindings to the Python DanmakuClient +/// for recording live chat (danmaku) alongside video streams. +pub struct DanmakuClient { + py_client: Arc>, +} + +impl DanmakuClient { + /// Creates a new DanmakuClient instance + pub fn new(py_client: Arc>) -> Self { + Self { py_client } + } +} + +#[async_trait] +impl Downloader for DanmakuClient { + /// Starts danmaku recording and manages lifecycle + async fn download( + &self, + _callback: Box, + ) -> AppResult { + let py_client = self.py_client.clone(); + tokio::task::spawn_blocking(move || { + Python::attach(|py| { + let py_client = py_client.bind(py); + py_client.call_method0("start")?; + Ok::<_, PyErr>(()) + }) + }) + .await + .change_context(AppError::Unknown)? + .change_context(AppError::Unknown)?; + + // Start the danmaku recording + // self.start() + // .map_err(|e| Box::new(e) as Box)?; + + // Return downloading status - the actual recording runs in the background + // The Python DanmakuClient handles the recording lifecycle internally + Ok(DownloadStatus::Downloading) + } + + /// Stops the danmaku recording + async fn stop(&self) -> Result<(), Box> { + let py_client = self.py_client.clone(); + // Call the DanmakuClient's stop method (not the trait method) + tokio::task::spawn_blocking(move || { + Python::attach(|py| { + let py_client = py_client.bind(py); + py_client.call_method0("stop")?; + Ok::<_, PyErr>(()) + }) + }) + .await??; + Ok(()) + } + + /// Saves current recording and starts new file (rolling) + fn rolling(&self, file_name: &str) -> Result<(), Box> { + // Forward to Python client.save() - this saves current recording + // and the Python client handles starting a new recording file + let py_client = self.py_client.clone(); + let file_name = file_name.to_string(); + Python::attach(|py| { + let py_client = py_client.bind(py); + py_client.call_method1("save", (file_name,))?; + Ok::<_, PyErr>(()) + })?; + Ok(()) + } +} diff --git a/crates/stream-gears/src/lib.rs b/crates/stream-gears/src/lib.rs index 21883c4..40e2a2a 100644 --- a/crates/stream-gears/src/lib.rs +++ b/crates/stream-gears/src/lib.rs @@ -1,9 +1,11 @@ -extern crate core; - +mod danmaku; mod login; mod server; mod uploader; +// Export DanmakuClient for external use +pub use danmaku::DanmakuClient; + use pyo3::prelude::*; use time::macros::format_description; use uploader::{PyCredit, StudioPre}; @@ -26,6 +28,8 @@ use tracing_subscriber::EnvFilter; use biliup::client::StatelessClient; use biliup::downloader::flv_parser::header; use biliup::downloader::httpflv::Connection; +use pyo3::exceptions::PyRuntimeError; +use tracing_appender::rolling::Rotation; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; @@ -39,7 +43,13 @@ fn init_tracing_with_rotation() { "[year]-[month]-[day] [hour]:[minute]:[second]" )); // 按日期滚动，每天创建新文件 - let file_appender = tracing_appender::rolling::daily("logs", "upload.log"); + let file_appender = tracing_appender::rolling::RollingFileAppender::builder() + .rotation(Rotation::DAILY) // rotate log files once every hour + .filename_prefix("biliup") // log file names will be prefixed with `myapp.` + .filename_suffix("log") // log file names will be suffixed with `.log` + .max_log_files(3) + .build("logs") // try to build an appender that stores log files in `/var/log` + .expect("initializing rolling file appender failed"); // 或者按小时滚动 // let file_appender = tracing_appender::rolling::hourly("logs", "upload.log"); @@ -446,6 +456,32 @@ fn upload( }) } +#[pyfunction] +pub fn main_loop(py: Python<'_>) -> PyResult<()> { + py.detach(|| server::_main().map_err(|e| PyRuntimeError::new_err(format!("{e:?}")))) + + // for item in plugin_list.iter() { + // // 每个元素都是一个类（type） + // let ty: &Bound = item.downcast()?; + // + // // 确认是否是 DownloadBase 的子类 + // if ty.is_subclass(download_base)? { + // let name: String = ty.getattr("__name__")?.extract()?; + // println!("发现插件类: {name}"); + // + // // 如果要实例化（若不是抽象类且构造器无参） + // // let obj = ty.call0()?; + // // 或者需要参数：ty.call1((arg1, arg2,))?; + // + // // 如果要调用方法（示例） + // // if let Ok(method) = obj.getattr("download") { + // // let result = method.call0()?; + // // println!("download() 返回: {}", result.repr()?.extract::()?); + // // } + // } + // } +} + /// A Python module implemented in Rust. #[pymodule] fn stream_gears(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -466,7 +502,7 @@ fn stream_gears(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(login_by_sms, m)?)?; m.add_function(wrap_pyfunction!(login_by_web_cookies, m)?)?; m.add_function(wrap_pyfunction!(login_by_web_qrcode, m)?)?; - m.add_function(wrap_pyfunction!(server::main_loop, m)?)?; + m.add_function(wrap_pyfunction!(main_loop, m)?)?; m.add_function(wrap_pyfunction!(server::config_bindings, m)?)?; m.add_class::()?; m.add_class::()?; diff --git a/crates/stream-gears/src/server.rs b/crates/stream-gears/src/server.rs index 929acd0..2e22819 100644 --- a/crates/stream-gears/src/server.rs +++ b/crates/stream-gears/src/server.rs @@ -1,8 +1,9 @@ -use crate::construct_headers; +use crate::{DanmakuClient, construct_headers}; use async_trait::async_trait; use biliup::downloader::util::Segmentable; use biliup_cli::server::app::ApplicationController; -use biliup_cli::server::config::Config; +use biliup_cli::server::common::util::{Recorder, media_ext_from_url, parse_time}; +use biliup_cli::server::config::{default_segment_time, Config}; use biliup_cli::server::core::download_manager::{ActorHandle, DownloadManager}; use biliup_cli::server::core::downloader::ffmpeg_downloader::FfmpegDownloader; use biliup_cli::server::core::downloader::stream_gears::StreamGears; @@ -10,10 +11,10 @@ use biliup_cli::server::core::downloader::{DownloadConfig, Downloader, Downloade use biliup_cli::server::core::plugin::{DownloadPlugin, StreamInfo, StreamStatus}; use biliup_cli::server::errors::{AppError, AppResult}; use biliup_cli::server::infrastructure::connection_pool::ConnectionManager; -use biliup_cli::server::infrastructure::context::Worker; +use biliup_cli::server::infrastructure::context::{Context, Worker}; use biliup_cli::server::infrastructure::repositories; +use biliup_cli::server::infrastructure::repositories::get_upload_config; use biliup_cli::server::infrastructure::service_register::ServiceRegister; -use biliup_cli::server::util::{Recorder, media_ext_from_url, parse_time}; use error_stack::{Report, ResultExt}; use fancy_regex::Regex; use pyo3::exceptions::PyRuntimeError; @@ -68,13 +69,20 @@ impl DownloadPlugin for PyPlugin { false } - async fn check_status(&self, url: &str) -> Result> { + async fn check_status(&self, ctx: &mut Context) -> Result> { info!("Checking status"); - match call_via_threads(self.plugin.clone(), url) + match call_via_threads(self.plugin.clone(), &ctx.worker.live_streamer.url) .await .change_context(AppError::Unknown)? { - Some(info) => Ok(StreamStatus::Live { stream_info: info }), + Some((info, danmaku)) => { + if let Some(danmaku) = danmaku { + let danmaku = + Arc::new(DanmakuClient::new(Arc::new(danmaku))) as Arc; + ctx.extension.insert(danmaku); + } + Ok(StreamStatus::Live { stream_info: info }) + } None => Ok(StreamStatus::Offline), } } @@ -82,45 +90,33 @@ impl DownloadPlugin for PyPlugin { async fn create_downloader( &self, stream_info: &StreamInfo, - worker: &Worker, - ) -> AppResult> { - let config = worker.get_config().await?; - let streamer = worker.get_streamer().await?; - // info!(stream_info=?stream_info, "Create downloader"); + // worker: &Worker, + config: Config, + recorder: Recorder, + ) -> Box { let raw_stream_url = &stream_info.raw_stream_url; - let suffix = streamer - .format - .unwrap_or_else(|| stream_info.suffix.clone()); - println!("suffix: {suffix}"); - let recorder = Recorder::new( - streamer.filename_prefix, - &streamer.remark, - &stream_info.title, - &suffix, - ); - match config.downloader { Some(DownloaderType::Ffmpeg) => { let config = DownloadConfig { - format: suffix, - segment_time: config.segment_time, + segment_time: config.segment_time.or_else(default_segment_time), file_size: Some(config.file_size), // 2GB headers: stream_info.stream_headers.clone(), - extra_args: vec![], - downloader_type: DownloaderType::FfmpegInternal, - filename_prefix: recorder.generate_filename(), + recorder: recorder, + // output_dir: PathBuf::from("./downloads") + output_dir: PathBuf::from("."), }; - Ok(Box::new(FfmpegDownloader::new( + Box::new(FfmpegDownloader::new( raw_stream_url, config, - PathBuf::from("./downloads"), - ))) + Vec::new(), + DownloaderType::FfmpegInternal, + )) } // Some(DownloaderType::StreamGears) => { // // }, - _ => Ok(Box::new(StreamGears::new( + _ => Box::new(StreamGears::new( raw_stream_url, construct_headers(&stream_info.stream_headers), recorder.filename_template(), @@ -129,20 +125,26 @@ impl DownloadPlugin for PyPlugin { Some(config.file_size), ), None, - ))), + )), } } + fn danmaku_init(&self) -> Option> { + todo!() + } + fn name(&self) -> &str { &self.name } } -async fn call_via_threads(obj: Arc>, url: &str) -> PyResult> { +async fn call_via_threads( + obj: Arc>, + url: &str, +) -> PyResult>)>> { let url = url.to_string(); - // obj. tokio::task::spawn_blocking(move || { - Python::attach(|py| -> PyResult> { + Python::attach(|py| -> PyResult>)>> { // 从 biliup.util 获取 loop（按你项目里真实的名字来取） let util = PyModule::import(py, "biliup.common.util")?; // 下面两行二选一（取决于 biliup.util 的 API）： @@ -211,7 +213,9 @@ pub fn from_py(actor_handle: Arc) -> PyResult> /// 从 Python 的 `self`（Bound）与 start_time / end_time 构造 StreamInfo /// end_time 的语义与 Python 中一致：若未提供或为“假值”，则使用 time.localtime() -pub fn stream_info_from_py(self_obj: &Bound<'_, PyAny>) -> PyResult { +pub fn stream_info_from_py( + self_obj: &Bound<'_, PyAny>, +) -> PyResult<(StreamInfo, Option>)> { // 从 self 上获取属性并抽取为 Rust 类型 let name: String = self_obj.getattr("fname")?.extract()?; let url: String = self_obj.getattr("url")?.extract()?; @@ -222,6 +226,15 @@ pub fn stream_info_from_py(self_obj: &Bound<'_, PyAny>) -> PyResult let platform: String = self_obj.getattr("platform")?.extract()?; let stream_headers: HashMap = self_obj.getattr("stream_headers")?.extract()?; self_obj.call_method1("update_headers", (&stream_headers,))?; + let danmaku_init = self_obj.call_method0("danmaku_init")?; + // let platform: Option = self_obj.getattr("danmaku")?.extract()?; + // danmaku 可能在条件下没有设置（比如 bilibili_danmaku 为 False） + let self_danmaku = self_obj.getattr("danmaku")?; + let danmaku = if !self_danmaku.is_none() { + Some(self_danmaku.unbind()) + } else { + None + }; // date 直接使用传入的 start_time（保留为 Python 对象） let date = OffsetDateTime::now_utc(); @@ -235,17 +248,20 @@ pub fn stream_info_from_py(self_obj: &Bound<'_, PyAny>) -> PyResult // } // };self.update_headers(self.stream_headers) - Ok(StreamInfo { - name, - url, - suffix: media_ext_from_url(&raw_stream_url).unwrap(), - raw_stream_url, - title, - date, - live_cover_path, - platform, - stream_headers, - }) + Ok(( + StreamInfo { + name, + url, + suffix: media_ext_from_url(&raw_stream_url).unwrap(), + raw_stream_url, + title, + date, + live_cover_path, + platform, + stream_headers, + }, + danmaku, + )) } #[pyclass] @@ -308,35 +324,8 @@ fn cfg_arc() -> &'static Arc> { &CONFIG } -#[pyfunction] -pub fn main_loop(py: Python<'_>) -> PyResult<()> { - py.detach(_main) - .map_err(|e| PyRuntimeError::new_err(format!("{e:?}"))) - - // for item in plugin_list.iter() { - // // 每个元素都是一个类（type） - // let ty: &Bound = item.downcast()?; - // - // // 确认是否是 DownloadBase 的子类 - // if ty.is_subclass(download_base)? { - // let name: String = ty.getattr("__name__")?.extract()?; - // println!("发现插件类: {name}"); - // - // // 如果要实例化（若不是抽象类且构造器无参） - // // let obj = ty.call0()?; - // // 或者需要参数：ty.call1((arg1, arg2,))?; - // - // // 如果要调用方法（示例） - // // if let Ok(method) = obj.getattr("download") { - // // let result = method.call0()?; - // // println!("download() 返回: {}", result.repr()?.extract::()?); - // // } - // } - // } -} - #[tokio::main] -async fn _main() -> AppResult<()> { +pub(crate) async fn _main() -> AppResult<()> { info!( "environment loaded and configuration parsed, initializing Postgres connection and running migrations..." ); @@ -359,8 +348,9 @@ async fn _main() -> AppResult<()> { // workers.push(Arc::new(Worker::new(streamer.id, service_register.pool.clone()))); if let Some(manager) = service_register.get_manager(&streamer.url) { let monitor = manager.ensure_monitor(); + let upload_config = get_upload_config(&service_register.pool, streamer.id).await?; let _ = service_register - .add_room(streamer.id, &streamer.url, monitor) + .add_room(monitor, streamer, upload_config) .await?; }; }