Refactor uploader logic and improve worker status handling

Moved upload and post-processing logic to a new common/upload.rs module and refactored UActor to use it. WorkerStatus now tracks the active downloader instance, with improved serialization and debug output. StreamGears downloader now supports cancellation via tokio-util's CancellationToken. API endpoints and service registration updated to match new worker and uploader logic. Various doc comments and error handling improved throughout the codebase.
This commit is contained in:
ForgQi committed 2025-09-24 01:26:11 +08:00
1 parent 267bda7360
commit 2a8c1d0bd6
21 files changed
+457 -304

No files matched your search

Generated
+3 -2
View File
@@ -427,6 +427,7 @@ dependencies = [
"thiserror 2.0.12",
"time",
"tokio",
"tokio-util",
"toml",
"tower",
"tower-http",
@@ -3971,9 +3972,9 @@ dependencies = [
[[package]]
name = "tokio-util"
version = "0.7.15"
version = "0.7.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "66a539a9ad6d5d281510d5bd368c973d636c02dbf8a67300bfb6b950696ad7df"
checksum = "14307c986784f72ef81c89db7d9e28d6ac26d16213b109ea501696195e6e3ce5"
dependencies = [
"bytes",
"futures-core",
+1
View File
@@ -56,6 +56,7 @@ tower-sessions = { version = "0.14.0", default-features = false, features = [
tower-sessions-sqlx-store = { version = "0.15.0", features = ["sqlite"] }
async-channel = { version = "2.5.0" }
nom = "8"
tokio-util = "0.7.16"
[features]
default = []
python-bridge = []
@@ -24,7 +24,7 @@ pub async fn archive_pre_endpoint(
.await
.change_context(AppError::Unknown)
.map_err(report_to_response)?;
// 尝试使用每个Cookie进行登录
for cookies in configurations {
if let Ok(bili) = login_by_cookies(cookies.value, None).await {
@@ -36,7 +36,7 @@ pub async fn archive_pre_endpoint(
));
}
}
// 没有可用的Cookie
Err(report_to_response(Report::from(AppError::Custom(
"无可用 cookie 文件".to_string(),
@@ -52,7 +52,7 @@ pub async fn get_myinfo_endpoint(
.await
.change_context(AppError::Custom(params["user"].to_string()))
.map_err(report_to_response)?;
// 获取用户信息
Ok(Json(
bili.my_info()
@@ -68,7 +68,7 @@ pub async fn get_proxy_endpoint(
Query(params): Query<HashMap<String, String>>,
) -> Result<Bytes, Response> {
info!(params = &params["url"]);
// 代理HTTP请求
client
.client
@@ -11,10 +11,10 @@ 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, debug_handler};
use biliup::credential::Credential;
use error_stack::{Report, ResultExt};
use ormlite::{Insert, Model};
@@ -44,7 +44,7 @@ pub async fn get_streamers_endpoint(
.iter()
.find_map(|worker| {
if worker.live_streamer.id == x.id {
Some(*worker.downloader_status.read().unwrap())
Some(format!("{:?}", *worker.downloader_status.read().unwrap()))
} else {
None
}
@@ -65,7 +65,6 @@ pub async fn post_streamers_endpoint(
info!("not supported url: {}", &payload.url);
return Err((StatusCode::BAD_REQUEST, "Not supported url").into_response());
};
let monitor = manager.ensure_monitor();
// You can insert the model directly.
let live_streamers = payload
@@ -77,14 +76,14 @@ pub async fn post_streamers_endpoint(
.await
.map_err(report_to_response)?;
service_register
.add_room(monitor, live_streamers.clone(), upload_config)
.add_room(&manager, live_streamers.clone(), upload_config)
.await
.map_err(report_to_response)?;
info!(workers=?service_register.workers, "successfully inserted new live streamers");
Ok(Json(live_streamers))
}
#[debug_handler(state = ServiceRegister)]
pub async fn put_streamers_endpoint(
State(service_register): State<ServiceRegister>,
State(pool): State<ConnectionPool>,
@@ -109,10 +108,8 @@ pub async fn put_streamers_endpoint(
.await
.map_err(report_to_response)?;
let monitor = manager.ensure_monitor();
service_register
.add_room(monitor, streamer.clone(), upload_config)
.add_room(&manager, streamer.clone(), upload_config)
.await
.map_err(report_to_response)?;
+1 -1
View File
@@ -76,7 +76,7 @@ impl ApplicationController {
.allow_methods(AllowMethods::any()),
)
.fallback(static_handler); // 静态文件处理回退
// 启动HTTP服务器
info!("routes initialized, listening on {}", addr);
let listener = tokio::net::TcpListener::bind(addr)
@@ -1,2 +1,3 @@
pub mod upload;
/// 通用工具函数
pub mod util;
@@ -0,0 +1,263 @@
use crate::server::core::download_manager::UActor;
use crate::server::core::downloader::SegmentEvent;
use crate::server::errors::{AppError, AppResult};
use crate::server::infrastructure::context::Worker;
use crate::server::infrastructure::models::UploadStreamer;
use crate::server::infrastructure::models::hook_step::process_video;
use async_channel::Receiver;
use biliup::bilibili::{BiliBili, ResponseData, Studio, Video};
use biliup::client::StatelessClient;
use biliup::credential::login_by_cookies;
use biliup::uploader::line::{Line, Probe};
use biliup::uploader::util::SubmitOption;
use biliup::uploader::{VideoFile, line};
use error_stack::ResultExt;
use futures::StreamExt;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::time::Instant;
use tracing::{error, info};
// 辅助结构体
struct UploadContext {
bilibili: BiliBili,
client: StatelessClient,
line: Line,
threads: usize,
upload_config: UploadStreamer,
}
#[derive(Default)]
struct UploadedVideos {
videos: Vec<Video>,
paths: Vec<PathBuf>,
}
pub async fn process_with_upload(
first_event: SegmentEvent,
rx: Receiver<SegmentEvent>,
worker: &Worker,
upload_config: UploadStreamer,
) -> AppResult<()> {
// 1. 初始化上传环境
let upload_context = initialize_upload_context(&worker, upload_config).await?;
// 2. 流水线处理视频上传
let uploaded_videos = pipeline_upload_videos(first_event, rx, &upload_context).await?;
// 3. 提交到B站
if !uploaded_videos.videos.is_empty() {
submit_to_bilibili(&upload_context, uploaded_videos.videos, &worker).await?;
}
// 4. 执行后处理
if !uploaded_videos.paths.is_empty() {
execute_postprocessor(uploaded_videos.paths, &worker).await?;
}
Ok(())
}
async fn initialize_upload_context(
worker: &Worker,
upload_config: UploadStreamer,
) -> AppResult<UploadContext> {
// 登录处理
let cookie_file = upload_config
.user_cookie
.clone()
.unwrap_or("cookies.json".to_string());
let bilibili = login_by_cookies(&cookie_file, None)
.await
.change_context(AppError::Unknown)?;
let config = worker.get_config();
// 获取上传线路
let line = get_upload_line(&worker.client, &config.lines).await?;
Ok(UploadContext {
bilibili,
client: worker.client.clone(),
line,
threads: worker.get_config().threads as usize,
upload_config,
})
}
async fn get_upload_line(client: &StatelessClient, line: &str) -> AppResult<Line> {
let line = match line {
"bda2" => line::bda2(),
"bda" => line::bda(),
"tx" => line::tx(),
"txa" => line::txa(),
"bldsa" => line::bldsa(),
"alia" => line::alia(),
_ => Probe::probe(&client.client).await.unwrap_or_default(),
};
Ok(line)
}
async fn pipeline_upload_videos(
first_event: SegmentEvent,
rx: Receiver<SegmentEvent>,
context: &UploadContext,
) -> AppResult<UploadedVideos> {
// let mut desc_v2 = Vec::new();
// for credit in desc_v2_credit {
// desc_v2.push(Credit {
// type_id: credit.type_id,
// raw_text: credit.raw_text,
// biz_id: credit.biz_id,
// });
// }
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?;
uploaded.videos.push(video);
uploaded.paths.push(event.file_path);
// 失败的文件不加入路径列表,避免后处理出错
}
Ok(uploaded)
}
async fn upload_single_file(file_path: &Path, context: &UploadContext) -> AppResult<Video> {
let bilibili = &context.bilibili;
let limit = context.threads;
let client = &context.client;
let line = &context.line;
let video_path = file_path;
println!(
"{:?}",
video_path
.canonicalize()
.change_context(AppError::Unknown)?
.to_str()
);
info!("{line:?}");
let video_file = VideoFile::new(video_path).change_context(AppError::Unknown)?;
let total_size = video_file.total_size;
let file_name = video_file.file_name.clone();
let uploader = line
.pre_upload(bilibili, video_file)
.await
.change_context(AppError::Unknown)?;
let instant = Instant::now();
let video = uploader
.upload(client.clone(), limit, |vs| {
vs.map(|vs| {
let chunk = vs?;
let len = chunk.len();
Ok((chunk, len))
})
})
.await
.change_context(AppError::Unknown)?;
let t = instant.elapsed().as_millis();
info!(
"Upload completed: {file_name} => cost {:.2}s, {:.2} MB/s.",
t as f64 / 1000.,
total_size as f64 / 1000. / t as f64
);
Ok(video)
}
async fn submit_to_bilibili(
context: &UploadContext,
videos: Vec<Video>,
worker: &Worker,
) -> AppResult<ResponseData> {
let studio = build_studio(context, videos).await?;
// let submit = match worker.config.read().unwrap().submit_api {
// Some(submit) => SubmitOption::from_str(&submit).unwrap_or(SubmitOption::App),
// _ => SubmitOption::App,
// };
// let submit_result = match submit {
// SubmitOption::BCutAndroid => {
// bilibili.submit_by_bcut_android(&studio, None).await
// }
// _ => bilibili.submit_by_app(&studio, None).await,
// };
let submit_option = worker
.config
.read()
.unwrap()
.submit_api
.as_deref()
.and_then(|s| SubmitOption::from_str(s).ok())
.unwrap_or(SubmitOption::App);
let result = match submit_option {
SubmitOption::BCutAndroid => context
.bilibili
.submit_by_bcut_android(&studio, None)
.await
.change_context(AppError::Unknown)?,
_ => context
.bilibili
.submit_by_app(&studio, None)
.await
.change_context(AppError::Unknown)?,
};
info!("Submit successful");
Ok(result)
}
async fn build_studio(context: &UploadContext, videos: Vec<Video>) -> AppResult<Studio> {
let upload_config = &context.upload_config;
// 使用 Builder 模式简化构建
let mut studio: Studio = Studio::builder()
.desc(upload_config.description.clone().unwrap_or_default())
.dtime(upload_config.dtime)
.copyright(upload_config.copyright.unwrap_or(2))
.cover(upload_config.cover_path.clone().unwrap_or_default())
.dynamic(upload_config.dynamic.clone().unwrap_or_default())
.source(upload_config.copyright_source.clone().unwrap_or_default())
.tag(upload_config.tags.join(","))
.tid(upload_config.tid.unwrap_or(171))
.title(upload_config.title.clone().unwrap_or_default())
.videos(videos)
.dolby(upload_config.dolby.unwrap_or_default())
// .lossless_music(upload_config.)
.no_reprint(upload_config.no_reprint.unwrap_or_default())
.charging_pay(upload_config.charging_pay.unwrap_or_default())
.up_close_reply(upload_config.up_close_reply.unwrap_or_default())
.up_selection_reply(upload_config.up_selection_reply.unwrap_or_default())
.up_close_danmu(upload_config.up_close_danmu.unwrap_or_default())
.desc_v2(None)
.extra_fields(
serde_json::from_str(&upload_config.extra_fields.clone().unwrap_or_default())
.unwrap_or_default(), // 处理额外字段
)
.build();
// 处理封面上传
if !studio.cover.is_empty()
&& let Ok(c) = &std::fs::read(&studio.cover).map_err(|e| error!(e=?e))
&& let Ok(url) = context.bilibili.cover_up(c).await.map_err(|e| error!(e=?e))
{
studio.cover = url;
};
Ok(studio)
}
pub async fn execute_postprocessor(video_paths: Vec<PathBuf>, worker: &Worker) -> AppResult<()> {
if let Some(processor) = worker.get_streamer().postprocessor {
let paths: Vec<&Path> = video_paths.iter().map(|p| p.as_path()).collect();
process_video(&paths, &processor).await?;
}
Ok(())
}
@@ -1,3 +1,4 @@
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::monitor::{Monitor, RoomsHandle};
@@ -35,7 +36,7 @@ pub struct DownloadManager {
impl DownloadManager {
/// 创建新的下载管理器实例
///
///
/// # 参数
/// * `plugin` - 下载插件实现
/// * `actor_handle` - Actor处理器
@@ -51,7 +52,7 @@ impl DownloadManager {
}
/// 确保监控器存在,如果不存在则创建新的
///
///
/// # 返回
/// 返回监控器的Arc引用
pub fn ensure_monitor(&self) -> Arc<Monitor> {
@@ -65,10 +66,10 @@ impl DownloadManager {
}
/// 检查URL是否匹配此下载管理器的插件
///
///
/// # 参数
/// * `url` - 要检查的URL
///
///
/// # 返回
/// 如果URL匹配返回true,否则返回false
pub fn matches(&self, url: &str) -> bool {
@@ -107,7 +108,7 @@ impl DActor {
}
/// 处理下载消息
///
///
/// # 参数
/// * `msg` - 要处理的下载消息
async fn handle_message(&mut self, msg: DownloaderMessage) -> AppResult<()> {
@@ -134,7 +135,7 @@ impl DActor {
// 更新工作器状态为工作中
ctx.worker
.change_status(Stage::Download, WorkerStatus::Working);
.change_status(Stage::Download, WorkerStatus::Working(downloader.clone()));
// 初始化弹幕客户端(如果存在)
let mut danmaku_client = None;
if let Some(danmaku) = ctx.extension.get::<Arc<dyn Downloader>>() {
@@ -228,177 +229,27 @@ impl UActor {
}
/// 处理上传消息
///
///
/// # 参数
/// * `msg` - 要处理的上传消息
async fn handle_message(&mut self, msg: UploaderMessage) {
match msg {
UploaderMessage::SegmentEvent(event, rx, worker) => {
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.inspect_err(|e| error!(e=?e)) else {
return;
};
let config = worker.get_config();
let client = worker.client.clone();
let line = config.lines;
let line = match line.as_str() {
"bda2" => line::bda2(),
"bda" => line::bda(),
"tx" => line::tx(),
"txa" => line::txa(),
"bldsa" => line::bldsa(),
"alia" => line::alia(),
_ => Probe::probe(&client.client).await.unwrap_or_default(),
};
let mut video_paths = Vec::new();
let mut videos = Vec::new();
if let Ok(video) = UActor::upload_file(
&bilibili,
config.threads as usize,
&client,
&line,
&event.file_path,
)
.await
.inspect_err(|e| error!(e=?e))
{
videos.push(video);
}
video_paths.push(event.file_path);
while let Ok(se) = rx.recv().await {
if let Ok(video) = UActor::upload_file(
&bilibili,
config.threads as usize,
&client,
&line,
&se.file_path,
)
.await
.map_err(|e| error!(e=?e))
{
videos.push(video);
let result = match worker.get_upload_config() {
Some(config) => process_with_upload(event, rx, &worker, config).await,
None => {
// 无上传配置时,直接执行后处理
execute_postprocessor(vec![event.file_path], &worker).await
}
video_paths.push(se.file_path);
}
// let mut desc_v2 = Vec::new();
// for credit in desc_v2_credit {
// desc_v2.push(Credit {
// type_id: credit.type_id,
// raw_text: credit.raw_text,
// biz_id: credit.biz_id,
// });
// }
let mut studio: Studio = Studio::builder()
.desc(upload_config.description.unwrap_or_default())
.dtime(upload_config.dtime)
.copyright(upload_config.copyright.unwrap_or(2))
.cover(upload_config.cover_path.unwrap_or_default())
.dynamic(upload_config.dynamic.unwrap_or_default())
.source(upload_config.copyright_source.unwrap_or_default())
.tag(upload_config.tags.join(","))
.tid(upload_config.tid.unwrap_or(171))
.title(upload_config.title.unwrap_or_default())
.videos(videos)
.dolby(upload_config.dolby.unwrap_or_default())
// .lossless_music(upload_config.)
.no_reprint(upload_config.no_reprint.unwrap_or_default())
.charging_pay(upload_config.charging_pay.unwrap_or_default())
.up_close_reply(upload_config.up_close_reply.unwrap_or_default())
.up_selection_reply(upload_config.up_selection_reply.unwrap_or_default())
.up_close_danmu(upload_config.up_close_danmu.unwrap_or_default())
.desc_v2(None)
.extra_fields(
serde_json::from_str(&upload_config.extra_fields.unwrap_or_default())
.unwrap_or_default(),
)
.build();
if !studio.cover.is_empty()
&& let Ok(c) = &std::fs::read(&studio.cover).map_err(|e| error!(e=?e))
&& let Ok(url) = bilibili.cover_up(c).await.map_err(|e| error!(e=?e))
{
studio.cover = url;
};
let submit = match config.submit_api {
Some(submit) => SubmitOption::from_str(&submit).unwrap_or(SubmitOption::App),
_ => SubmitOption::App,
};
let submit_result = match submit {
SubmitOption::BCutAndroid => {
bilibili.submit_by_bcut_android(&studio, None).await
}
_ => bilibili.submit_by_app(&studio, None).await,
};
info!(submit_result=?submit_result);
if submit_result.is_ok() {
info!("Submit successful");
let streamer = worker.get_streamer();
if let Some(post_processor) = streamer.postprocessor {
let _ = process_video(&(video_paths.iter().map(|p|p.as_path()).collect::<Vec<_>>()), &post_processor)
.await
.map_err(|e| error!(e=?e));
}
if let Err(e) = result {
error!("Process segment event failed: {}", e);
// 可以添加错误通知机制
}
}
}
}
async fn upload_file(
bilibili: &BiliBili,
limit: usize,
client: &StatelessClient,
line: &Line,
video_path: &Path,
) -> AppResult<Video> {
println!(
"{:?}",
video_path
.canonicalize()
.change_context(AppError::Unknown)?
.to_str()
);
info!("{line:?}");
let video_file = VideoFile::new(video_path).change_context(AppError::Unknown)?;
let total_size = video_file.total_size;
let file_name = video_file.file_name.clone();
let uploader = line
.pre_upload(bilibili, video_file)
.await
.change_context(AppError::Unknown)?;
let instant = Instant::now();
let video = uploader
.upload(client.clone(), limit, |vs| {
vs.map(|vs| {
let chunk = vs?;
let len = chunk.len();
Ok((chunk, len))
})
})
.await
.change_context(AppError::Unknown)?;
let t = instant.elapsed().as_millis();
info!(
"Upload completed: {file_name} => cost {:.2}s, {:.2} MB/s.",
t as f64 / 1000.,
total_size as f64 / 1000. / t as f64
);
Ok(video)
}
}
/// Actor处理器
@@ -420,7 +271,7 @@ pub struct ActorHandle {
impl ActorHandle {
/// 创建新的Actor处理器实例
///
///
/// # 参数
/// * `download_semaphore` - 下载Actor数量
/// * `update_semaphore` - 上传Actor数量
@@ -32,7 +32,7 @@ pub struct DownloadConfig {
impl DownloadConfig {
/// 生成输出文件名
///
///
/// # 返回
/// 返回完整的输出文件路径
fn generate_output_filename(&self) -> PathBuf {
@@ -98,10 +98,10 @@ pub enum DownloadStatus {
#[async_trait]
pub trait Downloader: Send + Sync {
/// 开始下载
///
///
/// # 参数
/// * `callback` - 分段完成时的回调函数
///
///
/// # 返回
/// 返回下载状态
async fn download(
@@ -113,7 +113,7 @@ pub trait Downloader: Send + Sync {
async fn stop(&self) -> Result<(), Box<dyn std::error::Error>>;
/// 滚动保存(用于弹幕等)
///
///
/// # 参数
/// * `file_name` - 文件名
fn rolling(&self, file_name: &str) -> Result<(), Box<dyn std::error::Error>> {
@@ -122,10 +122,10 @@ pub trait Downloader: Send + Sync {
}
/// 解析时长字符串 "HH:MM:SS" 为秒数
///
///
/// # 参数
/// * `duration` - 时长字符串,格式为"HH:MM:SS"
///
///
/// # 返回
/// 返回总秒数
fn parse_duration(duration: &str) -> u64 {
@@ -7,11 +7,12 @@ 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 error_stack::{ResultExt, bail};
use nom::Err;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::SystemTime;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info};
/// Stream-gears下载器实现
@@ -29,11 +30,13 @@ pub struct StreamGears {
proxy: Option<String>,
/// 下载状态
status: Arc<RwLock<DownloadStatus>>,
token: CancellationToken,
}
impl StreamGears {
/// 创建新的Stream-gears下载器实例
///
///
/// # 参数
/// * `url` - 流URL
/// * `header_map` - HTTP请求头
@@ -54,17 +57,11 @@ impl StreamGears {
segment,
proxy,
status: Arc::new(RwLock::new(DownloadStatus::Downloading)),
token: CancellationToken::new(),
}
}
}
#[async_trait]
impl Downloader for StreamGears {
/// 开始下载流
///
/// # 参数
/// * `callback` - 分段完成时的回调函数
async fn download(
async fn start_download(
&self,
callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
) -> AppResult<DownloadStatus> {
@@ -73,7 +70,7 @@ impl Downloader for StreamGears {
let _headers_in = self.header_map.clone();
let proxy = self.proxy.clone();
let _status = Arc::clone(&self.status);
// 创建HTTP客户端
let client = StatelessClient::new(self.header_map.clone(), proxy.as_deref());
// 获取可重试的响应
@@ -124,11 +121,31 @@ impl Downloader for StreamGears {
}
Ok(DownloadStatus::StreamEnded)
}
}
#[async_trait]
impl Downloader for StreamGears {
/// 开始下载流
///
/// # 参数
/// * `callback` - 分段完成时的回调函数
async fn download(
&self,
callback: Box<dyn Fn(SegmentEvent) + Send + Sync + 'static>,
) -> AppResult<DownloadStatus> {
tokio::select! {
_ = self.token.cancelled() => {
bail!(AppError::Custom("StreamGears token cancelled".into()))
}
res = self.start_download(callback) => {res}
}
}
/// 停止下载
async fn stop(&self) -> Result<(), Box<dyn std::error::Error>> {
// 仅发出取消信号并更新状态
// 如果底层下载函数不支持取消,这里不能真正中断正在进行的下载
todo!()
self.token.cancel();
Ok(())
}
}
+8 -8
View File
@@ -9,7 +9,7 @@ use tokio::task::JoinHandle;
use tracing::info;
/// 启动客户端监控循环
///
///
/// # 参数
/// * `rooms_handle` - 房间处理器
/// * `plugin` - 下载插件
@@ -73,7 +73,7 @@ pub struct Monitor {
impl Monitor {
/// 创建新的监控器实例
///
///
/// # 参数
/// * `plugin` - 下载插件
/// * `actor_handle` - Actor处理器
@@ -121,7 +121,7 @@ pub struct RoomsHandle {
impl RoomsHandle {
/// 创建新的房间处理器实例
///
///
/// # 参数
/// * `name` - 平台名称
pub fn new(name: &str) -> Self {
@@ -139,7 +139,7 @@ impl RoomsHandle {
}
/// 添加工作器到房间列表
///
///
/// # 参数
/// * `worker` - 要添加的工作器
pub async fn add(&self, worker: Arc<Worker>) {
@@ -148,10 +148,10 @@ impl RoomsHandle {
}
/// 删除指定ID的工作器
///
///
/// # 参数
/// * `id` - 要删除的工作器ID
///
///
/// # 返回
/// 返回剩余工作器数量
pub async fn del(&self, id: i64) -> usize {
@@ -168,7 +168,7 @@ impl RoomsHandle {
}
/// 获取下一个要处理的工作器
///
///
/// # 返回
/// 返回下一个工作器,如果没有则返回None
async fn next(&self) -> Option<Arc<Worker>> {
@@ -182,7 +182,7 @@ impl RoomsHandle {
}
/// 切换工作器状态(在活跃和等待列表之间)
///
///
/// # 参数
/// * `worker` - 要切换的工作器
pub async fn toggle(&self, worker: Arc<Worker>) {
+2 -1
View File
@@ -8,6 +8,7 @@ use error_stack::Report;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use time::OffsetDateTime;
/// 流信息结构
@@ -77,7 +78,7 @@ pub trait DownloadPlugin {
stream_info: &StreamInfo,
config: Config,
recorder: Recorder,
) -> Box<dyn Downloader>;
) -> Arc<dyn Downloader>;
/// 初始化弹幕客户端(可选)
fn danmaku_init(&self) -> Option<Box<dyn Downloader>>;
@@ -14,10 +14,10 @@ pub struct ConnectionManager;
impl ConnectionManager {
/// 创建新的数据库连接池
///
///
/// # 参数
/// * `path` - 数据库文件路径
///
///
/// # 返回
/// 返回配置好的SQLite连接池
pub async fn new_pool(path: &str) -> AppResult<ConnectionPool> {
@@ -27,16 +27,16 @@ impl ConnectionManager {
.change_context(AppError::Unknown)
.attach_with(|| path.to_string())?; // 创建 data/ 目录
}
let db_url = format!("sqlite://{path}");
// 创建数据库文件(如果不存在)
std::fs::OpenOptions::new()
.write(true)
.create(true)
.open(path)
.change_context(AppError::Unknown)?;
// 创建连接池,最大连接数设为2
let pool = SqlitePoolOptions::new()
.max_connections(2)
@@ -1,13 +1,15 @@
use crate::server::config::Config;
use crate::server::core::downloader::Downloader;
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 core::fmt;
use error_stack::ResultExt;
use ormlite::Model;
use serde::{Deserialize, Serialize};
use serde::{Deserialize, Serialize, Serializer};
use std::sync::{Arc, RwLock};
use tracing::info;
@@ -22,7 +24,7 @@ pub struct Context {
impl Context {
/// 创建新的上下文实例
///
///
/// # 参数
/// * `worker` - 工作器实例的Arc引用
pub fn new(worker: Arc<Worker>) -> Self {
@@ -52,7 +54,7 @@ pub struct Worker {
impl Worker {
/// 创建新的工作器实例
///
///
/// # 参数
/// * `live_streamer` - 直播主播信息
/// * `upload_streamer` - 上传配置(可选)
@@ -93,7 +95,7 @@ impl Worker {
}
/// 更改工作器状态
///
///
/// # 参数
/// * `stage` - 工作阶段(下载或上传)
/// * `status` - 新的工作状态
@@ -109,6 +111,10 @@ impl Worker {
}
}
pub fn find_worker(workers: &[Arc<Worker>], id: i64) -> Option<&Arc<Worker>> {
workers.iter().find(|worker| worker.live_streamer.id == id)
}
impl Drop for Worker {
/// 工作器销毁时的清理逻辑
fn drop(&mut self) {
@@ -135,13 +141,25 @@ pub enum Stage {
}
/// 工作器状态枚举
#[derive(Clone, Copy, Debug, Serialize, Deserialize, Default, PartialEq)]
#[derive(Default)]
pub enum WorkerStatus {
/// 正在工作
Working,
Working(Arc<dyn Downloader>),
/// 等待中
Pending,
/// 空闲状态(默认)
#[default]
Idle,
}
// 简单 Debug:打印状态名,忽略内部 downloader
impl fmt::Debug for WorkerStatus {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let name = match self {
WorkerStatus::Working(_) => "Working",
WorkerStatus::Pending => "Pending",
WorkerStatus::Idle => "Idle",
};
f.write_str(name)
}
}
@@ -4,12 +4,12 @@ use serde::{Deserialize, Serialize};
/// 直播主播响应数据传输对象
/// 包含主播信息和当前工作状态
#[derive(Serialize, Deserialize)]
#[derive(Serialize)]
pub struct LiveStreamerResponse {
/// 主播基本信息(展开到顶层)
#[serde(flatten)]
pub inner: LiveStreamer,
/// 当前工作状态
pub status: WorkerStatus,
pub status: String,
}
@@ -1,5 +1,7 @@
use std::borrow::Cow;
use crate::server::errors::{AppError, AppResult};
use error_stack::{ResultExt, bail};
use serde::{Deserialize, Serialize};
use std::borrow::Cow;
use std::path::Path;
use tokio::fs;
use tokio::io::AsyncWriteExt;
@@ -21,13 +23,13 @@ pub enum HookStep {
impl HookStep {
/// 执行钩子步骤操作
///
///
/// # 参数
/// * `video_paths` - 视频文件路径列表
///
///
/// # 返回
/// 执行成功返回Ok(()),失败返回错误信息
pub async fn execute(&self, video_paths: &[&Path]) -> Result<(), Box<dyn std::error::Error>> {
pub async fn execute(&self, video_paths: &[&Path]) -> AppResult<()> {
match self {
HookStep::Run { run } => {
// 执行自定义命令
@@ -43,35 +45,32 @@ impl HookStep {
}
HookStep::Remove(cmd) => {
// 未知命令,返回错误
return Err(format!("Unknown command: {}", cmd).into());
bail!(AppError::Custom(format!("Unknown command: {}", cmd)));
}
}
Ok(())
}
/// 执行自定义命令,将视频路径作为标准输入传入
///
///
/// # 参数
/// * `cmd` - 要执行的命令字符串
/// * `video_paths` - 视频文件路径列表
async fn execute_command(
&self,
cmd: &str,
video_paths: &[&Path],
) -> Result<(), Box<dyn std::error::Error>> {
async fn execute_command(&self, cmd: &str, video_paths: &[&Path]) -> AppResult<()> {
println!("Executing: {}", cmd);
// 解析命令和参数
let parts: Vec<&str> = cmd.split_whitespace().collect();
if parts.is_empty() {
return Err("Empty command".into());
bail!(AppError::Custom("Empty command".into()));
}
// 启动子进程,配置标准输入管道
let mut process = Command::new(parts[0])
.args(&parts[1..])
.stdin(std::process::Stdio::piped())
.spawn()?;
.spawn()
.change_context(AppError::Unknown)?;
// 将视频文件路径写入标准输入
if let Some(mut stdin) = process.stdin.take() {
@@ -80,38 +79,43 @@ impl HookStep {
.map(|p| p.to_string_lossy())
.reduce(|acc, e| Cow::from(acc.to_string() + "\n" + &*e))
.expect("At least one video path should exist");
stdin.write_all(paths_str.as_bytes()).await?;
stdin
.write_all(paths_str.as_bytes())
.await
.change_context(AppError::Unknown)?;
}
// 等待进程完成并检查退出状态
let status = process.wait().await?;
let status = process.wait().await.change_context(AppError::Unknown)?;
if !status.success() {
return Err(format!("Command failed with status: {}", status).into());
bail!(AppError::Custom(
format!("Command failed with status: {}", status).into()
));
}
Ok(())
}
/// 移动文件到指定目录
///
///
/// # 参数
/// * `video_paths` - 视频文件路径列表
/// * `target_dir` - 目标目录路径
async fn move_file(
&self,
video_paths: &[&Path],
target_dir: &str,
) -> Result<(), Box<dyn std::error::Error>> {
async fn move_file(&self, video_paths: &[&Path], target_dir: &str) -> AppResult<()> {
let target_path = Path::new(target_dir);
// 创建目标目录(如果不存在)
if !target_path.exists() {
fs::create_dir_all(target_path).await?;
fs::create_dir_all(target_path)
.await
.change_context(AppError::Unknown)?;
}
// 移动每个视频文件到目标目录
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(AppError::Custom("Invalid file name".to_string()))?;
let destination = target_path.join(file_name);
println!(
@@ -119,20 +123,24 @@ impl HookStep {
video_path.display(),
destination.display()
);
fs::rename(video_path, destination).await?;
fs::rename(video_path, destination)
.await
.change_context(AppError::Unknown)?;
}
Ok(())
}
/// 删除指定的视频文件
///
///
/// # 参数
/// * `video_paths` - 要删除的视频文件路径列表
async fn remove_file(&self, video_paths: &[&Path]) -> Result<(), Box<dyn std::error::Error>> {
async fn remove_file(&self, video_paths: &[&Path]) -> AppResult<()> {
// 逐个删除视频文件
for video_path in video_paths {
println!("Removing: {}", video_path.display());
fs::remove_file(video_path).await?;
fs::remove_file(video_path)
.await
.change_context(AppError::Unknown)?;
}
Ok(())
}
@@ -141,21 +149,18 @@ impl HookStep {
/// 处理所有后处理器步骤
/// 处理视频文件的主函数
/// 按顺序执行所有处理器步骤
///
///
/// # 参数
/// * `video_path` - 视频文件路径列表
/// * `processors` - 处理器步骤列表
pub async fn process_video(
video_path: &[&Path],
processors: &[HookStep],
) -> Result<(), Box<dyn std::error::Error>> {
pub async fn process_video(video_path: &[&Path], processors: &[HookStep]) -> AppResult<()> {
info!("Starting video processing...");
// 依次执行每个处理器步骤
for processor in processors {
processor.execute(video_path).await?;
}
info!("Video processing completed");
Ok(())
}
@@ -8,7 +8,7 @@ use error_stack::ResultExt;
use ormlite::{Insert, Model};
/// 根据ID获取直播主播信息
///
///
/// # 参数
/// * `pool` - 数据库连接池
/// * `id` - 主播ID
@@ -22,7 +22,7 @@ pub async fn get_streamer(pool: &ConnectionPool, id: i64) -> AppResult<LiveStrea
}
/// 获取主播的上传配置
///
///
/// # 参数
/// * `pool` - 数据库连接池
/// * `id` - 主播ID
@@ -43,11 +43,11 @@ pub async fn get_upload_config(
}
/// 删除指定的直播主播
///
///
/// # 参数
/// * `pool` - 数据库连接池
/// * `id` - 主播ID
///
///
/// # 返回
/// 返回被删除的主播信息
pub async fn del_streamer(pool: &ConnectionPool, id: i64) -> AppResult<LiveStreamer> {
@@ -60,7 +60,7 @@ pub async fn del_streamer(pool: &ConnectionPool, id: i64) -> AppResult<LiveStrea
Ok(streamer)
}
/// 获取所有直播主播信息
///
///
/// # 参数
/// * `pool` - 数据库连接池
pub async fn get_all_streamer(pool: &ConnectionPool) -> AppResult<Vec<LiveStreamer>> {
@@ -71,10 +71,10 @@ pub async fn get_all_streamer(pool: &ConnectionPool) -> AppResult<Vec<LiveStream
}
/// 从数据库获取全局配置
///
///
/// # 参数
/// * `pool` - 数据库连接池
///
///
/// # 返回
/// 返回全局配置,如果不存在则返回默认配置
pub async fn get_config(pool: &ConnectionPool) -> AppResult<Config> {
@@ -96,7 +96,7 @@ pub async fn get_config(pool: &ConnectionPool) -> AppResult<Config> {
}
/// 插入或更新全局配置到数据库
///
///
/// # 参数
/// * `pool` - 数据库连接池
/// * `config` - 要保存的配置
@@ -112,7 +112,7 @@ pub async fn insert_config(pool: &ConnectionPool, config: &Config) -> AppResult<
}
/// 获取所有上传配置
///
///
/// # 参数
/// * `pool` - 数据库连接池
pub async fn get_all_uploader(pool: &ConnectionPool) -> AppResult<Vec<UploadStreamer>> {
@@ -3,7 +3,7 @@ use crate::server::core::download_manager::{ActorHandle, DownloadManager};
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::context::{Context, Worker, WorkerStatus};
use crate::server::infrastructure::models::{LiveStreamer, UploadStreamer};
use axum::extract::FromRef;
use biliup::client::StatelessClient;
@@ -32,7 +32,7 @@ pub struct ServiceRegister {
/// 简单的服务容器,负责管理API端点通过axum扩展获取的各种服务
impl ServiceRegister {
/// 创建新的服务注册器实例
///
///
/// # 参数
/// * `pool` - 数据库连接池
/// * `config` - 全局配置
@@ -65,10 +65,10 @@ impl ServiceRegister {
}
/// 根据URL获取匹配的下载管理器
///
///
/// # 参数
/// * `url` - 直播流URL
///
///
/// # 返回
/// 返回匹配的下载管理器引用,如果没有匹配的则返回None
pub fn get_manager(&self, url: &str) -> Option<&DownloadManager> {
@@ -79,14 +79,14 @@ impl ServiceRegister {
}
/// 添加新的直播间到监控列表
///
///
/// # 参数
/// * `monitor` - 监控器实例
/// * `live_streamer` - 直播主播信息
/// * `upload_streamer` - 上传配置(可选)
pub async fn add_room(
&self,
monitor: Arc<Monitor>,
manager: &DownloadManager,
live_streamer: LiveStreamer,
upload_streamer: Option<UploadStreamer>,
) -> AppResult<Option<()>> {
@@ -98,6 +98,7 @@ impl ServiceRegister {
self.client.clone(),
));
// 将工作器添加到监控器和工作器列表中
let monitor = manager.ensure_monitor();
monitor.rooms_handle.add(worker.clone()).await;
self.workers.write().unwrap().push(worker.clone());
info!("add {worker:?} success");
@@ -105,7 +106,7 @@ impl ServiceRegister {
}
/// 删除指定ID的直播间
///
///
/// # 参数
/// * `id` - 要删除的直播间ID
pub async fn del_room(&self, id: i64) -> AppResult<()> {
@@ -136,7 +137,13 @@ impl ServiceRegister {
if len == 0 {
*manager.monitor.lock().unwrap() = None;
}
let downloader = match &*removed.downloader_status.read().unwrap() {
WorkerStatus::Working(downloader) => Some(downloader.clone()),
_ => None,
};
if let Some(downloader) = downloader {
let _ = downloader.stop().await;
}
Ok(())
}
}
@@ -71,7 +71,7 @@ impl Backend {
}
/// 检查是否存在用户
///
///
/// # 返回
/// 如果存在用户返回true,否则返回false
pub async fn exists(&self) -> AppResult<bool> {
@@ -85,10 +85,10 @@ impl Backend {
}
/// 创建新用户
///
///
/// # 参数
/// * `creds` - 用户凭据
///
///
/// # 返回
/// 返回创建的用户信息
pub async fn create_user(&mut self, creds: Credentials) -> AppResult<User> {
@@ -141,10 +141,10 @@ impl AuthnBackend for Backend {
type Error = Error;
/// 认证用户
///
///
/// # 参数
/// * `creds` - 用户凭据
///
///
/// # 返回
/// 如果认证成功返回用户信息,否则返回None
async fn authenticate(
@@ -165,7 +165,7 @@ impl AuthnBackend for Backend {
}
/// 根据用户ID获取用户信息
///
///
/// # 参数
/// * `user_id` - 用户ID
async fn get_user(&self, user_id: &UserId<Self>) -> Result<Option<Self::User>, Self::Error> {
+10 -18
View File
@@ -18,45 +18,37 @@ pub fn router(service_register: ServiceRegister) -> Router<()> {
// 主播管理相关路由
.route(
"/v1/streamers",
get(get_streamers_endpoint) // 获取主播列表
.post(post_streamers_endpoint) // 添加主播
.put(put_streamers_endpoint), // 更新主播
get(get_streamers_endpoint) // 获取主播列表
.post(post_streamers_endpoint) // 添加主播
.put(put_streamers_endpoint), // 更新主播
)
.route("/v1/streamers/{id}", delete(delete_streamers_endpoint)) // 删除主播
// 配置管理路由
.route(
"/v1/configuration",
get(get_configuration).put(put_configuration), // 获取/更新配置
)
// 主播信息路由
.route("/v1/streamer-info", get(get_streamer_info)) // 获取主播信息
// 上传模板管理路由
.route("/v1/upload/streamers", get(get_upload_streamers_endpoint)) // 获取上传模板列表
.route(
"/v1/upload/streamers/{id}",
delete(delete_template_endpoint) // 删除上传模板
delete(delete_template_endpoint) // 删除上传模板
.get(get_upload_streamer_endpoint), // 获取单个上传模板
)
.route("/v1/upload/streamers", post(add_upload_streamer_endpoint)) // 添加上传模板
// 用户管理路由
.route("/v1/users", get(get_users_endpoint).post(add_user_endpoint)) // 获取用户列表/添加用户
.route("/v1/users/{id}", delete(delete_user_endpoint)) // 删除用户
// B站API代理路由
.route("/bili/archive/pre", get(archive_pre_endpoint)) // 投稿预处理
.route("/bili/space/myinfo", get(get_myinfo_endpoint)) // 获取用户信息
.route("/bili/proxy", get(get_proxy_endpoint)) // 代理请求
.route("/bili/archive/pre", get(archive_pre_endpoint)) // 投稿预处理
.route("/bili/space/myinfo", get(get_myinfo_endpoint)) // 获取用户信息
.route("/bili/proxy", get(get_proxy_endpoint)) // 代理请求
// 认证相关路由
.route("/v1/get_qrcode", get(get_qrcode)) // 获取二维码
.route("/v1/login_by_qrcode", post(login_by_qrcode)) // 二维码登录
.route("/v1/get_qrcode", get(get_qrcode)) // 获取二维码
.route("/v1/login_by_qrcode", post(login_by_qrcode)) // 二维码登录
// 视频文件管理路由
.route("/v1/videos", get(get_videos)) // 获取视频列表
.route("/v1/videos", get(get_videos)) // 获取视频列表
.with_state(service_register) // 注入服务注册器状态
}
+5 -6
View File
@@ -3,7 +3,7 @@ use async_trait::async_trait;
use biliup::downloader::util::Segmentable;
use biliup_cli::server::app::ApplicationController;
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::config::{Config, default_segment_time};
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;
@@ -93,7 +93,7 @@ impl DownloadPlugin for PyPlugin {
// worker: &Worker,
config: Config,
recorder: Recorder,
) -> Box<dyn Downloader> {
) -> Arc<dyn Downloader> {
let raw_stream_url = &stream_info.raw_stream_url;
match config.downloader {
Some(DownloaderType::Ffmpeg) => {
@@ -106,7 +106,7 @@ impl DownloadPlugin for PyPlugin {
output_dir: PathBuf::from("."),
};
Box::new(FfmpegDownloader::new(
Arc::new(FfmpegDownloader::new(
raw_stream_url,
config,
Vec::new(),
@@ -116,7 +116,7 @@ impl DownloadPlugin for PyPlugin {
// Some(DownloaderType::StreamGears) => {
//
// },
_ => Box::new(StreamGears::new(
_ => Arc::new(StreamGears::new(
raw_stream_url,
construct_headers(&stream_info.stream_headers),
recorder.filename_template(),
@@ -347,10 +347,9 @@ pub(crate) async fn _main() -> AppResult<()> {
for streamer in all_streamer {
// 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(monitor, streamer, upload_config)
.add_room(&manager, streamer, upload_config)
.await?;
};
}