Skip to main content

aic_sdk/
processor_async.rs

1use crate::{AicError, Model, OtelConfig, Processor, ProcessorConfig, ProcessorContext};
2use async_lock::Mutex;
3use futures_channel::oneshot;
4use std::sync::{Arc, OnceLock};
5
6static RAYON_POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
7
8pub(crate) fn get_global_thread_pool() -> &'static rayon::ThreadPool {
9    RAYON_POOL.get_or_init(|| {
10        let num_threads = std::env::var("AIC_NUM_THREADS")
11            .ok()
12            .and_then(|s| s.parse::<usize>().ok())
13            .filter(|&n| n > 0)
14            .unwrap_or_else(|| {
15                std::thread::available_parallelism()
16                    .map(|n| n.get())
17                    .unwrap_or(1)
18            });
19
20        rayon::ThreadPoolBuilder::new()
21            .num_threads(num_threads)
22            .thread_name(|i| format!("aic-processing-thread-{i}"))
23            .build()
24            .expect("failed to build aic thread pool")
25    })
26}
27
28/// A wrapper around [`Processor`] for use in async contexts.
29///
30/// # Threading
31///
32/// Processing runs on a background thread pool shared across all
33/// [`ProcessorAsync`] instances. The pool defaults to one thread per logical
34/// CPU. Override with the `AIC_NUM_THREADS` environment variable, which is
35/// read once on first use.
36///
37/// # Example
38///
39/// ```rust,no_run
40/// use aic_sdk::{Model, ProcessorAsync, ProcessorConfig};
41/// #[tokio::main]
42/// async fn main() -> Result<(), aic_sdk::AicError> {
43///     let license_key = std::env::var("AIC_SDK_LICENSE").unwrap();
44///     let model = Model::from_file("/path/to/model.aicmodel")?;
45///     let config = ProcessorConfig::optimal(&model);
46///
47///     let processor = ProcessorAsync::new(&model, &license_key)?.with_config(&config).await?;
48///
49///     let mut audio = vec![0.0f32; config.block_size];
50///     let audio = processor.process(audio).await?;
51///     Ok(())
52/// }
53/// ```
54pub struct ProcessorAsync {
55    inner: Arc<Mutex<Processor<'static>>>,
56}
57
58impl ProcessorAsync {
59    /// Creates a new async audio enhancement processor instance.
60    ///
61    /// See [`Processor::new`] for details.
62    pub fn new(model: &Model<'static>, license_key: &str) -> Result<Self, AicError> {
63        let processor = Processor::new(model, license_key)?;
64        Ok(Self {
65            inner: Arc::new(Mutex::new(processor)),
66        })
67    }
68
69    /// Creates a new async audio enhancement processor instance with explicit
70    /// OpenTelemetry configuration.
71    ///
72    /// See [`Processor::with_otel_config`] for details.
73    pub fn with_otel_config(
74        model: &Model<'static>,
75        license_key: &str,
76        otel_config: &OtelConfig,
77    ) -> Result<Self, AicError> {
78        let processor = Processor::with_otel_config(model, license_key, otel_config)?;
79        Ok(Self {
80            inner: Arc::new(Mutex::new(processor)),
81        })
82    }
83
84    /// Initializes the async processor with the given configuration.
85    ///
86    /// This is a convenience method that calls [`ProcessorAsync::initialize`]
87    /// internally and returns `self`.
88    pub async fn with_config(self, config: &ProcessorConfig) -> Result<Self, AicError> {
89        self.initialize(config).await?;
90        Ok(self)
91    }
92
93    /// Initializes the processor with the given configuration.
94    ///
95    /// See [`Processor::initialize`] for details.
96    ///
97    /// # Warning
98    /// This allocates memory internally. Do not call from latency-sensitive paths.
99    pub async fn initialize(&self, config: &ProcessorConfig) -> Result<(), AicError> {
100        let config = config.clone();
101        let (tx, rx) = oneshot::channel();
102        let mut processor = self.inner.lock_arc().await;
103        get_global_thread_pool().spawn(move || {
104            let _ = tx.send(processor.initialize(&config));
105        });
106        rx.await.expect("Rayon worker dropped")
107    }
108
109    /// Processes mono audio.
110    ///
111    /// This method takes ownership of `audio`, moves it to a background processing
112    /// thread, and returns the processed audio block.
113    ///
114    /// See [`Processor::process`] for details.
115    pub async fn process(&self, mut audio: Vec<f32>) -> Result<Vec<f32>, AicError> {
116        let (tx, rx) = oneshot::channel();
117        let mut processor = self.inner.lock_arc().await;
118        get_global_thread_pool().spawn(move || {
119            let result = processor.process(&mut audio).map(|_| audio);
120            let _ = tx.send(result);
121        });
122        rx.await.expect("Rayon worker dropped")
123    }
124
125    /// Terminates the telemetry session associated with this processor.
126    ///
127    /// See [`Processor::terminate_session`] for details.
128    ///
129    /// # Warning
130    /// This may block until the session is terminated, so it runs on the background
131    /// thread pool rather than the calling task.
132    pub async fn terminate_session(&self) -> Result<(), AicError> {
133        let (tx, rx) = oneshot::channel();
134        let mut processor = self.inner.lock_arc().await;
135        get_global_thread_pool().spawn(move || {
136            let _ = tx.send(processor.terminate_session());
137        });
138        rx.await.expect("Rayon worker dropped")
139    }
140
141    /// Returns a [`ProcessorContext`] for real-time parameter control.
142    ///
143    /// See [`Processor::context`] for details.
144    pub async fn context(&self) -> ProcessorContext {
145        self.inner.lock().await.context()
146    }
147}