diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 22d2afaa7bc..73d5c370ece 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -2760,12 +2760,12 @@ mod tests { ) } - struct ReleasedExportSource { + struct CancellableExportSource { schema: GtSchemaRef, channels: std::sync::Mutex, oneshot::Receiver<()>)>>, } - impl DataSource for ReleasedExportSource { + impl DataSource for CancellableExportSource { fn get_stream( &self, _request: ScanRequest, @@ -2785,15 +2785,15 @@ mod tests { } #[tokio::test] - async fn test_metric_export_http_timeout_drains_ordinary_writer() { + async fn test_metric_export_http_timeout_cancels_ordinary_source() { let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout"); let (started_tx, started_rx) = oneshot::channel(); - let (release_tx, release_rx) = oneshot::channel(); + let (mut release_tx, release_rx) = oneshot::channel(); let info = test_table_info(1024, "source").unwrap(); let source = Arc::new(Table::new( Arc::new(info.clone()), FilterPushDownType::Unsupported, - Arc::new(ReleasedExportSource { + Arc::new(CancellableExportSource { schema: info.meta.schema.clone(), channels: std::sync::Mutex::new(Some((started_tx, release_rx))), }), @@ -2831,9 +2831,9 @@ mod tests { let server_task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap(); }); + let uri = reqwest::Url::from_directory_path(destination.path()).unwrap(); let sql = format!( - "COPY DATABASE greptime.public TO '{}/' WITH (experimental_metric_export='true')", - destination.path().display() + "COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')" ); let request = tokio::spawn(async move { reqwest::Client::builder() @@ -2852,21 +2852,10 @@ mod tests { .unwrap(); let response = request.await.unwrap(); assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT); - // A dropped ordinary stream would close this receiver before release. - release_tx.send(()).unwrap(); - let file = destination.path().join("source.parquet"); - tokio::time::timeout(Duration::from_secs(5), async { - loop { - if std::fs::read(&file) - .is_ok_and(|bytes| bytes.len() > 8 && bytes.ends_with(b"PAR1")) - { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .unwrap(); + tokio::time::timeout(Duration::from_secs(5), release_tx.closed()) + .await + .unwrap(); + assert!(!destination.path().join("source.parquet").exists()); server_task.abort(); } diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs index ec0d8f08b48..dbe18ba763a 100644 --- a/src/operator/src/statement/export_database.rs +++ b/src/operator/src/statement/export_database.rs @@ -348,7 +348,7 @@ mod tests { } .fail(); } - // Ordinary COPY continues its I/O even when Metric jobs cancel. + // Already-started ordinary I/O drains after cancellation. token.cancelled().await; started.send(10).unwrap(); finish_io.acquire().await.unwrap().forget();