Skip to content

Commit

Permalink
feat: 调整并重构视频音频流的选择逻辑，应该可以提升些许性能 (#212)
Browse files Browse the repository at this point in the history 
* feat: 调整并重构视频音频流的选择逻辑，应该可以提升些许性能

* test: 添加少量单元测试
  • Loading branch information
@amtoaer
amtoaer authored Jan 13, 2025
1 parent 05aa301 commit 7d9999d
Showing 1 changed file with 144 additions and 131 deletions.
275 changes: 144 additions & 131 deletions crates/bili_sync/src/bilibili/analyzer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ pub enum VideoQuality {
QualityDolby = 126,
Quality8k = 127,
}
#[derive(Debug, strum::FromRepr, PartialEq, PartialOrd, Serialize, Deserialize)]

#[derive(Debug, Clone, Copy, strum::FromRepr, PartialEq, Serialize, Deserialize)]
pub enum AudioQuality {
Quality64k = 30216,
Quality132k = 30232,
Expand All @@ -29,8 +30,25 @@ pub enum AudioQuality {
Quality192k = 30280,
}

impl AudioQuality {
#[inline]
pub fn as_sort_key(&self) -> isize {
match self {
// 这可以让 Dolby 和 Hi-RES 排在 192k 之后，且 Dolby 和 Hi-RES 之间的顺序不变
Self::QualityHiRES | Self::QualityDolby => (*self as isize) + 40,
_ => *self as isize,
}
}
}

impl PartialOrd<AudioQuality> for AudioQuality {
fn partial_cmp(&self, other: &AudioQuality) -> Option<std::cmp::Ordering> {
self.as_sort_key().partial_cmp(&other.as_sort_key())
}
}

#[allow(clippy::upper_case_acronyms)]
#[derive(Debug, strum::EnumString, strum::Display, PartialEq, PartialOrd, Serialize, Deserialize)]
#[derive(Debug, strum::EnumString, strum::Display, strum::AsRefStr, PartialEq, PartialOrd, Serialize, Deserialize)]
pub enum VideoCodecs {
#[strum(serialize = "hev")]
HEV,
Expand Down Expand Up @@ -115,26 +133,22 @@ impl PageAnalyzer {
}

fn is_flv_stream(&self) -> bool {
self.info.get("durl").is_some()
&& self.info["format"].is_string()
&& self.info["format"].as_str().unwrap().starts_with("flv")
self.info.get("durl").is_some() && self.info["format"].as_str().is_some_and(|f| f.starts_with("flv"))
}

fn is_html5_mp4_stream(&self) -> bool {
self.info.get("durl").is_some()
&& self.info["format"].is_string()
&& self.info["format"].as_str().unwrap().starts_with("mp4")
&& self.info["is_html5"].is_boolean()
&& self.info["is_html5"].as_bool().unwrap()
&& self.info["format"].as_str().is_some_and(|f| f.starts_with("mp4"))
&& self.info["is_html5"].as_bool().is_some_and(|b| b)
}

fn is_episode_try_mp4_stream(&self) -> bool {
self.info.get("durl").is_some()
&& self.info["format"].is_string()
&& self.info["format"].as_str().unwrap().starts_with("mp4")
&& !(self.info["is_html5"].is_boolean() && self.info["is_html5"].as_bool().unwrap())
&& self.info["format"].as_str().is_some_and(|f| f.starts_with("mp4"))
&& self.info["is_html5"].as_bool().is_none_or(|b| !b)
}

/// 获取所有的视频、音频流，并根据条件筛选
fn streams(&mut self, filter_option: &FilterOption) -> Result<Vec<Stream>> {
if self.is_flv_stream() {
return Ok(vec![Stream::Flv(
Expand All @@ -161,85 +175,78 @@ impl PageAnalyzer {
)]);
}
let mut streams: Vec<Stream> = Vec::new();
let videos_data = self.info["dash"]["video"].take();
let audios_data = self.info["dash"]["audio"].take();
let flac_data = self.info["dash"]["flac"].take();
let dolby_data = self.info["dash"]["dolby"].take();
for video_data in videos_data.as_array().ok_or(BiliError::RiskControlOccurred)?.iter() {
let video_stream_url = video_data["baseUrl"].as_str().unwrap().to_string();
let video_stream_quality = VideoQuality::from_repr(video_data["id"].as_u64().unwrap() as usize)
.ok_or(anyhow!("invalid video stream quality"))?;
if (video_stream_quality == VideoQuality::QualityHdr && filter_option.no_hdr)
|| (video_stream_quality == VideoQuality::QualityDolby && filter_option.no_dolby_video)
|| (video_stream_quality != VideoQuality::QualityDolby
&& video_stream_quality != VideoQuality::QualityHdr
&& (video_stream_quality < filter_option.video_min_quality
|| video_stream_quality > filter_option.video_max_quality))
// 此处过滤包含三种情况：
// 1. HDR 视频，但指定不需要 HDR
// 2. 杜比视界视频，但指定不需要杜比视界
// 3. 视频质量不在指定范围内
{
for video in self.info["dash"]["video"]
.as_array()
.ok_or(BiliError::RiskControlOccurred)?
.iter()
{
let (Some(url), Some(quality), Some(codecs)) = (
video["baseUrl"].as_str(),
video["id"].as_u64(),
video["codecs"].as_str(),
) else {
continue;
}

let video_codecs = video_data["codecs"].as_str().unwrap();
};
let quality = VideoQuality::from_repr(quality as usize).ok_or(anyhow!("invalid video stream quality"))?;
// 从视频流的 codecs 字段中获取编码格式，此处并非精确匹配而是判断包含，比如 codecs 是 av1.42c01e，需要匹配为 av1
let video_codecs = vec![VideoCodecs::HEV, VideoCodecs::AVC, VideoCodecs::AV1]
let codecs = [VideoCodecs::HEV, VideoCodecs::AVC, VideoCodecs::AV1]
.into_iter()
.find(|c| video_codecs.contains(c.to_string().as_str()));

let Some(video_codecs) = video_codecs else {
continue;
};
if !filter_option.codecs.contains(&video_codecs) {
.find(|c| codecs.contains(c.as_ref()))
.ok_or(anyhow!("invalid video stream codecs"))?;
if !filter_option.codecs.contains(&codecs)
|| quality < filter_option.video_min_quality
|| quality > filter_option.video_max_quality
|| (quality == VideoQuality::QualityHdr && filter_option.no_hdr)
|| (quality == VideoQuality::QualityDolby && filter_option.no_dolby_video)
{
continue;
}
streams.push(Stream::DashVideo {
url: video_stream_url,
quality: video_stream_quality,
codecs: video_codecs,
url: url.to_string(),
quality,
codecs,
});
}
if audios_data.is_array() {
for audio_data in audios_data.as_array().unwrap().iter() {
let audio_stream_url = audio_data["baseUrl"].as_str().unwrap().to_string();
let audio_stream_quality = AudioQuality::from_repr(audio_data["id"].as_u64().unwrap() as usize);
let Some(audio_stream_quality) = audio_stream_quality else {
if let Some(audios) = self.info["dash"]["audio"].as_array() {
for audio in audios.iter() {
let (Some(url), Some(quality)) = (audio["baseUrl"].as_str(), audio["id"].as_u64()) else {
continue;
};
if audio_stream_quality > filter_option.audio_max_quality
|| audio_stream_quality < filter_option.audio_min_quality
{
let quality =
AudioQuality::from_repr(quality as usize).ok_or(anyhow!("invalid audio stream quality"))?;
if quality < filter_option.audio_min_quality || quality > filter_option.audio_max_quality {
continue;
}
streams.push(Stream::DashAudio {
url: audio_stream_url,
quality: audio_stream_quality,
url: url.to_string(),
quality,
});
}
}
if !(filter_option.no_hires || flac_data["audio"].is_null()) {
// 允许 hires 且存在 flac 音频流才会进来
let flac_stream_url = flac_data["audio"]["baseUrl"].as_str().unwrap().to_string();
let flac_stream_quality =
AudioQuality::from_repr(flac_data["audio"]["id"].as_u64().unwrap() as usize).unwrap();
streams.push(Stream::DashAudio {
url: flac_stream_url,
quality: flac_stream_quality,
});
let flac = &self.info["dash"]["flac"]["audio"];
if !(filter_option.no_hires || flac.is_null()) {
let (Some(url), Some(quality)) = (flac["baseUrl"].as_str(), flac["id"].as_u64()) else {
bail!("invalid flac stream");
};
let quality = AudioQuality::from_repr(quality as usize).ok_or(anyhow!("invalid flac stream quality"))?;
if quality >= filter_option.audio_min_quality && quality <= filter_option.audio_max_quality {
streams.push(Stream::DashAudio {
url: url.to_string(),
quality,
});
}
}
if !(filter_option.no_dolby_audio || dolby_data["audio"].is_null()) {
// 同理，允许杜比音频且存在杜比音频流才会进来
let dolby_stream_data = dolby_data["audio"].as_array().and_then(|v| v.first());
if dolby_stream_data.is_some() {
let dolby_stream_data = dolby_stream_data.unwrap();
let dolby_stream_url = dolby_stream_data["baseUrl"].as_str().unwrap().to_string();
let dolby_stream_quality =
AudioQuality::from_repr(dolby_stream_data["id"].as_u64().unwrap() as usize).unwrap();
let dolby_audio = &self.info["dash"]["dolby"]["audio"][0];
if !(filter_option.no_dolby_audio || dolby_audio.is_null()) {
let (Some(url), Some(quality)) = (dolby_audio["baseUrl"].as_str(), dolby_audio["id"].as_u64()) else {
bail!("invalid dolby audio stream");
};
let quality =
AudioQuality::from_repr(quality as usize).ok_or(anyhow!("invalid dolby audio stream quality"))?;
if quality >= filter_option.audio_min_quality && quality <= filter_option.audio_max_quality {
streams.push(Stream::DashAudio {
url: dolby_stream_url,
quality: dolby_stream_quality,
url: url.to_string(),
quality,
});
}
}
Expand All @@ -250,68 +257,74 @@ impl PageAnalyzer {
let streams = self.streams(filter_option)?;
if self.is_flv_stream() || self.is_html5_mp4_stream() || self.is_episode_try_mp4_stream() {
// 按照 streams 中的假设，符合这三种情况的流只有一个，直接取
return Ok(BestStream::Mixed(streams.into_iter().next().unwrap()));
return Ok(BestStream::Mixed(
streams.into_iter().next().ok_or(anyhow!("no stream found"))?,
));
}
// 将视频流和音频流拆分，分别做排序
let (mut video_streams, mut audio_streams): (Vec<_>, Vec<_>) =
let (videos, audios): (Vec<Stream>, Vec<Stream>) =
streams.into_iter().partition(|s| matches!(s, Stream::DashVideo { .. }));
// 因为该处的排序与筛选选项有关，因此不能在外面实现 PartialOrd trait，只能在这里写闭包
video_streams.sort_by(|a, b| match (a, b) {
(
Stream::DashVideo {
quality: a_quality,
codecs: a_codecs,
..
},
Stream::DashVideo {
quality: b_quality,
codecs: b_codecs,
..
},
) => {
if a_quality == &VideoQuality::QualityDolby && !filter_option.no_dolby_video {
return std::cmp::Ordering::Greater;
}
if b_quality == &VideoQuality::QualityDolby && !filter_option.no_dolby_video {
return std::cmp::Ordering::Less;
}
if a_quality == &VideoQuality::QualityHdr && !filter_option.no_hdr {
return std::cmp::Ordering::Greater;
}
if b_quality == &VideoQuality::QualityHdr && !filter_option.no_hdr {
return std::cmp::Ordering::Less;
}
if a_quality != b_quality {
return a_quality.partial_cmp(b_quality).unwrap();
}
// 如果视频质量相同，按照偏好的编码优先级排序
filter_option
.codecs
.iter()
.position(|c| c == b_codecs)
.cmp(&filter_option.codecs.iter().position(|c| c == a_codecs))
}
_ => unreachable!(),
});
audio_streams.sort_by(|a, b| match (a, b) {
(Stream::DashAudio { quality: a_quality, .. }, Stream::DashAudio { quality: b_quality, .. }) => {
if a_quality == &AudioQuality::QualityDolby && !filter_option.no_dolby_audio {
return std::cmp::Ordering::Greater;
Ok(BestStream::VideoAudio {
video: Iterator::max_by(videos.into_iter(), |a, b| match (a, b) {
(
Stream::DashVideo {
quality: a_quality,
codecs: a_codecs,
..
},
Stream::DashVideo {
quality: b_quality,
codecs: b_codecs,
..
},
) => {
if a_quality != b_quality {
return a_quality.partial_cmp(b_quality).unwrap();
};
filter_option
.codecs
.iter()
.position(|c| c == b_codecs)
.cmp(&filter_option.codecs.iter().position(|c| c == a_codecs))
}
if b_quality == &AudioQuality::QualityDolby && !filter_option.no_dolby_audio {
return std::cmp::Ordering::Less;
_ => unreachable!(),
})
.ok_or(anyhow!("no video stream found"))?,
audio: Iterator::max_by(audios.into_iter(), |a, b| match (a, b) {
(Stream::DashAudio { quality: a_quality, .. }, Stream::DashAudio { quality: b_quality, .. }) => {
a_quality.partial_cmp(b_quality).unwrap()
}
a_quality.partial_cmp(b_quality).unwrap()
}
_ => unreachable!(),
});
if video_streams.is_empty() {
bail!("no video stream found");
}
Ok(BestStream::VideoAudio {
video: video_streams.remove(video_streams.len() - 1),
// 音频流可能为空，因此直接使用 pop 返回 Option
audio: audio_streams.pop(),
_ => unreachable!(),
}),
})
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_quality_order() {
assert!([
VideoQuality::Quality360p,
VideoQuality::Quality480p,
VideoQuality::Quality720p,
VideoQuality::Quality1080p,
VideoQuality::Quality1080pPLUS,
VideoQuality::Quality1080p60,
VideoQuality::Quality4k,
VideoQuality::QualityHdr,
VideoQuality::QualityDolby,
VideoQuality::Quality8k
]
.is_sorted());
assert!([
AudioQuality::Quality64k,
AudioQuality::Quality132k,
AudioQuality::Quality192k,
AudioQuality::QualityDolby,
AudioQuality::QualityHiRES,
]
.is_sorted());
}
}

0 comments on commit 7d9999d

Please sign in to comment.