GithubBackup/bili-sync
1
0
Fork 0
mirror of https://github.com/amtoaer/bili-sync.git synced 2026-05-06 20:42:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 尝试将任务处理部分重构为 stream 写法，增补注释

Browse Source
This commit is contained in:
amtoaer
2025-01-23 17:13:51 +08:00
parent b4177d4ffc
commit aa88f97eff
6 changed files with 127 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/bili_sync/src/adapter/collection.rs
View File

@@ -12,7 +12,7 @@ use sea_orm::{DatabaseConnection, TransactionTrait, Unchanged};
use crate::adapter::{helper, VideoListModel}; use crate::adapter::{helper, VideoListModel};
use crate::bilibili::{self, BiliClient, Collection, CollectionItem, CollectionType, VideoInfo}; use crate::bilibili::{self, BiliClient, Collection, CollectionItem, CollectionType, VideoInfo};
use crate::utils::status::Status; use crate::utils::status::STATUS_COMPLETED;
#[async_trait] #[async_trait]
impl VideoListModel for collection::Model { impl VideoListModel for collection::Model {
@@ -38,7 +38,7 @@ impl VideoListModel for collection::Model {
video::Column::CollectionId video::Column::CollectionId
.eq(self.id) .eq(self.id)
.and(video::Column::Valid.eq(true)) .and(video::Column::Valid.eq(true))
.and(video::Column::DownloadStatus.lt(Status::handled())) .and(video::Column::DownloadStatus.lt(STATUS_COMPLETED))
.and(video::Column::Category.eq(2)) .and(video::Column::Category.eq(2))
.and(video::Column::SinglePage.is_not_null()) .and(video::Column::SinglePage.is_not_null())
.into_condition(), .into_condition(),

4
crates/bili_sync/src/adapter/favorite.rs
View File

@@ -12,7 +12,7 @@ use sea_orm::{DatabaseConnection, TransactionTrait, Unchanged};
use crate::adapter::{helper, VideoListModel}; use crate::adapter::{helper, VideoListModel};
use crate::bilibili::{self, BiliClient, FavoriteList, VideoInfo}; use crate::bilibili::{self, BiliClient, FavoriteList, VideoInfo};
use crate::utils::status::Status; use crate::utils::status::STATUS_COMPLETED;
#[async_trait] #[async_trait]
impl VideoListModel for favorite::Model { impl VideoListModel for favorite::Model {
@@ -38,7 +38,7 @@ impl VideoListModel for favorite::Model {
video::Column::FavoriteId video::Column::FavoriteId
.eq(self.id) .eq(self.id)
.and(video::Column::Valid.eq(true)) .and(video::Column::Valid.eq(true))
.and(video::Column::DownloadStatus.lt(Status::handled())) .and(video::Column::DownloadStatus.lt(STATUS_COMPLETED))
.and(video::Column::Category.eq(2)) .and(video::Column::Category.eq(2))
.and(video::Column::SinglePage.is_not_null()) .and(video::Column::SinglePage.is_not_null())
.into_condition(), .into_condition(),

4
crates/bili_sync/src/adapter/submission.rs
View File

@@ -13,7 +13,7 @@ use sea_orm::{DatabaseConnection, TransactionTrait, Unchanged};
use crate::adapter::helper::video_with_path; use crate::adapter::helper::video_with_path;
use crate::adapter::{helper, VideoListModel}; use crate::adapter::{helper, VideoListModel};
use crate::bilibili::{self, BiliClient, Submission, VideoInfo}; use crate::bilibili::{self, BiliClient, Submission, VideoInfo};
use crate::utils::status::Status; use crate::utils::status::STATUS_COMPLETED;
#[async_trait] #[async_trait]
impl VideoListModel for submission::Model { impl VideoListModel for submission::Model {
@@ -39,7 +39,7 @@ impl VideoListModel for submission::Model {
video::Column::SubmissionId video::Column::SubmissionId
.eq(self.id) .eq(self.id)
.and(video::Column::Valid.eq(true)) .and(video::Column::Valid.eq(true))
.and(video::Column::DownloadStatus.lt(Status::handled())) .and(video::Column::DownloadStatus.lt(STATUS_COMPLETED))
.and(video::Column::Category.eq(2)) .and(video::Column::Category.eq(2))
.and(video::Column::SinglePage.is_not_null()) .and(video::Column::SinglePage.is_not_null())
.into_condition(), .into_condition(),

4
crates/bili_sync/src/adapter/watch_later.rs
View File

@@ -13,7 +13,7 @@ use sea_orm::{DatabaseConnection, TransactionTrait, Unchanged};
use crate::adapter::helper::video_with_path; use crate::adapter::helper::video_with_path;
use crate::adapter::{helper, VideoListModel}; use crate::adapter::{helper, VideoListModel};
use crate::bilibili::{self, BiliClient, VideoInfo, WatchLater}; use crate::bilibili::{self, BiliClient, VideoInfo, WatchLater};
use crate::utils::status::Status; use crate::utils::status::STATUS_COMPLETED;
#[async_trait] #[async_trait]
impl VideoListModel for watch_later::Model { impl VideoListModel for watch_later::Model {
@@ -39,7 +39,7 @@ impl VideoListModel for watch_later::Model {
video::Column::WatchLaterId video::Column::WatchLaterId
.eq(self.id) .eq(self.id)
.and(video::Column::Valid.eq(true)) .and(video::Column::Valid.eq(true))
.and(video::Column::DownloadStatus.lt(Status::handled())) .and(video::Column::DownloadStatus.lt(STATUS_COMPLETED))
.and(video::Column::Category.eq(2)) .and(video::Column::Category.eq(2))
.and(video::Column::SinglePage.is_not_null()) .and(video::Column::SinglePage.is_not_null())
.into_condition(), .into_condition(),

101
crates/bili_sync/src/utils/status.rs
View File

@@ -2,78 +2,85 @@ use anyhow::Result;
static STATUS_MAX_RETRY: u32 = 0b100; static STATUS_MAX_RETRY: u32 = 0b100;
static STATUS_OK: u32 = 0b111; static STATUS_OK: u32 = 0b111;
pub static STATUS_COMPLETED: u32 = 1 << 31;
/// 用来表示下载的状态，不想写太多列了，所以仅使用一个 u32 表示。 /// 用来表示下载的状态，不想写太多列了，所以仅使用一个 u32 表示。
/// 从低位开始，固定每三位表示一种数据的状态，从 0b000 开始，每失败一次加一，最多 0b100（即重试 4 次）， /// 从低位开始，固定每三位表示一种子任务的状态。
/// 如果成功，将对应的三位设置为 0b111。 /// 子任务状态从 0b000 开始，每执行失败一次将状态加一，最多 0b100（即允许重试 4 次），该值定义为 STATUS_MAX_RETRY。
/// 当所有任务都成功或者由于尝试次数过多失败，为 status 最高位打上标记 1，将来不再继续尝试。 /// 如果子任务执行成功，将状态设置为 0b111，该值定义为 STATUS_OK。
/// 子任务达到最大失败次数或者执行成功时，认为该子任务已经完成。
/// 当所有子任务都已经完成时，为最高位打上标记 1，表示整个下载任务已经完成。
#[derive(Clone)] #[derive(Clone)]
pub struct Status(u32); pub struct Status(u32);
impl Status { impl Status {
/// 如果 status 整体大于等于 1 << 31，则表示任务已经被处理过，不再需要重试。
/// 数据库可以使用 status < Status::handled() 来筛选需要处理的内容。
pub const fn handled() -> u32 {
1 << 31
}
fn new(status: u32) -> Self { fn new(status: u32) -> Self {
Self(status) Self(status)
} }
/// 一般仅需要被内部调用，用来设置最高位的标记 /// 设置最高位的完成标记
fn set_flag(&mut self, handled: bool) { fn set_completed(&mut self, completed: bool) {
if handled { if completed {
self.0 |= 1 << 31; self.0 |= 1 << 31;
} else { } else {
self.0 &= !(1 << 31); self.0 &= !(1 << 31);
} }
} }
/// 从低到高检查状态，如果该位置的任务应该继续尝试执行，则返回 true，否则返回 false // 获取最高位的完成标记
fn should_run(&self, size: usize) -> Vec<bool> { fn get_completed(&self) -> bool {
(0..size).map(|x| self.check_continue(x)).collect() self.0 >> 31 == 1
} }
/// 如果任务的执行次数小于 STATUS_MAX_RETRY，说明可以继续运行 /// 获取某个子任务的状态
fn get_status(&self, offset: usize) -> u32 {
let helper = !0u32;
(self.0 & (helper << (offset * 3)) & (helper >> (32 - 3 * offset - 3))) >> (offset * 3)
}
// 将某个子任务的状态加一（在任务失败时使用）
fn plus_one(&mut self, offset: usize) {
self.0 += 1 << (3 * offset);
}
// 设置某个子任务的状态为 STATUS_OK（在任务成功时使用）
fn set_ok(&mut self, offset: usize) {
self.0 |= STATUS_OK << (3 * offset);
}
/// 检查某个子任务是否还应该继续执行，实际是检查该子任务的状态是否小于 STATUS_MAX_RETRY
fn check_continue(&self, offset: usize) -> bool { fn check_continue(&self, offset: usize) -> bool {
self.get_status(offset) < STATUS_MAX_RETRY self.get_status(offset) < STATUS_MAX_RETRY
} }
/// 根据任务结果更新状态，如果任务成功，设置为 STATUS_OK，否则加一 /// 依次检查所有子任务是否还应该继续执行，返回一个 bool 数组
fn should_run(&self, size: usize) -> Vec<bool> {
(0..size).map(|x| self.check_continue(x)).collect()
}
/// 根据子任务执行结果更新子任务的状态
/// 如果 Result 是 Ok，那么认为任务执行成功，将状态设置为 STATUS_OK
/// 如果 Result 是 Err，那么认为任务执行失败，将状态加一
fn set_result(&mut self, result: &Result<()>, offset: usize) {
if self.get_status(offset) < STATUS_MAX_RETRY {
match result {
Ok(_) => self.set_ok(offset),
Err(_) => self.plus_one(offset),
}
}
}
/// 根据任务结果更新状态，任务结果是一个 Result 数组，需要与子任务一一对应
/// 如果所有子任务都已经完成，那么打上最高位的完成标记
fn update_status(&mut self, result: &[Result<()>]) { fn update_status(&mut self, result: &[Result<()>]) {
for (i, res) in result.iter().enumerate() { for (i, res) in result.iter().enumerate() {
self.set_result(res, i); self.set_result(res, i);
} }
if self.should_run(result.len()).iter().all(|x| !x) { if self.should_run(result.len()).iter().all(|x| !x) {
// 所有任务都成功或者由于尝试次数过多失败，为 status 最高位打上标记，将来不再重试 // 所有任务都成功或者由于尝试次数过多失败，为 status 最高位打上标记，将来不再重试
self.set_flag(true) self.set_completed(true)
} }
} }
fn set_result(&mut self, result: &Result<()>, offset: usize) {
if result.is_ok() {
// 如果任务已经执行到最大次数，那么此时 Result 也是 Ok，此时不应该更新状态
if self.get_status(offset) < STATUS_MAX_RETRY {
self.set_ok(offset);
}
} else {
self.plus_one(offset);
}
}
fn plus_one(&mut self, offset: usize) {
self.0 += 1 << (3 * offset);
}
fn set_ok(&mut self, offset: usize) {
self.0 |= STATUS_OK << (3 * offset);
}
fn get_status(&self, offset: usize) -> u32 {
let helper = !0u32;
(self.0 & (helper << (offset * 3)) & (helper >> (32 - 3 * offset - 3))) >> (offset * 3)
}
} }
impl From<Status> for u32 { impl From<Status> for u32 {
@@ -82,7 +89,7 @@ impl From<Status> for u32 {
} }
} }
/// 从前到后分别表示：视频封面、视频信息、Up 主头像、Up 主信息、分 P 下载 /// 包含五个子任务，从前到后依次是：视频封面、视频信息、Up 主头像、Up 主信息、分 P 下载
#[derive(Clone)] #[derive(Clone)]
pub struct VideoStatus(Status); pub struct VideoStatus(Status);
@@ -107,7 +114,7 @@ impl From<VideoStatus> for u32 {
} }
} }
/// 从前到后分别表示：视频封面、视频内容、视频信息 /// 包含四个子任务，从前到后分别是：视频封面、视频内容、视频信息、视频弹幕
#[derive(Clone)] #[derive(Clone)]
pub struct PageStatus(Status); pub struct PageStatus(Status);
@@ -124,6 +131,10 @@ impl PageStatus {
assert!(result.len() == 4, "PageStatus should have 4 status"); assert!(result.len() == 4, "PageStatus should have 4 status");
self.0.update_status(result) self.0.update_status(result)
} }
pub fn get_completed(&self) -> bool {
self.0.get_completed()
}
} }
impl From<PageStatus> for u32 { impl From<PageStatus> for u32 {
@@ -149,6 +160,6 @@ mod test {
} }
status.update_status(&[Err(anyhow!("")), Ok(()), Ok(())]); status.update_status(&[Err(anyhow!("")), Ok(()), Ok(())]);
assert_eq!(status.should_run(3), vec![false, false, false]); assert_eq!(status.should_run(3), vec![false, false, false]);
assert_eq!(u32::from(status), 0b111_111_100 | Status::handled()); assert_eq!(u32::from(status), 0b111_111_100 | STATUS_COMPLETED);
} }
} }

129
crates/bili_sync/src/workflow.rs
View File

@@ -102,16 +102,16 @@ pub async fn download_unprocessed_videos(
connection: &DatabaseConnection, connection: &DatabaseConnection,
) -> Result<()> { ) -> Result<()> {
video_list_model.log_download_video_start(); video_list_model.log_download_video_start();
let unhandled_videos_pages = video_list_model.unhandled_video_pages(connection).await?;
let semaphore = Semaphore::new(CONFIG.concurrent_limit.video); let semaphore = Semaphore::new(CONFIG.concurrent_limit.video);
let downloader = Downloader::new(bili_client.client.clone()); let downloader = Downloader::new(bili_client.client.clone());
let mut uppers_mutex: HashMap<i64, (Mutex<()>, Mutex<()>)> = HashMap::new(); let mut uppers_mutex: HashMap<i64, (Mutex<()>, Mutex<()>)> = HashMap::new();
let unhandled_videos_pages = video_list_model.unhandled_video_pages(connection).await?;
for (video_model, _) in &unhandled_videos_pages { for (video_model, _) in &unhandled_videos_pages {
uppers_mutex uppers_mutex
.entry(video_model.upper_id) .entry(video_model.upper_id)
.or_insert_with(|| (Mutex::new(()), Mutex::new(()))); .or_insert_with(|| (Mutex::new(()), Mutex::new(())));
} }
let mut tasks = unhandled_videos_pages let tasks = unhandled_videos_pages
.into_iter() .into_iter()
.map(|(video_model, pages_model)| { .map(|(video_model, pages_model)| {
let upper_mutex = uppers_mutex.get(&video_model.upper_id).expect("upper mutex not found"); let upper_mutex = uppers_mutex.get(&video_model.upper_id).expect("upper mutex not found");
@@ -127,32 +127,32 @@ pub async fn download_unprocessed_videos(
) )
}) })
.collect::<FuturesUnordered<_>>(); .collect::<FuturesUnordered<_>>();
let mut models = Vec::with_capacity(10); let mut download_aborted = false;
while let Some(res) = tasks.next().await { let mut stream = tasks
match res { // 触发风控时设置 download_aborted 标记并终止流
Ok(model) => { .take_while(|res| {
models.push(model); if res
.as_ref()
.is_err_and(|e| e.downcast_ref::<DownloadAbortError>().is_some())
{
download_aborted = true;
} }
Err(e) => { futures::future::ready(!download_aborted)
if e.downcast_ref::<DownloadAbortError>().is_some() { })
error!("下载视频时触发风控，将终止收藏夹下所有下载任务，等待下一轮执行"); // 过滤掉没有触发风控的普通 Err，只保留正确返回的 Model
break; .filter_map(|res| futures::future::ready(res.ok()))
} // 将成功返回的 Model 按十个一组合并
} .chunks(10);
} while let Some(models) = stream.next().await {
// 满十个就写入数据库
if models.len() == 10 {
update_videos_model(std::mem::replace(&mut models, Vec::with_capacity(10)), connection).await?;
}
}
if !models.is_empty() {
update_videos_model(models, connection).await?; update_videos_model(models, connection).await?;
} }
if download_aborted {
error!("下载视频时触发风控，终止收藏夹下所有下载任务，等待下一轮执行");
}
video_list_model.log_download_video_end(); video_list_model.log_download_video_end();
Ok(()) Ok(())
} }
/// 暂时这样做，后面提取成上下文
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
pub async fn download_video_pages( pub async fn download_video_pages(
bili_client: &BiliClient, bili_client: &BiliClient,
@@ -232,7 +232,7 @@ pub async fn download_video_pages(
&video_model.bvid, &video_model.name, task_name, e &video_model.bvid, &video_model.name, task_name, e
), ),
}); });
if let Err(e) = results.into_iter().nth(4).expect("not enough results") { if let Err(e) = results.into_iter().nth(4).context("page download result not found")? {
if e.downcast_ref::<DownloadAbortError>().is_some() { if e.downcast_ref::<DownloadAbortError>().is_some() {
return Err(e); return Err(e);
} }
@@ -242,6 +242,7 @@ pub async fn download_video_pages(
Ok(video_active_model) Ok(video_active_model)
} }
/// 分发并执行分页下载任务，当且仅当所有分页成功下载或达到最大重试次数时返回 Ok，否则根据失败原因返回对应的错误
pub async fn dispatch_download_page( pub async fn dispatch_download_page(
should_run: bool, should_run: bool,
bili_client: &BiliClient, bili_client: &BiliClient,
@@ -254,57 +255,56 @@ pub async fn dispatch_download_page(
return Ok(()); return Ok(());
} }
let child_semaphore = Semaphore::new(CONFIG.concurrent_limit.page); let child_semaphore = Semaphore::new(CONFIG.concurrent_limit.page);
let mut tasks = pages let tasks = pages
.into_iter() .into_iter()
.map(|page_model| download_page(bili_client, video_model, page_model, &child_semaphore, downloader)) .map(|page_model| download_page(bili_client, video_model, page_model, &child_semaphore, downloader))
.collect::<FuturesUnordered<_>>(); .collect::<FuturesUnordered<_>>();
let mut models = Vec::with_capacity(10); let (mut download_aborted, mut error_occurred) = (false, false);
let (mut should_error, mut is_break) = (false, false); let mut stream = tasks
while let Some(res) = tasks.next().await { .take_while(|res| {
match res { match res {
Ok(model) => { Ok(model) => {
if let Set(status) = model.download_status { // 当前函数返回的是所有分页的下载状态，只要有任何一个分页返回新的下载状态标识位是 false，当前函数就应该认为是失败的
let status = PageStatus::new(status); if model
if status.should_run().iter().any(|v| *v) { .download_status
// 有一个分页没变成终止状态（即下载成功或者重试次数达到限制），就应该向上层传递 Error .try_as_ref()
should_error = true; .is_none_or(|status| !PageStatus::new(*status).get_completed())
{
error_occurred = true;
} }
} }
models.push(model); Err(e) => {
} if e.downcast_ref::<DownloadAbortError>().is_some() {
Err(e) => { download_aborted = true;
if e.downcast_ref::<DownloadAbortError>().is_some() { }
should_error = true;
is_break = true;
break;
} }
} }
} // 仅在发生风控时终止流，其它情况继续执行
if models.len() == 10 { futures::future::ready(!download_aborted)
update_pages_model(std::mem::replace(&mut models, Vec::with_capacity(10)), connection).await?; })
} .filter_map(|res| futures::future::ready(res.ok()))
} .chunks(10);
if !models.is_empty() { while let Some(models) = stream.next().await {
update_pages_model(models, connection).await?; update_pages_model(models, connection).await?;
} }
if should_error { if download_aborted {
if is_break { error!(
error!( "下载视频 {} - {} 的分页时触发风控，将异常向上传递...",
"下载视频 {} - {} 的分页时触发风控，将异常向上传递...", &video_model.bvid, &video_model.name
&video_model.bvid, &video_model.name );
); bail!(DownloadAbortError());
bail!(DownloadAbortError()); }
} else { if error_occurred {
error!( error!(
"下载视频 {} - {} 的分页时出现了错误，将在下一轮尝试重新处理", "下载视频 {} - {} 的分页时出现了错误，将在下一轮尝试重新处理",
&video_model.bvid, &video_model.name &video_model.bvid, &video_model.name
); );
bail!(ProcessPageError()); bail!(ProcessPageError());
}
} }
Ok(()) Ok(())
} }
/// 下载某个分页，未发生风控且正常运行时返回 Ok(Page::ActiveModel)，其中 status 字段存储了新的下载状态，发生风控时返回 DownloadAbortError
pub async fn download_page( pub async fn download_page(
bili_client: &BiliClient, bili_client: &BiliClient,
video_model: &video::Model, video_model: &video::Model,
@@ -312,10 +312,7 @@ pub async fn download_page(
semaphore: &Semaphore, semaphore: &Semaphore,
downloader: &Downloader, downloader: &Downloader,
) -> Result<page::ActiveModel> { ) -> Result<page::ActiveModel> {
let permit = semaphore.acquire().await; let _permit = semaphore.acquire().await.context("acquire semaphore failed")?;
if let Err(e) = permit {
return Err(e.into());
}
let mut status = PageStatus::new(page_model.download_status); let mut status = PageStatus::new(page_model.download_status);
let seprate_status = status.should_run(); let seprate_status = status.should_run();
let is_single_page = video_model.single_page.context("single_page is null")?; let is_single_page = video_model.single_page.context("single_page is null")?;
@@ -415,8 +412,8 @@ pub async fn download_page(
&video_model.bvid, &video_model.name, page_model.pid, task_name, e &video_model.bvid, &video_model.name, page_model.pid, task_name, e
), ),
}); });
// 查看下载视频的状态，该状态会影响上层是否 break // 如果下载视频时触发风控，直接返回 DownloadAbortError
if let Err(e) = results.into_iter().nth(1).expect("not enough results") { if let Err(e) = results.into_iter().nth(1).context("video download result not found")? {
if let Ok(BiliError::RiskControlOccurred) = e.downcast::<BiliError>() { if let Ok(BiliError::RiskControlOccurred) = e.downcast::<BiliError>() {
bail!(DownloadAbortError()); bail!(DownloadAbortError());
} }