fix: lagged output in unified_exec buffer (#4992)

Handle `Lagged` error when the broadcast buffer of the unified_exec is
full
This commit is contained in:
jif-oai
2025-10-09 17:06:07 +01:00
committed by GitHub
parent ca6a0358de
commit 3ddd4d47d0
2 changed files with 141 additions and 5 deletions

View File

@@ -110,11 +110,22 @@ impl ManagedUnifiedExecSession {
let buffer_clone = Arc::clone(&output_buffer);
let notify_clone = Arc::clone(&output_notify);
let output_task = tokio::spawn(async move {
while let Ok(chunk) = receiver.recv().await {
let mut guard = buffer_clone.lock().await;
guard.push_chunk(chunk);
drop(guard);
notify_clone.notify_waiters();
loop {
match receiver.recv().await {
Ok(chunk) => {
let mut guard = buffer_clone.lock().await;
guard.push_chunk(chunk);
drop(guard);
notify_clone.notify_waiters();
}
// If we lag behind the broadcast buffer, skip missed
// messages but keep the task alive to continue streaming.
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
continue;
}
// When the sender closes, exit the task.
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
}
}
});