xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(sql-file): throttle progress events

This commit is contained in:
t8y2 2026-07-22 01:18:26 +08:00
parent 25e799c34b
commit 2188db3a26
4 changed files with 247 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

185
crates/dbx-core/src/sql_file_import.rs
View File

@ -1,5 +1,5 @@
use std::path::Path;
use std::time::Instant;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, BufReader};
use tokio_util::sync::CancellationToken;
@ -34,6 +34,77 @@ struct StatementErrorDecision {
const SQL_FILE_READ_CHUNK_BYTES: usize = 256 * 1024;
const SQL_FILE_STATEMENT_BATCH_SIZE: usize = 256;
const SQL_FILE_PREVIEW_ENCODING_SAMPLE_BYTES: usize = 1024 * 1024;
const SQL_FILE_PROGRESS_EMIT_INTERVAL: Duration = Duration::from_millis(100);
pub struct SqlFileProgressEmitter<F, C = fn() -> Instant> {
emit: F,
now: C,
last_regular_emit_at: Option<Instant>,
pending_regular: Option<SqlFileProgress>,
}
impl<F> SqlFileProgressEmitter<F>
where
F: FnMut(SqlFileProgress),
{
pub fn new(emit: F) -> Self {
Self::with_clock(emit, Instant::now)
}
}
impl<F, C> SqlFileProgressEmitter<F, C>
where
F: FnMut(SqlFileProgress),
C: FnMut() -> Instant,
{
fn with_clock(emit: F, now: C) -> Self {
Self { emit, now, last_regular_emit_at: None, pending_regular: None }
}
pub fn emit(&mut self, progress: SqlFileProgress) {
if sql_file_progress_is_immediate(progress.status) {
// Preserve ordering and final counters before terminal or failure events.
self.flush_pending();
(self.emit)(progress);
return;
}
self.pending_regular = Some(progress);
let now = (self.now)();
if self
.last_regular_emit_at
.is_none_or(|last_emit_at| now.duration_since(last_emit_at) >= SQL_FILE_PROGRESS_EMIT_INTERVAL)
{
self.flush_pending_at(now);
}
}
fn flush_pending(&mut self) {
if self.pending_regular.is_none() {
return;
}
let now = (self.now)();
self.flush_pending_at(now);
}
fn flush_pending_at(&mut self, now: Instant) {
if let Some(progress) = self.pending_regular.take() {
self.last_regular_emit_at = Some(now);
(self.emit)(progress);
}
}
}
fn sql_file_progress_is_immediate(status: SqlFileStatus) -> bool {
matches!(
status,
SqlFileStatus::Started
| SqlFileStatus::StatementFailed
| SqlFileStatus::Done
| SqlFileStatus::Error
| SqlFileStatus::Cancelled
)
}
struct SqlFileExecutionProgress {
statement_index: usize,
@ -1249,6 +1320,7 @@ fn statement_error_decision(
mod tests {
use super::*;
use crate::models::connection::DatabaseType;
use std::cell::Cell;
use std::sync::atomic::{AtomicU64, Ordering};
static TEMP_SQL_FILE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
@ -1263,6 +1335,117 @@ mod tests {
path
}
fn test_progress(status: SqlFileStatus, statement_index: usize) -> SqlFileProgress {
SqlFileProgress {
execution_id: "test-execution".to_string(),
status,
statement_index,
success_count: statement_index,
failure_count: 0,
affected_rows: statement_index as u64,
elapsed_ms: statement_index as u128,
statement_summary: format!("statement {statement_index}"),
error: None,
}
}
#[test]
fn progress_emitter_compresses_high_frequency_regular_events() {
let base = Instant::now();
let elapsed = Cell::new(Duration::ZERO);
let mut emitted = Vec::new();
{
let mut emitter =
SqlFileProgressEmitter::with_clock(|progress| emitted.push(progress), || base + elapsed.get());
for statement_index in 1..=1_000 {
elapsed.set(Duration::from_millis((statement_index - 1) as u64));
emitter.emit(test_progress(SqlFileStatus::Running, statement_index));
emitter.emit(test_progress(SqlFileStatus::StatementDone, statement_index));
}
emitter.emit(test_progress(SqlFileStatus::Done, 1_000));
}
let regular_count = emitted
.iter()
.filter(|progress| matches!(progress.status, SqlFileStatus::Running | SqlFileStatus::StatementDone))
.count();
assert_eq!(regular_count, 11);
assert_eq!(emitted.last().unwrap().status, SqlFileStatus::Done);
assert_eq!(emitted[emitted.len() - 2].statement_index, 1_000);
}
#[test]
fn progress_emitter_sends_key_events_immediately() {
let base = Instant::now();
let elapsed = Cell::new(Duration::ZERO);
let mut emitted = Vec::new();
{
let mut emitter =
SqlFileProgressEmitter::with_clock(|progress| emitted.push(progress), || base + elapsed.get());
for status in [
SqlFileStatus::Started,
SqlFileStatus::StatementFailed,
SqlFileStatus::Error,
SqlFileStatus::Cancelled,
SqlFileStatus::Done,
] {
emitter.emit(test_progress(status, 1));
}
}
assert_eq!(
emitted.iter().map(|progress| progress.status).collect::<Vec<_>>(),
vec![
SqlFileStatus::Started,
SqlFileStatus::StatementFailed,
SqlFileStatus::Error,
SqlFileStatus::Cancelled,
SqlFileStatus::Done,
]
);
}
#[test]
fn progress_emitter_flushes_latest_counters_before_terminal_event() {
let base = Instant::now();
let elapsed = Cell::new(Duration::ZERO);
let mut emitted = Vec::new();
{
let mut emitter =
SqlFileProgressEmitter::with_clock(|progress| emitted.push(progress), || base + elapsed.get());
emitter.emit(test_progress(SqlFileStatus::Running, 1));
elapsed.set(Duration::from_millis(10));
emitter.emit(test_progress(SqlFileStatus::StatementDone, 2));
emitter.emit(test_progress(SqlFileStatus::Done, 2));
}
assert_eq!(emitted.len(), 3);
assert_eq!(emitted[1].status, SqlFileStatus::StatementDone);
assert_eq!(emitted[1].statement_index, 2);
assert_eq!(emitted[2].status, SqlFileStatus::Done);
}
#[test]
fn progress_emitter_keeps_small_file_progress_timely() {
let base = Instant::now();
let mut emitted = Vec::new();
{
let mut emitter = SqlFileProgressEmitter::with_clock(|progress| emitted.push(progress), || base);
emitter.emit(test_progress(SqlFileStatus::Started, 0));
emitter.emit(test_progress(SqlFileStatus::Running, 1));
emitter.emit(test_progress(SqlFileStatus::StatementDone, 1));
emitter.emit(test_progress(SqlFileStatus::Done, 1));
}
assert_eq!(
emitted.iter().map(|progress| progress.status).collect::<Vec<_>>(),
vec![SqlFileStatus::Started, SqlFileStatus::Running, SqlFileStatus::StatementDone, SqlFileStatus::Done,]
);
}
#[test]
fn sqlserver_sql_file_splits_go_batches_without_sending_delimiters() {
let statements = split_sql_file_import_statements(

41
crates/dbx-web/src/routes/sql_file.rs
View File

@ -8,6 +8,7 @@ use dbx_core::sql;
use dbx_core::sql::{SqlFileProgress, SqlFileRequest, SqlFileStatus};
use dbx_core::sql_file_import::{
execute_sql_file_content, sql_file_error_progress, sql_file_progress as build_sql_file_progress,
SqlFileProgressEmitter,
};
use futures::stream::Stream;
use serde::Deserialize;
@ -92,21 +93,32 @@ pub async fn execute_sql_file(
tokio::spawn(async move {
let started_at = std::time::Instant::now();
let mut progress_emitter = SqlFileProgressEmitter::new(|progress| {
send_sql_file_progress(&tx, progress);
});
progress_emitter.emit(build_sql_file_progress(
&req.execution_id,
SqlFileStatus::Started,
0,
0,
0,
0,
started_at,
"",
None,
));
match std::fs::metadata(&file_path) {
Ok(meta) if meta.len() > 200 * 1024 * 1024 => {
send_sql_file_progress(
&tx,
sql_file_error_progress(
&req.execution_id,
started_at,
format!("File too large: {} bytes (max {} bytes)", meta.len(), 200 * 1024 * 1024),
),
);
progress_emitter.emit(sql_file_error_progress(
&req.execution_id,
started_at,
format!("File too large: {} bytes (max {} bytes)", meta.len(), 200 * 1024 * 1024),
));
cleanup_sql_file_execution(&state_clone, &req.execution_id).await;
return;
}
Err(e) => {
send_sql_file_progress(&tx, sql_file_error_progress(&req.execution_id, started_at, e.to_string()));
progress_emitter.emit(sql_file_error_progress(&req.execution_id, started_at, e.to_string()));
cleanup_sql_file_execution(&state_clone, &req.execution_id).await;
return;
}
@ -119,19 +131,14 @@ pub async fn execute_sql_file(
}) {
Ok(content) => content,
Err(e) => {
send_sql_file_progress(&tx, sql_file_error_progress(&req.execution_id, started_at, e.to_string()));
progress_emitter.emit(sql_file_error_progress(&req.execution_id, started_at, e.to_string()));
cleanup_sql_file_execution(&state_clone, &req.execution_id).await;
return;
}
};
send_sql_file_progress(
&tx,
build_sql_file_progress(&req.execution_id, SqlFileStatus::Started, 0, 0, 0, 0, started_at, "", None),
);
let _ = execute_sql_file_content(&app, &req, &file_content, token, started_at, |progress| {
send_sql_file_progress(&tx, progress);
progress_emitter.emit(progress);
})
.await;
@ -182,7 +189,7 @@ pub async fn sql_file_progress(
let tx = channels.get(&execution_id).ok_or_else(|| AppError("Execution not found".to_string()))?;
let rx = tx.subscribe();
drop(channels);
Ok(crate::sse::sse_from_channel(rx))
Ok(crate::sse::sse_from_lossy_channel(rx))
}
pub async fn cancel_sql_file(

26
crates/dbx-web/src/sse.rs
View File

@ -1,13 +1,33 @@
use axum::response::sse::{Event, KeepAlive, Sse};
use futures::stream::Stream;
use tokio::sync::broadcast;
use tokio::sync::broadcast::{self, error::RecvError};
pub fn sse_from_channel(
rx: broadcast::Receiver<String>,
) -> Sse<impl Stream<Item = Result<Event, std::convert::Infallible>>> {
sse_from_channel_with_lag_policy(rx, false)
}
pub fn sse_from_lossy_channel(
rx: broadcast::Receiver<String>,
) -> Sse<impl Stream<Item = Result<Event, std::convert::Infallible>>> {
sse_from_channel_with_lag_policy(rx, true)
}
fn sse_from_channel_with_lag_policy(
mut rx: broadcast::Receiver<String>,
recover_from_lag: bool,
) -> Sse<impl Stream<Item = Result<Event, std::convert::Infallible>>> {
let stream = async_stream::stream! {
while let Ok(data) = rx.recv().await {
yield Ok(Event::default().data(data));
loop {
match rx.recv().await {
Ok(data) => yield Ok(Event::default().data(data)),
// Only cumulative progress streams may skip stale snapshots; token and
// data streams retain the previous fail-closed behavior on message loss.
Err(RecvError::Lagged(_)) if recover_from_lag => continue,
Err(RecvError::Lagged(_)) => break,
Err(RecvError::Closed) => break,
}
}
};
Sse::new(stream).keep_alive(KeepAlive::default())

49
src-tauri/src/commands/sql_file.rs
View File

@ -10,7 +10,7 @@ use tokio_util::sync::CancellationToken;
use crate::commands::connection::{ensure_connection_writable, AppState};
use dbx_core::sql_file_import::{
execute_sql_file_path, mysql_like_sql_file_can_execute_without_selected_database, read_sql_file_preview,
sql_file_progress,
sql_file_progress, SqlFileProgressEmitter,
};
pub use dbx_core::sql::{SqlFilePreview, SqlFileRequest, SqlFileStatus};
@ -62,8 +62,6 @@ pub async fn execute_sql_file(
}
let started_at = Instant::now();
emit_progress(&app, &request.execution_id, SqlFileStatus::Started, 0, 0, 0, 0, started_at, "", None);
let result = execute_sql_file_inner(&app, &state, &request, token, started_at).await;
{
let mut executions = sql_file_executions().write().await;
@ -90,6 +88,20 @@ async fn execute_sql_file_inner(
token: CancellationToken,
started_at: Instant,
) -> Result<(), String> {
let mut progress_emitter = SqlFileProgressEmitter::new(|progress| {
let _ = app.emit("sql-file-progress", progress);
});
progress_emitter.emit(sql_file_progress(
&request.execution_id,
SqlFileStatus::Started,
0,
0,
0,
0,
started_at,
"",
None,
));
execute_sql_file_path(
state.inner().as_ref(),
request,
@ -97,7 +109,7 @@ async fn execute_sql_file_inner(
token,
started_at,
|progress| {
let _ = app.emit("sql-file-progress", progress);
progress_emitter.emit(progress);
},
)
.await
@ -120,35 +132,6 @@ fn remove_sql_file_execution(executions: &mut HashMap<String, CancellationToken>
executions.remove(execution_id);
}
#[allow(clippy::too_many_arguments)]
fn emit_progress(
app: &AppHandle,
execution_id: &str,
status: SqlFileStatus,
statement_index: usize,
success_count: usize,
failure_count: usize,
affected_rows: u64,
started_at: Instant,
statement_summary: &str,
error: Option<String>,
) {
let _ = app.emit(
"sql-file-progress",
sql_file_progress(
execution_id,
status,
statement_index,
success_count,
failure_count,
affected_rows,
started_at,
statement_summary,
error,
),
);
}
#[cfg(test)]
async fn run_statements_for_test(
statements: Vec<String>,