add frame-size-ms to AudioStream (#658)

Author: theomonnom

livekit-ffi/protocol/audio_frame.proto

@@ -29,6 +29,7 @@ message NewAudioStreamRequest {
   optional uint32 num_channels = 4;
   optional string audio_filter_module_id = 5; // Unique identifier passed in LoadAudioFilterPluginRequest
   optional string audio_filter_options = 6;
+  optional uint32 frame_size_ms = 7;
 }
 message NewAudioStreamResponse { required OwnedAudioStream stream = 1; }
 
@@ -40,6 +41,7 @@ message AudioStreamFromParticipantRequest {
   optional uint32 num_channels = 6;
   optional string audio_filter_module_id = 7;
   optional string audio_filter_options = 8;
+  optional uint32 frame_size_ms = 9;
 }
 
 message AudioStreamFromParticipantResponse { required OwnedAudioStream stream = 1; }

livekit-ffi/src/livekit.proto.rs

@@ -4126,6 +4126,8 @@ pub struct NewAudioStreamRequest {
     pub audio_filter_module_id: ::core::option::Option<::prost::alloc::string::String>,
     #[prost(string, optional, tag="6")]
     pub audio_filter_options: ::core::option::Option<::prost::alloc::string::String>,
+    #[prost(uint32, optional, tag="7")]
+    pub frame_size_ms: ::core::option::Option<u32>,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
@@ -4150,6 +4152,8 @@ pub struct AudioStreamFromParticipantRequest {
     pub audio_filter_module_id: ::core::option::Option<::prost::alloc::string::String>,
     #[prost(string, optional, tag="8")]
     pub audio_filter_options: ::core::option::Option<::prost::alloc::string::String>,
+    #[prost(uint32, optional, tag="9")]
+    pub frame_size_ms: ::core::option::Option<u32>,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]

livekit-ffi/src/server/audio_stream.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::borrow::Cow;
 use std::time::Duration;
 
 use futures_util::StreamExt;
@@ -135,6 +136,9 @@ impl FfiAudioStream {
                     server.watch_handle_dropped(new_stream.track_handle),
                     true,
                     info,
+                    new_stream.frame_size_ms,
+                    sample_rate.try_into().unwrap(),
+                    num_channels.try_into().unwrap(),
                 ));
                 server.watch_panic(handle);
                 Ok::<FfiAudioStream, FfiError>(audio_stream)
@@ -312,6 +316,9 @@ impl FfiAudioStream {
                         handle_dropped_rx,
                         false,
                         info,
+                        request.frame_size_ms,
+                        sample_rate.try_into().unwrap(),
+                        num_channels.try_into().unwrap(),
                     )
                     .await;
                     let _ = done_tx.send(());
@@ -348,7 +355,14 @@ impl FfiAudioStream {
         mut handle_dropped_rx: oneshot::Receiver<()>,
         send_eos: bool,
         mut filter_info: Option<AudioFilterInfo>,
+        frame_size_ms: Option<u32>,
+        sample_rate: u32,
+        num_channels: u32,
     ) {
+        let mut buf = Vec::new();
+        let target_samples = frame_size_ms
+            .map(|ms| sample_rate as usize * ms as usize / 1000 * num_channels as usize);
+
         loop {
             tokio::select! {
                 _ = &mut self_dropped_rx => {
@@ -377,26 +391,59 @@ impl FfiAudioStream {
                         }
                     }
 
-                    let handle_id = server.next_id();
-                    let buffer_info = proto::AudioFrameBufferInfo::from(&frame);
-                    server.store_handle(handle_id, frame);
-
-                    if let Err(err) = server.send_event(proto::ffi_event::Message::AudioStreamEvent(
-                        proto::AudioStreamEvent {
-                            stream_handle: stream_handle_id,
-                            message: Some(proto::audio_stream_event::Message::FrameReceived(
-                                proto::AudioFrameReceived {
-                                    frame: proto::OwnedAudioFrameBuffer {
-                                        handle: proto::FfiOwnedHandle { id: handle_id },
-                                        info: buffer_info,
-                                    },
+                    if let Some(target) = target_samples {
+                        buf.extend_from_slice(&frame.data);
+                        while buf.len() >= target {
+                            let data = buf.split_off(target);
+                            let mut frame_data = std::mem::replace(&mut buf, data);
+                            let new_frame = AudioFrame {
+                                data: Cow::Owned(frame_data),
+                                sample_rate,
+                                num_channels,
+                                samples_per_channel: target as u32 / num_channels,
+                            };
+                            let handle_id = server.next_id();
+                            let buffer_info = proto::AudioFrameBufferInfo::from(&new_frame);
+                            server.store_handle(handle_id, new_frame);
+                            if let Err(err) = server.send_event(proto::ffi_event::Message::AudioStreamEvent(
+                                proto::AudioStreamEvent {
+                                    stream_handle: stream_handle_id,
+                                    message: Some(proto::audio_stream_event::Message::FrameReceived(
+                                        proto::AudioFrameReceived {
+                                            frame: proto::OwnedAudioFrameBuffer {
+                                                handle: proto::FfiOwnedHandle { id: handle_id },
+                                                info: buffer_info,
+                                            },
+                                        },
+                                    )),
                                 },
-                            )),
-                        },
-                    )) {
-                        server.drop_handle(handle_id);
-                        log::warn!("failed to send audio frame: {}", err);
+                            )) {
+                                server.drop_handle(handle_id);
+                                log::warn!("failed to send audio frame: {}", err);
+                            }
+                        }
+                    } else {
+                        let handle_id = server.next_id();
+                        let buffer_info = proto::AudioFrameBufferInfo::from(&frame);
+                        server.store_handle(handle_id, frame);
+                        if let Err(err) = server.send_event(proto::ffi_event::Message::AudioStreamEvent(
+                            proto::AudioStreamEvent {
+                                stream_handle: stream_handle_id,
+                                message: Some(proto::audio_stream_event::Message::FrameReceived(
+                                    proto::AudioFrameReceived {
+                                        frame: proto::OwnedAudioFrameBuffer {
+                                            handle: proto::FfiOwnedHandle { id: handle_id },
+                                            info: buffer_info,
+                                        },
+                                    },
+                                )),
+                            },
+                        )) {
+                            server.drop_handle(handle_id);
+                            log::warn!("failed to send audio frame: {}", err);
+                        }
                     }
+
                 }
             }
         }