tenant.openagents/omega

1use std::{
2    num::NonZero,
3    sync::{
4        Arc, Mutex,
5        atomic::{AtomicBool, Ordering},
6    },
7    time::Duration,
8};
9
10use crossbeam::queue::ArrayQueue;
11use log::warn;
12use rodio::{
13    ChannelCount, Sample, SampleRate, Source, conversions::SampleRateConverter, nz,
14    source::UniformSourceIterator,
15};
16
17const MAX_CHANNELS: usize = 8;
18
19#[derive(Debug, thiserror::Error)]
20#[error("Replay duration is too short must be >= 100ms")]
21pub struct ReplayDurationTooShort;
22
23// These all require constant sources (so the span is infinitely long)
24// this is not guaranteed by rodio however we know it to be true in all our
25// applications. Rodio desperately needs a constant source concept.
26pub trait RodioExt: Source + Sized {
27    fn process_buffer<const N: usize, F>(self, callback: F) -> ProcessBuffer<N, Self, F>
28    where
29        F: FnMut(&mut [Sample; N]);
30    fn inspect_buffer<const N: usize, F>(self, callback: F) -> InspectBuffer<N, Self, F>
31    where
32        F: FnMut(&[Sample; N]);
33    fn replayable(
34        self,
35        duration: Duration,
36    ) -> Result<(Replay, Replayable<Self>), ReplayDurationTooShort>;
37    fn take_samples(self, n: usize) -> TakeSamples<Self>;
38    fn constant_params(
39        self,
40        channel_count: ChannelCount,
41        sample_rate: SampleRate,
42    ) -> UniformSourceIterator<Self>;
43    fn constant_samplerate(self, sample_rate: SampleRate) -> ConstantSampleRate<Self>;
44    fn possibly_disconnected_channels_to_mono(self) -> ToMono<Self>;
45}
46
47impl<S: Source> RodioExt for S {
48    fn process_buffer<const N: usize, F>(self, callback: F) -> ProcessBuffer<N, Self, F>
49    where
50        F: FnMut(&mut [Sample; N]),
51    {
52        ProcessBuffer {
53            inner: self,
54            callback,
55            buffer: [0.0; N],
56            next: N,
57        }
58    }
59    fn inspect_buffer<const N: usize, F>(self, callback: F) -> InspectBuffer<N, Self, F>
60    where
61        F: FnMut(&[Sample; N]),
62    {
63        InspectBuffer {
64            inner: self,
65            callback,
66            buffer: [0.0; N],
67            free: 0,
68        }
69    }
70    /// Maintains a live replay with a history of at least `duration` seconds.
71    ///
72    /// Note:
73    /// History can be 100ms longer if the source drops before or while the
74    /// replay is being read
75    ///
76    /// # Errors
77    /// If duration is smaller than 100ms
78    fn replayable(
79        self,
80        duration: Duration,
81    ) -> Result<(Replay, Replayable<Self>), ReplayDurationTooShort> {
82        if duration < Duration::from_millis(100) {
83            return Err(ReplayDurationTooShort);
84        }
85
86        let samples_per_second = self.sample_rate().get() as usize * self.channels().get() as usize;
87        let samples_to_queue = duration.as_secs_f64() * samples_per_second as f64;
88        let samples_to_queue =
89            (samples_to_queue as usize).next_multiple_of(self.channels().get().into());
90
91        let chunk_size =
92            (samples_per_second.div_ceil(10)).next_multiple_of(self.channels().get() as usize);
93        let chunks_to_queue = samples_to_queue.div_ceil(chunk_size);
94
95        let is_active = Arc::new(AtomicBool::new(true));
96        let queue = Arc::new(ReplayQueue::new(chunks_to_queue, chunk_size));
97        Ok((
98            Replay {
99                rx: Arc::clone(&queue),
100                buffer: Vec::new().into_iter(),
101                sleep_duration: duration / 2,
102                sample_rate: self.sample_rate(),
103                channel_count: self.channels(),
104                source_is_active: is_active.clone(),
105            },
106            Replayable {
107                tx: queue,
108                inner: self,
109                buffer: Vec::with_capacity(chunk_size),
110                chunk_size,
111                is_active,
112            },
113        ))
114    }
115    fn take_samples(self, n: usize) -> TakeSamples<S> {
116        TakeSamples {
117            inner: self,
118            left_to_take: n,
119        }
120    }
121    fn constant_params(
122        self,
123        channel_count: ChannelCount,
124        sample_rate: SampleRate,
125    ) -> UniformSourceIterator<Self> {
126        UniformSourceIterator::new(self, channel_count, sample_rate)
127    }
128    fn constant_samplerate(self, sample_rate: SampleRate) -> ConstantSampleRate<Self> {
129        ConstantSampleRate::new(self, sample_rate)
130    }
131    fn possibly_disconnected_channels_to_mono(self) -> ToMono<Self> {
132        ToMono::new(self)
133    }
134}
135
136pub struct ConstantSampleRate<S: Source> {
137    inner: SampleRateConverter<S>,
138    channels: ChannelCount,
139    sample_rate: SampleRate,
140}
141
142impl<S: Source> ConstantSampleRate<S> {
143    fn new(source: S, target_rate: SampleRate) -> Self {
144        let input_sample_rate = source.sample_rate();
145        let channels = source.channels();
146        let inner = SampleRateConverter::new(source, input_sample_rate, target_rate, channels);
147        Self {
148            inner,
149            channels,
150            sample_rate: target_rate,
151        }
152    }
153}
154
155impl<S: Source> Iterator for ConstantSampleRate<S> {
156    type Item = rodio::Sample;
157
158    fn next(&mut self) -> Option<Self::Item> {
159        self.inner.next()
160    }
161
162    fn size_hint(&self) -> (usize, Option<usize>) {
163        self.inner.size_hint()
164    }
165}
166
167impl<S: Source> Source for ConstantSampleRate<S> {
168    fn current_span_len(&self) -> Option<usize> {
169        None
170    }
171
172    fn channels(&self) -> ChannelCount {
173        self.channels
174    }
175
176    fn sample_rate(&self) -> SampleRate {
177        self.sample_rate
178    }
179
180    fn total_duration(&self) -> Option<Duration> {
181        None // not supported (not used by us)
182    }
183}
184
185const TYPICAL_NOISE_FLOOR: Sample = 1e-3;
186
187/// constant source, only works on a single span
188pub struct ToMono<S> {
189    inner: S,
190    input_channel_count: ChannelCount,
191    connected_channels: ChannelCount,
192    /// running mean of second channel 'volume'
193    means: [f32; MAX_CHANNELS],
194}
195impl<S: Source> ToMono<S> {
196    fn new(input: S) -> Self {
197        let channels = input
198            .channels()
199            .min(const { NonZero::<u16>::new(MAX_CHANNELS as u16).unwrap() });
200        if channels < input.channels() {
201            warn!("Ignoring input channels {}..", channels.get());
202        }
203
204        Self {
205            connected_channels: channels,
206            input_channel_count: channels,
207            inner: input,
208            means: [TYPICAL_NOISE_FLOOR; MAX_CHANNELS],
209        }
210    }
211}
212
213impl<S: Source> Source for ToMono<S> {
214    fn current_span_len(&self) -> Option<usize> {
215        None
216    }
217
218    fn channels(&self) -> ChannelCount {
219        rodio::nz!(1)
220    }
221
222    fn sample_rate(&self) -> SampleRate {
223        self.inner.sample_rate()
224    }
225
226    fn total_duration(&self) -> Option<Duration> {
227        self.inner.total_duration()
228    }
229}
230
231fn update_mean(mean: &mut f32, sample: Sample) {
232    const HISTORY: f32 = 500.0;
233    *mean *= (HISTORY - 1.0) / HISTORY;
234    *mean += sample.abs() / HISTORY;
235}
236
237impl<S: Source> Iterator for ToMono<S> {
238    type Item = Sample;
239
240    fn next(&mut self) -> Option<Self::Item> {
241        let mut mono_sample = 0f32;
242        let mut active_channels = 0;
243        for channel in 0..self.input_channel_count.get() as usize {
244            let sample = self.inner.next()?;
245            mono_sample += sample;
246
247            update_mean(&mut self.means[channel], sample);
248            if self.means[channel] > TYPICAL_NOISE_FLOOR / 10.0 {
249                active_channels += 1;
250            }
251        }
252        mono_sample /= self.connected_channels.get() as f32;
253        self.connected_channels = NonZero::new(active_channels).unwrap_or(nz!(1));
254
255        Some(mono_sample)
256    }
257}
258
259/// constant source, only works on a single span
260pub struct TakeSamples<S> {
261    inner: S,
262    left_to_take: usize,
263}
264
265impl<S: Source> Iterator for TakeSamples<S> {
266    type Item = Sample;
267
268    fn next(&mut self) -> Option<Self::Item> {
269        if self.left_to_take == 0 {
270            None
271        } else {
272            self.left_to_take -= 1;
273            self.inner.next()
274        }
275    }
276
277    fn size_hint(&self) -> (usize, Option<usize>) {
278        (0, Some(self.left_to_take))
279    }
280}
281
282impl<S: Source> Source for TakeSamples<S> {
283    fn current_span_len(&self) -> Option<usize> {
284        None // does not support spans
285    }
286
287    fn channels(&self) -> ChannelCount {
288        self.inner.channels()
289    }
290
291    fn sample_rate(&self) -> SampleRate {
292        self.inner.sample_rate()
293    }
294
295    fn total_duration(&self) -> Option<Duration> {
296        Some(Duration::from_secs_f64(
297            self.left_to_take as f64
298                / self.sample_rate().get() as f64
299                / self.channels().get() as f64,
300        ))
301    }
302}
303
304/// constant source, only works on a single span
305#[derive(Debug)]
306struct ReplayQueue {
307    inner: ArrayQueue<Vec<Sample>>,
308    normal_chunk_len: usize,
309    /// The last chunk in the queue may be smaller than
310    /// the normal chunk size. This is always equal to the
311    /// size of the last element in the queue.
312    /// (so normally chunk_size)
313    last_chunk: Mutex<Vec<Sample>>,
314}
315
316impl ReplayQueue {
317    fn new(queue_len: usize, chunk_size: usize) -> Self {
318        Self {
319            inner: ArrayQueue::new(queue_len),
320            normal_chunk_len: chunk_size,
321            last_chunk: Mutex::new(Vec::new()),
322        }
323    }
324    /// Returns the length in samples
325    fn len(&self) -> usize {
326        self.inner.len().saturating_sub(1) * self.normal_chunk_len
327            + self
328                .last_chunk
329                .lock()
330                .expect("Self::push_last can not poison this lock")
331                .len()
332    }
333
334    fn pop(&self) -> Option<Vec<Sample>> {
335        self.inner.pop() // removes element that was inserted first
336    }
337
338    fn push_last(&self, mut samples: Vec<Sample>) {
339        let mut last_chunk = self
340            .last_chunk
341            .lock()
342            .expect("Self::len can not poison this lock");
343        std::mem::swap(&mut *last_chunk, &mut samples);
344    }
345
346    fn push_normal(&self, samples: Vec<Sample>) {
347        let _pushed_out_of_ringbuf = self.inner.force_push(samples);
348    }
349}
350
351/// constant source, only works on a single span
352pub struct ProcessBuffer<const N: usize, S, F>
353where
354    S: Source + Sized,
355    F: FnMut(&mut [Sample; N]),
356{
357    inner: S,
358    callback: F,
359    /// Buffer used for both input and output.
360    buffer: [Sample; N],
361    /// Next already processed sample is at this index
362    /// in buffer.
363    ///
364    /// If this is equal to the length of the buffer we have no more samples and
365    /// we must get new ones and process them
366    next: usize,
367}
368
369impl<const N: usize, S, F> Iterator for ProcessBuffer<N, S, F>
370where
371    S: Source + Sized,
372    F: FnMut(&mut [Sample; N]),
373{
374    type Item = Sample;
375
376    fn next(&mut self) -> Option<Self::Item> {
377        self.next += 1;
378        if self.next < self.buffer.len() {
379            let sample = self.buffer[self.next];
380            return Some(sample);
381        }
382
383        for sample in &mut self.buffer {
384            *sample = self.inner.next()?
385        }
386        (self.callback)(&mut self.buffer);
387
388        self.next = 0;
389        Some(self.buffer[0])
390    }
391
392    fn size_hint(&self) -> (usize, Option<usize>) {
393        self.inner.size_hint()
394    }
395}
396
397impl<const N: usize, S, F> Source for ProcessBuffer<N, S, F>
398where
399    S: Source + Sized,
400    F: FnMut(&mut [Sample; N]),
401{
402    fn current_span_len(&self) -> Option<usize> {
403        None
404    }
405
406    fn channels(&self) -> rodio::ChannelCount {
407        self.inner.channels()
408    }
409
410    fn sample_rate(&self) -> rodio::SampleRate {
411        self.inner.sample_rate()
412    }
413
414    fn total_duration(&self) -> Option<std::time::Duration> {
415        self.inner.total_duration()
416    }
417}
418
419/// constant source, only works on a single span
420pub struct InspectBuffer<const N: usize, S, F>
421where
422    S: Source + Sized,
423    F: FnMut(&[Sample; N]),
424{
425    inner: S,
426    callback: F,
427    /// Stores already emitted samples, once its full we call the callback.
428    buffer: [Sample; N],
429    /// Next free element in buffer. If this is equal to the buffer length
430    /// we have no more free elements.
431    free: usize,
432}
433
434impl<const N: usize, S, F> Iterator for InspectBuffer<N, S, F>
435where
436    S: Source + Sized,
437    F: FnMut(&[Sample; N]),
438{
439    type Item = Sample;
440
441    fn next(&mut self) -> Option<Self::Item> {
442        let Some(sample) = self.inner.next() else {
443            return None;
444        };
445
446        self.buffer[self.free] = sample;
447        self.free += 1;
448
449        if self.free == self.buffer.len() {
450            (self.callback)(&self.buffer);
451            self.free = 0
452        }
453
454        Some(sample)
455    }
456
457    fn size_hint(&self) -> (usize, Option<usize>) {
458        self.inner.size_hint()
459    }
460}
461
462impl<const N: usize, S, F> Source for InspectBuffer<N, S, F>
463where
464    S: Source + Sized,
465    F: FnMut(&[Sample; N]),
466{
467    fn current_span_len(&self) -> Option<usize> {
468        None
469    }
470
471    fn channels(&self) -> rodio::ChannelCount {
472        self.inner.channels()
473    }
474
475    fn sample_rate(&self) -> rodio::SampleRate {
476        self.inner.sample_rate()
477    }
478
479    fn total_duration(&self) -> Option<std::time::Duration> {
480        self.inner.total_duration()
481    }
482}
483
484/// constant source, only works on a single span
485#[derive(Debug)]
486pub struct Replayable<S: Source> {
487    inner: S,
488    buffer: Vec<Sample>,
489    chunk_size: usize,
490    tx: Arc<ReplayQueue>,
491    is_active: Arc<AtomicBool>,
492}
493
494impl<S: Source> Iterator for Replayable<S> {
495    type Item = Sample;
496
497    fn next(&mut self) -> Option<Self::Item> {
498        if let Some(sample) = self.inner.next() {
499            self.buffer.push(sample);
500            // If the buffer is full send it
501            if self.buffer.len() == self.chunk_size {
502                self.tx.push_normal(std::mem::take(&mut self.buffer));
503            }
504            Some(sample)
505        } else {
506            let last_chunk = std::mem::take(&mut self.buffer);
507            self.tx.push_last(last_chunk);
508            self.is_active.store(false, Ordering::Relaxed);
509            None
510        }
511    }
512
513    fn size_hint(&self) -> (usize, Option<usize>) {
514        self.inner.size_hint()
515    }
516}
517
518impl<S: Source> Source for Replayable<S> {
519    fn current_span_len(&self) -> Option<usize> {
520        self.inner.current_span_len()
521    }
522
523    fn channels(&self) -> ChannelCount {
524        self.inner.channels()
525    }
526
527    fn sample_rate(&self) -> SampleRate {
528        self.inner.sample_rate()
529    }
530
531    fn total_duration(&self) -> Option<Duration> {
532        self.inner.total_duration()
533    }
534}
535
536/// constant source, only works on a single span
537#[derive(Debug)]
538pub struct Replay {
539    rx: Arc<ReplayQueue>,
540    buffer: std::vec::IntoIter<Sample>,
541    sleep_duration: Duration,
542    sample_rate: SampleRate,
543    channel_count: ChannelCount,
544    source_is_active: Arc<AtomicBool>,
545}
546
547impl Replay {
548    pub fn source_is_active(&self) -> bool {
549        // - source could return None and not drop
550        // - source could be dropped before returning None
551        self.source_is_active.load(Ordering::Relaxed) && Arc::strong_count(&self.rx) < 2
552    }
553
554    /// Duration of what is in the buffer and can be returned without blocking.
555    pub fn duration_ready(&self) -> Duration {
556        let samples_per_second = self.channels().get() as u32 * self.sample_rate().get();
557
558        let seconds_queued = self.samples_ready() as f64 / samples_per_second as f64;
559        Duration::from_secs_f64(seconds_queued)
560    }
561
562    /// Number of samples in the buffer and can be returned without blocking.
563    pub fn samples_ready(&self) -> usize {
564        self.rx.len() + self.buffer.len()
565    }
566}
567
568impl Iterator for Replay {
569    type Item = Sample;
570
571    fn next(&mut self) -> Option<Self::Item> {
572        if let Some(sample) = self.buffer.next() {
573            return Some(sample);
574        }
575
576        loop {
577            if let Some(new_buffer) = self.rx.pop() {
578                self.buffer = new_buffer.into_iter();
579                return self.buffer.next();
580            }
581
582            if !self.source_is_active() {
583                return None;
584            }
585
586            // The queue does not support blocking on a next item. We want this queue as it
587            // is quite fast and provides a fixed size. We know how many samples are in a
588            // buffer so if we do not get one now we must be getting one after `sleep_duration`.
589            std::thread::sleep(self.sleep_duration);
590        }
591    }
592
593    fn size_hint(&self) -> (usize, Option<usize>) {
594        ((self.rx.len() + self.buffer.len()), None)
595    }
596}
597
598impl Source for Replay {
599    fn current_span_len(&self) -> Option<usize> {
600        None // source is not compatible with spans
601    }
602
603    fn channels(&self) -> ChannelCount {
604        self.channel_count
605    }
606
607    fn sample_rate(&self) -> SampleRate {
608        self.sample_rate
609    }
610
611    fn total_duration(&self) -> Option<Duration> {
612        None
613    }
614}
615
616#[cfg(test)]
617mod tests {
618    use rodio::{nz, static_buffer::StaticSamplesBuffer};
619
620    use super::*;
621
622    const SAMPLES: [Sample; 5] = [0.0, 1.0, 2.0, 3.0, 4.0];
623
624    fn test_source() -> StaticSamplesBuffer {
625        StaticSamplesBuffer::new(nz!(1), nz!(1), &SAMPLES)
626    }
627
628    mod process_buffer {
629        use super::*;
630
631        #[test]
632        fn callback_gets_all_samples() {
633            let input = test_source();
634
635            let _ = input
636                .process_buffer::<{ SAMPLES.len() }, _>(|buffer| assert_eq!(*buffer, SAMPLES))
637                .count();
638        }
639        #[test]
640        fn callback_modifies_yielded() {
641            let input = test_source();
642
643            let yielded: Vec<_> = input
644                .process_buffer::<{ SAMPLES.len() }, _>(|buffer| {
645                    for sample in buffer {
646                        *sample += 1.0;
647                    }
648                })
649                .collect();
650            assert_eq!(
651                yielded,
652                SAMPLES.into_iter().map(|s| s + 1.0).collect::<Vec<_>>()
653            )
654        }
655        #[test]
656        fn source_truncates_to_whole_buffers() {
657            let input = test_source();
658
659            let yielded = input
660                .process_buffer::<3, _>(|buffer| assert_eq!(buffer, &SAMPLES[..3]))
661                .count();
662            assert_eq!(yielded, 3)
663        }
664    }
665
666    mod inspect_buffer {
667        use super::*;
668
669        #[test]
670        fn callback_gets_all_samples() {
671            let input = test_source();
672
673            let _ = input
674                .inspect_buffer::<{ SAMPLES.len() }, _>(|buffer| assert_eq!(*buffer, SAMPLES))
675                .count();
676        }
677        #[test]
678        fn source_does_not_truncate() {
679            let input = test_source();
680
681            let yielded = input
682                .inspect_buffer::<3, _>(|buffer| assert_eq!(buffer, &SAMPLES[..3]))
683                .count();
684            assert_eq!(yielded, SAMPLES.len())
685        }
686    }
687
688    mod instant_replay {
689        use super::*;
690
691        #[test]
692        fn continues_after_history() {
693            let input = test_source();
694
695            let (mut replay, mut source) = input
696                .replayable(Duration::from_secs(3))
697                .expect("longer than 100ms");
698
699            source.by_ref().take(3).count();
700            let yielded: Vec<Sample> = replay.by_ref().take(3).collect();
701            assert_eq!(&yielded, &SAMPLES[0..3],);
702
703            source.count();
704            let yielded: Vec<Sample> = replay.collect();
705            assert_eq!(&yielded, &SAMPLES[3..5],);
706        }
707
708        #[test]
709        fn keeps_only_latest() {
710            let input = test_source();
711
712            let (mut replay, mut source) = input
713                .replayable(Duration::from_secs(2))
714                .expect("longer than 100ms");
715
716            source.by_ref().take(5).count(); // get all items but do not end the source
717            let yielded: Vec<Sample> = replay.by_ref().take(2).collect();
718            assert_eq!(&yielded, &SAMPLES[3..5]);
719            source.count(); // exhaust source
720            assert_eq!(replay.next(), None);
721        }
722
723        #[test]
724        fn keeps_correct_amount_of_seconds() {
725            let input = StaticSamplesBuffer::new(nz!(1), nz!(16_000), &[0.0; 40_000]);
726
727            let (replay, mut source) = input
728                .replayable(Duration::from_secs(2))
729                .expect("longer than 100ms");
730
731            // exhaust but do not yet end source
732            source.by_ref().take(40_000).count();
733
734            // take all samples we can without blocking
735            let ready = replay.samples_ready();
736            let n_yielded = replay.take_samples(ready).count();
737
738            let max = source.sample_rate().get() * source.channels().get() as u32 * 2;
739            let margin = 16_000 / 10; // 100ms
740            assert!(n_yielded as u32 >= max - margin);
741        }
742
743        #[test]
744        fn samples_ready() {
745            let input = StaticSamplesBuffer::new(nz!(1), nz!(16_000), &[0.0; 40_000]);
746            let (mut replay, source) = input
747                .replayable(Duration::from_secs(2))
748                .expect("longer than 100ms");
749            assert_eq!(replay.by_ref().samples_ready(), 0);
750
751            source.take(8000).count(); // half a second
752            let margin = 16_000 / 10; // 100ms
753            let ready = replay.samples_ready();
754            assert!(ready >= 8000 - margin);
755        }
756    }
757}
758