use super::stream::{StreamState, stream_registry};
use core::ffi::c_void;
use std::sync::atomic::Ordering;
use tokio::sync::mpsc;
use wslcsdk_sys::types::WslcProcessIOHandle;
pub(super) fn record_dropped(state: &StreamState, is_stdout: bool, bytes: u64) {
let (counter, direction) = if is_stdout {
(&state.stdout_dropped_bytes, "标准输出")
} else {
(&state.stderr_dropped_bytes, "标准错误")
};
let previous = counter.fetch_add(bytes, Ordering::Relaxed);
if previous == 0 {
log::warn!(
"流式通道消费不及时,已开始丢弃数据以维持背压,流 ID: {},方向: {},本次丢弃 {} 字节",
state.id,
direction,
bytes
);
}
}
pub(super) unsafe extern "system" fn stream_io_trampoline(
io_handle: WslcProcessIOHandle,
data: *const u8,
data_bytes: u32,
context: *mut c_void,
) {
if context.is_null() || data.is_null() || data_bytes == 0 {
return;
}
let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let stream_id = context as usize as u64;
let state_arc = stream_registry()
.get(&stream_id)
.and_then(|weak| weak.upgrade());
let Some(state) = state_arc else {
return;
};
let is_stdout = io_handle == WslcProcessIOHandle::Stdout;
let sender_slot = if is_stdout {
&state.stdout_tx
} else if io_handle == WslcProcessIOHandle::Stderr {
&state.stderr_tx
} else {
return;
};
let Ok(guard) = sender_slot.lock() else {
return;
};
let Some(sender) = guard.as_ref() else {
return;
};
let slice = unsafe { std::slice::from_raw_parts(data, data_bytes as usize) };
let chunk = slice.to_vec();
if let Err(mpsc::error::TrySendError::Full(dropped)) = sender.try_send(chunk) {
record_dropped(&state, is_stdout, dropped.len() as u64);
}
drop(guard);
}));
}
pub(super) unsafe extern "system" fn stream_exit_trampoline(exit_code: i32, context: *mut c_void) {
if context.is_null() {
return;
}
let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let stream_id = context as usize as u64;
let state_arc = stream_registry()
.get(&stream_id)
.and_then(|weak| weak.upgrade());
if let Some(state) = state_arc {
state.close_io_channels();
if let Ok(mut guard) = state.exit_tx.lock()
&& let Some(tx) = guard.take()
{
let _ = tx.send(exit_code);
}
}
}));
}