rtl_sdr_rs/
async_read.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5#![cfg_attr(test, allow(dead_code))]
6
7//! Continuous IQ streaming backed by nusb's endpoint transfer queue.
8
9use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
10use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
11use std::sync::Arc;
12use std::thread::{self, JoinHandle};
13use std::time::Duration;
14
15use nusb::transfer::{Buffer, Bulk, In, TransferError};
16use nusb::Endpoint;
17
18use crate::error::{Result, RtlsdrError};
19use crate::{RtlSdr, TunerGain, DEFAULT_ASYNC_BUF_NUMBER, DEFAULT_BUF_LENGTH};
20
21const CANCEL_POLL: Duration = Duration::from_millis(100);
22
23/// A configuration change applied to an [`AsyncReadHandle`] stream.
24#[derive(Debug, Clone, PartialEq, Eq)]
25#[non_exhaustive]
26pub enum AsyncReadConfigChange {
27    CenterFrequency(u32),
28    TunerGain(TunerGain),
29    SampleRate(u32),
30}
31
32/// An item produced by an owned asynchronous reader.
33#[derive(Debug, PartialEq, Eq)]
34#[non_exhaustive]
35pub enum AsyncReadEvent {
36    /// IQ bytes captured under the indicated configuration generation.
37    Samples { generation: u64, data: Vec<u8> },
38    /// An exact stream boundary following a successful configuration change.
39    /// All subsequent samples carry this generation until the next change.
40    Reconfigured {
41        generation: u64,
42        change: AsyncReadConfigChange,
43    },
44}
45
46struct ControlCommand {
47    change: AsyncReadConfigChange,
48    reply: Sender<Result<u64>>,
49}
50
51/// Cloneable control access to an owned asynchronous reader.
52#[derive(Clone)]
53pub struct AsyncReadControlHandle {
54    command_tx: Sender<ControlCommand>,
55    stop: Arc<AtomicBool>,
56    dropped: Arc<AtomicU64>,
57}
58
59impl AsyncReadControlHandle {
60    /// Retunes the receiver and returns the new stream generation.
61    pub fn set_center_freq(&self, frequency: u32) -> Result<u64> {
62        self.send(AsyncReadConfigChange::CenterFrequency(frequency))
63    }
64
65    /// Alias for [`Self::set_center_freq`].
66    pub fn tune(&self, frequency: u32) -> Result<u64> {
67        self.set_center_freq(frequency)
68    }
69
70    /// Changes tuner gain and returns the new stream generation.
71    pub fn set_tuner_gain(&self, gain: TunerGain) -> Result<u64> {
72        self.send(AsyncReadConfigChange::TunerGain(gain))
73    }
74
75    /// Changes the sample rate and returns the new stream generation.
76    pub fn set_sample_rate(&self, sample_rate: u32) -> Result<u64> {
77        self.send(AsyncReadConfigChange::SampleRate(sample_rate))
78    }
79
80    /// Requests streaming shutdown.
81    pub fn stop(&self) {
82        self.stop.store(true, Ordering::Release);
83    }
84
85    /// Number of sample chunks dropped because the consumer queue was full.
86    pub fn dropped_chunks(&self) -> u64 {
87        self.dropped.load(Ordering::Acquire)
88    }
89
90    fn send(&self, change: AsyncReadConfigChange) -> Result<u64> {
91        if self.stop.load(Ordering::Acquire) {
92            return Err(control_error("async reader has stopped"));
93        }
94
95        let (reply, response) = mpsc::channel();
96        self.command_tx
97            .send(ControlCommand { change, reply })
98            .map_err(|_| control_error("async control channel is closed"))?;
99        response
100            .recv()
101            .map_err(|_| control_error("async reader stopped before acknowledging control"))?
102    }
103}
104
105/// An owned continuous IQ stream.
106///
107/// Dropping this handle stops the stream and joins its worker thread. Use
108/// [`Self::control_handle`] to adjust the device from another thread.
109pub struct AsyncReadHandle {
110    event_rx: Option<Receiver<Result<AsyncReadEvent>>>,
111    queued_samples: Arc<AtomicUsize>,
112    control: AsyncReadControlHandle,
113    worker: Option<JoinHandle<()>>,
114}
115
116impl AsyncReadHandle {
117    pub fn control_handle(&self) -> AsyncReadControlHandle {
118        self.control.clone()
119    }
120
121    pub fn recv(&self) -> Option<Result<AsyncReadEvent>> {
122        let event = self.event_rx.as_ref()?.recv().ok()?;
123        self.release_sample_slot(&event);
124        Some(event)
125    }
126
127    pub fn try_recv(&self) -> Option<Result<AsyncReadEvent>> {
128        let event = self.event_rx.as_ref()?.try_recv().ok()?;
129        self.release_sample_slot(&event);
130        Some(event)
131    }
132
133    pub fn stop(&self) {
134        self.control.stop();
135    }
136
137    pub fn dropped_chunks(&self) -> u64 {
138        self.control.dropped_chunks()
139    }
140
141    fn release_sample_slot(&self, event: &Result<AsyncReadEvent>) {
142        if matches!(event, Ok(AsyncReadEvent::Samples { .. })) {
143            self.queued_samples.fetch_sub(1, Ordering::AcqRel);
144        }
145    }
146}
147
148impl Iterator for AsyncReadHandle {
149    type Item = Result<AsyncReadEvent>;
150
151    fn next(&mut self) -> Option<Self::Item> {
152        self.recv()
153    }
154}
155
156impl Drop for AsyncReadHandle {
157    fn drop(&mut self) {
158        self.stop();
159        // Disconnect before joining so the worker notices that no events can
160        // be delivered while it winds down.
161        drop(self.event_rx.take());
162        if let Some(worker) = self.worker.take() {
163            let _ = worker.join();
164        }
165    }
166}
167
168/// One-shot tear-down signal for [`crate::RtlSdr::read_async`].
169///
170/// Once cancelled, a handle remains cancelled. Create a new handle for each
171/// subsequent streaming session.
172#[derive(Clone, Default, Debug)]
173pub struct CancelHandle {
174    flag: Arc<AtomicBool>,
175}
176
177impl CancelHandle {
178    pub fn new() -> Self {
179        Self::default()
180    }
181
182    pub fn cancel(&self) {
183        self.flag.store(true, Ordering::Release);
184    }
185
186    pub fn is_cancelled(&self) -> bool {
187        self.flag.load(Ordering::Acquire)
188    }
189}
190
191pub(crate) fn read_async_blocking<F>(
192    endpoint: &mut Endpoint<Bulk, In>,
193    buf_num: usize,
194    buf_len: usize,
195    cancel: &CancelHandle,
196    mut callback: F,
197) -> Result<()>
198where
199    F: FnMut(&[u8]),
200{
201    let buf_num = if buf_num == 0 {
202        DEFAULT_ASYNC_BUF_NUMBER
203    } else {
204        buf_num
205    };
206    let buf_len = if buf_len == 0 {
207        DEFAULT_BUF_LENGTH
208    } else {
209        buf_len
210    };
211
212    if buf_len % 512 != 0 {
213        return Err(RtlsdrError::RtlsdrErr(format!(
214            "Invalid async buffer length {buf_len} (must be multiple of 512)"
215        )));
216    }
217
218    if cancel.is_cancelled() {
219        return Ok(());
220    }
221
222    for _ in 0..buf_num {
223        endpoint.submit(Buffer::new(buf_len));
224    }
225
226    while endpoint.pending() > 0 {
227        if cancel.is_cancelled() {
228            endpoint.cancel_all();
229        }
230
231        let Some(completion) = endpoint.wait_next_complete(CANCEL_POLL) else {
232            continue;
233        };
234
235        let should_resubmit = !cancel.is_cancelled() && completion.status.is_ok();
236        if completion.actual_len > 0 {
237            callback(&completion.buffer[..completion.actual_len]);
238        }
239
240        match completion.status {
241            Ok(()) if should_resubmit => endpoint.submit(completion.buffer),
242            Ok(()) => {}
243            Err(TransferError::Cancelled) if cancel.is_cancelled() => {}
244            Err(e) => {
245                endpoint.cancel_all();
246                drain(endpoint);
247                return Err(RtlsdrError::from_usb_transfer(e));
248            }
249        }
250    }
251
252    Ok(())
253}
254
255pub(crate) fn start_async_reader(
256    sdr: RtlSdr,
257    buf_num: usize,
258    buf_len: usize,
259) -> Result<AsyncReadHandle> {
260    let (buf_num, buf_len) = normalize_buffer_config(buf_num, buf_len)?;
261    let endpoint = sdr.sdr.async_endpoint()?;
262    let queue_len = buf_num
263        .checked_mul(2)
264        .ok_or_else(|| control_error("async buffer count is too large"))?;
265    let (event_tx, event_rx) = mpsc::channel();
266    let (command_tx, command_rx) = mpsc::channel();
267    let stop = Arc::new(AtomicBool::new(false));
268    let dropped = Arc::new(AtomicU64::new(0));
269    let queued_samples = Arc::new(AtomicUsize::new(0));
270
271    let control = AsyncReadControlHandle {
272        command_tx,
273        stop: Arc::clone(&stop),
274        dropped: Arc::clone(&dropped),
275    };
276    let worker_queued_samples = Arc::clone(&queued_samples);
277    let worker = thread::spawn(move || {
278        OwnedReaderWorker {
279            sdr,
280            endpoint,
281            buf_num,
282            buf_len,
283            queue_len,
284            event_tx,
285            command_rx,
286            stop,
287            dropped,
288            queued_samples: worker_queued_samples,
289        }
290        .run()
291    });
292
293    Ok(AsyncReadHandle {
294        event_rx: Some(event_rx),
295        queued_samples,
296        control,
297        worker: Some(worker),
298    })
299}
300
301struct OwnedReaderWorker {
302    sdr: RtlSdr,
303    endpoint: Endpoint<Bulk, In>,
304    buf_num: usize,
305    buf_len: usize,
306    queue_len: usize,
307    event_tx: Sender<Result<AsyncReadEvent>>,
308    command_rx: Receiver<ControlCommand>,
309    stop: Arc<AtomicBool>,
310    dropped: Arc<AtomicU64>,
311    queued_samples: Arc<AtomicUsize>,
312}
313
314impl OwnedReaderWorker {
315    fn run(mut self) {
316        for _ in 0..self.buf_num {
317            self.endpoint.submit(Buffer::new(self.buf_len));
318        }
319
320        let mut generation = 0;
321        while self.endpoint.pending() > 0 && !self.stop.load(Ordering::Acquire) {
322            match self.command_rx.try_recv() {
323                Ok(command) => {
324                    let buffers = cancel_and_collect(&mut self.endpoint);
325                    match apply_change(&mut self.sdr, &command.change) {
326                        Ok(()) => {
327                            if let Err(error) = self.sdr.reset_buffer() {
328                                let message = format!(
329                                    "failed to reset buffer after reconfiguration: {error}"
330                                );
331                                let _ = command.reply.send(Err(control_error(&message)));
332                                break;
333                            }
334
335                            generation += 1;
336                            let _ = command.reply.send(Ok(generation));
337                            let event = AsyncReadEvent::Reconfigured {
338                                generation,
339                                change: command.change,
340                            };
341                            if self.event_tx.send(Ok(event)).is_err() {
342                                break;
343                            }
344                        }
345                        Err(error) => {
346                            let _ = command.reply.send(Err(error));
347                        }
348                    }
349
350                    if self.stop.load(Ordering::Acquire) {
351                        break;
352                    }
353                    for buffer in buffers {
354                        self.endpoint.submit(buffer);
355                    }
356                    continue;
357                }
358                Err(TryRecvError::Disconnected | TryRecvError::Empty) => {}
359            }
360
361            let Some(completion) = self.endpoint.wait_next_complete(CANCEL_POLL) else {
362                continue;
363            };
364
365            if let Err(error) = completion.status {
366                self.endpoint.cancel_all();
367                drain(&mut self.endpoint);
368                let _ = self
369                    .event_tx
370                    .send(Err(RtlsdrError::from_usb_transfer(error)));
371                break;
372            }
373
374            if completion.actual_len > 0 {
375                if reserve_sample_slot(&self.queued_samples, self.queue_len) {
376                    let event = AsyncReadEvent::Samples {
377                        generation,
378                        data: completion.buffer[..completion.actual_len].to_vec(),
379                    };
380                    if self.event_tx.send(Ok(event)).is_err() {
381                        self.queued_samples.fetch_sub(1, Ordering::AcqRel);
382                        break;
383                    }
384                } else {
385                    self.dropped.fetch_add(1, Ordering::AcqRel);
386                }
387            }
388            self.endpoint.submit(completion.buffer);
389        }
390
391        self.endpoint.cancel_all();
392        drain(&mut self.endpoint);
393        drop(self.endpoint);
394        self.stop.store(true, Ordering::Release);
395        let _ = self.sdr.close();
396    }
397}
398
399fn reserve_sample_slot(queued: &AtomicUsize, limit: usize) -> bool {
400    queued
401        .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
402            (count < limit).then_some(count + 1)
403        })
404        .is_ok()
405}
406
407fn apply_change(sdr: &mut RtlSdr, change: &AsyncReadConfigChange) -> Result<()> {
408    match change {
409        AsyncReadConfigChange::CenterFrequency(frequency) => sdr.set_center_freq(*frequency),
410        AsyncReadConfigChange::TunerGain(gain) => sdr.set_tuner_gain(gain.clone()),
411        AsyncReadConfigChange::SampleRate(sample_rate) => sdr.set_sample_rate(*sample_rate),
412    }
413}
414
415fn cancel_and_collect(endpoint: &mut Endpoint<Bulk, In>) -> Vec<Buffer> {
416    endpoint.cancel_all();
417    let mut buffers = Vec::with_capacity(endpoint.pending());
418    while endpoint.pending() > 0 {
419        if let Some(completion) = endpoint.wait_next_complete(CANCEL_POLL) {
420            buffers.push(completion.buffer);
421        }
422    }
423    buffers
424}
425
426fn normalize_buffer_config(buf_num: usize, buf_len: usize) -> Result<(usize, usize)> {
427    let buf_num = if buf_num == 0 {
428        DEFAULT_ASYNC_BUF_NUMBER
429    } else {
430        buf_num
431    };
432    let buf_len = if buf_len == 0 {
433        DEFAULT_BUF_LENGTH
434    } else {
435        buf_len
436    };
437    if buf_len % 512 != 0 {
438        return Err(control_error(&format!(
439            "Invalid async buffer length {buf_len} (must be multiple of 512)"
440        )));
441    }
442    Ok((buf_num, buf_len))
443}
444
445fn control_error(message: &str) -> RtlsdrError {
446    RtlsdrError::RtlsdrErr(message.to_string())
447}
448
449fn drain(endpoint: &mut Endpoint<Bulk, In>) {
450    while endpoint.pending() > 0 {
451        let _ = endpoint.wait_next_complete(CANCEL_POLL);
452    }
453}
454
455#[cfg(test)]
456mod tests {
457    use super::{
458        normalize_buffer_config, reserve_sample_slot, AsyncReadConfigChange,
459        AsyncReadControlHandle, CancelHandle,
460    };
461    use crate::{DEFAULT_ASYNC_BUF_NUMBER, DEFAULT_BUF_LENGTH};
462    use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
463    use std::sync::{mpsc, Arc};
464    use std::thread;
465
466    #[test]
467    fn cancellation_is_shared_and_sticky() {
468        let cancel = CancelHandle::new();
469        let clone = cancel.clone();
470
471        assert!(!cancel.is_cancelled());
472        clone.cancel();
473        assert!(cancel.is_cancelled());
474        assert!(clone.is_cancelled());
475    }
476
477    #[test]
478    fn owned_reader_buffer_config_uses_defaults_and_validates_alignment() {
479        assert_eq!(
480            normalize_buffer_config(0, 0).unwrap(),
481            (DEFAULT_ASYNC_BUF_NUMBER, DEFAULT_BUF_LENGTH)
482        );
483        assert!(normalize_buffer_config(1, 511).is_err());
484        assert_eq!(normalize_buffer_config(3, 1024).unwrap(), (3, 1024));
485    }
486
487    #[test]
488    fn control_calls_wait_for_the_worker_acknowledgement() {
489        let (command_tx, command_rx) = mpsc::channel();
490        let control = AsyncReadControlHandle {
491            command_tx,
492            stop: Arc::new(AtomicBool::new(false)),
493            dropped: Arc::new(AtomicU64::new(0)),
494        };
495
496        let caller = thread::spawn(move || control.set_center_freq(101_100_000));
497        let command = command_rx.recv().unwrap();
498        assert_eq!(
499            command.change,
500            AsyncReadConfigChange::CenterFrequency(101_100_000)
501        );
502        command.reply.send(Ok(7)).unwrap();
503        assert_eq!(caller.join().unwrap().unwrap(), 7);
504    }
505
506    #[test]
507    fn stopped_control_handle_rejects_new_commands() {
508        let (command_tx, _command_rx) = mpsc::channel();
509        let control = AsyncReadControlHandle {
510            command_tx,
511            stop: Arc::new(AtomicBool::new(true)),
512            dropped: Arc::new(AtomicU64::new(0)),
513        };
514
515        assert!(control.set_sample_rate(1_920_000).is_err());
516    }
517
518    #[test]
519    fn sample_queue_reservation_stops_at_the_limit() {
520        let queued = AtomicUsize::new(0);
521        assert!(reserve_sample_slot(&queued, 2));
522        assert!(reserve_sample_slot(&queued, 2));
523        assert!(!reserve_sample_slot(&queued, 2));
524        assert_eq!(queued.load(Ordering::Acquire), 2);
525    }
526}
rtl_sdr_rs/async_read.rs

rtl_sdr_rs/
async_read.rs
