fix(core,protocol): bound disconnect shutdown

This commit is contained in:
Edison Jwa
2026-06-08 20:15:50 +09:00
parent 8487acf167
commit 7c341d42e5
5 changed files with 674 additions and 388 deletions
+171 -46
View File
@@ -46,6 +46,9 @@ use crate::ProtocolError;
const SPEAKING_ACTIVITY_WINDOW: Duration = Duration::from_millis(750);
const INBOUND_VOICE_SEND_TIMEOUT: Duration = Duration::from_millis(40);
const PROFILE_REFRESH_RESULT_TIMEOUT: Duration = Duration::from_secs(3);
const OUTBOUND_VOICE_PACKETS_PER_TICK: usize = 8;
const DISCONNECT_REPLY_TIMEOUT: Duration = Duration::from_secs(1);
const DISCONNECT_EVENT_DRAIN_TIMEOUT: Duration = Duration::from_millis(500);
type PendingMoves = HashMap<
MessageHandle,
@@ -84,6 +87,30 @@ async fn send_with_timeout<T: Send>(
}
}
fn drain_voice_packets_for_tick<T, E>(
voice_out_rx: &mut mpsc::Receiver<T>,
max_packets: usize,
mut send: impl FnMut(T) -> Result<(), E>,
) -> usize {
let mut drained = 0;
for _ in 0..max_packets {
let packet = match voice_out_rx.try_recv() {
Ok(packet) => packet,
Err(_) => break,
};
let _ = send(packet);
drained += 1;
}
drained
}
async fn bounded_drain_stream<S>(stream: S, timeout_duration: Duration)
where
S: futures::Stream,
{
let _ = tokio::time::timeout(timeout_duration, stream.for_each(|_| future::ready(()))).await;
}
/// Pick the TeamSpeak `client_version`/platform/signature triple
/// (sourced from `ReSpeak/tsdeclarations/Versions.csv`, baked into
/// `tsproto-types` at vendor-time) that best matches the *runtime*
@@ -341,8 +368,21 @@ impl ProtocolClient {
/// Disconnect cleanly. Blocks until the task exits.
pub async fn disconnect(self) {
let (tx, rx) = oneshot::channel();
if self.tx.send(Request::Disconnect(tx)).await.is_ok() {
let _ = rx.await;
let request_path = async {
if self.tx.send(Request::Disconnect(tx)).await.is_ok() {
let _ = rx.await;
}
};
if tokio::time::timeout(DISCONNECT_REPLY_TIMEOUT, request_path)
.await
.is_err()
{
warn!(
target: "chanora_protocol",
timeout_ms = DISCONNECT_REPLY_TIMEOUT.as_millis() as u64,
"disconnect request did not complete before timeout"
);
}
}
@@ -676,12 +716,14 @@ async fn connection_task(
// Main loop: pump events, service requests, forward voice.
loop {
// 1. Drain any outbound voice packets first — they're time-sensitive.
while let Ok(pkt) = voice_out_rx.try_recv() {
// 1. Send a bounded batch of outbound voice packets first — they're
// time-sensitive, but control requests must still make progress.
drain_voice_packets_for_tick(&mut voice_out_rx, OUTBOUND_VOICE_PACKETS_PER_TICK, |pkt| {
if let Err(e) = con.send_audio(pkt) {
warn!(target: "chanora_protocol", error = %e, "send_audio failed");
}
}
Ok::<(), ()>(())
});
// 2. Advance event stream by at most one event with a small timeout.
let pump = async {
@@ -689,21 +731,19 @@ async fn connection_task(
tokio::time::timeout(Duration::from_millis(20), ev_stream.next()).await
};
match pump.await {
Ok(Some(Ok(item))) => {
match item {
StreamItem::Audio(buf) => {
handle_audio_stream_item(&channels.voice_in, &mut voice_activity, buf).await;
}
other => handle_non_audio_stream_item(
&con,
other,
&channels.chat,
&channels.activity,
&channels.delta,
&mut pending_moves,
),
Ok(Some(Ok(item))) => match item {
StreamItem::Audio(buf) => {
handle_audio_stream_item(&channels.voice_in, &mut voice_activity, buf).await;
}
}
other => handle_non_audio_stream_item(
&con,
other,
&channels.chat,
&channels.activity,
&channels.delta,
&mut pending_moves,
),
},
Ok(Some(Err(e))) => {
warn!(target: "chanora_protocol", error = %e, "event error");
// Some errors are transient; treat persistent ones
@@ -807,7 +847,7 @@ async fn connection_task(
}
Ok(Request::Disconnect(reply)) => {
let _ = con.disconnect(DisconnectOptions::new());
con.events().for_each(|_| future::ready(())).await;
bounded_drain_stream(con.events(), DISCONNECT_EVENT_DRAIN_TIMEOUT).await;
let _ = reply.send(());
info!(target: "chanora_protocol", "clean disconnect");
exit!(DisconnectReason::UserRequested);
@@ -815,7 +855,7 @@ async fn connection_task(
Err(mpsc::error::TryRecvError::Empty) => {}
Err(mpsc::error::TryRecvError::Disconnected) => {
let _ = con.disconnect(DisconnectOptions::new());
con.events().for_each(|_| future::ready(())).await;
bounded_drain_stream(con.events(), DISCONNECT_EVENT_DRAIN_TIMEOUT).await;
info!(target: "chanora_protocol", "handle dropped; implicit disconnect");
exit!(DisconnectReason::UserRequested);
}
@@ -927,7 +967,9 @@ fn handle_non_audio_stream_item(
let mapped = match target {
tsclientlib::MessageTarget::Server => MessageTarget::Server,
tsclientlib::MessageTarget::Channel => MessageTarget::Channel,
tsclientlib::MessageTarget::Client(id) => MessageTarget::Client(id.0 as u64),
tsclientlib::MessageTarget::Client(id) => {
MessageTarget::Client(id.0 as u64)
}
tsclientlib::MessageTarget::Poke(id) => MessageTarget::Poke(id.0 as u64),
};
let _ = chat_tx.try_send(ChatMessage {
@@ -1125,7 +1167,15 @@ async fn fetch_client_profile(
) -> Result<ClientProfile, ProtocolError> {
let target_id = TsClientId(client_id as u16);
let (database_id, uid_b64, has_optional, has_connection, is_own, needs_server_groups, needs_channel_groups) = {
let (
database_id,
uid_b64,
has_optional,
has_connection,
is_own,
needs_server_groups,
needs_channel_groups,
) = {
let state = con
.get_state()
.map_err(|e| ProtocolError::Backend(format!("get_state: {e}")))?;
@@ -1214,15 +1264,9 @@ async fn fetch_client_profile(
}
let db_info = if refresh_plan.needs_client_db_info {
request_client_db_info(
con,
database_id,
channels,
pending_moves,
voice_activity,
)
.await
.ok()
request_client_db_info(con, database_id, channels, pending_moves, voice_activity)
.await
.ok()
} else {
None
};
@@ -1289,10 +1333,18 @@ async fn fetch_client_profile(
.or_else(|| db_info.as_ref().map(|info| info.created.unix_timestamp())),
last_connected_unix_seconds: optional
.map(|info| info.last_connected.unix_timestamp())
.or_else(|| db_info.as_ref().map(|info| info.last_connected.unix_timestamp())),
.or_else(|| {
db_info
.as_ref()
.map(|info| info.last_connected.unix_timestamp())
}),
connections_total: optional
.map(|info| u64::from(info.connections_total))
.or_else(|| db_info.as_ref().map(|info| u64::from(info.connections_total))),
.or_else(|| {
db_info
.as_ref()
.map(|info| u64::from(info.connections_total))
}),
online_seconds: connection
.and_then(|info| info.connected_time.map(|duration| duration.whole_seconds())),
idle_milliseconds: connection.map(|info| duration_millis(info.idle_time)),
@@ -1325,14 +1377,10 @@ async fn fetch_client_profile(
.or_else(|| db_info.as_ref().map(|info| info.bytes_uploaded_total)),
packet_loss_client_to_server_total: net_stats
.map(|s| s.get_packetloss())
.or_else(|| {
connection.map(|info| info.client_to_server_packetloss_total)
}),
.or_else(|| connection.map(|info| info.client_to_server_packetloss_total)),
packet_loss_server_to_client_total: net_stats
.map(|s| s.get_packetloss_s2c_total())
.or_else(|| {
connection.and_then(|info| info.server_to_client_packetloss_total)
}),
.or_else(|| connection.and_then(|info| info.server_to_client_packetloss_total)),
})
}
@@ -1876,10 +1924,12 @@ const _: () = {
#[cfg(test)]
mod tests {
use super::{
client_profile_refresh_plan, is_server_query_client_type, send_with_timeout,
server_socket_from_config, sort_channels_tree_by, std_duration_millis,
ConnectConfig, SendTimeoutError,
bounded_drain_stream, client_profile_refresh_plan, drain_voice_packets_for_tick,
is_server_query_client_type, send_with_timeout, server_socket_from_config,
sort_channels_tree_by, std_duration_millis, ConnectConfig, ProtocolClient, Request,
SendTimeoutError, DISCONNECT_REPLY_TIMEOUT,
};
use futures::stream;
use std::time::Duration;
use tokio::sync::mpsc;
use tsproto_types::ClientType;
@@ -2122,6 +2172,83 @@ mod tests {
assert_eq!(result, Err(SendTimeoutError::Timeout(2)));
}
#[tokio::test]
async fn disconnect_request_send_is_bounded_when_request_channel_is_full() {
let (tx, _rx) = mpsc::channel(1);
let (reply_tx, _reply_rx) = tokio::sync::oneshot::channel();
tx.send(Request::Snapshot(reply_tx))
.await
.expect("seed first request");
let (disconnect_tx, _disconnect_rx) = tokio::sync::oneshot::channel();
let result = send_with_timeout(
&tx,
Request::Disconnect(disconnect_tx),
Duration::from_millis(10),
)
.await;
assert!(matches!(result, Err(SendTimeoutError::Timeout(_))));
}
#[tokio::test]
async fn protocol_client_disconnect_returns_when_request_channel_is_full() {
let (tx, _rx) = mpsc::channel(1);
let (snapshot_tx, _snapshot_rx) = tokio::sync::oneshot::channel();
tx.send(Request::Snapshot(snapshot_tx))
.await
.expect("seed first request");
let (voice_out_tx, _voice_out_rx) = mpsc::channel(1);
let (_voice_in_tx, voice_in_rx) = mpsc::channel(1);
let (_lost_tx, lost_rx) = tokio::sync::oneshot::channel();
let (_chat_tx, chat_rx) = mpsc::channel(1);
let (_activity_tx, activity_rx) = mpsc::channel(1);
let (_delta_tx, delta_rx) = mpsc::channel(1);
let client = ProtocolClient {
tx,
voice_out_tx,
voice_in_rx: std::sync::Mutex::new(Some(voice_in_rx)),
lost_rx: std::sync::Mutex::new(Some(lost_rx)),
chat_rx: std::sync::Mutex::new(Some(chat_rx)),
activity_rx: std::sync::Mutex::new(Some(activity_rx)),
delta_rx: std::sync::Mutex::new(Some(delta_rx)),
};
tokio::time::timeout(
DISCONNECT_REPLY_TIMEOUT + Duration::from_millis(100),
client.disconnect(),
)
.await
.expect("disconnect should not wait indefinitely for request channel capacity");
}
#[tokio::test]
async fn voice_drain_stops_at_per_tick_budget() {
let (tx, mut rx) = mpsc::channel(8);
for value in 0_u8..5 {
tx.send(value).await.expect("seed voice packet");
}
let mut sent = Vec::new();
let drained = drain_voice_packets_for_tick(&mut rx, 2, |value| {
sent.push(value);
Ok::<(), ()>(())
});
assert_eq!(drained, 2);
assert_eq!(sent, vec![0, 1]);
assert_eq!(rx.len(), 3);
}
#[tokio::test]
async fn disconnect_stream_drain_returns_after_timeout() {
let start = tokio::time::Instant::now();
bounded_drain_stream(stream::pending::<()>(), Duration::from_millis(10)).await;
assert!(start.elapsed() < Duration::from_millis(100));
}
}
fn forward_delta(
@@ -2204,9 +2331,7 @@ fn forward_delta(
old: PropertyValue::Channel(channel),
..
} => {
let _ = delta_tx.try_send(ProtocolDelta::ChannelRemoved {
id: channel.id.0,
});
let _ = delta_tx.try_send(ProtocolDelta::ChannelRemoved { id: channel.id.0 });
}
Event::PropertyChanged {
id: PropertyId::Channel(channel_id),