diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 54b46fc11..711a73033 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -3,6 +3,8 @@ //! Handles to operations a server may run asynchronously. +use std::future::Future; +use std::pin::Pin; use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; use arrow_array::RecordBatch; @@ -250,6 +252,18 @@ where pub(crate) fn spawned(task: JoinHandle>) -> Self { Self::new_typed(Box::new(SpawnedJob::new(task))) } + + pub(crate) fn spawned_with_cancellation(task: JoinHandle>, cancel: F) -> Self + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let cancellation: Cancellation = Arc::new(move || Box::pin(cancel())); + Self::new_typed(Box::new(SpawnedJob::new_with_cancellation( + task, + cancellation, + ))) + } } impl Job @@ -501,10 +515,28 @@ impl Outcome { struct SpawnedJob { outcome: watch::Receiver>, abort: AbortHandle, + cancellation: Option, } +type CancellationFuture = Pin> + Send>>; +type Cancellation = Arc CancellationFuture + Send + Sync>; + impl SpawnedJob { fn new(task: JoinHandle>) -> Self + where + T: Serialize + Send + 'static, + { + Self::new_inner(task, None) + } + + fn new_with_cancellation(task: JoinHandle>, cancellation: Cancellation) -> Self + where + T: Serialize + Send + 'static, + { + Self::new_inner(task, Some(cancellation)) + } + + fn new_inner(task: JoinHandle>, cancellation: Option) -> Self where T: Serialize + Send + 'static, { @@ -526,7 +558,11 @@ impl SpawnedJob { }; let _ = tx.send(Some(outcome)); }); - Self { outcome, abort } + Self { + outcome, + abort, + cancellation, + } } } @@ -556,6 +592,19 @@ impl JobHandle for SpawnedJob { } async fn cancel(&self) -> Result<()> { + if let Some(cancellation) = &self.cancellation { + let mut outcome = self.outcome.clone(); + tokio::select! { + biased; + settled = outcome.wait_for(Option::is_some) => { + settled.map_err(|_| Error::Runtime { + message: "job outcome was dropped before cancellation completed".to_string(), + })?; + return Ok(()); + } + result = cancellation() => result?, + } + } self.abort.abort(); Ok(()) } diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 42e4c3206..2aff95934 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -880,12 +880,7 @@ impl Database for RemoteDatabase { quote_sql_identifier(&request.name), request.query ); - Ok(Job::spawned(tokio::spawn(async move { - let query = client.submit(&statement, &namespace).await?; - let mut reader = query.reader().await?; - while reader.try_next().await?.is_some() {} - Ok(()) - }))) + Ok(client.submit_as_job(statement, namespace, || async { Ok(()) })) } async fn drop_materialized_view_async( diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 489158922..e8ddcf6ff 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::fs; +use std::future::{Future, pending}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::{Duration, Instant}; @@ -20,12 +21,14 @@ use arrow_schema::{Schema, SchemaRef}; use futures::TryStreamExt; use http::header::{HeaderMap, HeaderName, HeaderValue}; use prost::Message; -use tokio::sync::{Mutex, Notify, OnceCell, mpsc}; +use serde::{Serialize, de::DeserializeOwned}; +use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch}; use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity}; use uuid::Uuid; use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; use crate::error::{Error, Result}; +use crate::job::Job; use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; @@ -209,6 +212,42 @@ impl SqlClient { .await } + pub(super) fn submit_as_job( + &self, + statement: String, + namespace: Vec, + finish: F, + ) -> Job + where + T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, + F: FnOnce() -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + { + let client = self.clone(); + let (query_tx, query_rx) = watch::channel(None::>); + let task = tokio::spawn(async move { + let query = Arc::new(client.submit(&statement, &namespace).await?); + query_tx.send_replace(Some(query.clone())); + let mut reader = query.reader().await?; + while reader.try_next().await?.is_some() {} + finish().await + }); + Job::spawned_with_cancellation(task, move || { + let mut query_rx = query_rx.clone(); + async move { + loop { + let query = query_rx.borrow().clone(); + if let Some(query) = query { + return query.cancel().await; + } + if query_rx.changed().await.is_err() { + pending::<()>().await; + } + } + } + }) + } + pub(super) async fn describe(&self, query_id: Uuid) -> Result { let query = self .queries diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index fbc43234a..0d7efc9c4 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -21,7 +21,10 @@ use super::*; use crate::database::Database; use crate::remote::RemoteCatalogOptions; use crate::remote::client::HeaderProvider; +use crate::remote::client::test_utils::client_with_handler; use crate::remote::db::RemoteDatabase; +use crate::remote::table::RemoteTable; +use crate::table::BaseTable; #[derive(Debug, Default)] struct DelayedHeaderProvider { @@ -73,6 +76,7 @@ struct TestSqlService { first_continuation_count: Arc, transient_poll_failures: Arc, headers: Arc>>, + gate_refresh_release: Arc, result: RecordBatch, large_result: RecordBatch, dictionary_result: RecordBatch, @@ -121,6 +125,7 @@ impl Default for TestSqlService { first_continuation_count: Arc::new(AtomicUsize::new(0)), transient_poll_failures: Arc::new(AtomicUsize::new(0)), headers: Arc::new(std::sync::Mutex::new(Vec::new())), + gate_refresh_release: Arc::new(Notify::new()), result, large_result, dictionary_result, @@ -279,6 +284,7 @@ impl FlightService for TestSqlService { ) -> std::result::Result::DoGetStream>, Status> { self.do_get_count.fetch_add(1, Ordering::SeqCst); let ticket = request.get_ref().ticket.as_ref(); + let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; let large = ticket == b"SELECT large message"; @@ -290,7 +296,11 @@ impl FlightService for TestSqlService { self.result.clone() }; let schema = result.schema(); + let gate_refresh_release = self.gate_refresh_release.clone(); let input = futures::stream::once(async move { + if gate_refresh { + gate_refresh_release.notified().await; + } if slow { tokio::time::sleep(Duration::from_millis(250)).await; } @@ -448,10 +458,6 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { #[tokio::test] async fn remote_refresh_twice_keeps_latest_selector() { - use crate::remote::client::test_utils::client_with_handler; - use crate::remote::table::RemoteTable; - use crate::table::BaseTable; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let incoming = futures::stream::try_unfold(listener, |listener| async { @@ -531,6 +537,66 @@ async fn remote_refresh_twice_keeps_latest_selector() { server.await.unwrap().unwrap(); } +#[tokio::test] +async fn cancelling_remote_refresh_cancels_sql_query() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let incoming = futures::stream::try_unfold(listener, |listener| async { + let (socket, _) = listener.accept().await?; + Ok::<_, std::io::Error>(Some((socket, listener))) + }); + let service = TestSqlService::default(); + let cancel_count = service.cancel_count.clone(); + let do_get_count = service.do_get_count.clone(); + let gate_refresh_release = service.gate_refresh_release.clone(); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn( + tonic::transport::Server::builder() + .add_service(FlightServiceServer::new(service)) + .serve_with_incoming_shutdown(incoming, async { + let _ = shutdown_rx.await; + }), + ); + let sql_client = SqlClient::new( + "analytics".into(), + None, + "test-key".into(), + None, + Some(format!("grpc://{address}")), + ClientConfig::default(), + ); + let client = client_with_handler(|request| -> http::Response { + panic!( + "cancelled refresh must not read result metadata: {}", + request.url().path() + ); + }); + let table = RemoteTable::new_with_sql_client( + client, + "cancel_view".into(), + vec![], + "cancel_view".into(), + Default::default(), + Some(sql_client), + ); + let job = table.refresh_materialized_view_async(None).await.unwrap(); + tokio::time::timeout(Duration::from_secs(10), async { + while do_get_count.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + job.cancel().await.unwrap(); + assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. }))); + assert_eq!(cancel_count.load(Ordering::SeqCst), 1); + + gate_refresh_release.notify_one(); + shutdown_tx.send(()).unwrap(); + server.await.unwrap().unwrap(); +} + #[tokio::test] async fn submits_polls_fetches_cancels_and_reuses_client() { let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 79511d55d..4ba5e7e76 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2157,33 +2157,32 @@ impl BaseTable for RemoteTable { let server_version = self.server_version.clone(); let freshness = self.freshness.clone(); let freshness_request = self.snapshot_freshness_headers(); - Ok(Job::spawned(tokio::spawn(async move { - let query = sql_client.submit(&statement, &namespace).await?; - let mut reader = query.reader().await?; - while reader.try_next().await?.is_some() {} - - let table = Self::new_with_sql_client( - client, - name, - namespace, - identifier, - server_version, - Some(sql_client), - ); - table.checkout_latest().await?; - let version = table.version().await?; - let rows_written = - u64::try_from(table.count_rows(None).await?).map_err(|_| Error::Runtime { - message: "materialized-view row count exceeds u64".to_string(), - })?; - freshness_request.observe_version(&freshness, version); - Ok(RefreshMaterializedViewResult { - mode: crate::materialized_view::RefreshMode::Rebuild, - rows_written, - source_version: 0, - version, - }) - }))) + let job_sql_client = sql_client.clone(); + Ok( + sql_client.submit_as_job(statement, namespace.clone(), move || async move { + let table = Self::new_with_sql_client( + client, + name, + namespace, + identifier, + server_version, + Some(job_sql_client), + ); + table.checkout_latest().await?; + let version = table.version().await?; + let rows_written = + u64::try_from(table.count_rows(None).await?).map_err(|_| Error::Runtime { + message: "materialized-view row count exceeds u64".to_string(), + })?; + freshness_request.observe_version(&freshness, version); + Ok(RefreshMaterializedViewResult { + mode: crate::materialized_view::RefreshMode::Rebuild, + rows_written, + source_version: 0, + version, + }) + }), + ) } async fn query_snapshot(&self) -> Result> { let description = self.describe().await?;