From 2188db3a26061fbfea2137ac0a55ca07b4bad584 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 22 Jul 2026 01:18:26 +0800 Subject: [PATCH] fix(sql-file): throttle progress events --- crates/dbx-core/src/sql_file_import.rs | 185 ++++++++++++++++++++++++- crates/dbx-web/src/routes/sql_file.rs | 41 +++--- crates/dbx-web/src/sse.rs | 26 +++- src-tauri/src/commands/sql_file.rs | 49 +++---- 4 files changed, 247 insertions(+), 54 deletions(-) diff --git a/crates/dbx-core/src/sql_file_import.rs b/crates/dbx-core/src/sql_file_import.rs index 27409bffa..718636d32 100644 --- a/crates/dbx-core/src/sql_file_import.rs +++ b/crates/dbx-core/src/sql_file_import.rs @@ -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 Instant> { + emit: F, + now: C, + last_regular_emit_at: Option, + pending_regular: Option, +} + +impl SqlFileProgressEmitter +where + F: FnMut(SqlFileProgress), +{ + pub fn new(emit: F) -> Self { + Self::with_clock(emit, Instant::now) + } +} + +impl SqlFileProgressEmitter +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![ + 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![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( diff --git a/crates/dbx-web/src/routes/sql_file.rs b/crates/dbx-web/src/routes/sql_file.rs index 4b416adbd..e80aaa606 100644 --- a/crates/dbx-web/src/routes/sql_file.rs +++ b/crates/dbx-web/src/routes/sql_file.rs @@ -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( diff --git a/crates/dbx-web/src/sse.rs b/crates/dbx-web/src/sse.rs index ac46d5eca..fa7ed3550 100644 --- a/crates/dbx-web/src/sse.rs +++ b/crates/dbx-web/src/sse.rs @@ -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, +) -> Sse>> { + sse_from_channel_with_lag_policy(rx, false) +} + +pub fn sse_from_lossy_channel( + rx: broadcast::Receiver, +) -> Sse>> { + sse_from_channel_with_lag_policy(rx, true) +} + +fn sse_from_channel_with_lag_policy( mut rx: broadcast::Receiver, + recover_from_lag: bool, ) -> Sse>> { 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()) diff --git a/src-tauri/src/commands/sql_file.rs b/src-tauri/src/commands/sql_file.rs index a4c5aba19..f61b85cac 100644 --- a/src-tauri/src/commands/sql_file.rs +++ b/src-tauri/src/commands/sql_file.rs @@ -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 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, -) { - 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,