okooooooo/biliup
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor monitor and download manager for improved cleanup

Browse Source 
Refactored the monitor system by merging RoomsHandle into Monitor, centralizing plugin and worker management, and improving monitor task spawning. Updated DownloadManager to use the new Monitor structure, added a cleanup method for graceful shutdown, and ensured worker statuses are reset on exit. Adjusted related modules and API endpoints to accommodate these changes, and improved code clarity and error handling throughout.
This commit is contained in:
ForgQi committed 2025-11-17 03:18:01 +08:00
1 parent 2810ac9a88
commit 5de3dcaa6d
9 files changed
+274 -263
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/biliup-cli/src/server/api/endpoints.rs
+1 -1
View File
@@ -151,8 +151,8 @@ pub async fn pause_streamers_endpoint(
let worker_status = w.downloader_status.read().unwrap().clone();
match worker_status {
WorkerStatus::Working(_) => {
managers.make_waker(id).await;
w.change_status(Stage::Download, WorkerStatus::Pause).await;
managers.make_waker(id).await;
info!(workers=?&w.live_streamer.url, "successfully pause live streamers");
}
WorkerStatus::Pause => {
crates/biliup-cli/src/server/app.rs
+5 -3
View File
@@ -13,6 +13,7 @@ use axum_login::{AuthManagerLayerBuilder, login_required};
use error_stack::ResultExt;
use std::net::SocketAddr;
use time::Duration;
use tokio::io::AsyncWriteExt;
use tokio::signal;
use tokio::task::AbortHandle;
use tower_http::cors::{AllowMethods, CorsLayer};
@@ -63,7 +64,7 @@ impl ApplicationController {
// 构建应用程序路由
// 是否启用登录保护
let mut app = server::router::router(service_register);
let mut app = server::router::router(service_register.clone());
if enable_login_guard {
app = app
.route_layer(login_required!(Backend)) // 添加登录验证中间件
@@ -92,7 +93,7 @@ impl ApplicationController {
.change_context(AppError::Unknown)?;
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
.with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle(), service_register))
.await
.change_context(AppError::Unknown)
.attach("error while starting API server")?;
@@ -122,7 +123,7 @@ impl ApplicationController {
}
/// 优雅关闭信号处理
async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
async fn shutdown_signal(deletion_task_abort_handle: AbortHandle, service_register: ServiceRegister) {
// 监听Ctrl+C信号
let ctrl_c = async {
signal::ctrl_c()
@@ -148,4 +149,5 @@ async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
_ = ctrl_c => { deletion_task_abort_handle.abort() },
_ = terminate => { deletion_task_abort_handle.abort() },
}
service_register.cleanup().await;
}
crates/biliup-cli/src/server/common/download.rs
+11 -11
View File
@@ -4,7 +4,7 @@ use crate::server::core::downloader::{
DanmakuClient, DownloadConfig, DownloadStatus, DownloaderRuntime, DownloaderType, SegmentEvent,
SegmentInfo,
};
use crate::server::core::monitor::RoomsHandle;
use crate::server::core::monitor::Monitor;
use crate::server::core::plugin::{DownloadBase, StreamStatus};
use crate::server::errors::{AppError, AppResult};
use crate::server::infrastructure::context::{Context, Stage, WorkerStatus};
@@ -190,7 +190,7 @@ impl DownloadTask {
mut ctx: Context,
mut processor: SegmentEventProcessor,
mut plugin: Box<dyn DownloadBase>,
) -> AppResult<DownloadStatus> {
) -> AppResult<()> {
let worker = ctx.worker.clone();
// 重试配置
let mut retry_count = 0;
@@ -239,7 +239,7 @@ impl DownloadTask {
}
if self.token.is_cancelled() {
info!("用户手动停止");
info!("task is cancelled");
break components;
}
@@ -268,10 +268,12 @@ impl DownloadTask {
// 清理资源
// 确保状态更新和资源清理
worker.change_status(Stage::Download, WorkerStatus::Idle).await;
worker
.change_status(Stage::Download, WorkerStatus::Idle)
.await;
self.done_notify.notify_one();
info!("Download task completed: {:?}", result);
result
Ok(())
}
async fn initialize_components(
@@ -381,7 +383,7 @@ pub struct DActor {
/// 上传消息发送器
sender: Sender<UploaderMessage>,
rooms_handle: Arc<RoomsHandle>,
rooms_handle: Arc<Monitor>,
}
impl DActor {
@@ -389,7 +391,7 @@ impl DActor {
pub fn new(
receiver: Receiver<DownloaderMessage>,
sender: Sender<UploaderMessage>,
rooms_handle: Arc<RoomsHandle>,
rooms_handle: Arc<Monitor>,
) -> Self {
Self {
receiver,
@@ -425,14 +427,12 @@ impl DActor {
// 初始化组件
let processor = SegmentEventProcessor::new(self.sender.clone(), ctx.clone());
// 更新工作器状态为工作中
// rooms_handle
// .toggle(worker.clone(), WorkerStatus::Working(task.clone()))
// .await;
worker.change_status(Stage::Download, WorkerStatus::Working(task.clone())).await;
process(&[], &ctx.worker.get_streamer().preprocessor).await;
let option = &ctx.worker.get_streamer().downloaded_processor;
let result = task.execute(ctx, processor, downloader).await;
let _ = task.execute(ctx, processor, downloader).await;
process(&[], option).await;
crates/biliup-cli/src/server/core/download_manager.rs
+32 -60
View File
@@ -1,9 +1,9 @@
use crate::server::common::download::{DActor, DownloaderMessage};
use crate::server::common::upload::UActor;
use crate::server::core::monitor::{Monitor, RoomsHandle};
use crate::server::core::monitor::Monitor;
use crate::server::core::plugin::DownloadPlugin;
use crate::server::infrastructure::connection_pool::ConnectionPool;
use crate::server::infrastructure::context::Worker;
use crate::server::infrastructure::context::{Stage, Worker, WorkerStatus};
use async_channel::{Sender, bounded};
use std::collections::HashMap;
use std::collections::hash_map::Entry;
@@ -16,11 +16,7 @@ use tracing::info;
pub struct DownloadManager {
/// 下载插件
// plugins: Vec<Arc<dyn DownloadPlugin + Send + Sync>>,
rooms_handle: Arc<RoomsHandle>,
monitor: Monitor,
monitors: RwLock<HashMap<String, JoinHandle<()>>>,
rooms_handle: Arc<Monitor>,
/// 下载信号量数量
download_semaphore: u32,
@@ -47,7 +43,7 @@ impl DownloadManager {
let mut d_kills = Vec::new();
let mut u_kills = Vec::new();
let rooms_handle = Arc::new(RoomsHandle::new());
let rooms_handle = Arc::new(Monitor::new(down_tx.clone(), pool.clone()));
// 创建下载Actor
for _ in 0..download_semaphore {
let mut d_actor = DActor::new(down_rx.clone(), up_tx.clone(), rooms_handle.clone());
@@ -61,8 +57,6 @@ impl DownloadManager {
u_kills.push(u_kill)
}
let monitor = Monitor::new(rooms_handle.clone(), down_tx.clone(), pool.clone());
Self {
download_semaphore,
update_semaphore,
@@ -70,8 +64,6 @@ impl DownloadManager {
d_kills,
u_kills,
rooms_handle,
monitor,
monitors: Default::default(),
}
}
@@ -80,26 +72,13 @@ impl DownloadManager {
tokio::spawn(async move {
let name = plugin.name().to_string();
rooms_handle.add_plugin(plugin).await;
info!("Added plugin[{}] to", name);
info!("Added plugin[{}]", name);
});
}
pub async fn add_room(&self, worker: Worker) -> Option<()> {
let arc = Arc::new(worker);
let dp = self.rooms_handle.add(arc.clone()).await?;
let platform_name = dp.name().to_owned();
match self.monitors.write().unwrap().entry(platform_name.clone()) {
Entry::Occupied(entry) => {}
Entry::Vacant(entry) => {
let monitor = self.monitor.clone();
entry.insert(tokio::spawn(async move {
monitor.start_monitor(&platform_name, dp).await;
}));
}
}
self.rooms_handle.add(arc.clone()).await?;
Some(())
}
@@ -110,13 +89,14 @@ impl DownloadManager {
pub async fn get_rooms(&self) -> Vec<Arc<Worker>> {
self.rooms_handle.get_all().await
}
/// 移出工作队列
pub async fn make_waker(&self, id: i64) {
self.rooms_handle.make_waker(id).await
}
pub async fn wake_waker(&self, id: i64) {
self.rooms_handle.wake_waker(id).await
self.rooms_handle.wake_waker(id).await;
}
pub async fn get_room_by_id(&self, id: i64) -> Option<Arc<Worker>> {
@@ -124,37 +104,29 @@ impl DownloadManager {
.get_all()
.await
.iter()
.find(|worker| worker.id() == id).cloned()
.find(|worker| worker.id() == id)
.cloned()
}
pub async fn cleanup(&self) {
let vec = self.rooms_handle.get_all().await;
for worker in vec {
worker.change_status(Stage::Download, WorkerStatus::Idle).await;
}
info!("Cleanup complete");
}
}
// /// Actor处理器
// /// 管理下载和上传Actor的生命周期
// pub struct ActorHandle {
//
//
// }
//
// impl ActorHandle {
// /// 创建新的Actor处理器实例
// ///
// /// # 参数
// /// * `download_semaphore` - 下载Actor数量
// /// * `update_semaphore` - 上传Actor数量
// pub fn new() -> Self {
//
// }
// }
//
// impl Drop for ActorHandle {
// fn drop(&mut self) {
// // 发送端随 ActorHandle 一起被 drop，会关闭通道（如果没有其他 sender 克隆）。
// // 为避免 tokio 任务在后台“挂着”，这里直接 abort。
// for h in &self.d_kills {
// h.abort();
// }
// for h in &self.u_kills {
// h.abort();
// }
// }
// }
impl Drop for DownloadManager {
fn drop(&mut self) {
// 发送端随 ActorHandle 一起被 drop，会关闭通道（如果没有其他 sender 克隆）。
// 为避免 tokio 任务在后台“挂着”，这里直接 abort。
for h in &self.d_kills {
h.abort();
}
for h in &self.u_kills {
h.abort();
}
info!("exit download manager");
}
}
crates/biliup-cli/src/server/core/downloader/ytdlp.rs
+15 -12
View File
@@ -150,9 +150,10 @@ impl YouTubeDownloader {
s.push_str(&format!("[vcodec~='^({})']", v));
}
if !self.cfg.is_live
&& let Some(f) = &self.cfg.max_filesize {
s.push_str(&format!("[filesize<{}]", f));
}
&& let Some(f) = &self.cfg.max_filesize
{
s.push_str(&format!("[filesize<{}]", f));
}
if let Some(h) = self.cfg.max_height {
s.push_str(&format!("[height<={}]", h));
}
@@ -183,9 +184,10 @@ impl YouTubeDownloader {
cmd.arg("--proxy").arg(proxy);
}
if !self.cfg.is_live
&& let Some(archive) = &self.cfg.download_archive {
cmd.arg("--download-archive").arg(archive);
}
&& let Some(archive) = &self.cfg.download_archive
{
cmd.arg("--download-archive").arg(archive);
}
// 自定义附加参数
for a in &self.cfg.extra_ytdlp_args {
@@ -241,12 +243,13 @@ impl YouTubeDownloader {
// 清理临时目录
if self.cfg.use_temp_dir_for_ytdlp
&& let Err(e) = fs::remove_dir_all(&download_dir).await {
warn!(
"清理残留文件失败，请手动删除: {}，原因: {e}",
download_dir.display()
);
}
&& let Err(e) = fs::remove_dir_all(&download_dir).await
{
warn!(
"清理残留文件失败，请手动删除: {}，原因: {e}",
download_dir.display()
);
}
Ok(())
}
crates/biliup-cli/src/server/core/monitor.rs
+185 -144
View File
@@ -10,132 +10,27 @@ use ormlite::Model;
use ormlite::model::ModelBuilder;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use tracing::{debug, error, info};
#[derive(Debug, Clone)]
pub struct Monitor {
rooms_handle: Arc<RoomsHandle>,
/// 下载消息发送器
down_sender: Sender<DownloaderMessage>,
pool: ConnectionPool,
}
impl Monitor {
pub fn new(
rooms_handle: Arc<RoomsHandle>,
down_sender: Sender<DownloaderMessage>,
pool: ConnectionPool,
) -> Self {
Self {
rooms_handle,
down_sender,
pool,
}
}
/// 启动客户端监控循环
///
/// # 参数
/// * `rooms_handle` - 房间处理器
/// * `plugin` - 下载插件
/// * `actor_handle` - Actor处理器
/// * `interval` - 监控间隔（秒）
pub(crate) async fn start_monitor(
&self,
platform_name: &str,
plugin: Arc<dyn DownloadPlugin + Send + Sync>,
) {
info!("start -> [{platform_name}]");
loop {
let mut interval = 30;
// 获取下一个要检查的房间
if let Some(room) = self.rooms_handle.next(platform_name).await {
let url = room.get_streamer().url;
interval = room.get_config().event_loop_interval;
let mut ctx = Context::new(room.clone(), self.pool.clone());
// 检查直播状态
let mut downloader = plugin.create_downloader(&mut ctx);
match downloader.check_stream().await {
Ok(StreamStatus::Live { mut stream_info }) => {
let sql_no_id = &stream_info.streamer_info;
let insert = match StreamerInfo::builder()
.url(sql_no_id.url.clone())
.name(room.live_streamer.remark.clone())
.title(sql_no_id.title.clone())
.date(sql_no_id.date)
.live_cover_path(sql_no_id.live_cover_path.clone())
.insert(&ctx.pool)
.await
{
Ok(insert) => insert,
Err(e) => {
error!(e=?e, "插入数据库失败");
continue;
}
};
info!(url = url, "room: is live -> 开播了");
// 更新状态为等待中
room.change_status(Stage::Download, WorkerStatus::Pending)
.await;
stream_info.streamer_info = insert;
let streamer = room.get_streamer();
// 确定文件格式后缀
let suffix = streamer
.format
.unwrap_or_else(|| stream_info.suffix.clone());
// 创建录制器
let recorder = Recorder::new(
streamer
.filename_prefix
.or(room.get_config().filename_prefix.clone()),
stream_info.streamer_info.clone(),
&suffix,
);
// 修改 ctx
ctx.stream_info = *stream_info;
ctx.recorder = recorder;
// 发送下载开始消息
if self
.down_sender
.send(DownloaderMessage::Start(downloader, ctx))
.await
.is_ok()
{
info!("成功开始录制 {}", url)
}
}
Ok(StreamStatus::Offline) => {
self.rooms_handle.wake_waker(room.id()).await;
debug!(url = ctx.worker.live_streamer.url, "未开播")
}
Err(e) => {
self.rooms_handle.wake_waker(room.id()).await;
error!(e=?e, ctx=ctx.worker.live_streamer.url,"检查直播间出错")
}
};
}
// 等待下一次检查
tokio::time::sleep(Duration::from_secs(interval)).await;
}
info!("exit -> [{platform_name}]")
}
}
use tracing::{debug, error, info, trace};
/// 房间处理器
/// 管理多个直播间的状态和操作
#[derive(Debug)]
pub struct RoomsHandle {
pub struct Monitor {
/// 消息发送器
sender: tokio::sync::mpsc::Sender<ActorMessage>,
/// Actor任务句柄
kill: JoinHandle<()>,
pool: ConnectionPool,
/// 下载消息发送器
down_sender: Sender<DownloaderMessage>,
monitors: RwLock<HashMap<String, JoinHandle<()>>>,
}
impl Drop for RoomsHandle {
impl Drop for Monitor {
/// 监控器销毁时的清理逻辑
fn drop(&mut self) {
let sender = self.sender.clone();
@@ -150,36 +45,129 @@ impl Drop for RoomsHandle {
}
}
impl Default for RoomsHandle {
fn default() -> Self {
Self::new()
}
}
impl RoomsHandle {
impl Monitor {
/// 创建新的房间处理器实例
///
/// # 参数
/// * `name` - 平台名称
pub fn new() -> Self {
pub fn new(down_sender: Sender<DownloaderMessage>, pool: ConnectionPool) -> Self {
// 创建消息通道
let (sender, receiver) = tokio::sync::mpsc::channel(1);
let mut actor = RoomsActor::new(receiver);
// 启动Actor任务
let kill = tokio::spawn(async move { actor.run().await });
Self { sender, kill }
Self {
sender,
kill,
pool,
down_sender,
monitors: Default::default(),
}
}
/// 启动客户端监控循环
///
/// # 参数
/// * `rooms_handle` - 房间处理器
/// * `plugin` - 下载插件
/// * `actor_handle` - Actor处理器
/// * `interval` - 监控间隔（秒）
pub(crate) async fn start_monitor(
self: &Arc<Self>,
platform_name: &str,
plugin: Arc<dyn DownloadPlugin + Send + Sync>,
) {
info!("start -> [{platform_name}]");
// 获取下一个要检查的房间
while let Some(room) = self.next(platform_name).await {
let url = room.get_streamer().url;
let interval = room.get_config().event_loop_interval;
let mut ctx = Context::new(room.clone(), self.pool.clone());
// 检查直播状态
let mut downloader = plugin.create_downloader(&mut ctx);
match downloader.check_stream().await {
Ok(StreamStatus::Live { mut stream_info }) => {
let sql_no_id = &stream_info.streamer_info;
let insert = match StreamerInfo::builder()
.url(sql_no_id.url.clone())
.name(room.live_streamer.remark.clone())
.title(sql_no_id.title.clone())
.date(sql_no_id.date)
.live_cover_path(sql_no_id.live_cover_path.clone())
.insert(&ctx.pool)
.await
{
Ok(insert) => insert,
Err(e) => {
error!(e=?e, "插入数据库失败");
continue;
}
};
info!(url = url, "room: is live -> 开播了");
// 更新状态为等待中
room.change_status(Stage::Download, WorkerStatus::Pending)
.await;
stream_info.streamer_info = insert;
let streamer = room.get_streamer();
// 确定文件格式后缀
let suffix = streamer
.format
.unwrap_or_else(|| stream_info.suffix.clone());
// 创建录制器
let recorder = Recorder::new(
streamer
.filename_prefix
.or(room.get_config().filename_prefix.clone()),
stream_info.streamer_info.clone(),
&suffix,
);
// 修改 ctx
ctx.stream_info = *stream_info;
ctx.recorder = recorder;
// 发送下载开始消息
if self
.down_sender
.send(DownloaderMessage::Start(downloader, ctx))
.await
.is_ok()
{
info!("成功开始录制 {}", url)
}
}
Ok(StreamStatus::Offline) => {
self.wake_waker(room.id()).await;
*room.downloader_status.write().unwrap() = WorkerStatus::Idle;
debug!(url = ctx.worker.live_streamer.url, "未开播")
}
Err(e) => {
self.wake_waker(room.id()).await;
*room.downloader_status.write().unwrap() = WorkerStatus::Idle;
error!(e=?e, ctx=ctx.worker.live_streamer.url,"检查直播间出错")
}
};
// 等待下一次检查
tokio::time::sleep(Duration::from_secs(interval)).await;
}
info!("exit -> [{platform_name}]")
}
/// 添加工作器到房间列表
///
/// # 参数
/// * `worker` - 要添加的工作器
pub async fn add(&self, worker: Arc<Worker>) -> Option<Arc<dyn DownloadPlugin + Send + Sync>> {
pub async fn add(
self: &Arc<Self>,
worker: Arc<Worker>,
) -> Option<Arc<dyn DownloadPlugin + Send + Sync>> {
let (send, recv) = oneshot::channel();
let msg = ActorMessage::Add(send, worker);
let msg = ActorMessage::Add(send, worker.clone());
let _ = self.sender.send(msg).await;
recv.await.expect("Actor task has been killed")
let plugin = recv.await.expect("Actor task has been killed")?;
self.rooms_handle_pool(plugin.clone());
Some(plugin)
}
/// 添加工作器到房间列表
@@ -254,7 +242,7 @@ impl RoomsHandle {
///
/// # 返回
/// 返回下一个工作器，如果没有则返回None
async fn next(&self, platform_name: &str) -> Option<Arc<Worker>> {
async fn next(self: &Arc<Self>, platform_name: &str) -> Option<Arc<Worker>> {
let (send, recv) = oneshot::channel();
let msg = ActorMessage::NextRoom {
respond_to: send,
@@ -267,21 +255,26 @@ impl RoomsHandle {
recv.await.expect("Actor task has been killed")
}
/// 切换工作器状态（在活跃和等待列表之间）
/// 放回工作队列
///
/// # 参数
/// * `worker` - 要切换的工作器
pub async fn wake_waker(&self, id: i64) {
pub async fn wake_waker(
self: &Arc<Self>,
id: i64,
) -> Option<Arc<dyn DownloadPlugin + Send + Sync>> {
let (send, recv) = oneshot::channel();
let msg = ActorMessage::WakeWaker(send, id);
// 忽略发送错误
let _ = self.sender.send(msg).await;
recv.await.expect("Actor task has been killed")
let plugin = recv.await.expect("Actor task has been killed")?;
self.rooms_handle_pool(plugin.clone());
Some(plugin)
}
/// 切换工作器状态（在活跃和等待列表之间）
/// 移出工作队列
///
/// # 参数
/// * `worker` - 要切换的工作器
@@ -294,6 +287,45 @@ impl RoomsHandle {
let _ = self.sender.send(msg).await;
recv.await.expect("Actor task has been killed")
}
fn spawn_monitor_task(
this: Arc<Self>,
plugin: Arc<dyn DownloadPlugin + Send + Sync>,
platform_name: String,
) -> JoinHandle<()> {
tokio::spawn(async move {
this.start_monitor(&platform_name, plugin).await;
})
}
fn rooms_handle_pool(self: &Arc<Self>, plugin: Arc<dyn DownloadPlugin + Send + Sync>) {
let platform_name = plugin.name().to_owned();
match self.monitors.write().unwrap().entry(platform_name.clone()) {
Entry::Occupied(mut entry) => {
// 已经有一个任务了，检查是否结束
if entry.get().is_finished() {
// 旧任务已经结束，重新 spawn 一个
let handle = Self::spawn_monitor_task(
Arc::clone(self),
plugin.clone(),
platform_name.clone(),
);
entry.insert(handle); // 替换旧的 JoinHandle
} else {
// 任务还在跑，不做任何事
}
}
Entry::Vacant(entry) => {
// 没有任务，正常 spawn
let handle = Self::spawn_monitor_task(
Arc::clone(self),
plugin.clone(),
platform_name.clone(),
);
entry.insert(handle);
}
}
}
}
/// Actor消息枚举
@@ -331,8 +363,11 @@ enum ActorMessage {
platform_name: String,
},
/// 放回工作队列
WakeWaker(oneshot::Sender<()>, i64),
/// 切换工作器状态
WakeWaker(
oneshot::Sender<Option<Arc<dyn DownloadPlugin + Send + Sync>>>,
i64,
),
/// 移出工作队列
MakeWaker(oneshot::Sender<()>, i64),
Shutdown,
}
@@ -392,9 +427,8 @@ impl RoomsActor {
let _ = respond_to.send(());
}
ActorMessage::WakeWaker(sender, id) => {
self.push_back(id);
// `let _ =` 忽略发送时的任何错误
let _ = sender.send(());
let _ = sender.send(self.push_back(id));
}
ActorMessage::Shutdown => {
return;
@@ -445,18 +479,20 @@ impl RoomsActor {
entry.insert(VecDeque::from([worker.clone()])); // 插入新值
}
}
info!("Added room [{}]", worker.live_streamer.url);
debug!("Added room [{}]", worker.live_streamer.url);
Some(plugin)
}
fn add_plugin(&mut self, plugin: Arc<dyn DownloadPlugin + Send + Sync>) {
self.plugins.push(plugin);
debug!("Added plugin size[{}]", self.plugins.len());
}
fn get_worker(&mut self, id: i64) -> Option<Arc<Worker>> {
self.all_workers
.iter()
.find(|worker| worker.id() == id).cloned()
.find(|worker| worker.id() == id)
.cloned()
}
fn get_by_platform(&mut self, platform_name: &str) -> Vec<Arc<Worker>> {
@@ -484,19 +520,18 @@ impl RoomsActor {
}
/// 放回工作队列
fn push_back(&mut self, id: i64) -> Option<()> {
fn push_back(&mut self, id: i64) -> Option<Arc<dyn DownloadPlugin + Send + Sync>> {
// 如果内部Vec是空的，迭代结束（虽然是循环迭代器，但空集合无法产生任何值）
let worker = self.get_worker(id)?;
let plugin = self.matches(&worker.live_streamer.url)?;
self.platforms
.get_mut(plugin.name())?
.push_back(worker.clone());
*worker.downloader_status.write().unwrap() = WorkerStatus::Idle;
Some(())
Some(plugin)
}
/// 移除工作队列
/// 移出工作队列
fn pop(&mut self, id: i64) {
for (_name, queue) in self.platforms.iter_mut() {
if let Some(pos) = queue.iter().position(|w| w.id() == id) {
@@ -531,9 +566,15 @@ impl RoomsActor {
/// # 返回
/// 如果URL匹配返回true，否则返回false
pub fn matches(&self, url: &str) -> Option<Arc<dyn DownloadPlugin + Send + Sync>> {
if let Some(plugin) = self.plugins.iter().next() {
plugin.matches(url);
return Some(plugin.clone());
for plugin in &self.plugins {
trace!(
platform_name = plugin.name(),
url = url,
"Found plugin for URL"
);
if plugin.matches(url) {
return Some(plugin.clone());
}
}
None
}
crates/biliup-cli/src/server/core/plugin/twitch.rs
+15 -15
View File
@@ -1,9 +1,7 @@
use crate::server::common::util::media_ext_from_url;
use crate::server::core::downloader::ffmpeg_downloader::FfmpegDownloader;
use crate::server::core::downloader::stream_gears::StreamGears;
use crate::server::core::downloader::streamlink::{
Platform, Streamlink, StreamlinkDownloader,
};
use crate::server::core::downloader::streamlink::{Platform, Streamlink, StreamlinkDownloader};
use crate::server::core::downloader::{DanmakuClient, DownloaderRuntime, DownloaderType};
use crate::server::core::plugin::{DownloadBase, DownloadPlugin, StreamInfoExt, StreamStatus};
use crate::server::errors::{AppError, AppResult};
@@ -261,10 +259,11 @@ impl TwitchDownloader {
for (index, gql) in gql_list.iter().enumerate() {
if let Some(data) = &gql.data
&& let Some(user) = &data.user
&& let Some(stream) = &user.stream
&& stream.stream_type.as_deref() == Some("live") {
live_urls.push(check_urls[index].clone());
}
&& let Some(stream) = &user.stream
&& stream.stream_type.as_deref() == Some("live")
{
live_urls.push(check_urls[index].clone());
}
}
Ok(live_urls)
@@ -313,14 +312,15 @@ impl TwitchUtils {
// 第二次重试时不使用 auth_token（因为已经失效）
if retry == 0
&& let Some(auth_token) = Self::get_auth_token() {
headers.insert(
"Authorization",
format!("OAuth {}", auth_token)
.parse()
.change_context(AppError::Unknown)?,
);
}
&& let Some(auth_token) = Self::get_auth_token()
{
headers.insert(
"Authorization",
format!("OAuth {}", auth_token)
.parse()
.change_context(AppError::Unknown)?,
);
}
let client = Client::new();
let resp = client
crates/biliup-cli/src/server/infrastructure/service_register.rs
+4 -15
View File
@@ -69,6 +69,10 @@ impl ServiceRegister {
self.client.clone(),
)
}
pub async fn cleanup(&self) {
self.managers.cleanup().await;
}
}
// impl FromRef<ServiceRegister> for ConnectionPool {
@@ -76,18 +80,3 @@ impl ServiceRegister {
// app_state.pool.clone()
// }
// }
// pub(crate) async fn add_room(&self, id: i64, pool: ConnectionPool) -> Arc<Worker> {
// let arc = self.ensure_monitor().rooms_handle.clone();
// let worker = Arc::new(Worker::new(id, pool, arc.clone()));
// arc.add(worker.clone()).await;
// worker
// }
//
// pub async fn del_room(&self, id: i64) {
// let len = self.ensure_monitor().rooms_handle.del(id).await;
// info!("{id} removed, remained len {len}");
// if len == 0 {
// *self.monitor.lock().unwrap() = None;
// }
// }
crates/stream-gears/src/server.rs
+6 -2
View File
@@ -33,7 +33,7 @@ use std::net::ToSocketAddrs;
use std::ops::Deref;
use std::sync::{Arc, LazyLock, RwLock};
use time::macros::format_description;
use tracing::{debug, info};
use tracing::{debug, info, warn};
use tracing_appender::rolling::Rotation;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::layer::SubscriberExt;
@@ -193,6 +193,7 @@ impl DownloadPlugin for PyPlugin {
}
false
}
fn create_downloader(&self, ctx: &mut Context) -> Box<dyn DownloadBase> {
let url = ctx.worker.live_streamer.url.to_string();
let remark = ctx.worker.live_streamer.remark.to_string();
@@ -468,13 +469,16 @@ pub(crate) async fn _main(args: &[String]) -> AppResult<()> {
for streamer in all_streamer {
// workers.push(Arc::new(Worker::new(streamer.id, service_register.pool.clone())));
let upload_config = get_upload_config(&conn_pool, streamer.id).await?;
let url = streamer.url.clone();
let worker = Worker::new(
streamer,
upload_config,
CONFIG.clone(),
service_register.client.clone(),
);
service_register.managers.add_room(worker).await;
if service_register.managers.add_room(worker).await.is_none() {
warn!(url = url, "Could not add room to manager");
}
}
info!("migrations successfully ran, initializing axum server...");