mirror of
https://github.com/GuillemCastro/spotify-dl.git
synced 2024-11-23 18:50:24 +01:00
async encoding
This commit is contained in:
parent
c2b636236d
commit
fbe55b73fd
5 changed files with 113 additions and 64 deletions
|
@ -1,6 +1,6 @@
|
|||
use std::fmt::Write;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use futures::StreamExt;
|
||||
|
@ -100,6 +100,7 @@ impl Downloader {
|
|||
);
|
||||
|
||||
let pb = self.progress_bar.add(ProgressBar::new(file_size as u64));
|
||||
pb.enable_steady_tick(Duration::from_millis(100));
|
||||
pb.set_style(ProgressStyle::with_template("{spinner:.green} {msg} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")?
|
||||
.with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
|
||||
.progress_chars("#>-"));
|
||||
|
@ -122,24 +123,23 @@ impl Downloader {
|
|||
samples.append(&mut content);
|
||||
}
|
||||
SinkEvent::Finished => {
|
||||
pb.finish_with_message(format!("Downloaded {}", &file_name));
|
||||
tracing::info!("Finished downloading track: {:?}", file_name);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!("Downloaded track: {:?}", file_name);
|
||||
|
||||
tracing::info!("Encoding track: {:?}", file_name);
|
||||
pb.set_message(format!("Encoding {}", &file_name));
|
||||
let samples = Samples::new(samples, 44100, 2, 16);
|
||||
let encoder = crate::encoder::get_encoder(options.format)?;
|
||||
let stream = encoder.encode(samples).await?;
|
||||
|
||||
let format = options.format.clone();
|
||||
tokio_rayon::spawn(move || {
|
||||
tracing::info!("Encoding track: {:?}", file_name);
|
||||
let encoder = crate::encoder::get_encoder(format).unwrap();
|
||||
let stream = encoder.encode(samples).unwrap();
|
||||
tracing::info!("Writing track: {:?} to file: {}", file_name, &path);
|
||||
stream.write_to_file(&path).unwrap();
|
||||
}).await;
|
||||
pb.set_message(format!("Writing {}", &file_name));
|
||||
tracing::info!("Writing track: {:?} to file: {}", file_name, &path);
|
||||
stream.write_to_file(&path).await?;
|
||||
|
||||
pb.finish_with_message(format!("Downloaded {}", &file_name));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
|
@ -1,28 +1,24 @@
|
|||
use flacenc::bitsink::ByteSink;
|
||||
use flacenc::component::BitRepr;
|
||||
use flacenc::error::Verified;
|
||||
use flacenc::error::Verify;
|
||||
|
||||
use super::execute_with_result;
|
||||
use super::EncodedStream;
|
||||
use super::Encoder;
|
||||
use super::Samples;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct FlacEncoder {
|
||||
config: Verified<flacenc::config::Encoder>,
|
||||
}
|
||||
pub struct FlacEncoder;
|
||||
|
||||
impl FlacEncoder {
|
||||
pub fn new() -> anyhow::Result<Self> {
|
||||
let config = flacenc::config::Encoder::default()
|
||||
.into_verified()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to verify encoder config: {:?}", e))?;
|
||||
Ok(FlacEncoder { config })
|
||||
Ok(Self)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Encoder for FlacEncoder {
|
||||
fn encode(&self, samples: Samples) -> anyhow::Result<EncodedStream> {
|
||||
async fn encode(&self, samples: Samples) -> anyhow::Result<EncodedStream> {
|
||||
let source = flacenc::source::MemSource::from_samples(
|
||||
&samples.samples,
|
||||
samples.channels as usize,
|
||||
|
@ -30,15 +26,33 @@ impl Encoder for FlacEncoder {
|
|||
samples.sample_rate as usize,
|
||||
);
|
||||
|
||||
let flac_stream =
|
||||
flacenc::encode_with_fixed_block_size(&self.config, source, self.config.block_size)
|
||||
let config = flacenc::config::Encoder::default()
|
||||
.into_verified()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to verify encoder config: {:?}", e))?;
|
||||
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
|
||||
rayon::spawn(execute_with_result(
|
||||
move || {
|
||||
let flac_stream = flacenc::encode_with_fixed_block_size(
|
||||
&config,
|
||||
source,
|
||||
config.block_size,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to encode flac: {:?}", e))?;
|
||||
|
||||
let mut byte_sink = ByteSink::new();
|
||||
flac_stream
|
||||
.write(&mut byte_sink)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to write flac stream: {:?}", e))?;
|
||||
let mut byte_sink = ByteSink::new();
|
||||
flac_stream
|
||||
.write(&mut byte_sink)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to write flac stream: {:?}", e))?;
|
||||
|
||||
Ok(EncodedStream::new(byte_sink.into_inner()))
|
||||
Ok(byte_sink.into_inner())
|
||||
},
|
||||
tx,
|
||||
));
|
||||
|
||||
let byte_sink: Vec<u8> = rx.await??;
|
||||
|
||||
Ok(EncodedStream::new(byte_sink))
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,10 +2,11 @@ mod flac;
|
|||
#[cfg(feature = "mp3")]
|
||||
mod mp3;
|
||||
|
||||
use std::str::FromStr;
|
||||
use std::{path::Path, str::FromStr};
|
||||
|
||||
use anyhow::Result;
|
||||
use lazy_static::lazy_static;
|
||||
use tokio::sync::oneshot::Sender;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
|
||||
pub enum Format {
|
||||
|
@ -54,8 +55,9 @@ pub fn get_encoder(format: Format) -> anyhow::Result<&'static Box<dyn Encoder +
|
|||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
pub trait Encoder {
|
||||
fn encode(&self, samples: Samples) -> Result<EncodedStream>;
|
||||
async fn encode(&self, samples: Samples) -> Result<EncodedStream>;
|
||||
}
|
||||
|
||||
pub struct Samples {
|
||||
|
@ -85,15 +87,27 @@ impl EncodedStream {
|
|||
EncodedStream { stream }
|
||||
}
|
||||
|
||||
pub fn write_to_file<P: std::convert::AsRef<std::path::Path>>(&self, path: P) -> Result<()> {
|
||||
pub async fn write_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
|
||||
if !path.as_ref().exists() {
|
||||
std::fs::create_dir_all(
|
||||
tokio::fs::create_dir_all(
|
||||
path.as_ref()
|
||||
.parent()
|
||||
.ok_or(anyhow::anyhow!("Could not create path"))?,
|
||||
)?;
|
||||
).await?;
|
||||
}
|
||||
std::fs::write(path, &self.stream)?;
|
||||
tokio::fs::write(path, &self.stream).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn execute_with_result<F, T>(func: F, tx: Sender<anyhow::Result<T>>) -> impl FnOnce()
|
||||
where
|
||||
F: FnOnce() -> anyhow::Result<T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
move || {
|
||||
let result = func();
|
||||
// Ignore the error if the receiver has been dropped
|
||||
let _ = tx.send(result);
|
||||
}
|
||||
}
|
|
@ -4,48 +4,69 @@ use mp3lame_encoder::Builder;
|
|||
use mp3lame_encoder::FlushNoGap;
|
||||
use mp3lame_encoder::InterleavedPcm;
|
||||
|
||||
use super::Samples;
|
||||
use super::Encoder;
|
||||
use super::execute_with_result;
|
||||
use super::EncodedStream;
|
||||
use super::Encoder;
|
||||
use super::Samples;
|
||||
|
||||
pub struct Mp3Encoder;
|
||||
|
||||
impl Mp3Encoder {
|
||||
fn build_encoder(&self, sample_rate: u32, channels: u32) -> anyhow::Result<mp3lame_encoder::Encoder> {
|
||||
let mut builder = Builder::new().
|
||||
ok_or(anyhow::anyhow!("Failed to create mp3 encoder"))?;
|
||||
fn build_encoder(
|
||||
&self,
|
||||
sample_rate: u32,
|
||||
channels: u32,
|
||||
) -> anyhow::Result<mp3lame_encoder::Encoder> {
|
||||
let mut builder = Builder::new().ok_or(anyhow::anyhow!("Failed to create mp3 encoder"))?;
|
||||
|
||||
builder.set_sample_rate(sample_rate)
|
||||
builder
|
||||
.set_sample_rate(sample_rate)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to set sample rate for mp3 encoder: {}", e))?;
|
||||
builder.set_num_channels(channels as u8)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to set number of channels for mp3 encoder: {}", e))?;
|
||||
builder.set_brate(mp3lame_encoder::Birtate::Kbps160)
|
||||
builder.set_num_channels(channels as u8).map_err(|e| {
|
||||
anyhow::anyhow!("Failed to set number of channels for mp3 encoder: {}", e)
|
||||
})?;
|
||||
builder
|
||||
.set_brate(mp3lame_encoder::Birtate::Kbps160)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to set bitrate for mp3 encoder: {}", e))?;
|
||||
|
||||
builder.build()
|
||||
builder
|
||||
.build()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to build mp3 encoder: {}", e))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Encoder for Mp3Encoder {
|
||||
fn encode(&self, samples: Samples) -> anyhow::Result<EncodedStream> {
|
||||
async fn encode(&self, samples: Samples) -> anyhow::Result<EncodedStream> {
|
||||
let mut mp3_encoder = self.build_encoder(samples.sample_rate, samples.channels)?;
|
||||
|
||||
let samples: Vec<i16> = samples.samples.iter().map(|&x| x as i16).collect();
|
||||
let input = InterleavedPcm(samples.as_slice());
|
||||
let mut mp3_out_buffer = Vec::new();
|
||||
mp3_out_buffer.reserve(mp3lame_encoder::max_required_buffer_size(samples.len()));
|
||||
let encoded_size = mp3_encoder.encode(input, mp3_out_buffer.spare_capacity_mut())
|
||||
.map_err(|e| anyhow!("Failed to encode mp3: {}", e))?;
|
||||
unsafe {
|
||||
mp3_out_buffer.set_len(mp3_out_buffer.len().wrapping_add(encoded_size));
|
||||
}
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
|
||||
let encoded_size = mp3_encoder.flush::<FlushNoGap>(mp3_out_buffer.spare_capacity_mut())
|
||||
.map_err(|e| anyhow!("Failed to flush mp3 encoder: {}", e))?;
|
||||
unsafe {
|
||||
mp3_out_buffer.set_len(mp3_out_buffer.len().wrapping_add(encoded_size));
|
||||
}
|
||||
rayon::spawn(execute_with_result(
|
||||
move || {
|
||||
let samples: Vec<i16> = samples.samples.iter().map(|&x| x as i16).collect();
|
||||
let input = InterleavedPcm(samples.as_slice());
|
||||
let mut mp3_out_buffer = Vec::new();
|
||||
mp3_out_buffer.reserve(mp3lame_encoder::max_required_buffer_size(samples.len()));
|
||||
let encoded_size = mp3_encoder
|
||||
.encode(input, mp3_out_buffer.spare_capacity_mut())
|
||||
.map_err(|e| anyhow!("Failed to encode mp3: {}", e))?;
|
||||
unsafe {
|
||||
mp3_out_buffer.set_len(mp3_out_buffer.len().wrapping_add(encoded_size));
|
||||
}
|
||||
|
||||
let encoded_size = mp3_encoder
|
||||
.flush::<FlushNoGap>(mp3_out_buffer.spare_capacity_mut())
|
||||
.map_err(|e| anyhow!("Failed to flush mp3 encoder: {}", e))?;
|
||||
unsafe {
|
||||
mp3_out_buffer.set_len(mp3_out_buffer.len().wrapping_add(encoded_size));
|
||||
}
|
||||
Ok(mp3_out_buffer)
|
||||
},
|
||||
tx,
|
||||
));
|
||||
|
||||
let mp3_out_buffer = rx.await??;
|
||||
|
||||
Ok(EncodedStream::new(mp3_out_buffer))
|
||||
}
|
||||
|
|
|
@ -1,5 +1,5 @@
|
|||
pub mod download;
|
||||
pub mod file_sink;
|
||||
// pub mod file_sink;
|
||||
pub mod channel_sink;
|
||||
pub mod session;
|
||||
pub mod track;
|
||||
|
|
Loading…
Reference in a new issue