501 lines
16 KiB
Rust
501 lines
16 KiB
Rust
//! DiskSink - Écrit le flux audio dans un fichier
|
|
//!
|
|
//! Ce module fournit un sink qui écrit les chunks audio sur disque,
|
|
//! avec support de la dérivation automatique du nom de fichier depuis la source.
|
|
|
|
use crate::{events::SourceNameUpdateEvent, nodes::AudioError, AudioChunk};
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
use tokio::fs::File;
|
|
use tokio::io::AsyncWriteExt;
|
|
use tokio::sync::{mpsc, RwLock};
|
|
|
|
/// Configuration pour le DiskSink
|
|
#[derive(Debug, Clone)]
|
|
pub struct DiskSinkConfig {
|
|
/// Chemin racine où écrire les fichiers
|
|
pub output_dir: PathBuf,
|
|
|
|
/// Nom de fichier explicite (optionnel)
|
|
/// Si None, sera dérivé du nom de la source
|
|
pub filename: Option<String>,
|
|
|
|
/// Format d'écriture
|
|
pub format: AudioFileFormat,
|
|
|
|
/// Taille du buffer d'écriture (en chunks)
|
|
pub buffer_size: usize,
|
|
}
|
|
|
|
impl Default for DiskSinkConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
output_dir: PathBuf::from("."),
|
|
filename: None,
|
|
format: AudioFileFormat::Wav,
|
|
buffer_size: 100,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Formats de fichiers audio supportés
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub enum AudioFileFormat {
|
|
/// Format WAV (non compressé)
|
|
Wav,
|
|
/// Format FLAC (compressé sans perte)
|
|
Flac,
|
|
/// Format brut PCM
|
|
Raw,
|
|
}
|
|
|
|
impl AudioFileFormat {
|
|
/// Retourne l'extension de fichier appropriée
|
|
pub fn extension(&self) -> &str {
|
|
match self {
|
|
AudioFileFormat::Wav => "wav",
|
|
AudioFileFormat::Flac => "flac",
|
|
AudioFileFormat::Raw => "pcm",
|
|
}
|
|
}
|
|
}
|
|
|
|
/// DiskSink - Écrit le flux audio dans un fichier sur disque
|
|
///
|
|
/// Ce sink consomme les chunks audio et les écrit dans un fichier.
|
|
/// Le nom du fichier peut être dérivé automatiquement du nom de la source
|
|
/// via les événements `SourceNameUpdateEvent`.
|
|
///
|
|
/// # Caractéristiques
|
|
///
|
|
/// - Écriture asynchrone avec buffer
|
|
/// - Dérivation automatique du nom de fichier depuis la source
|
|
/// - Support de plusieurs formats (WAV, FLAC, PCM brut)
|
|
/// - Gestion du gain : applique le gain avant l'écriture
|
|
///
|
|
/// # Exemples
|
|
///
|
|
/// ```no_run
|
|
/// use pmoaudio::{DiskSink, DiskSinkConfig};
|
|
/// use std::path::PathBuf;
|
|
///
|
|
/// #[tokio::main]
|
|
/// async fn main() {
|
|
/// let config = DiskSinkConfig {
|
|
/// output_dir: PathBuf::from("/tmp/audio"),
|
|
/// filename: Some("output.wav".to_string()),
|
|
/// ..Default::default()
|
|
/// };
|
|
///
|
|
/// let (sink, sink_tx) = DiskSink::new("disk1".to_string(), config, 10);
|
|
///
|
|
/// tokio::spawn(async move {
|
|
/// sink.run().await.unwrap()
|
|
/// });
|
|
/// }
|
|
/// ```
|
|
pub struct DiskSink {
|
|
/// Identifiant du sink
|
|
node_id: String,
|
|
|
|
/// Channel pour recevoir les chunks audio
|
|
rx: mpsc::Receiver<Arc<AudioChunk>>,
|
|
|
|
/// Configuration
|
|
config: DiskSinkConfig,
|
|
|
|
/// Nom de fichier résolu (partagé)
|
|
resolved_filename: Arc<RwLock<Option<PathBuf>>>,
|
|
|
|
/// Receiver pour les événements de nom de source (optionnel)
|
|
source_name_rx: Option<mpsc::Receiver<SourceNameUpdateEvent>>,
|
|
|
|
/// Writer pour le fichier
|
|
writer: Option<AudioFileWriter>,
|
|
}
|
|
|
|
impl DiskSink {
|
|
/// Crée un nouveau DiskSink
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `node_id` - Identifiant unique du sink
|
|
/// * `config` - Configuration du sink
|
|
/// * `channel_size` - Taille du buffer du channel
|
|
pub fn new(
|
|
node_id: String,
|
|
config: DiskSinkConfig,
|
|
channel_size: usize,
|
|
) -> (Self, mpsc::Sender<Arc<AudioChunk>>) {
|
|
let (tx, rx) = mpsc::channel(channel_size);
|
|
|
|
let sink = Self {
|
|
node_id,
|
|
rx,
|
|
config,
|
|
resolved_filename: Arc::new(RwLock::new(None)),
|
|
source_name_rx: None,
|
|
writer: None,
|
|
};
|
|
|
|
(sink, tx)
|
|
}
|
|
|
|
/// Configure la source des événements de nom de source
|
|
pub fn set_source_name_source(&mut self, rx: mpsc::Receiver<SourceNameUpdateEvent>) {
|
|
self.source_name_rx = Some(rx);
|
|
}
|
|
|
|
/// Résout le nom du fichier de sortie
|
|
///
|
|
/// Si un filename explicite est fourni dans la config, l'utilise.
|
|
/// Sinon, utilise le source_name avec l'extension appropriée.
|
|
fn resolve_filename(&self, source_name: Option<&str>) -> PathBuf {
|
|
let filename = if let Some(ref explicit_name) = self.config.filename {
|
|
explicit_name.clone()
|
|
} else if let Some(name) = source_name {
|
|
// Nettoyer le nom de la source pour en faire un nom de fichier valide
|
|
let clean_name = name
|
|
.chars()
|
|
.map(|c| {
|
|
if c.is_alphanumeric() || c == '_' || c == '-' {
|
|
c
|
|
} else {
|
|
'_'
|
|
}
|
|
})
|
|
.collect::<String>();
|
|
|
|
format!("{}.{}", clean_name, self.config.format.extension())
|
|
} else {
|
|
// Fallback sur un nom par défaut
|
|
format!("{}.{}", self.node_id, self.config.format.extension())
|
|
};
|
|
|
|
self.config.output_dir.join(filename)
|
|
}
|
|
|
|
/// Initialise le writer pour le fichier de sortie
|
|
async fn initialize_writer(&mut self, source_name: Option<&str>) -> Result<(), AudioError> {
|
|
let path = self.resolve_filename(source_name);
|
|
*self.resolved_filename.write().await = Some(path.clone());
|
|
|
|
// Créer le répertoire parent si nécessaire
|
|
if let Some(parent) = path.parent() {
|
|
tokio::fs::create_dir_all(parent).await.map_err(|e| {
|
|
AudioError::ProcessingError(format!("Failed to create directory: {}", e))
|
|
})?;
|
|
}
|
|
|
|
// Créer le writer approprié selon le format
|
|
let writer = match self.config.format {
|
|
AudioFileFormat::Wav => AudioFileWriter::new_wav(path).await?,
|
|
AudioFileFormat::Flac => {
|
|
// FLAC nécessiterait une bibliothèque externe, pour l'instant utiliser WAV
|
|
AudioFileWriter::new_wav(path).await?
|
|
}
|
|
AudioFileFormat::Raw => AudioFileWriter::new_raw(path).await?,
|
|
};
|
|
|
|
self.writer = Some(writer);
|
|
Ok(())
|
|
}
|
|
|
|
/// Démarre la boucle de traitement du DiskSink
|
|
pub async fn run(mut self) -> Result<DiskSinkStats, AudioError> {
|
|
let mut stats = DiskSinkStats::new(self.node_id.clone());
|
|
let mut source_name: Option<String> = None;
|
|
let mut initialized = false;
|
|
|
|
loop {
|
|
tokio::select! {
|
|
// Recevoir les chunks audio
|
|
chunk_opt = self.rx.recv() => {
|
|
match chunk_opt {
|
|
Some(chunk) => {
|
|
// Initialiser le writer à la réception du premier chunk
|
|
if !initialized {
|
|
self.initialize_writer(source_name.as_deref()).await?;
|
|
initialized = true;
|
|
}
|
|
|
|
// Appliquer le gain avant l'écriture
|
|
let chunk_with_gain = if chunk.gain_db().abs() > f64::EPSILON {
|
|
Arc::clone(&chunk).apply_gain()
|
|
} else {
|
|
Arc::clone(&chunk)
|
|
};
|
|
|
|
// Écrire le chunk
|
|
if let Some(ref mut writer) = self.writer {
|
|
writer.write_chunk(&chunk_with_gain).await?;
|
|
stats.record_chunk(&chunk_with_gain);
|
|
}
|
|
}
|
|
None => {
|
|
// Channel fermé, terminer
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Recevoir les mises à jour du nom de source
|
|
source_event_opt = async {
|
|
if let Some(ref mut rx) = self.source_name_rx {
|
|
rx.recv().await
|
|
} else {
|
|
std::future::pending().await
|
|
}
|
|
} => {
|
|
if let Some(event) = source_event_opt {
|
|
source_name = Some(event.source_name.clone());
|
|
|
|
// Si on n'a pas encore initialisé, le nom sera utilisé plus tard
|
|
// Sinon, on pourrait décider de fermer le fichier actuel et d'en créer un nouveau
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fermer le fichier proprement
|
|
if let Some(writer) = self.writer {
|
|
writer.close().await?;
|
|
}
|
|
|
|
stats.finalize();
|
|
Ok(stats)
|
|
}
|
|
}
|
|
|
|
/// Writer pour fichiers audio
|
|
struct AudioFileWriter {
|
|
file: File,
|
|
format: AudioFileFormat,
|
|
sample_rate: Option<u32>,
|
|
total_samples: usize,
|
|
}
|
|
|
|
impl AudioFileWriter {
|
|
/// Crée un writer WAV
|
|
async fn new_wav(path: PathBuf) -> Result<Self, AudioError> {
|
|
let file = File::create(path)
|
|
.await
|
|
.map_err(|e| AudioError::ProcessingError(format!("Failed to create file: {}", e)))?;
|
|
|
|
Ok(Self {
|
|
file,
|
|
format: AudioFileFormat::Wav,
|
|
sample_rate: None,
|
|
total_samples: 0,
|
|
})
|
|
}
|
|
|
|
/// Crée un writer pour PCM brut
|
|
async fn new_raw(path: PathBuf) -> Result<Self, AudioError> {
|
|
let file = File::create(path)
|
|
.await
|
|
.map_err(|e| AudioError::ProcessingError(format!("Failed to create file: {}", e)))?;
|
|
|
|
Ok(Self {
|
|
file,
|
|
format: AudioFileFormat::Raw,
|
|
sample_rate: None,
|
|
total_samples: 0,
|
|
})
|
|
}
|
|
|
|
/// Écrit un chunk audio
|
|
async fn write_chunk(&mut self, chunk: &AudioChunk) -> Result<(), AudioError> {
|
|
// Enregistrer le sample rate du premier chunk
|
|
if self.sample_rate.is_none() {
|
|
let sr = chunk.sample_rate();
|
|
self.sample_rate = Some(sr);
|
|
|
|
// Pour WAV, écrire l'en-tête (simplifié)
|
|
if matches!(self.format, AudioFileFormat::Wav) {
|
|
self.write_wav_header(sr).await?;
|
|
}
|
|
}
|
|
|
|
// Convertir en bytes (little-endian 16-bit PCM)
|
|
let mut bytes = Vec::with_capacity(chunk.len() * 4);
|
|
let max_val = chunk.bit_depth().max_value();
|
|
for frame in chunk.frames() {
|
|
let left = (frame[0] as f32 / max_val).clamp(-1.0, 1.0);
|
|
let right = (frame[1] as f32 / max_val).clamp(-1.0, 1.0);
|
|
let sample_i16 = (left * 32767.0) as i16;
|
|
bytes.extend_from_slice(&sample_i16.to_le_bytes());
|
|
let sample_r16 = (right * 32767.0) as i16;
|
|
bytes.extend_from_slice(&sample_r16.to_le_bytes());
|
|
}
|
|
|
|
self.file.write_all(&bytes).await.map_err(|e| {
|
|
AudioError::ProcessingError(format!("Failed to write audio data: {}", e))
|
|
})?;
|
|
|
|
self.total_samples += chunk.len();
|
|
Ok(())
|
|
}
|
|
|
|
/// Écrit un en-tête WAV simplifié
|
|
async fn write_wav_header(&mut self, sample_rate: u32) -> Result<(), AudioError> {
|
|
// En-tête WAV basique (sera mis à jour à la fermeture)
|
|
let mut header = Vec::new();
|
|
|
|
// RIFF chunk
|
|
header.extend_from_slice(b"RIFF");
|
|
header.extend_from_slice(&0u32.to_le_bytes()); // Taille (à mettre à jour)
|
|
header.extend_from_slice(b"WAVE");
|
|
|
|
// fmt chunk
|
|
header.extend_from_slice(b"fmt ");
|
|
header.extend_from_slice(&16u32.to_le_bytes()); // Taille du fmt chunk
|
|
header.extend_from_slice(&1u16.to_le_bytes()); // Format PCM
|
|
header.extend_from_slice(&2u16.to_le_bytes()); // 2 canaux (stéréo)
|
|
header.extend_from_slice(&sample_rate.to_le_bytes());
|
|
header.extend_from_slice(&(sample_rate * 4).to_le_bytes()); // Byte rate
|
|
header.extend_from_slice(&4u16.to_le_bytes()); // Block align
|
|
header.extend_from_slice(&16u16.to_le_bytes()); // Bits per sample
|
|
|
|
// data chunk header
|
|
header.extend_from_slice(b"data");
|
|
header.extend_from_slice(&0u32.to_le_bytes()); // Taille des données (à mettre à jour)
|
|
|
|
self.file.write_all(&header).await.map_err(|e| {
|
|
AudioError::ProcessingError(format!("Failed to write WAV header: {}", e))
|
|
})?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Ferme le fichier et met à jour l'en-tête si nécessaire
|
|
async fn close(mut self) -> Result<(), AudioError> {
|
|
if matches!(self.format, AudioFileFormat::Wav) {
|
|
// Mettre à jour les tailles dans l'en-tête WAV
|
|
let data_size = (self.total_samples * 4) as u32; // 2 bytes per sample * 2 channels
|
|
let file_size = data_size + 36;
|
|
|
|
// Positionner au début et réécrire les tailles
|
|
use tokio::io::AsyncSeekExt;
|
|
self.file
|
|
.seek(std::io::SeekFrom::Start(4))
|
|
.await
|
|
.map_err(|e| {
|
|
AudioError::ProcessingError(format!("Failed to seek in file: {}", e))
|
|
})?;
|
|
self.file
|
|
.write_all(&file_size.to_le_bytes())
|
|
.await
|
|
.map_err(|e| {
|
|
AudioError::ProcessingError(format!("Failed to update file size: {}", e))
|
|
})?;
|
|
|
|
self.file
|
|
.seek(std::io::SeekFrom::Start(40))
|
|
.await
|
|
.map_err(|e| {
|
|
AudioError::ProcessingError(format!("Failed to seek in file: {}", e))
|
|
})?;
|
|
self.file
|
|
.write_all(&data_size.to_le_bytes())
|
|
.await
|
|
.map_err(|e| {
|
|
AudioError::ProcessingError(format!("Failed to update data size: {}", e))
|
|
})?;
|
|
}
|
|
|
|
self.file
|
|
.flush()
|
|
.await
|
|
.map_err(|e| AudioError::ProcessingError(format!("Failed to flush file: {}", e)))?;
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Statistiques du DiskSink
|
|
#[derive(Debug, Clone)]
|
|
pub struct DiskSinkStats {
|
|
pub node_id: String,
|
|
pub chunks_written: u64,
|
|
pub total_samples: u64,
|
|
pub total_duration_sec: f64,
|
|
}
|
|
|
|
impl DiskSinkStats {
|
|
pub fn new(node_id: String) -> Self {
|
|
Self {
|
|
node_id,
|
|
chunks_written: 0,
|
|
total_samples: 0,
|
|
total_duration_sec: 0.0,
|
|
}
|
|
}
|
|
|
|
pub fn record_chunk(&mut self, chunk: &AudioChunk) {
|
|
self.chunks_written += 1;
|
|
self.total_samples += chunk.len() as u64;
|
|
self.total_duration_sec += chunk.len() as f64 / chunk.sample_rate() as f64;
|
|
}
|
|
|
|
pub fn finalize(&mut self) {
|
|
// Pourrait effectuer des calculs finaux ici
|
|
}
|
|
|
|
pub fn display(&self) {
|
|
println!("\n=== DiskSink Statistics: {} ===", self.node_id);
|
|
println!("Chunks written: {}", self.chunks_written);
|
|
println!("Total samples: {}", self.total_samples);
|
|
println!("Total duration: {:.3} sec", self.total_duration_sec);
|
|
println!("============================\n");
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::BitDepth;
|
|
|
|
#[tokio::test]
|
|
async fn test_disk_sink_basic() {
|
|
let temp_dir = std::env::temp_dir().join("pmoaudio_test");
|
|
tokio::fs::create_dir_all(&temp_dir).await.unwrap();
|
|
|
|
let config = DiskSinkConfig {
|
|
output_dir: temp_dir.clone(),
|
|
filename: Some("test_output.wav".to_string()),
|
|
format: AudioFileFormat::Wav,
|
|
buffer_size: 10,
|
|
};
|
|
|
|
let (sink, tx) = DiskSink::new("test".to_string(), config, 10);
|
|
|
|
let handle = tokio::spawn(async move { sink.run().await });
|
|
|
|
// Envoyer quelques chunks
|
|
for i in 0..5 {
|
|
let chunk = AudioChunk::from_channels_f32(
|
|
i,
|
|
vec![0.5; 1000],
|
|
vec![0.5; 1000],
|
|
48000,
|
|
BitDepth::B24,
|
|
);
|
|
tx.send(chunk).await.unwrap();
|
|
}
|
|
|
|
drop(tx);
|
|
|
|
let stats = handle.await.unwrap().unwrap();
|
|
assert_eq!(stats.chunks_written, 5);
|
|
|
|
// Vérifier que le fichier existe
|
|
let output_path = temp_dir.join("test_output.wav");
|
|
assert!(output_path.exists());
|
|
|
|
// Nettoyage
|
|
tokio::fs::remove_file(output_path).await.ok();
|
|
tokio::fs::remove_dir(temp_dir).await.ok();
|
|
}
|
|
}
|