jacob/songbird
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Driver: Implement audio scheduler (#179)

Browse Source 
This PR implements a custom scheduler for audio threads, which reduces thread use and (often) memory consumption.

To save threads and memory (e.g., packet buffer allocations), Songbird parks Mixer tasks which do not have any live Tracks.
These are now all co-located on a single async 'Idle' task.
This task is responsible for managing UDP keepalive messages for each task, maintaining event state, and executing any Mixer task messages.
Whenever any message arrives which adds a `Track`, the mixer task is moved to a live thread.
The Idle task inspects task counts and execution time on each thread, choosing the first live thread with room, and creating a new one if needed.

Each live thread is responsible for running as many live mixers as it can in a single tick every 20ms: this currently defaults to 16 mixers per thread, but is user-configurable.
A live thread also stores RTP packet blocks to be written into by each sub-task.
Each live thread has a conservative limit of 18ms that it will aim to stay under: if all work takes longer than this, it will offload the task with the highest mixing cost once per tick onto another (possibly new) live worker thread.
This commit is contained in:
Kyle Simpson
2023-05-21 13:50:54 +01:00
parent a5f7d3f488
commit 3daf11f5d1
34 changed files with 2828 additions and 561 deletions
Download Patch File Download Diff File Expand all files Collapse all files

468
src/driver/tasks/mixer/mod.rs
View File

@@ -34,11 +34,12 @@ use discortp::{
rtp::{MutableRtpPacket, RtpPacket},
MutablePacket,
};
use flume::{Receiver, Sender, TryRecvError};
use flume::{Receiver, SendError, Sender, TryRecvError};
use rand::random;
use rubato::{FftFixedOut, Resampler};
use std::{
io::Write,
result::Result as StdResult,
sync::Arc,
time::{Duration, Instant},
};
@@ -51,11 +52,11 @@ use symphonia_core::{
units::Time,
};
use tokio::runtime::Handle;
use tracing::{debug, error, instrument, warn};
use tracing::error;
use xsalsa20poly1305::TAG_SIZE;
#[cfg(test)]
use crate::driver::test_config::{OutputMessage, OutputMode, TickStyle};
use crate::driver::test_config::{OutputMessage, OutputMode};
#[cfg(test)]
use discortp::Packet as _;
@@ -70,10 +71,9 @@ pub struct Mixer {
pub interconnect: Interconnect,
pub mix_rx: Receiver<MixerMessage>,
pub muted: bool,
pub packet: [u8; VOICE_PACKET_MAX],
// pub packet: [u8; VOICE_PACKET_MAX],
pub prevent_events: bool,
pub silence_frames: u8,
pub skip_sleep: bool,
pub soft_clip: SoftClip,
thread_pool: BlockyTaskPool,
pub ws: Option<Sender<WsMessage>>,
@@ -89,7 +89,10 @@ pub struct Mixer {
resample_scratch: AudioBuffer<f32>,
#[cfg(test)]
remaining_loops: Option<u64>,
pub remaining_loops: Option<u64>,
#[cfg(test)]
raw_msg: Option<OutputMessage>,
}
fn new_encoder(bitrate: Bitrate, mix_mode: MixMode) -> Result<OpusEncoder> {
@@ -111,18 +114,8 @@ impl Mixer {
.expect("Failed to create encoder in mixing thread with known-good values.");
let soft_clip = SoftClip::new(config.mix_mode.to_opus());
let mut packet = [0u8; VOICE_PACKET_MAX];
let keepalive_packet = [0u8; MutableKeepalivePacket::minimum_packet_size()];
let mut rtp = MutableRtpPacket::new(&mut packet[..]).expect(
"FATAL: Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
rtp.set_version(RTP_VERSION);
rtp.set_payload_type(RTP_PROFILE_TYPE);
rtp.set_sequence(random::<u16>().into());
rtp.set_timestamp(random::<u32>().into());
let tracks = Vec::with_capacity(1.max(config.preallocated_tracks));
let track_handles = Vec::with_capacity(1.max(config.preallocated_tracks));
@@ -165,10 +158,8 @@ impl Mixer {
interconnect,
mix_rx,
muted: false,
packet,
prevent_events: false,
silence_frames: 0,
skip_sleep: false,
soft_clip,
thread_pool,
ws: None,
@@ -185,110 +176,55 @@ impl Mixer {
#[cfg(test)]
remaining_loops: None,
#[cfg(test)]
raw_msg: None,
}
}
fn run(&mut self) {
let mut events_failure = false;
let mut conn_failure = false;
fn set_bitrate(&mut self, bitrate: Bitrate) -> Result<()> {
self.encoder.set_bitrate(bitrate).map_err(Into::into)
}
#[cfg(test)]
let ignore_check = self.config.override_connection.is_some();
#[cfg(not(test))]
let ignore_check = false;
'runner: loop {
if self.conn_active.is_some() || ignore_check {
loop {
match self.mix_rx.try_recv() {
Ok(m) => {
let (events, conn, should_exit) = self.handle_message(m);
events_failure |= events;
conn_failure |= conn;
if should_exit {
break 'runner;
}
},
Err(TryRecvError::Disconnected) => {
break 'runner;
},
Err(TryRecvError::Empty) => {
break;
},
};
}
// The above action may have invalidated the connection; need to re-check!
// Also, if we're in a test mode we should unconditionally run packet mixing code.
if self.conn_active.is_some() || ignore_check {
if let Err(e) = self
.cycle()
.and_then(|_| self.audio_commands_events())
.and_then(|_| {
self.check_and_send_keepalive()
.or_else(Error::disarm_would_block)
})
{
events_failure |= e.should_trigger_interconnect_rebuild();
conn_failure |= e.should_trigger_connect();
debug!("Mixer thread cycle: {:?}", e);
}
}
} else {
match self.mix_rx.recv() {
Ok(m) => {
let (events, conn, should_exit) = self.handle_message(m);
events_failure |= events;
conn_failure |= conn;
if should_exit {
break 'runner;
}
},
Err(_) => {
break 'runner;
},
}
}
// event failure? rebuild interconnect.
// ws or udp failure? full connect
// (soft reconnect is covered by the ws task.)
//
// in both cases, send failure is fatal,
// but will only occur on disconnect.
// expecting this is fairly noisy, so exit silently.
if events_failure {
self.prevent_events = true;
let sent = self
.interconnect
.core
.send(CoreMessage::RebuildInterconnect);
events_failure = false;
if sent.is_err() {
break;
}
}
if conn_failure {
self.conn_active = None;
let sent = self.interconnect.core.send(CoreMessage::FullReconnect);
conn_failure = false;
if sent.is_err() {
break;
}
}
pub(crate) fn do_rebuilds(
&mut self,
event_failure: bool,
conn_failure: bool,
) -> StdResult<(), SendError<CoreMessage>> {
// event failure? rebuild interconnect.
// ws or udp failure? full connect
// (soft reconnect is covered by the ws task.)
//
// in both cases, send failure is fatal,
// but will only occur on disconnect.
if event_failure {
self.rebuild_interconnect()?;
}
if conn_failure {
self.full_reconnect_gateway()?;
}
Ok(())
}
pub(crate) fn rebuild_interconnect(&mut self) -> StdResult<(), SendError<CoreMessage>> {
self.prevent_events = true;
self.interconnect
.core
.send(CoreMessage::RebuildInterconnect)
}
pub(crate) fn full_reconnect_gateway(&mut self) -> StdResult<(), SendError<CoreMessage>> {
self.conn_active = None;
self.interconnect.core.send(CoreMessage::FullReconnect)
}
#[inline]
fn handle_message(&mut self, msg: MixerMessage) -> (bool, bool, bool) {
pub(crate) fn handle_message(
&mut self,
msg: MixerMessage,
packet: &mut [u8],
) -> (bool, bool, bool) {
let mut events_failure = false;
let mut conn_failure = false;
let mut should_exit = false;
@@ -323,7 +259,7 @@ impl Mixer {
},
MixerMessage::SetConn(conn, ssrc) => {
self.conn_active = Some(conn);
let mut rtp = MutableRtpPacket::new(&mut self.packet[..]).expect(
let mut rtp = MutableRtpPacket::new(packet).expect(
"Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
@@ -332,10 +268,7 @@ impl Mixer {
rtp.set_timestamp(random::<u32>().into());
self.deadline = Instant::now();
let mut ka = MutableKeepalivePacket::new(&mut self.keepalive_packet[..])
.expect("FATAL: Insufficient bytes given to keepalive packet.");
ka.set_ssrc(ssrc);
self.keepalive_deadline = self.deadline + UDP_KEEPALIVE_GAP;
self.update_keepalive(ssrc);
Ok(())
},
MixerMessage::DropConn => {
@@ -420,6 +353,9 @@ impl Mixer {
},
MixerMessage::Ws(new_ws_handle) => {
self.ws = new_ws_handle;
if let Err(e) = self.send_gateway_speaking() {
conn_failure |= e.should_trigger_connect();
}
Ok(())
},
MixerMessage::Poison => {
@@ -436,8 +372,15 @@ impl Mixer {
(events_failure, conn_failure, should_exit)
}
pub(crate) fn update_keepalive(&mut self, ssrc: u32) {
let mut ka = MutableKeepalivePacket::new(&mut self.keepalive_packet[..])
.expect("FATAL: Insufficient bytes given to keepalive packet.");
ka.set_ssrc(ssrc);
self.keepalive_deadline = self.deadline + UDP_KEEPALIVE_GAP;
}
#[inline]
fn fire_event(&self, event: EventMessage) -> Result<()> {
pub(crate) fn fire_event(&self, event: EventMessage) -> Result<()> {
// As this task is responsible for noticing the potential death of an event context,
// it's responsible for not forcibly recreating said context repeatedly.
if !self.prevent_events {
@@ -476,7 +419,7 @@ impl Mixer {
}
#[inline]
fn audio_commands_events(&mut self) -> Result<()> {
pub(crate) fn audio_commands_events(&mut self) -> Result<()> {
// Apply user commands.
for (i, track) in self.tracks.iter_mut().enumerate() {
// This causes fallible event system changes,
@@ -558,48 +501,19 @@ impl Mixer {
}
#[cfg(test)]
fn _march_deadline(&mut self) {
match &self.config.tick_style {
TickStyle::Timed => {
std::thread::sleep(self.deadline.saturating_duration_since(Instant::now()));
self.deadline += TIMESTEP_LENGTH;
},
TickStyle::UntimedWithExecLimit(rx) => {
if self.remaining_loops.is_none() {
if let Ok(new_val) = rx.recv() {
self.remaining_loops = Some(new_val.wrapping_sub(1));
}
}
if let Some(cnt) = self.remaining_loops.as_mut() {
if *cnt == 0 {
self.remaining_loops = None;
} else {
*cnt = cnt.wrapping_sub(1);
}
}
},
#[inline]
pub(crate) fn test_signal_empty_tick(&self) {
match &self.config.override_connection {
Some(OutputMode::Raw(tx)) =>
drop(tx.send(crate::driver::test_config::TickMessage::NoEl)),
Some(OutputMode::Rtp(tx)) =>
drop(tx.send(crate::driver::test_config::TickMessage::NoEl)),
None => {},
}
}
#[cfg(not(test))]
#[inline(always)]
#[allow(clippy::inline_always)] // Justified, this is a very very hot path
fn _march_deadline(&mut self) {
std::thread::sleep(self.deadline.saturating_duration_since(Instant::now()));
self.deadline += TIMESTEP_LENGTH;
}
#[inline]
fn march_deadline(&mut self) {
if self.skip_sleep {
return;
}
self._march_deadline();
}
pub fn cycle(&mut self) -> Result<()> {
pub fn mix_and_build_packet(&mut self, packet: &mut [u8]) -> Result<usize> {
// symph_mix is an `AudioBuffer` (planar format), we need to convert this
// later into an interleaved `SampleBuffer` for libopus.
self.symph_mix.clear();
@@ -609,7 +523,7 @@ impl Mixer {
// Walk over all the audio files, combining into one audio frame according
// to volume, play state, etc.
let mut mix_len = {
let out = self.mix_tracks();
let out = self.mix_tracks(packet);
self.sample_buffer.copy_interleaved_typed(&self.symph_mix);
@@ -626,7 +540,7 @@ impl Mixer {
if mix_len == MixType::MixedPcm(0) {
if self.silence_frames > 0 {
self.silence_frames -= 1;
let mut rtp = MutableRtpPacket::new(&mut self.packet[..]).expect(
let mut rtp = MutableRtpPacket::new(packet).expect(
"FATAL: Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
@@ -638,28 +552,7 @@ impl Mixer {
mix_len = MixType::Passthrough(SILENT_FRAME.len());
} else {
// Per official guidelines, send 5x silence BEFORE we stop speaking.
if let Some(ws) = &self.ws {
// NOTE: this explicit `drop` should prevent a catastrophic thread pileup.
// A full reconnect might cause an inner closed connection.
// It's safer to leave the central task to clean this up and
// pass the mixer a new channel.
drop(ws.send(WsMessage::Speaking(false)));
}
self.march_deadline();
#[cfg(test)]
match &self.config.override_connection {
Some(OutputMode::Raw(tx)) =>
drop(tx.send(crate::driver::test_config::TickMessage::NoEl)),
Some(OutputMode::Rtp(tx)) =>
drop(tx.send(crate::driver::test_config::TickMessage::NoEl)),
None => {},
}
self.advance_rtp_timestamp();
return Ok(());
return Ok(0);
}
} else {
self.silence_frames = 5;
@@ -676,20 +569,14 @@ impl Mixer {
}
}
if let Some(ws) = &self.ws {
ws.send(WsMessage::Speaking(true))?;
}
// Wait till the right time to send this packet:
// usually a 20ms tick, in test modes this is either a finite number of runs or user input.
self.march_deadline();
// For the benefit of test cases, send the raw un-RTP'd data.
#[cfg(test)]
let send_status = if let Some(OutputMode::Raw(tx)) = &self.config.override_connection {
let out = if let Some(OutputMode::Raw(_)) = &self.config.override_connection {
// This case has been handled before buffer clearing above.
let msg = match mix_len {
MixType::Passthrough(len) if len == SILENT_FRAME.len() => OutputMessage::Silent,
MixType::Passthrough(len) => {
let rtp = RtpPacket::new(&self.packet[..]).expect(
let rtp = RtpPacket::new(&packet).expect(
"FATAL: Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
@@ -704,19 +591,15 @@ impl Mixer {
),
};
drop(tx.send(msg.into()));
self.raw_msg = Some(msg);
Ok(())
Ok(1)
} else {
self.prep_and_send_packet(mix_len)
self.prep_packet(mix_len, packet)
};
#[cfg(not(test))]
let send_status = self.prep_and_send_packet(mix_len);
send_status.or_else(Error::disarm_would_block)?;
self.advance_rtp_counters();
let out = self.prep_packet(mix_len, packet);
// Zero out all planes of the mix buffer if any audio was written.
if matches!(mix_len, MixType::MixedPcm(a) if a > 0) {
@@ -725,15 +608,11 @@ impl Mixer {
}
}
Ok(())
}
fn set_bitrate(&mut self, bitrate: Bitrate) -> Result<()> {
self.encoder.set_bitrate(bitrate).map_err(Into::into)
out
}
#[inline]
fn prep_and_send_packet(&mut self, mix_len: MixType) -> Result<()> {
fn prep_packet(&mut self, mix_len: MixType, packet: &mut [u8]) -> Result<usize> {
let send_buffer = self.sample_buffer.samples();
let conn = self
@@ -741,92 +620,97 @@ impl Mixer {
.as_mut()
.expect("Shouldn't be mixing packets without access to a cipher + UDP dest.");
let index = {
let mut rtp = MutableRtpPacket::new(&mut self.packet[..]).expect(
"FATAL: Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
let mut rtp = MutableRtpPacket::new(packet).expect(
"FATAL: Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
let payload = rtp.payload_mut();
let crypto_mode = conn.crypto_state.kind();
let payload = rtp.payload_mut();
let crypto_mode = conn.crypto_state.kind();
// If passthrough, Opus payload in place already.
// Else encode into buffer with space for AEAD encryption headers.
let payload_len = match mix_len {
MixType::Passthrough(opus_len) => opus_len,
MixType::MixedPcm(_samples) => {
let total_payload_space = payload.len() - crypto_mode.payload_suffix_len();
self.encoder.encode_float(
&send_buffer[..self.config.mix_mode.sample_count_in_frame()],
&mut payload[TAG_SIZE..total_payload_space],
)?
},
};
let final_payload_size = conn
.crypto_state
.write_packet_nonce(&mut rtp, TAG_SIZE + payload_len);
// Packet encryption ignored in test modes.
#[cfg(not(test))]
let encrypt = true;
#[cfg(test)]
let encrypt = self.config.override_connection.is_none();
if encrypt {
conn.crypto_state.kind().encrypt_in_place(
&mut rtp,
&conn.cipher,
final_payload_size,
)?;
}
RtpPacket::minimum_packet_size() + final_payload_size
// If passthrough, Opus payload in place already.
// Else encode into buffer with space for AEAD encryption headers.
let payload_len = match mix_len {
MixType::Passthrough(opus_len) => opus_len,
MixType::MixedPcm(_samples) => {
let total_payload_space = payload.len() - crypto_mode.payload_suffix_len();
self.encoder.encode_float(
&send_buffer[..self.config.mix_mode.sample_count_in_frame()],
&mut payload[TAG_SIZE..total_payload_space],
)?
},
};
let final_payload_size = conn
.crypto_state
.write_packet_nonce(&mut rtp, TAG_SIZE + payload_len);
// Packet encryption ignored in test modes.
#[cfg(not(test))]
let encrypt = true;
#[cfg(test)]
let encrypt = self.config.override_connection.is_none();
if encrypt {
conn.crypto_state.kind().encrypt_in_place(
&mut rtp,
&conn.cipher,
final_payload_size,
)?;
}
Ok(RtpPacket::minimum_packet_size() + final_payload_size)
}
#[inline]
pub(crate) fn send_packet(&self, packet: &[u8]) -> Result<()> {
#[cfg(test)]
let send_status = if let Some(OutputMode::Raw(tx)) = &self.config.override_connection {
// This case has been handled before buffer clearing in `mix_and_build_packet`.
drop(tx.send(self.raw_msg.clone().unwrap().into()));
Ok(())
} else {
self._send_packet(packet)
};
#[cfg(not(test))]
let send_status = self._send_packet(packet);
send_status.or_else(Error::disarm_would_block)?;
Ok(())
}
#[inline]
fn _send_packet(&self, packet: &[u8]) -> Result<()> {
let conn = self
.conn_active
.as_ref()
.expect("Shouldn't be mixing packets without access to a cipher + UDP dest.");
#[cfg(test)]
if let Some(OutputMode::Rtp(tx)) = &self.config.override_connection {
// Test mode: send unencrypted (compressed) packets to local receiver.
drop(tx.send(self.packet[..index].to_vec().into()));
drop(tx.send(packet.to_vec().into()));
} else {
conn.udp_tx.send(&self.packet[..index])?;
conn.udp_tx.send(packet)?;
}
#[cfg(not(test))]
{
// Normal operation: send encrypted payload to UDP Tx task.
conn.udp_tx.send(&self.packet[..index])?;
conn.udp_tx.send(packet)?;
}
Ok(())
}
#[inline]
fn advance_rtp_counters(&mut self) {
let mut rtp = MutableRtpPacket::new(&mut self.packet[..]).expect(
"FATAL: Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
rtp.set_sequence(rtp.get_sequence() + 1);
rtp.set_timestamp(rtp.get_timestamp() + MONO_FRAME_SIZE as u32);
}
#[inline]
// Even if we don't send a packet, we *do* need to keep advancing the timestamp
// to make it easier for a receiver to reorder packets and compute jitter measures
// wrt. our clock difference vs. theirs.
fn advance_rtp_timestamp(&mut self) {
let mut rtp = MutableRtpPacket::new(&mut self.packet[..]).expect(
"FATAL: Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
rtp.set_timestamp(rtp.get_timestamp() + MONO_FRAME_SIZE as u32);
}
#[inline]
fn check_and_send_keepalive(&mut self) -> Result<()> {
pub(crate) fn check_and_send_keepalive(&mut self, now: Option<Instant>) -> Result<()> {
if let Some(conn) = self.conn_active.as_mut() {
if Instant::now() >= self.keepalive_deadline {
let now = now.unwrap_or_else(Instant::now);
if now >= self.keepalive_deadline {
conn.udp_tx.send(&self.keepalive_packet)?;
self.keepalive_deadline += UDP_KEEPALIVE_GAP;
}
@@ -836,9 +720,29 @@ impl Mixer {
}
#[inline]
fn mix_tracks(&mut self) -> MixType {
pub(crate) fn send_gateway_speaking(&self) -> Result<()> {
if let Some(ws) = &self.ws {
ws.send(WsMessage::Speaking(true))?;
}
Ok(())
}
#[inline]
pub(crate) fn send_gateway_not_speaking(&self) {
if let Some(ws) = &self.ws {
// NOTE: this explicit `drop` should prevent a catastrophic thread pileup.
// A full reconnect might cause an inner closed connection.
// It's safer to leave the central task to clean this up and
// pass the mixer a new channel.
drop(ws.send(WsMessage::Speaking(false)));
}
}
#[inline]
fn mix_tracks(&mut self, packet: &mut [u8]) -> MixType {
// Get a slice of bytes to write in data for Opus packet passthrough.
let mut rtp = MutableRtpPacket::new(&mut self.packet[..]).expect(
let mut rtp = MutableRtpPacket::new(packet).expect(
"FATAL: Too few bytes in self.packet for RTP header.\
(Blame: VOICE_PACKET_MAX?)",
);
@@ -962,17 +866,3 @@ impl Mixer {
MixType::MixedPcm(len)
}
}
/// The mixing thread is a synchronous context due to its compute-bound nature.
///
/// We pass in an async handle for the benefit of some Input classes (e.g., restartables)
/// who need to run their restart code elsewhere and return blank data until such time.
#[instrument(skip(interconnect, mix_rx, async_handle))]
pub(crate) fn runner(
interconnect: Interconnect,
mix_rx: Receiver<MixerMessage>,
async_handle: Handle,
config: Config,
) {
Mixer::new(mix_rx, async_handle, interconnect, config).run();
}