ReplayGain optimizations/cleanup

This commit is contained in:
Agatha Lovelace 2024-03-29 23:21:40 +01:00
parent 8771669a47
commit 9dff83a52f
Signed by: sorceress
GPG Key ID: 01D0B3AB10CED4F8
2 changed files with 240 additions and 124 deletions

View File

@ -1,13 +1,14 @@
use std::{ffi::OsStr, fs::File, hash::Hasher, path::Path};
use adler::Adler32;
use lofty::{AudioFile, ItemKey, TaggedFileExt};
use lofty::{AudioFile, TaggedFileExt};
use miette::{miette, IntoDiagnostic, Result};
use rayon::prelude::*;
use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, QuerySelect, Set};
use symphonia::{
core::{
formats::{FormatOptions, FormatReader},
errors::Error as SymphoniaError,
formats::{FormatOptions, FormatReader, Packet},
io::{MediaSourceStream, MediaSourceStreamOptions},
meta::{Limit, MetadataOptions},
probe::Hint,
@ -15,10 +16,10 @@ use symphonia::{
default::get_probe,
};
use time::OffsetDateTime;
use tracing::debug;
use tracing::{debug, warn};
use walkdir::{DirEntry, WalkDir};
use crate::backend::replaygain::{format_gain, ReplayGain};
use crate::backend::replaygain::{ReplayGainResult, ReplayGainState};
use super::{
config::Source,
@ -49,14 +50,87 @@ fn get_packets(path: &Path) -> Result<Box<dyn FormatReader>> {
.map(|v| v.format)
}
fn hash_packets(data: &mut Box<dyn FormatReader>) -> u64 {
let mut adler = Adler32::new();
struct FormatReaderIter {
inner: Box<dyn FormatReader>,
error: Option<SymphoniaError>,
hash: Adler32,
rg: ReplayGainState,
}
while let Ok(packet) = data.next_packet() {
adler.write(&packet.data);
impl FormatReaderIter {
/// Initialize a new iterator over `FormatReader`.
/// Track ReplayGain values will be computed if `rg` is None.
fn new(
inner: Box<dyn FormatReader>,
rg: Option<ReplayGainResult>,
) -> Result<Self, EleanorError> {
let track = inner
.default_track()
.ok_or_else(|| miette!("No default track was found"))?;
let params = &track.codec_params;
let (sample_rate, channels) = (params.sample_rate, params.channels);
// Only stereo is supported.
// In future we could try duplicating mono tracks and downmixing anything higher.
let rg = (!channels.is_some_and(|x| x.count() == 2))
.then_some(ReplayGainState::Failed(
miette!("Unsupported channel configuration: {:?}", channels).into(),
))
.unwrap_or_else(|| ReplayGainState::new_with_override(sample_rate, track, rg));
Ok(Self {
inner,
error: None,
hash: Adler32::new(),
rg,
})
}
adler.finish()
fn process(mut self) -> (u64, Result<ReplayGainResult, EleanorError>) {
// Loop over all packets
while let Some(packet) = (&mut self).next() {
// Hash the packet
self.hash.write(&packet.data);
// Copy into replaygain
self.rg.handle_packet(&packet);
}
let hash = self.hash.finish();
// Check for errors during iteration
if let Some(error) = self.error {
return (hash, Err(error.into()));
}
let rg = match self.rg {
ReplayGainState::Finished(rg_res) => Ok(rg_res),
ReplayGainState::Computing(rg) => Ok(rg.finish()),
ReplayGainState::Failed(e) => Err(e),
};
(hash, rg)
}
}
impl Iterator for &mut FormatReaderIter {
type Item = Packet;
fn next(&mut self) -> Option<Self::Item> {
let res_packet = self.inner.next_packet();
match res_packet {
Err(symphonia::core::errors::Error::IoError(ref packet_error))
if packet_error.kind() == std::io::ErrorKind::UnexpectedEof =>
{
None
}
Err(e) => {
self.error = Some(e);
None
}
Ok(packet) => Some(packet),
}
}
}
fn index_song(
@ -65,6 +139,8 @@ fn index_song(
force: bool,
indexed_ts: OffsetDateTime,
) -> Result<Option<library::ActiveModel>, EleanorError> {
let filename = file.path().display().to_string();
// Re-index previously indexed files
if !force {
let modified = file
@ -75,7 +151,7 @@ fn index_song(
.is_ok();
if modified {
debug!("Skipping file {}", file.path().display());
debug!("Skipping file {filename}");
return Ok(None);
}
}
@ -87,42 +163,31 @@ fn index_song(
let tags = tagged_file.primary_tag().or(tagged_file.first_tag());
let properties = tagged_file.properties();
// Hash audio packets
let mut packets = get_packets(file.path())?;
let hash = hash_packets(&mut packets);
// Hash audio packets and calculate replaygain
let (hash, rg) = FormatReaderIter::new(
get_packets(file.path())?,
// Fall back to metadata
if let rg @ ReplayGainResult {
track_gain: Some(_),
track_peak: Some(_),
..
} = ReplayGainResult::from(tags)
{
Some(rg)
} else {
None
},
)?
.process();
let rg_track_gain = tags
.and_then(|t| t.get_string(&ItemKey::ReplayGainTrackGain))
.and_then(|v| format_gain(v).ok());
let rg_track_peak = tags
.and_then(|t| t.get_string(&ItemKey::ReplayGainTrackPeak))
.and_then(|v| format_gain(v).ok());
let rg_album_gain = tags
.and_then(|t| t.get_string(&ItemKey::ReplayGainAlbumGain))
.and_then(|v| format_gain(v).ok());
let rg_album_peak = tags
.and_then(|t| t.get_string(&ItemKey::ReplayGainAlbumPeak))
.and_then(|v| format_gain(v).ok());
// Check for existing ReplayGain tags.
let mut rg = if let (Some(track_gain), Some(track_peak)) = (rg_track_gain, rg_track_peak) {
Ok(ReplayGain {
track_gain,
track_peak,
album_gain: None,
album_peak: None,
})
} else {
// Calculate replaygain values for the audio track
let mut packets = get_packets(file.path())?;
ReplayGain::try_calculate(&mut packets)
};
// Set album gain and peak, if present in metadata.
if let Ok(rg) = &mut rg {
rg.album_gain = rg_album_gain;
rg.album_peak = rg_album_peak;
// Ignore replaygain errors
if let Err(e) = &rg {
warn!(
track = filename,
"Encountered error while calculating ReplayGain: {}", e
);
}
let rg = rg.ok();
let song: library::ActiveModel = library::ActiveModel {
path: Set(file
@ -151,18 +216,10 @@ fn index_song(
disc: Set(tags.and_then(lofty::Accessor::disk).map(|t| t as i32)),
year: Set(tags.and_then(lofty::Accessor::year).map(|t| t as i32)),
duration: Set(properties.duration().as_millis().try_into()?),
rg_track_gain: Set(rg.as_ref().map(|v| f64::from(v.track_gain)).ok()),
rg_track_peak: Set(rg.as_ref().map(|v| f64::from(v.track_peak)).ok()),
rg_album_gain: Set(rg
.as_ref()
.map(|v| v.album_gain.map(f64::from))
.ok()
.flatten()),
rg_album_peak: Set(rg
.as_ref()
.map(|v| v.album_peak.map(f64::from))
.ok()
.flatten()),
rg_track_gain: Set(rg.as_ref().and_then(|rg| rg.track_gain).map(f64::from)),
rg_track_peak: Set(rg.as_ref().and_then(|rg| rg.track_peak).map(f64::from)),
rg_album_gain: Set(rg.as_ref().and_then(|rg| rg.album_gain).map(f64::from)),
rg_album_peak: Set(rg.as_ref().and_then(|rg| rg.album_peak).map(f64::from)),
..Default::default()
};

View File

@ -1,105 +1,164 @@
use lofty::{ItemKey, Tag};
use miette::miette;
use symphonia::core::{audio::SampleBuffer, codecs::DecoderOptions, formats::FormatReader};
use symphonia::core::{
audio::SampleBuffer,
codecs::{Decoder, DecoderOptions},
formats::{Packet, Track},
};
use tracing::warn;
use super::error::EleanorError;
#[derive(Debug, Clone)]
/// A wrapper around the data needed for ReplayGain calculation.
pub struct ReplayGain {
pub track_gain: f32,
pub track_peak: f32,
pub album_gain: Option<f32>,
pub album_peak: Option<f32>,
data: Vec<f32>,
rg: replaygain::ReplayGain,
decoder: Box<dyn Decoder>,
track_id: u32,
sample_buf: Option<SampleBuffer<f32>>,
}
impl ReplayGain {
pub(crate) fn try_calculate(audio: &mut Box<dyn FormatReader>) -> Result<Self, EleanorError> {
let track = audio
.default_track()
.ok_or_else(|| miette!("No default track was found"))?;
pub fn init(sample_rate: usize, track: &Track) -> Result<Self, EleanorError> {
let rg = replaygain::ReplayGain::new(sample_rate)
.ok_or(miette!("Unsupported sample rate: {}", sample_rate))?;
let params = &track.codec_params;
let (sample_rate, channels) = (params.sample_rate, params.channels);
let Some(sample_rate) = sample_rate else {
return Err(miette!("Sample rate must be known").into());
};
// Only stereo is supported.
if !channels.is_some_and(|x| x.count() == 2) {
return Err(miette!("Unsupported channel configuration: {:?}", channels).into());
}
// Only 44.1kHz and 48kHz are supported.
let Some(mut rg) = replaygain::ReplayGain::new(sample_rate as usize) else {
return Err(miette!("Unsupported sample rate: {:?}", sample_rate).into());
};
let mut decoder =
symphonia::default::get_codecs().make(params, &DecoderOptions::default())?;
let track_id = track.id;
let mut samples: Vec<f32> = vec![];
let mut sample_buf = None;
loop {
let packet = match audio.next_packet() {
Ok(packet) => packet,
Err(symphonia::core::errors::Error::IoError(ref packet_error))
if packet_error.kind() == std::io::ErrorKind::UnexpectedEof =>
{
// End of audio stream
break;
}
Err(e) => {
return Err(e.into());
}
};
let decoder = symphonia::default::get_codecs().make(params, &DecoderOptions::default())?;
// Skip packets belonging to other audio tracks
if packet.track_id() != track_id {
continue;
}
Ok(Self {
rg,
data: Vec::new(),
decoder,
track_id,
sample_buf: None,
})
}
match decoder.decode(&packet) {
pub fn handle_packet(&mut self, packet: &Packet) -> Result<(), EleanorError> {
if packet.track_id() == self.track_id {
match self.decoder.decode(packet) {
Ok(buffer) => {
if sample_buf.is_none() {
if self.sample_buf.is_none() {
let spec = *buffer.spec();
let duration = buffer.capacity() as u64;
sample_buf = Some(SampleBuffer::<f32>::new(duration, spec));
self.sample_buf = Some(SampleBuffer::<f32>::new(duration, spec));
}
if let Some(target) = &mut sample_buf {
if let Some(target) = self.sample_buf.as_mut() {
target.copy_interleaved_ref(buffer);
}
}
Err(symphonia::core::errors::Error::DecodeError(_)) => (),
Err(_) => break,
Err(e) => return Err(e.into()),
}
if let Some(buf) = &mut sample_buf {
samples.extend(buf.samples());
if let Some(buf) = self.sample_buf.as_mut() {
self.data.extend(buf.samples());
}
}
Ok(())
}
if samples.is_empty() {
return Err(miette!("No samples were decoded from input audio").into());
pub fn finish(mut self) -> ReplayGainResult {
let (track_gain, track_peak) = if self.data.is_empty() {
warn!("No samples were decoded from input audio");
(None, None)
} else {
self.rg.process_samples(&self.data);
let (gain, peak) = self.rg.finish();
(Some(gain), Some(peak))
};
rg.process_samples(&samples);
let (track_gain, track_peak) = rg.finish();
Ok(Self {
ReplayGainResult {
track_gain,
track_peak,
album_gain: None,
album_peak: None,
})
}
}
}
pub(crate) fn format_gain(gain: &str) -> Result<f32, EleanorError> {
/// A representation of the possible calculation states
pub enum ReplayGainState {
Finished(ReplayGainResult),
Computing(Box<ReplayGain>),
Failed(EleanorError),
}
impl ReplayGainState {
pub fn handle_packet(&mut self, packet: &Packet) {
if let ReplayGainState::Computing(rg) = self {
let result = rg.handle_packet(packet);
if let Err(e) = result {
*self = ReplayGainState::Failed(e);
}
}
}
pub fn new_with_override(
sample_rate: Option<u32>,
track: &Track,
r#override: Option<ReplayGainResult>,
) -> Self {
let Some(sample_rate) = sample_rate else {
return Self::Failed(miette!("Sample rate must be known").into());
};
match r#override {
Some(rg_res) => Self::Finished(rg_res),
None => match ReplayGain::init(sample_rate as usize, track)
.map(Box::new)
.map(Self::Computing)
{
Ok(rg) => rg,
Err(e) => Self::Failed(e),
},
}
}
}
/// Finished ReplayGain calculation results.
///
/// Calculation errors may happen, but are non-critical
/// and should be ignored by the player.
#[derive(Debug, Clone)]
pub struct ReplayGainResult {
pub track_gain: Option<f32>,
pub track_peak: Option<f32>,
pub album_gain: Option<f32>,
pub album_peak: Option<f32>,
}
impl From<Option<&Tag>> for ReplayGainResult {
fn from(tags: Option<&Tag>) -> Self {
let track_gain = tags
.and_then(|t| t.get_string(&ItemKey::ReplayGainTrackGain))
.and_then(|v| parse_gain(v).ok());
let track_peak = tags
.and_then(|t| t.get_string(&ItemKey::ReplayGainTrackPeak))
.and_then(|v| parse_gain(v).ok());
let album_gain = tags
.and_then(|t| t.get_string(&ItemKey::ReplayGainAlbumGain))
.and_then(|v| parse_gain(v).ok());
let album_peak = tags
.and_then(|t| t.get_string(&ItemKey::ReplayGainAlbumPeak))
.and_then(|v| parse_gain(v).ok());
Self {
track_gain,
track_peak,
album_gain,
album_peak,
}
}
}
pub fn parse_gain(gain: &str) -> Result<f32, EleanorError> {
gain.chars()
.filter(|c| c.is_numeric() || matches!(*c, '-' | '+' | '.'))
.collect::<String>()
@ -110,7 +169,7 @@ pub(crate) fn format_gain(gain: &str) -> Result<f32, EleanorError> {
#[test]
fn rg_parse_expected() {
for (input, expected) in [("-8.97 dB", -8.97), ("12.75 dB", 12.75), ("0.00 dB", 0.0)] {
assert_eq!(format_gain(input).unwrap(), expected)
assert_eq!(parse_gain(input).unwrap(), expected)
}
}
@ -124,13 +183,13 @@ fn rg_parse_malformed() {
("0.93dB", 0.93),
("24.12 deutscheBahn", 24.12),
] {
assert_eq!(format_gain(input).unwrap(), expected)
assert_eq!(parse_gain(input).unwrap(), expected)
}
}
#[test]
fn rg_parse_invalid() {
for input in ["", "akjdfnhlkjfh", "🥺", " DB "] {
assert!(format_gain(input).is_err())
assert!(parse_gain(input).is_err())
}
}