diff --git a/crates/pyspark-rs/src/functions.rs b/crates/pyspark-rs/src/functions.rs index 74da2f4f..14ac379f 100644 --- a/crates/pyspark-rs/src/functions.rs +++ b/crates/pyspark-rs/src/functions.rs @@ -333,9 +333,11 @@ fn numpy_ndarray_ndim(obj: &Bound<'_, PyAny>) -> PyResult> { /// The literal of a NumPy scalar, typed as PySpark's `LiteralExpression._infer_type` types it: a /// bool, a signed int of 8 to 64 bits (tinyint to bigint), or a float32 or float64. A str or bytes /// scalar is `None`, for the string and binary conversions to take. Any other NumPy scalar (an -/// unsigned int, float16, longdouble, complex, datetime64 or timedelta64) is -/// UNSUPPORTED_DATA_TYPE, named by its Python type. NumPy is looked up in `sys.modules`, as a -/// NumPy scalar exists only when NumPy is imported. +/// unsigned int, float16, complex, datetime64 or timedelta64) is UNSUPPORTED_DATA_TYPE, +/// named by its Python type. A supported dtype with a scalar class outside PySpark's accepted +/// classes raises a bare `AssertionError`, such as a distinct `longlong`, or an eight-byte +/// `longdouble` on macOS arm64 and Windows. Wider `longdouble` dtypes are unsupported. +/// NumPy is looked up in `sys.modules`, as a NumPy scalar exists only when NumPy is imported. fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult> { let py = obj.py(); let modules = py.import("sys")?.getattr("modules")?; @@ -352,6 +354,36 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult &["int8"], + ("i", 2) => &["int16", "int8"], + ("i", 4) => &["int32", "int16", "int8"], + ("i", 8) => &["int64", "int32", "int16", "int8"], + ("f", 4) => &["float32"], + ("f", 8) => &["float64", "float32"], + _ => &[], + }; + if !scalar_classes.is_empty() { + let mut accepted = match kind.as_str() { + "i" => obj.is_instance_of::(), + "f" => obj.is_instance_of::(), + _ => false, + }; + if !accepted { + for class in scalar_classes { + if obj.is_instance(&numpy.getattr(*class)?)? { + accepted = true; + break; + } + } + } + if !accepted { + return Err(pyo3::exceptions::PyAssertionError::new_err(())); + } + } let literal = match (kind.as_str(), itemsize) { ("U", _) | ("S", _) => return Ok(None), ("i", 1) => LiteralExpression::Byte(obj.extract::()? as i32), diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 94aa92dd..a4515a80 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -1799,11 +1799,12 @@ impl PySparkSession { /// Stop this Spark session. Best-effort and idempotent, mirroring pyspark: a failure /// to release the session on the server is logged, not raised. Also clears the - /// active session and the cached `SQLContext` if they belong to this session. + /// active session and the cached `SQLContext` if they belong to this session. Listener + /// callbacks are not joined; the current event's listener snapshot may finish afterward. fn stop(&self, py: Python<'_>) -> PyResult<()> { crate::progress::forget_session(&self.session); crate::profiler::forget_session(&self.session); - if let Err(e) = py.detach(|| self.session.stop()) { + if let Err(e) = py.detach(|| self.session.stop().map_err(Box::new)) { let logger = py .import("logging")? .call_method1("getLogger", ("pyspark",))?; @@ -2075,8 +2076,9 @@ impl PySparkSession { } /// Close the session without releasing it on the server, for `client.close`. - fn _close(&self) { - self.session.close(); + /// Do not join listener callbacks; the current event's listener snapshot may finish afterward. + fn _close(&self, py: Python<'_>) { + py.detach(|| self.session.close()); } /// The client of the session, as PySpark Connect's `SparkSession._client`, which the diff --git a/crates/pyspark-rs/src/streaming.rs b/crates/pyspark-rs/src/streaming.rs index 58a3c9c4..3ce62698 100644 --- a/crates/pyspark-rs/src/streaming.rs +++ b/crates/pyspark-rs/src/streaming.rs @@ -1244,6 +1244,8 @@ impl PyStreamingQueryManager { /// it. The id of each registration is appended to the listener's `_rust_listener_ids`, as a /// listener added twice is in the bus twice. Should the server fail to add its listener, this /// warns, as PySpark Connect's does. + /// A terminal event-stream end or error clears the registrations; add listeners again + /// to open a new subscription. fn addListener(&self, py: Python<'_>, listener: Py) -> PyResult<()> { let session = crate::session::PySparkSession::new(self.session.clone()); listener @@ -1277,6 +1279,9 @@ impl PyStreamingQueryManager { /// the bus of the session, the first one that the bus has (as the bus removes the first /// occurrence), and the others, which are of other buses, stay. Should the server fail to /// remove its listener, this warns, and the listener stays, as PySpark Connect's does. + /// A callback coalesced with another caller's in-progress last-listener removal or close + /// keeps its ids so a failed operation can be retried. Only a completed removal consumes + /// one matching id from this session's bus. fn removeListener(&self, py: Python<'_>, listener: Py) -> PyResult<()> { let bound = listener.bind(py); let Ok(ids) = bound.getattr("_rust_listener_ids") else { @@ -1319,9 +1324,21 @@ impl PyStreamingQueryManager { } /// Remove all client-side listeners and stop the dispatch thread. Mirrors - /// `StreamingQueryManager.close`. + /// `StreamingQueryManager.close`. A server removal failure warns and keeps the + /// listeners registered, as PySpark Connect's does, so close can be retried. fn close(&self, py: Python<'_>) -> PyResult<()> { - crate::dataframe::detached(py, || self.inner.close().map_err(Box::new)) + if let Err(error) = crate::dataframe::detached(py, || self.inner.close().map_err(Box::new)) + { + return warn( + py, + format!( + "Failed to remove the listener because of exception: {}\n\ + The listener is not removed, please remove it again.", + error.value(py) + ), + ); + } + Ok(()) } fn streamListenerEvents(&self, py: Python<'_>) -> PyResult { diff --git a/crates/spark-connect/src/progress.rs b/crates/spark-connect/src/progress.rs index 5677f7c3..5a3a1130 100644 --- a/crates/spark-connect/src/progress.rs +++ b/crates/spark-connect/src/progress.rs @@ -2,11 +2,14 @@ //! `pyspark.sql.connect.shell.progress`: the progress handlers of the session are called with //! each progress of the server that has tasks, and once more, as done, when the execution ends. +use std::sync::Arc; + use spark_connect_core::client::ReattachableResponseStream; use spark_connect_core::error::Result; +use spark_connect_core::runtime::CallbackGuard; use spark_connect_proto as proto; -use crate::session::SparkSession; +use crate::session::{ProgressHandler, ProgressHandlers, SparkSession}; /// The progress of a stage of an execution, as PySpark's `StageInfo`. #[derive(Debug, Clone, PartialEq, Eq)] @@ -44,7 +47,7 @@ pub struct ProgressUpdate { /// The progress of one execution, as PySpark's `Progress` of the execution of a request. pub(crate) struct ProgressTracker { - session: SparkSession, + handlers: ProgressHandlers, operation_id: Option, stages: Vec, inflight_tasks: i64, @@ -53,7 +56,7 @@ pub(crate) struct ProgressTracker { impl ProgressTracker { pub(crate) fn new(session: &SparkSession, operation_id: Option) -> Self { ProgressTracker { - session: session.clone(), + handlers: session.progress_handlers(), operation_id, stages: Vec::new(), inflight_tasks: 0, @@ -80,12 +83,29 @@ impl ProgressTracker { } fn notify(&self, done: bool) { - self.session.notify_progress(|| ProgressUpdate { + // Clone handlers out of the registry lock so callbacks can add or remove handlers. + let handlers: Vec = self + .handlers + .lock() + .unwrap() + .iter() + .map(|(_, handler)| Arc::clone(handler)) + .collect(); + if handlers.is_empty() { + return; + } + let update = ProgressUpdate { stages: self.stages.clone(), inflight_tasks: self.inflight_tasks, operation_id: self.operation_id.clone(), done, - }); + }; + // A handler may run a query whose progress calls the handlers again. Preserve the + // execution callback-depth guard without retaining the session or its listener bus. + let _callback = CallbackGuard::enter(); + for handler in handlers { + handler(&update); + } } } @@ -124,6 +144,7 @@ impl ExecutionStream { #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use super::*; @@ -186,4 +207,50 @@ mod tests { ); session.remove_progress_handler(id); } + + #[test] + fn detached_progress_registry_preserves_the_execution_recursion_bound() { + use spark_connect_core::error::SparkErrorKind; + use spark_connect_core::runtime::{check_callback_depth, MAX_CALLBACK_DEPTH}; + + let session = SparkSession::builder() + .remote("sc://localhost:15002") + .get_or_create() + .expect("session"); + let handlers = Arc::downgrade(&session.progress_handlers()); + let calls = Arc::new(AtomicUsize::new(0)); + let failures = Arc::new(AtomicUsize::new(0)); + let seen_calls = Arc::clone(&calls); + let seen_failures = Arc::clone(&failures); + session.register_progress_handler(move |update| { + if update.done { + return; + } + let call = seen_calls.fetch_add(1, Ordering::SeqCst); + match check_callback_depth() { + Err(error) => { + assert_eq!(error.kind, SparkErrorKind::RecursionError); + seen_failures.fetch_add(1, Ordering::SeqCst); + } + Ok(()) => { + // Bound this test even if the guard regresses; no network execution or + // unbounded native recursion is needed to observe the missing guard. + assert!(call < MAX_CALLBACK_DEPTH); + let nested = ProgressTracker { + handlers: handlers.upgrade().expect("registry is alive"), + operation_id: Some("nested".to_string()), + stages: Vec::new(), + inflight_tasks: 0, + }; + nested.notify(false); + } + } + }); + let tracker = ProgressTracker::new(&session, None); + tracker.notify(false); + assert_eq!(calls.load(Ordering::SeqCst), MAX_CALLBACK_DEPTH); + assert_eq!(failures.load(Ordering::SeqCst), 1); + assert!(check_callback_depth().is_ok()); + session.clear_progress_handlers(); + } } diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index 6dc8d152..67944769 100644 --- a/crates/spark-connect/src/session.rs +++ b/crates/spark-connect/src/session.rs @@ -8,7 +8,7 @@ use std::sync::{Arc, Mutex}; use spark_connect_core::channel::ChannelBuilder; use spark_connect_core::client::{ExecutePlanHook, SparkConnectClient}; use spark_connect_core::error::{Result, SparkError}; -use spark_connect_core::runtime::{block_on, check_callback_depth, CallbackGuard}; +use spark_connect_core::runtime::{block_on, check_callback_depth}; use spark_connect_core::thread_entries::forget_when_thread_ends; use spark_connect_proto as proto; @@ -24,6 +24,9 @@ use crate::types::DataType; /// execution ends. Mirrors `SparkSession.registerProgressHandler`. pub type ProgressHandler = Arc; +/// The shared registry of progress handlers and their ids, without retaining a session. +pub(crate) type ProgressHandlers = Arc>>; + /// Metrics captured from the last executed query. Mirrors the data behind /// `pyspark.sql.DataFrame.executionInfo` / `client.core.ExecutionInfo`. #[derive(Debug, Clone, Default)] @@ -140,7 +143,7 @@ pub struct SparkSession { last_execution: Arc>>, /// Registered progress handlers, keyed by id (for removal). Invoked with each progress of /// an execution, and when it ends. Mirrors `registerProgressHandler`. - progress_handlers: Arc>>, + progress_handlers: ProgressHandlers, /// Monotonic id source for `register_progress_handler`. progress_handler_id: Arc, /// The handler of the updates of the Python accumulators of its executions. @@ -268,29 +271,9 @@ impl SparkSession { self.progress_handlers.lock().unwrap().clear(); } - /// Invoke all registered progress handlers with the progress of an execution, which - /// `update` builds only when there is a handler. Called by the progress of executions; not - /// part of the public API. - pub(crate) fn notify_progress(&self, update: impl FnOnce() -> ProgressUpdate) { - // Clone the handler `Arc`s out of the lock so a handler can't deadlock by - // touching the session's handler list. - let handlers: Vec = { - self.progress_handlers - .lock() - .unwrap() - .iter() - .map(|(_, h)| Arc::clone(h)) - .collect() - }; - if handlers.is_empty() { - return; - } - let update = update(); - // A handler may run a query of its own, whose progress calls the handlers again. - let _callback = CallbackGuard::enter(); - for h in handlers { - h(&update); - } + /// Share the progress-handler registry without retaining the session or its listener bus. + pub(crate) fn progress_handlers(&self) -> ProgressHandlers { + Arc::clone(&self.progress_handlers) } /// Execute `request` on the server, with the progress of the execution for the progress @@ -1141,13 +1124,20 @@ impl SparkSession { /// /// Idempotent, mirroring pyspark: the session is marked stopped even if releasing it /// on the server fails, and stopping an already stopped session is a no-op. + /// + /// Cancel local listener registration/dispatch first, release the server session, then + /// close the client even if release fails. As PySpark's stop, this does not make a + /// separate listener-removal RPC: ReleaseSession releases the server's listeners too. + /// Teardown does not wait for listener callbacks; the current event's listener snapshot + /// may finish after stop returns, while no later event is dispatched. pub fn stop(&self) -> Result<()> { if self.stopped.swap(true, Ordering::SeqCst) { return Ok(()); } - // As PySpark's `stop()`, the client is closed after the session is released, so that - // its requests raise `NO_ACTIVE_SESSION` from now on. + self.streams().cancel_listener_bus(); let released = block_on(self.client.release_session()); + // As PySpark's `stop()`, close the client even if server cleanup fails, so its + // future requests raise `NO_ACTIVE_SESSION`. self.client.close(); released } @@ -1159,9 +1149,12 @@ impl SparkSession { } /// Close this session without releasing it on the server, as `SparkConnectClient.close` - /// closes its channel: it is stopped, and `stop()` does nothing more. + /// closes its channel: it is stopped, and `stop()` does nothing more. Cancel local listener + /// registration/dispatch before closing the client. Callbacks are not joined; the current + /// event's listener snapshot may finish after close returns. pub fn close(&self) { self.stopped.store(true, Ordering::SeqCst); + self.streams().cancel_listener_bus(); self.client.close(); } diff --git a/crates/spark-connect/src/streaming.rs b/crates/spark-connect/src/streaming.rs index 77ee15e2..20f45903 100644 --- a/crates/spark-connect/src/streaming.rs +++ b/crates/spark-connect/src/streaming.rs @@ -4,10 +4,10 @@ //! for building and executing streaming workloads. use std::collections::HashMap; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Condvar, Mutex, MutexGuard, Weak}; use uuid; -use spark_connect_core::error::{Result, SparkError}; +use spark_connect_core::error::{Result, SparkError, SparkErrorKind}; use spark_connect_core::runtime::block_on; use spark_connect_proto as proto; @@ -755,15 +755,49 @@ pub struct ListenerEventStream { stream: crate::progress::ExecutionStream, buffered_events: std::vec::IntoIter<(i32, String)>, done: bool, + cancel: Option>, } impl ListenerEventStream { + /// Read the next response unless canceled, checking cancellation before a ready response. + /// Cancellation consumes its oneshot receiver and ends the stream like EOF; once done, + /// the transport is never polled again. + fn message(&mut self) -> Result> { + if self.done { + return Ok(None); + } + let response = block_on(async { + if let Some(cancel) = &mut self.cancel { + tokio::select! { + biased; + _ = cancel => None, + response = self.stream.message() => Some(response), + } + } else { + Some(self.stream.message().await) + } + }); + match response { + None | Some(Ok(None)) => { + // A completed oneshot must never be polled again. + self.cancel.take(); + self.done = true; + Ok(None) + } + Some(Err(error)) => { + self.done = true; + Err(error) + } + Some(response) => response, + } + } + /// Read the responses of the server until the one that tells that it added the listener of /// the session (`listener_bus_listener_added`), keeping the events of them, which the /// iterator gives first, as PySpark Connect's `_register_server_side_listener` waits for it. fn wait_until_listener_added(&mut self) -> Result<()> { loop { - match block_on(self.stream.message()) { + match self.message() { Ok(Some(resp)) => { if let Some( proto::execute_plan_response::ResponseType::StreamingQueryListenerEventsResult(res), @@ -785,6 +819,9 @@ impl ListenerEventStream { } } Ok(None) => { + // Terminal cancellation ends this wait like EOF. Teardown marks the + // session stopped first; add_listener's post-wait check rejects this + // outcome instead of publishing a worker on the canceled stream. self.done = true; return Ok(()); } @@ -811,7 +848,7 @@ impl Iterator for ListenerEventStream { } loop { - match block_on(self.stream.message()) { + match self.message() { Ok(Some(resp)) => { if let Some( proto::execute_plan_response::ResponseType::StreamingQueryListenerEventsResult(res), @@ -880,26 +917,44 @@ pub trait StreamingQueryListener: Send + Sync { } } -/// The client-side listener bus of the streaming queries of a session, which all of its -/// `StreamingQueryManager`s share, as PySpark Connect's `StreamingQueryListenerBus` of -/// `spark.streams`: the registered listeners, in their order, and the thread that posts the -/// events of the server to them, from the first listener to the removal of the last. `ops` -/// serializes the additions and removals, which talk to the server, while the thread takes -/// `state` only to read the listeners, and calls them without any lock. +/// The client-side listener bus of one session. Handles of the same session share it; +/// a new server session has a separate bus. It never owns a session or a manager. #[derive(Default)] pub(crate) struct ListenerBus { - ops: Mutex<()>, state: Mutex, + changed: Condvar, } #[derive(Default)] struct ListenerBusState { listeners: Vec<(String, Arc)>, thread: Option>, + // Dropping the bus wakes a blocked reader, so its worker does not retain a session. + cancel: Option>, + phase: ListenerBusPhase, + worker_id: Option, + generation: u64, +} + +#[derive(Debug, Default, PartialEq, Eq)] +enum ListenerBusPhase { + #[default] + Idle, + Starting, + Stopping, } impl ListenerBus { - /// Post `event` to every listener, as PySpark Connect's `post_to_all`. + /// Clear only the registrations of the worker that has finished this generation. + fn finish_generation(&self, generation: u64) { + let mut state = self.state.lock().unwrap(); + if state.generation == generation { + state.listeners.clear(); + state.cancel.take(); + } + } + + /// Post a synchronous start event without holding the state lock across callbacks. fn post_to_all(&self, event: &StreamingQueryListenerEvent) { let listeners: Vec> = self .state @@ -909,29 +964,49 @@ impl ListenerBus { .iter() .map(|(_, listener)| Arc::clone(listener)) .collect(); - for listener in &listeners { + for listener in listeners { listener.on_event(event); } } } -/// The background dispatch loop: posts each event of the stream of the server, on which the -/// server has added the listener of the session, to the listeners. It ends when the stream does, -/// as when the server was asked to remove its listener (`RemoveListenerBusListener`). An error -/// of the stream removes all the listeners, as PySpark Connect's handler thread does. -fn run_listener_event_loop(stream: ListenerEventStream, bus: Arc) { - for item in stream { - match item { - Ok((event_type, event_json)) => bus.post_to_all(&StreamingQueryListenerEvent { - event_type, - event_json, - }), - Err(_) => { - bus.state.lock().unwrap().listeners.clear(); - return; +/// The worker retains neither the session nor its bus. The stream owns only transport +/// and progress-handler state; dropping the last session/manager cancels the reader. +fn run_listener_event_loop( + mut stream: ListenerEventStream, + bus: Weak, + generation: u64, +) { + for item in &mut stream { + let Ok((event_type, event_json)) = item else { + break; + }; + let Some(bus) = bus.upgrade() else { + break; + }; + let listeners: Vec> = { + let state = bus.state.lock().unwrap(); + if state.generation != generation || state.listeners.is_empty() { + break; } + state + .listeners + .iter() + .map(|(_, listener)| Arc::clone(listener)) + .collect() + }; + drop(bus); + let event = StreamingQueryListenerEvent { + event_type, + event_json, + }; + for listener in listeners { + listener.on_event(&event); } } + if let Some(bus) = bus.upgrade() { + bus.finish_generation(generation); + } } /// Manager for active streaming queries. @@ -945,61 +1020,130 @@ pub struct StreamingQueryManager { } impl StreamingQueryManager { - /// Create a new StreamingQueryManager, of the listener bus of `session`. + /// Create a handle to this session's shared StreamingQueryManager state. pub(crate) fn new(session: SparkSession) -> Self { let bus = session.listener_bus(); StreamingQueryManager { session, bus } } - /// Register a client-side listener. Returns an id that can be passed to - /// [`StreamingQueryManager::remove_listener`]. Mirrors `StreamingQueryManager.addListener`: - /// for the first listener of the session, the server adds its listener before this returns, - /// so that the queries started then send their started events, and a thread posts the events - /// of the server to the listeners. Should the server fail to add it, the listener is not - /// added. + /// Wait for another caller's start/stop operation without holding a lock needed by + /// callbacks. A callback during shutdown may remove/close the bus as a no-op, but + /// cannot wait for its own worker to finish in order to start a new subscription. + fn idle_state(&self, removing: bool) -> Result>> { + let mut state = self.bus.state.lock().unwrap(); + while state.phase != ListenerBusPhase::Idle { + if state.worker_id == Some(std::thread::current().id()) { + if removing && state.phase == ListenerBusPhase::Stopping { + return Ok(None); + } + return Err(SparkError::connect_msg( + "cannot change the listener bus from its own callback while another \ + lifecycle operation is in progress", + )); + } + state = self.bus.changed.wait(state).unwrap(); + } + Ok(Some(state)) + } + + /// The error of registering a listener on a stopped session, as PySpark Connect's client. + fn no_active_session() -> SparkError { + SparkError::classed(SparkErrorKind::Connect, "NO_ACTIVE_SESSION", &[]) + } + + /// Register a listener, waiting for the server's acknowledgement when opening the bus. + /// The returned id can be removed through any manager handle of this session. + /// A terminal event-stream error or end clears all registrations; add the listeners + /// again to open a new subscription. Terminal teardown cancels acknowledgement waits + /// once the RPC stream is established. Establishing the stream itself follows the + /// transport's deadline and retry policy; a stopped session then reports NO_ACTIVE_SESSION. pub fn add_listener(&self, listener: Arc) -> Result { + let mut state = self.idle_state(false)?.expect("adding obtains the state"); + if self.session.is_stopped() { + return Err(Self::no_active_session()); + } let id = uuid::Uuid::new_v4().to_string(); - let _ops = self.bus.ops.lock().unwrap(); - let first = { - let mut state = self.bus.state.lock().unwrap(); + if !state.listeners.is_empty() { state.listeners.push((id.clone(), listener)); - state.listeners.len() == 1 - }; - if first { - let registered = self.listener_event_stream().and_then(|mut stream| { - stream.wait_until_listener_added()?; - Ok(stream) - }); - let stream = match registered { - Ok(stream) => stream, - Err(e) => { - self.bus.state.lock().unwrap().listeners.clear(); - return Err(e); - } - }; - let bus = Arc::clone(&self.bus); - let thread = std::thread::spawn(move || run_listener_event_loop(stream, bus)); - self.bus.state.lock().unwrap().thread = Some(thread); + return Ok(id); + } + state.phase = ListenerBusPhase::Starting; + // An earlier worker must not clear this registration's cancellation sender. + state.generation = state.generation.wrapping_add(1); + let generation = state.generation; + let (cancel, receiver) = tokio::sync::oneshot::channel(); + state.cancel = Some(cancel); + let previous = state.thread.take(); + drop(state); + Self::join_listener_thread(previous); + let stream = self.listener_event_stream().and_then(|mut stream| { + // Cancel a pending acknowledgement once the RPC stream exists. Establishing + // the stream above still follows the transport's deadline and retry policy. + stream.cancel = Some(receiver); + stream.wait_until_listener_added()?; + Ok(stream) + }); + let mut state = self.bus.state.lock().unwrap(); + state.phase = ListenerBusPhase::Idle; + self.bus.changed.notify_all(); + // Terminal teardown sets stopped before canceling the stream. Its acknowledgement + // wait treats cancellation like EOF, so this check rejects a canceled registration. + if self.session.is_stopped() { + state.cancel.take(); + return Err(Self::no_active_session()); + } + match stream { + Ok(stream) => { + state.listeners.push((id.clone(), listener)); + let bus = Arc::downgrade(&self.bus); + let worker = std::thread::spawn(move || { + run_listener_event_loop(stream, bus, generation); + }); + state.worker_id = Some(worker.thread().id()); + state.thread = Some(worker); + Ok(id) + } + Err(error) => { + state.cancel.take(); + Err(error) + } } - Ok(id) } - /// Remove a client-side listener by id, which does nothing for an id of no listener. Mirrors - /// `StreamingQueryManager.removeListener`: for the last listener, the server is asked to - /// remove its listener, which ends its stream, and this returns once the thread has posted - /// the events left. Should that request fail, the listener stays, and the error is returned. + /// Remove a listener by id. Unknown ids do nothing. Removing the last listener drains + /// the server event stream before returning, as PySpark's removeListener does. + /// A callback during another caller's shutdown returns without waiting for its own + /// worker; that caller owns the final result and may retry a failed removal. pub fn remove_listener(&self, id: &str) -> Result<()> { self.remove_listener_found(id).map(|_| ()) } - /// As [`remove_listener`](Self::remove_listener), which returns whether the listener of `id` - /// was in this bus (and so is removed now), for a client that keeps the ids of its listeners - /// per bus, as the ids of another session's bus are not in this one. + /// As [`remove_listener`](Self::remove_listener), returning whether a registration of `id` + /// was removed from this bus. An absent id or a callback coalesced with an in-progress + /// shutdown returns false, so a wrapper retains its token if that shutdown later fails. + /// An id from another session's bus is absent here and leaves that registration untouched. pub fn remove_listener_found(&self, id: &str) -> Result { - let _ops = self.bus.ops.lock().unwrap(); - { - let mut state = self.bus.state.lock().unwrap(); - let Some(index) = state.listeners.iter().position(|(lid, _)| lid == id) else { + self.remove_listeners(Some(id)) + } + + /// Remove all listeners and stop this session's shared dispatch thread. + /// A server removal failure retains the registrations and worker so closing can be retried. + pub fn close(&self) -> Result<()> { + self.remove_listeners(None).map(|_| ()) + } + + fn remove_listeners(&self, id: Option<&str>) -> Result { + let Some(mut state) = self.idle_state(true)? else { + // Another caller waits for this callback. It owns the shutdown result; keep + // the registration token until that caller has actually removed the listener. + return Ok(false); + }; + if let Some(id) = id { + let Some(index) = state + .listeners + .iter() + .position(|(listener_id, _)| listener_id == id) + else { return Ok(false); }; if state.listeners.len() > 1 { @@ -1007,32 +1151,44 @@ impl StreamingQueryManager { return Ok(true); } } - self.remove_server_listener()?; - self.join_listener_thread(); - self.bus - .state - .lock() - .unwrap() - .listeners - .retain(|(lid, _)| lid != id); - Ok(true) + let had_listeners = !state.listeners.is_empty(); + let generation = state.generation; + state.phase = ListenerBusPhase::Stopping; + let worker = state.thread.take(); + drop(state); + let removal = if had_listeners { + self.remove_server_listener() + } else { + Ok(()) + }; + self.finish_listener_removal(generation, worker, removal) + .map(|_| had_listeners) } - /// Remove all client-side listeners and stop the dispatch thread. Mirrors - /// `StreamingQueryManager.close`. Should the server fail to remove its listener, the thread - /// is left to end with its stream. - pub fn close(&self) -> Result<()> { - let _ops = self.bus.ops.lock().unwrap(); - if self.bus.state.lock().unwrap().listeners.is_empty() { - return Ok(()); - } - if self.remove_server_listener().is_ok() { - self.join_listener_thread(); - } + /// Reconcile an unlocked removal RPC, joining canceled generations outside the state lock. + fn finish_listener_removal( + &self, + generation: u64, + worker: Option>, + removal: Result<()>, + ) -> Result<()> { let mut state = self.bus.state.lock().unwrap(); - state.listeners.clear(); - state.thread = None; - Ok(()) + if removal.is_err() && state.generation == generation { + // A failed ordinary removal keeps its registrations and worker for a retry. + state.thread = worker; + } else { + // Keep worker_id and Stopping until after the join: callbacks must still be + // able to recognize their worker and coalesce a concurrent remove/close. + drop(state); + Self::join_listener_thread(worker); + state = self.bus.state.lock().unwrap(); + state.listeners.clear(); + state.cancel.take(); + state.worker_id = None; + } + state.phase = ListenerBusPhase::Idle; + self.bus.changed.notify_all(); + removal } /// The client-side listeners of the session, in the order that they were added, as the @@ -1053,28 +1209,46 @@ impl StreamingQueryManager { self.bus.post_to_all(event); } - /// Ask the server to remove the listener of the session, which ends its stream of events. + /// Locally cancel registration and dispatch of a stopped/closed session before or + /// without server cleanup. Do not wait for callbacks: the current event's listener + /// snapshot may finish, but the generation change prevents another event dispatch. + /// Never wait for a lifecycle phase: this can be called by its worker callback while + /// another caller is already draining that worker. That caller keeps its join handle. + pub(crate) fn cancel_listener_bus(&self) { + let worker = { + let mut state = self.bus.state.lock().unwrap(); + state.listeners.clear(); + state.cancel.take(); + state.generation = state.generation.wrapping_add(1); + state.thread.take() + }; + // A callback may wait for an RPC ended by ReleaseSession, or for a lock held + // by the terminal caller. Detach instead of making teardown wait for it. + drop(worker); + } + + /// Join the event worker unless called from that worker itself. + fn join_listener_thread(handle: Option>) { + if let Some(handle) = handle { + // A listener may remove itself from its callback. It cannot join itself; + // after the callback returns the stopped stream lets its worker exit. + if handle.thread().id() != std::thread::current().id() { + let _ = handle.join(); + } + } + } + + /// Send RemoveListenerBusListener so the server ends the event stream. fn remove_server_listener(&self) -> Result<()> { - let mut lb = proto::StreamingQueryListenerBusCommand::default(); - lb.command = Some( + let mut command = proto::StreamingQueryListenerBusCommand::default(); + command.command = Some( proto::streaming_query_listener_bus_command::Command::RemoveListenerBusListener(true), ); crate::dataframe::execute_command_collect( &self.session, - proto::command::CommandType::StreamingQueryListenerBusCommand(lb), - ) - .map(|_| ()) - } - - /// Wait for the dispatch thread to post the events left of its ended stream, unless this is - /// that thread, as when a listener removes the last listener. - fn join_listener_thread(&self) { - let thread = self.bus.state.lock().unwrap().thread.take(); - if let Some(thread) = thread { - if thread.thread().id() != std::thread::current().id() { - let _ = thread.join(); - } - } + proto::command::CommandType::StreamingQueryListenerBusCommand(command), + )?; + Ok(()) } /// Get all active streaming queries. @@ -1262,6 +1436,7 @@ impl StreamingQueryManager { stream: response_stream, buffered_events: vec![].into_iter(), done: false, + cancel: None, }) } @@ -1307,6 +1482,167 @@ mod tests { .expect("failed to build session") } + #[test] + fn progress_tracker_does_not_keep_listener_bus_alive() { + let session = session(); + let bus = Arc::downgrade(&session.listener_bus()); + let tracker = crate::progress::ProgressTracker::new(&session, None); + drop(session); + assert!(bus.upgrade().is_none()); + drop(tracker); + } + + #[test] + fn a_callback_coalesced_with_shutdown_does_not_consume_its_registration() { + struct Listener; + impl StreamingQueryListener for Listener { + fn on_event(&self, _: &StreamingQueryListenerEvent) {} + } + let spark = session(); + let manager = spark.streams(); + { + let mut state = manager.bus.state.lock().unwrap(); + state + .listeners + .push(("registered".to_string(), Arc::new(Listener))); + state.phase = ListenerBusPhase::Stopping; + state.worker_id = Some(std::thread::current().id()); + } + assert!(!manager.remove_listener_found("registered").unwrap()); + assert_eq!(manager.listeners().len(), 1); + let mut state = manager.bus.state.lock().unwrap(); + state.phase = ListenerBusPhase::Idle; + state.worker_id = None; + } + + #[test] + fn removing_a_duplicate_registration_preserves_other_and_foreign_ids() { + struct Listener; + impl StreamingQueryListener for Listener { + fn on_event(&self, _: &StreamingQueryListenerEvent) {} + } + let spark = session(); + let manager = spark.streams(); + let listener: Arc = Arc::new(Listener); + { + let mut state = manager.bus.state.lock().unwrap(); + state + .listeners + .push(("first".to_string(), Arc::clone(&listener))); + state + .listeners + .push(("second".to_string(), Arc::clone(&listener))); + } + assert!(!manager.remove_listener_found("foreign").unwrap()); + assert_eq!(manager.listeners().len(), 2); + assert!(manager.remove_listener_found("first").unwrap()); + let state = manager.bus.state.lock().unwrap(); + assert_eq!(state.listeners.len(), 1); + assert_eq!(state.listeners[0].0, "second"); + assert!(Arc::ptr_eq(&state.listeners[0].1, &listener)); + } + + /// Install a synthetic quiet worker to check cancellation sender/receiver wiring. + /// This bypasses run_listener_event_loop and does not establish its Weak bus ownership. + fn waiting_listener_worker(session: &SparkSession) -> std::sync::mpsc::Receiver<()> { + struct Listener; + impl StreamingQueryListener for Listener { + fn on_event(&self, _: &StreamingQueryListenerEvent) {} + } + + let (cancel, receiver) = tokio::sync::oneshot::channel(); + let (finished, observed) = std::sync::mpsc::channel(); + let worker = std::thread::spawn(move || { + let _ = block_on(receiver); + let _ = finished.send(()); + }); + let bus = session.listener_bus(); + let mut state = bus.state.lock().unwrap(); + state + .listeners + .push(("listener".to_string(), Arc::new(Listener))); + state.cancel = Some(cancel); + state.worker_id = Some(worker.thread().id()); + state.thread = Some(worker); + observed + } + + #[test] + fn dropping_the_last_native_session_cancels_its_listener_worker() { + let session = session(); + let finished = waiting_listener_worker(&session); + drop(session); + finished + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("dropping the bus cancels its quiet worker"); + } + + #[test] + fn closing_a_session_eventually_cancels_its_listener_worker() { + let session = session(); + let finished = waiting_listener_worker(&session); + session.close(); + assert!(session.is_stopped()); + assert!(session.streams().listeners().is_empty()); + finished + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("close eventually cancels the quiet worker without joining it"); + } + + #[test] + fn a_failed_removal_joins_its_terminally_canceled_worker() { + let session = session(); + let finished = waiting_listener_worker(&session); + let manager = session.streams(); + let (generation, worker) = { + let mut state = manager.bus.state.lock().unwrap(); + state.phase = ListenerBusPhase::Stopping; + (state.generation, state.thread.take()) + }; + // The in-flight removal owns the join handle, so terminal close only cancels it. + session.close(); + let error = manager + .finish_listener_removal( + generation, + worker, + Err(SparkError::connect_msg("listener removal failed")), + ) + .expect_err("the caller still receives its removal failure"); + assert_eq!(error.to_string(), "listener removal failed"); + finished + .try_recv() + .expect("failed removal joins a generation canceled during its RPC"); + let state = manager.bus.state.lock().unwrap(); + assert_eq!(state.phase, ListenerBusPhase::Idle); + assert!(state.worker_id.is_none()); + assert!(state.thread.is_none()); + assert!(state.listeners.is_empty()); + } + + #[test] + fn an_old_generation_cannot_clear_new_listeners_or_cancel_their_worker() { + let session = session(); + let finished = waiting_listener_worker(&session); + let bus = session.listener_bus(); + let old_generation = { + let mut state = bus.state.lock().unwrap(); + let old_generation = state.generation; + state.generation = state.generation.wrapping_add(1); + old_generation + }; + bus.finish_generation(old_generation); + { + let state = bus.state.lock().unwrap(); + assert_eq!(state.listeners.len(), 1); + assert!(state.cancel.is_some()); + assert!(!state.thread.as_ref().unwrap().is_finished()); + } + session.close(); + finished + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("only terminal close eventually cancels the current worker"); + } + #[test] fn stream_reader_format_option() { let spark = session(); diff --git a/crates/spark-connect/tests/e2e_streaming.rs b/crates/spark-connect/tests/e2e_streaming.rs index 8092457b..3486e1be 100644 --- a/crates/spark-connect/tests/e2e_streaming.rs +++ b/crates/spark-connect/tests/e2e_streaming.rs @@ -238,3 +238,75 @@ fn streaming_listeners_are_those_of_the_session() { assert!(events.contains(&QUERY_PROGRESS_EVENT), "{events:?}"); assert_eq!(events.len(), seen, "{events:?}"); } + +/// Exercise the real event reader before dropping every native session/manager owner. +/// After the probe query terminates the subscription is quiet; its worker must not keep +/// the bus, its listener, or its cancellation sender alive by retaining a strong bus Arc. +#[test] +fn streaming_last_native_owner_drop_cancels_a_quiet_listener_worker() { + if !should_run() { + return; + } + use spark_connect::streaming::{ + StreamingQueryListener, StreamingQueryListenerEvent, QUERY_TERMINATED_EVENT, + }; + use std::sync::{mpsc, Arc}; + use std::time::Duration; + + struct DropObservedListener { + terminated: mpsc::Sender<()>, + dropped: mpsc::Sender<()>, + } + impl StreamingQueryListener for DropObservedListener { + fn on_event(&self, event: &StreamingQueryListenerEvent) { + if event.event_type == QUERY_TERMINATED_EVENT { + let _ = self.terminated.send(()); + } + } + } + impl Drop for DropObservedListener { + fn drop(&mut self) { + let _ = self.dropped.send(()); + } + } + + let spark = session(); + let manager = spark.streams(); + let (terminated, observed_termination) = mpsc::channel(); + let (dropped, observed_drop) = mpsc::channel(); + let listener = Arc::new(DropObservedListener { + terminated, + dropped, + }); + manager + .add_listener(listener.clone()) + .expect("register the native ownership listener"); + let query = spark + .read_stream() + .format("rate") + .option("rowsPerSecond", "1") + .load(None) + .write_stream() + .format("noop") + .query_name("e2e_last_native_listener_owner") + .start("") + .expect("start the ownership probe query"); + query.stop().expect("stop the ownership probe query"); + // A started callback is synchronous. This asynchronous bus event proves that + // the real worker entered its loop before the test drops any of its owners. + observed_termination + .recv_timeout(Duration::from_secs(30)) + .expect("the real listener worker receives the query termination"); + assert_eq!(manager.listeners().len(), 1); + + // The listener stores only channels, so it cannot retain a session or manager. + // Do not remove the listener or close/stop the session: cancellation must follow + // the last owner's drop, rather than explicit cleanup. + drop(query); + drop(listener); + drop(manager); + drop(spark); + observed_drop + .recv_timeout(Duration::from_secs(5)) + .expect("dropping every native owner releases the listener of the quiet real worker"); +} diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index 0d4beee0..44dd2e33 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -21,7 +21,8 @@ import re from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from threading import Event +from queue import Empty, Queue +from threading import Condition, Event # The drop-in package does not ship generated service protobuf classes, so these @@ -258,6 +259,10 @@ class Service: # Each AnalyzePlanRequest, recorded. analyze_requests = [] + def __init__(self): + self.listener_condition = Condition() + self.listener_epoch = 0 + def reset(self): # Back to the state of a new server: the tests of a module share one. self.calls = 0 @@ -265,7 +270,6 @@ def reset(self): self.data_type = struct_schema() self.failure = None self.block_response = False - self.response_gate.clear() self.block_execute = False self.analyze_requests = [] self.execute_requests = [] @@ -284,16 +288,48 @@ def reset(self): self.clone_requests = [] # The DDL of the JSON of each JsonToDDL. self.ddl_string = "" + with self.listener_condition: + # Late completion of an old RPC must not count towards the next test. + self.listener_epoch += 1 + # Release old waiters without opening the next test's response gate. + self.response_gate.set() + self.response_gate = Event() + self.execute_started = 0 + self.execute_gate_timeouts = 0 + self.release_session_unblocks_execute = False + for queue in getattr(self, "listener_queues", {}).values(): + queue.put(None) + self.listener_queues = {} + self.listener_subscriptions = 0 + self.listener_removals = 0 + self.listener_streams_finished = 0 + self.listener_condition.notify_all() + for name in ("listener_register_gate", "listener_remove_gate"): + gate = getattr(self, name, None) + if gate is not None: + gate.set() + self.listener_bus_enabled = False + self.listener_register_failure = False + self.listener_register_blocked = False + self.listener_register_gate = Event() + self.listener_remove_failure = False + self.listener_remove_blocked = False + self.listener_remove_gate = Event() + self.release_session_failure = False + self.release_session_calls = 0 def analyze(self, request, context): self.calls += 1 self.analyze_requests.append(request) + with self.listener_condition: + response_gate = self.response_gate + blocked = self.block_response # base.proto: AnalyzePlanRequest.is_streaming=8; AnalyzePlanResponse.is_streaming=6, # IsStreaming.is_streaming=1, of `streaming`, which a test sets. if 8 in parse_message(request): streaming = getattr(self, "streaming", False) return message_field(6, varint_field(1, 1) if streaming else b"") - if self.block_response and not self.response_gate.wait(5): + if blocked and not response_gate.wait(5): context.abort( grpc.StatusCode.DEADLINE_EXCEEDED, "Python thread could not run while metadata lookup waited", @@ -364,14 +400,144 @@ def analyze(self, request, context): read_confs = [] set_confs = [] - def execute(self, request, context): - self.execute_calls += 1 - self.execute_requests.append(request) - if self.block_execute and not self.response_gate.wait(5): - context.abort( - grpc.StatusCode.DEADLINE_EXCEEDED, - "Python thread could not run while the statement ran", + def release_session(self, request, context): + with self.listener_condition: + self.release_session_calls += 1 + # Opt in only for the terminal callback-RPC regression. Ordinary tests + # retain independent response and session-release behavior. + if self.release_session_unblocks_execute: + self.response_gate.set() + if self.release_session_failure: + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "session release failed") + return b"" + + def finish_listener_stream(self, epoch, session_id, queue): + with self.listener_condition: + if epoch != self.listener_epoch: + return + if self.listener_queues.get(session_id) is queue: + self.listener_queues.pop(session_id) + self.listener_streams_finished += 1 + self.listener_condition.notify_all() + + def wait_listener_calls(self, counter, count, timeout): + with self.listener_condition: + return self.listener_condition.wait_for( + lambda: getattr(self, counter) >= count, timeout=timeout ) + + def wait_listener_streams_finished(self, count, timeout): + return self.wait_listener_calls("listener_streams_finished", count, timeout) + + def wait_execute_started(self, count, timeout): + with self.listener_condition: + epoch = self.listener_epoch + reached = self.listener_condition.wait_for( + lambda: epoch != self.listener_epoch or self.execute_started >= count, + timeout=timeout, + ) + return reached and epoch == self.listener_epoch and self.execute_started >= count + + def execute(self, request, context): + with self.listener_condition: + self.execute_calls += 1 + self.execute_requests.append(request) + execute_epoch = self.listener_epoch + response_gate = self.response_gate + blocked = self.block_execute + fields = parse_message(request) + plan = parse_message(fields.get(3, [b""])[0]) + command = parse_message(plan.get(2, [b""])[0]) + listener = parse_message(command.get(11, [b""])[0]) + if self.listener_bus_enabled and listener: + session_id = fields[1][0].decode() + header = message_field(12, fields.get(6, [b""])[0]) + if 1 in listener: + queue = Queue() + with self.listener_condition: + epoch = self.listener_epoch + self.listener_subscriptions += 1 + registration_failed = self.listener_register_failure + if not registration_failed: + self.listener_queues[session_id] = queue + self.listener_condition.notify_all() + if registration_failed: + context.abort( + grpc.StatusCode.INVALID_ARGUMENT, "listener registration failed" + ) + counted = False + try: + if self.listener_register_blocked: + gate = self.listener_register_gate + # Force an established-stream/pre-acknowledgement window with + # early headers, which real servers seldom expose. An RPC still + # being established follows its transport deadline/retry policy. + context.send_initial_metadata(()) + context.add_callback(gate.set) + if not gate.wait(15): + context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + "listener registration gate was not released", + ) + if not context.is_active(): + return + yield (header + message_field(13, "registered") + + message_field(16, varint_field(2, 1))) + sequence = 0 + while context.is_active(): + try: + event = queue.get(timeout=0.1) + except Empty: + continue + if event is None: + break + event_type, event_json = event + if event_type is None: + context.abort(grpc.StatusCode.INVALID_ARGUMENT, event_json) + sequence += 1 + payload = message_field(1, event_json) + varint_field(2, event_type) + yield (header + message_field(13, f"event-{sequence}") + + message_field(16, message_field(1, payload))) + # The client returns at ResultComplete, before server-side EOF/finally. + self.finish_listener_stream(epoch, session_id, queue) + counted = True + yield header + message_field(13, "finished") + message_field(14, b"") + finally: + if not counted: + self.finish_listener_stream(epoch, session_id, queue) + return + with self.listener_condition: + self.listener_removals += 1 + # Keep an old, gated removal from closing a queue created after reset. + queue = self.listener_queues.get(session_id) + gate = self.listener_remove_gate + blocked = self.listener_remove_blocked + removal_failed = self.listener_remove_failure + self.listener_condition.notify_all() + if blocked and not gate.wait(15): + context.abort( + grpc.StatusCode.INVALID_ARGUMENT, "listener removal gate was not released" + ) + if removal_failed: + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "listener removal failed") + if queue is not None: + queue.put(None) + yield header + message_field(13, "removed") + message_field(14, b"") + return + if blocked: + with self.listener_condition: + if execute_epoch == self.listener_epoch: + self.execute_started += 1 + self.listener_condition.notify_all() + if not response_gate.wait(5): + with self.listener_condition: + if execute_epoch == self.listener_epoch: + self.execute_gate_timeouts += 1 + self.listener_condition.notify_all() + context.abort( + grpc.StatusCode.DEADLINE_EXCEEDED, + "Python thread could not run while the statement ran", + ) if self.execute_failure: code, message = ( self.execute_failure @@ -476,11 +642,12 @@ def config(self, request, context): ) service = Service() + service.reset() with ThreadPoolExecutor(max_workers=4) as executor: server = grpc.server(executor) handlers = { "AnalyzePlan": grpc.unary_unary_rpc_method_handler(service.analyze), - "ReleaseSession": grpc.unary_unary_rpc_method_handler(lambda request, context: b""), + "ReleaseSession": grpc.unary_unary_rpc_method_handler(service.release_session), "ExecutePlan": grpc.unary_stream_rpc_method_handler(service.execute), "ReattachExecute": grpc.unary_stream_rpc_method_handler( lambda request, context: iter([message_field(14, b"")]) @@ -512,6 +679,36 @@ def config(self, request, context): elif command[0] == "reset": service.reset() connection.send(None) + elif command[0] == "listener_event": + _, session_id, event_type, event_json = command + with service.listener_condition: + queue = service.listener_queues.get(session_id) + if queue is not None: + queue.put((event_type, event_json)) + connection.send(None) + elif command[0] == "end_listener_stream": + _, session_id, failure = command + with service.listener_condition: + queue = service.listener_queues.get(session_id) + if queue is not None: + queue.put((None, "listener event stream failed") if failure else None) + connection.send(queue is not None) + elif command[0] == "wait_listener_removal": + connection.send(service.wait_listener_calls("listener_removals", *command[1:])) + elif command[0] == "wait_listener_registration": + connection.send(service.wait_listener_calls( + "listener_subscriptions", *command[1:] + )) + elif command[0] == "wait_listener_streams_finished": + connection.send(service.wait_listener_streams_finished(*command[1:])) + elif command[0] == "wait_execute_started": + connection.send(service.wait_execute_started(*command[1:])) + elif command[0] == "release_listener_registration": + service.listener_register_gate.set() + connection.send(None) + elif command[0] == "release_listener_removal": + service.listener_remove_gate.set() + connection.send(None) finally: server.stop(0).wait() @@ -540,6 +737,46 @@ def reset(self): assert self.connection.poll(10), "Mock server did not respond" self.connection.recv() + def listener_event(self, session_id, event_type, event_json): + self.connection.send(("listener_event", session_id, event_type, json.dumps(event_json))) + assert self.connection.poll(10), "Mock server did not respond" + self.connection.recv() + + def end_listener_stream(self, session_id, failure=False): + self.connection.send(("end_listener_stream", session_id, failure)) + assert self.connection.poll(10), "Mock server did not respond" + return self.connection.recv() + + def wait_listener_removal(self, count=1, timeout=5): + self.connection.send(("wait_listener_removal", count, timeout)) + assert self.connection.poll(timeout + 5), "Mock server did not respond" + return self.connection.recv() + + def wait_listener_registration(self, count=1, timeout=5): + self.connection.send(("wait_listener_registration", count, timeout)) + assert self.connection.poll(timeout + 5), "Mock server did not respond" + return self.connection.recv() + + def wait_listener_streams_finished(self, count, timeout=5): + self.connection.send(("wait_listener_streams_finished", count, timeout)) + assert self.connection.poll(timeout + 5), "Mock server did not respond" + return self.connection.recv() + + def wait_execute_started(self, count=1, timeout=5): + self.connection.send(("wait_execute_started", count, timeout)) + assert self.connection.poll(timeout + 5), "Mock server did not respond" + return self.connection.recv() + + def release_listener_registration(self): + self.connection.send(("release_listener_registration",)) + assert self.connection.poll(10), "Mock server did not respond" + self.connection.recv() + + def release_listener_removal(self): + self.connection.send(("release_listener_removal",)) + assert self.connection.poll(10), "Mock server did not respond" + self.connection.recv() + @contextmanager def mock_connect_session(): diff --git a/python/tests/streaming_listener_cases.py b/python/tests/streaming_listener_cases.py new file mode 100644 index 00000000..a5750312 --- /dev/null +++ b/python/tests/streaming_listener_cases.py @@ -0,0 +1,782 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Streaming listener lifecycle tests against the separately hosted mock service.""" + +import subprocess +import sys +import tempfile +import threading +import time +import warnings +from pathlib import Path + +import pytest + +from pyspark.sql.streaming.listener import StreamingQueryListener + + +class IdleListener(StreamingQueryListener): + def __init__(self): + super().__init__() + self.timestamps = [] + self.condition = threading.Condition() + + def onQueryStarted(self, event): + pass + + def onQueryProgress(self, event): + pass + + def onQueryIdle(self, event): + with self.condition: + self.timestamps.append(event.timestamp) + self.condition.notify_all() + + def onQueryTerminated(self, event): + pass + + def wait_for(self, timestamp): + with self.condition: + assert self.condition.wait_for(lambda: timestamp in self.timestamps, timeout=5) + + +def emit_idle(service, session, timestamp): + service.listener_event(session.session_id, 3, { + "id": "00000000-0000-4000-8000-000000000001", + "runId": "00000000-0000-4000-8000-000000000002", + "timestamp": timestamp, + }) + + +def run_listener_case(*arguments): + # A deadlocked native callback must fail without hanging the test runner. A file + # keeps a killed child's mock-service process from holding a capture pipe open. + with tempfile.TemporaryFile(mode="w+") as output: + try: + result = subprocess.run( + [sys.executable, str(Path(__file__)), *arguments], + stdout=output, stderr=subprocess.STDOUT, timeout=30, + ) + except subprocess.TimeoutExpired: + output.seek(0) + pytest.fail("Concurrent listener shutdown timed out:\n" + output.read()) + output.seek(0) + assert result.returncode == 0, output.read() + + +def listener_warnings(caught): + prefixes = ("Failed to add the listener", "Failed to remove the listener") + return [str(item.message) for item in caught + if issubclass(item.category, UserWarning) and str(item.message).startswith(prefixes)] + + +class TestStreamingListenerLifecycle: + @pytest.mark.parametrize("remove_failure, callback_action", [ + (False, "remove"), (True, "remove"), (False, "stop"), (True, "stop"), (False, "add"), + (False, "remove_duplicate"), (True, "remove_duplicate"), + ]) + def test_streaming_listener_callback_lifecycle_calls_during_external_close( + self, remove_failure, callback_action, + ): + run_listener_case( + "--concurrent-close-case", str(int(remove_failure)), callback_action, + ) + + @pytest.mark.parametrize("method", ["stop", "client_close"]) + def test_streaming_listener_terminal_teardown_cancels_established_stream_registration( + self, method, + ): + run_listener_case("--pending-registration-case", method) + + def test_streaming_listener_stop_releases_a_session_with_a_callback_rpc(self): + run_listener_case("--callback-rpc-case") + + @pytest.mark.parametrize("method", ["stop", "client_close"]) + @pytest.mark.parametrize("blocker", ["event", "lock"]) + def test_streaming_listener_terminal_teardown_does_not_wait_for_a_callback( + self, method, blocker, + ): + run_listener_case("--blocked-callback-case", method, blocker) + + def test_streaming_listener_started_waits_require_the_requested_rpc(self): + run_listener_case("--second-rpc-case") + + @pytest.mark.parametrize("termination", ["end", "error"]) + def test_streaming_listener_ended_stream_clears_registry_and_reopens(self, termination): + run_listener_case("--ended-stream-case", termination) + + def test_streaming_listener_managers_share_a_session_registry(self, mock_connect): + spark, service = mock_connect + service.listener_bus_enabled = True + first, second = IdleListener(), IdleListener() + assert spark.streams is spark.streams + alias = spark.range(1).sparkSession.streams + try: + spark.streams.addListener(first) + spark.streams.addListener(second) + assert service.listener_subscriptions == 1 + emit_idle(service, spark, "both") + first.wait_for("both") + second.wait_for("both") + + # A handle obtained through a DataFrame can remove a listener registered + # through the original session, without stopping the other listener. + alias.removeListener(first) + assert service.listener_removals == 0 + emit_idle(service, spark, "second") + second.wait_for("second") + assert first.timestamps == ["both"] + spark.streams.removeListener(first) + assert service.listener_removals == 0 + spark.streams.removeListener(second) + assert service.listener_removals == 1 + assert service.wait_listener_streams_finished(1) + spark.streams.close() + assert service.listener_removals == 1 + finally: + spark.streams.close() + + def test_streaming_listener_close_reopens_the_same_manager(self, mock_connect): + spark, service = mock_connect + service.listener_bus_enabled = True + manager = spark.streams + first, second = IdleListener(), IdleListener() + try: + manager.addListener(first) + spark.streams.close() + assert service.listener_removals == 1 + assert service.wait_listener_streams_finished(1) + spark.streams.addListener(second) + assert service.listener_subscriptions == 2 + emit_idle(service, spark, "reopened") + second.wait_for("reopened") + assert first.timestamps == [] + finally: + manager.close() + assert service.listener_removals == 2 + assert service.wait_listener_streams_finished(2) + + def test_streaming_listener_new_sessions_have_separate_buses(self, mock_connect): + spark, service = mock_connect + service.listener_bus_enabled = True + other = spark.newSession() + first, second = IdleListener(), IdleListener() + try: + spark.streams.addListener(first) + other.streams.addListener(second) + assert service.listener_subscriptions == 2 + other.streams.close() + emit_idle(service, spark, "original") + first.wait_for("original") + assert second.timestamps == [] + assert service.listener_removals == 1 + finally: + other.stop() + spark.streams.close() + assert service.listener_removals == 2 + + def test_streaming_listener_duplicate_ids_are_removed_only_from_the_matching_bus( + self, mock_connect, + ): + spark, service = mock_connect + service.listener_bus_enabled = True + other = spark.newSession() + listener = IdleListener() + try: + spark.streams.addListener(listener) + spark.streams.addListener(listener) + first, second = listener._rust_listener_ids + other.streams.addListener(listener) + third = listener._rust_listener_ids[-1] + assert len({first, second, third}) == 3 + assert spark.streams._sqlb._listener_bus == [listener, listener] + assert other.streams._sqlb._listener_bus == [listener] + assert service.listener_subscriptions == 2 + + other.streams.removeListener(listener) + assert listener._rust_listener_ids == [first, second] + assert other.streams._sqlb._listener_bus == [] + assert spark.streams._sqlb._listener_bus == [listener, listener] + assert service.listener_removals == 1 + assert service.wait_listener_streams_finished(1) + + spark.streams.removeListener(listener) + assert listener._rust_listener_ids == [second] + assert spark.streams._sqlb._listener_bus == [listener] + assert service.listener_removals == 1 + spark.streams.removeListener(listener) + assert not hasattr(listener, "_rust_listener_ids") + assert spark.streams._sqlb._listener_bus == [] + assert service.listener_removals == 2 + assert service.wait_listener_streams_finished(2) + finally: + other.stop() + spark.streams.close() + + def test_streaming_listener_registration_failure_can_be_retried(self, mock_connect): + spark, service = mock_connect + service.listener_bus_enabled = True + service.listener_register_failure = True + listener = IdleListener() + with pytest.warns(UserWarning, match="listener registration failed") as caught: + assert spark.streams.addListener(listener) is None + assert "The listener is not added, please add it again." in str(caught[0].message) + assert not hasattr(listener, "_rust_listener_ids") + service.listener_register_failure = False + try: + spark.streams.addListener(listener) + emit_idle(service, spark, "retried") + listener.wait_for("retried") + assert service.listener_subscriptions == 2 + finally: + spark.streams.close() + + @pytest.mark.parametrize("method", ["removeListener", "close"]) + def test_streaming_listener_removal_failure_can_be_retried(self, mock_connect, method): + spark, service = mock_connect + service.listener_bus_enabled = True + listener = IdleListener() + spark.streams.addListener(listener) + listener_ids = list(listener._rust_listener_ids) + + def remove(): + if method == "removeListener": + return spark.streams.removeListener(listener) + return spark.streams.close() + + service.listener_remove_failure = True + try: + with pytest.warns(UserWarning, match="listener removal failed") as caught: + assert remove() is None + assert ( + "The listener is not removed, please remove it again." + in str(caught[0].message) + ) + assert listener._rust_listener_ids == listener_ids + assert service.listener_streams_finished == 0 + emit_idle(service, spark, "still-registered") + listener.wait_for("still-registered") + + service.listener_remove_failure = False + assert remove() is None + assert service.listener_removals == 2 + assert service.wait_listener_streams_finished(1) + # A completed remove consumes its id; manager close leaves stale private ids, + # which must not prevent registering and removing the same object again. + remaining_ids = listener_ids if method == "close" else [] + assert getattr(listener, "_rust_listener_ids", []) == remaining_ids + # A fresh registration after removal opens a new server subscription. + spark.streams.addListener(listener) + assert listener._rust_listener_ids[:-1] == remaining_ids + assert listener._rust_listener_ids[-1] not in listener_ids + assert service.listener_subscriptions == 2 + emit_idle(service, spark, "registered-again") + listener.wait_for("registered-again") + spark.streams.removeListener(listener) + assert getattr(listener, "_rust_listener_ids", []) == remaining_ids + assert service.listener_removals == 3 + assert service.wait_listener_streams_finished(2) + finally: + service.listener_remove_failure = False + spark.streams.close() + + @pytest.mark.parametrize("release_failure", [False, True]) + def test_streaming_listener_stop_releases_without_a_listener_removal_rpc( + self, mock_connect, caplog, release_failure, + ): + spark, service = mock_connect + service.listener_bus_enabled = True + spark.streams.addListener(IdleListener()) + # An ordinary close would block/fail here; terminal stop must use only ReleaseSession. + service.listener_remove_blocked = True + service.listener_remove_failure = True + service.release_session_failure = release_failure + try: + with caplog.at_level("WARNING", logger="pyspark"): + assert spark.stop() is None + assert spark.is_stopped + assert service.listener_removals == 0 + assert service.release_session_calls == 1 + assert spark.streams._sqlb._listener_bus == [] + assert service.wait_listener_streams_finished(1) + messages = [record.message for record in caplog.records if record.name == "pyspark"] + if release_failure: + assert len(messages) == 1 + assert messages[0].startswith( + "session.stop(): Session could not be released. Error: " + ) + assert messages[0].endswith("session release failed") + else: + assert messages == [] + spark.stop() + assert service.listener_removals == 0 + assert service.release_session_calls == 1 + finally: + service.listener_remove_failure = False + service.listener_remove_blocked = False + service.release_session_failure = False + spark.streams.close() + + def test_streaming_listener_client_close_cancels_without_releasing(self, mock_connect): + spark, service = mock_connect + service.listener_bus_enabled = True + spark.streams.addListener(IdleListener()) + spark.client.close() + assert spark.is_stopped + assert spark.streams._sqlb._listener_bus == [] + assert service.wait_listener_streams_finished(1) + assert service.listener_removals == 0 + assert service.release_session_calls == 0 + spark.stop() + assert service.release_session_calls == 0 + + +def concurrent_close_case(remove_failure, callback_action): + from _mock_connect_server import mock_connect_session + + with mock_connect_session() as (spark, service): + service.listener_bus_enabled = True + service.listener_remove_blocked = True + service.listener_remove_failure = remove_failure + entered, release, removed = (threading.Event() for _ in range(3)) + callback_threads, errors = [], [] + registrations = 2 if callback_action == "remove_duplicate" else 1 + caller_thread = threading.get_ident() + attempted = IdleListener() + + class RemovingListener(IdleListener): + def onQueryIdle(self, event): + if event.timestamp != "closing": + return super().onQueryIdle(event) + callback_threads.append(threading.get_ident()) + entered.set() + assert release.wait(5) + if callback_action == "stop": + spark.stop() + elif callback_action == "add": + spark.streams.addListener(attempted) + else: + spark.streams.removeListener(self) + if len(callback_threads) == registrations: + removed.set() + + def close(): + try: + assert spark.streams.close() is None + except Exception as error: + errors.append(error) + + listener = RemovingListener() + for _ in range(registrations): + spark.streams.addListener(listener) + listener_ids = list(listener._rust_listener_ids) + closing = threading.Thread(target=close, daemon=True) + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + emit_idle(service, spark, "closing") + assert entered.wait(5) + closing.start() + try: + assert service.wait_listener_removal() + release.set() + # The RPC stays gated until the callback's remove/stop has returned. + assert removed.wait(5) + assert len(callback_threads) == registrations + assert all(thread not in (caller_thread, closing.ident) + for thread in callback_threads) + assert closing.is_alive() + assert listener._rust_listener_ids == listener_ids + finally: + release.set() + service.release_listener_removal() + closing.join(5) + assert not closing.is_alive() + assert errors == [] + relevant = listener_warnings(caught) + add_warnings = [message for message in relevant + if message.startswith("Failed to add the listener")] + close_warnings = [message for message in relevant + if message.startswith("Failed to remove the listener")] + if callback_action == "add": + assert len(add_warnings) == 1 + assert "cannot change the listener bus from its own callback" in add_warnings[0] + assert "The listener is not added, please add it again." in add_warnings[0] + assert not hasattr(attempted, "_rust_listener_ids") + assert attempted.timestamps == [] + else: + assert add_warnings == [] + if remove_failure: + assert len(close_warnings) == 1 + assert "listener removal failed" in close_warnings[0] + else: + assert close_warnings == [] + if callback_action == "stop": + assert spark.is_stopped + assert spark.streams._sqlb._listener_bus == [] + assert service.release_session_calls == 1 + assert service.listener_removals == 1 + assert service.wait_listener_streams_finished(1) + return + if remove_failure: + assert spark.streams._sqlb._listener_bus == [listener] * registrations + assert service.listener_streams_finished == 0 + emit_idle(service, spark, "still-registered") + listener.wait_for("still-registered") + service.listener_remove_failure = False + service.listener_remove_blocked = False + for remaining in range(registrations - 1, -1, -1): + spark.streams.removeListener(listener) + assert spark.streams._sqlb._listener_bus == [listener] * remaining + if remaining: + assert listener._rust_listener_ids == listener_ids[-remaining:] + else: + assert not hasattr(listener, "_rust_listener_ids") + assert service.listener_removals == 2 + else: + assert service.listener_removals == 1 + assert service.wait_listener_streams_finished(1) + # Successful external close leaves stale ids. A fresh add appends a new id + # and opens exactly one subscription; failed-close retries consumed their ids. + remaining_ids = list(getattr(listener, "_rust_listener_ids", [])) + spark.streams.addListener(listener) + assert listener._rust_listener_ids[:-1] == remaining_ids + assert listener._rust_listener_ids[-1] not in listener_ids + assert service.listener_subscriptions == 2 + emit_idle(service, spark, "registered-again") + listener.wait_for("registered-again") + finally: + release.set() + service.release_listener_removal() + if closing.ident is not None: + closing.join(5) + service.listener_remove_failure = False + service.listener_remove_blocked = False + spark.streams.close() + + +def callback_rpc_case(): + from _mock_connect_server import mock_connect_session + + with mock_connect_session() as (spark, service): + service.listener_bus_enabled = True + finished, returned = threading.Event(), threading.Event() + callback_errors, stopping_errors = [], [] + + class RpcListener(IdleListener): + def onQueryIdle(self, event): + try: + spark.range(1).collect() + except Exception as error: + callback_errors.append(error) + finally: + finished.set() + + def stop(): + try: + assert spark.stop() is None + returned.set() + except Exception as error: + stopping_errors.append(error) + + spark.streams.addListener(RpcListener()) + service.block_execute = True + service.release_session_unblocks_execute = True + stopping = threading.Thread(target=stop, daemon=True) + try: + emit_idle(service, spark, "callback-rpc") + assert service.wait_execute_started(1) + assert not finished.is_set() + stopping.start() + # The mock's five-second fallback must not rescue a join-before-release deadlock. + assert returned.wait(2), "stop waited for the callback RPC before releasing its session" + # No explicit response release has run: only ReleaseSession can open the gate. + assert finished.wait(2), "ReleaseSession did not unblock the callback RPC" + assert service.release_session_calls == 1 + assert service.execute_gate_timeouts == 0 + assert spark.is_stopped + assert spark.streams._sqlb._listener_bus == [] + assert service.listener_removals == 0 + finally: + service.release_response() + if stopping.ident is not None: + stopping.join(5) + assert finished.wait(5), "callback RPC did not finish after its response was released" + assert not stopping.is_alive() + assert stopping_errors == [] + # A callback already in flight can finish normally or observe terminal session closure. + assert all("NO_ACTIVE_SESSION" in str(error) for error in callback_errors) + assert service.execute_gate_timeouts == 0 + assert service.wait_listener_streams_finished(1) + + +def blocked_callback_case(method, blocker): + from _mock_connect_server import mock_connect_session + + with mock_connect_session() as (spark, service): + service.listener_bus_enabled = True + caller_ready, requested, returned = (threading.Event() for _ in range(3)) + entered, finished = threading.Event(), threading.Event() + release_caller, release_callback = threading.Event(), threading.Event() + callback_lock, errors = threading.Lock(), [] + + class BlockingListener(IdleListener): + def onQueryIdle(self, event): + try: + if event.timestamp == "blocked": + entered.set() + if blocker == "event": + assert release_callback.wait(5), "callback Event was not released" + else: + assert callback_lock.acquire(timeout=5), "callback lock timed out" + callback_lock.release() + super().onQueryIdle(event) + except Exception as error: + errors.append(error) + finally: + if event.timestamp == "blocked": + finished.set() + + def teardown(): + # The same caller holds the lock and invokes terminal teardown. Joining its + # callback would wait for a lock that this caller cannot release until return. + with callback_lock: + caller_ready.set() + try: + assert requested.wait(5), "terminal teardown was not requested" + if method == "stop": + assert spark.stop() is None + else: + assert spark.client.close() is None + returned.set() + assert release_caller.wait(5), "terminal caller was not released" + except Exception as error: + errors.append(error) + + listener = BlockingListener() + spark.streams.addListener(listener) + stopping = threading.Thread(target=teardown, daemon=True) + try: + stopping.start() + assert caller_ready.wait(5) + emit_idle(service, spark, "blocked") + assert entered.wait(5) + # Queue another event while the callback still owns the worker, before + # cancellation can retire the server queue. It must never start a dispatch. + emit_idle(service, spark, "queued-before-terminal") + requested.set() + assert returned.wait(2), "terminal teardown waited for an in-progress callback" + assert not finished.is_set() + assert listener.timestamps == [] + assert spark.is_stopped + assert spark.streams._sqlb._listener_bus == [] + assert service.listener_removals == 0 + assert service.release_session_calls == (1 if method == "stop" else 0) + # The current event's listener snapshot may finish after terminal return. + # A fresh event must not start another dispatch from the canceled stream. + emit_idle(service, spark, "after-terminal") + finally: + requested.set() + release_callback.set() + release_caller.set() + stopping.join(5) + if entered.is_set(): + assert finished.wait(5), "in-progress callback did not finish during cleanup" + assert not stopping.is_alive() + assert errors == [] + assert service.wait_listener_streams_finished(1) + assert listener.timestamps == ["blocked"] + + +def pending_registration_case(method): + from _mock_connect_server import mock_connect_session + + with mock_connect_session() as (spark, service): + service.listener_bus_enabled = True + # Early mock headers establish a stream before its acknowledgement. Establishing + # the RPC itself remains subject to the transport's deadline/retry policy. + service.listener_register_blocked = True + listener, errors = IdleListener(), [] + + def add(): + try: + assert spark.streams.addListener(listener) is None + except Exception as error: + errors.append(error) + + def teardown(): + try: + if method == "stop": + assert spark.stop() is None + else: + assert spark.client.close() is None + except Exception as error: + errors.append(error) + + starting = threading.Thread(target=add, daemon=True) + stopping = threading.Thread(target=teardown, daemon=True) + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + # Unrelated process-wide warnings must not change lifecycle assertions. + warnings.warn("unrelated warning used to check filtering", DeprecationWarning) + starting.start() + assert service.wait_listener_registration() + stopping.start() + stopping.join(5) + assert not stopping.is_alive(), ( + "terminal teardown waited for stream acknowledgement" + ) + starting.join(5) + assert not starting.is_alive(), "registration was not canceled" + assert errors == [] + relevant = listener_warnings(caught) + assert len(relevant) == 1 + assert "NO_ACTIVE_SESSION" in relevant[0] + assert "The listener is not added, please add it again." in relevant[0] + assert not hasattr(listener, "_rust_listener_ids") + assert spark.is_stopped + assert spark.streams._sqlb._listener_bus == [] + assert service.listener_removals == 0 + assert service.release_session_calls == (1 if method == "stop" else 0) + assert service.wait_listener_streams_finished(1) + finally: + service.release_listener_registration() + starting.join(5) + if stopping.ident is not None: + stopping.join(5) + + +def ended_stream_case(termination): + from _mock_connect_server import mock_connect_session + + with mock_connect_session() as (spark, service): + service.listener_bus_enabled = True + manager = spark.streams + listener = IdleListener() + try: + manager.addListener(listener) + stale_ids = list(listener._rust_listener_ids) + emit_idle(service, spark, "before-end") + listener.wait_for("before-end") + # End the subscription without a client removal RPC, either at ResultComplete + # or with a nonretryable server error. Server completion precedes native cleanup. + assert service.end_listener_stream(spark.session_id, failure=termination == "error") + assert service.wait_listener_streams_finished(1) + deadline = time.monotonic() + 5 + while manager._sqlb._listener_bus: + assert time.monotonic() < deadline, "ended stream kept its registrations" + time.sleep(0.01) + assert listener._rust_listener_ids == stale_ids + manager.removeListener(listener) + assert listener._rust_listener_ids == stale_ids + assert service.listener_removals == 0 + + # Reuse the same manager and object: stale tokens must not match the new + # generation or prevent removal of its one active registration. + manager.addListener(listener) + assert listener._rust_listener_ids[:-1] == stale_ids + assert listener._rust_listener_ids[-1] not in stale_ids + assert service.listener_subscriptions == 2 + assert manager._sqlb._listener_bus == [listener] + emit_idle(service, spark, "after-end") + listener.wait_for("after-end") + assert listener.timestamps == ["before-end", "after-end"] + manager.removeListener(listener) + assert manager._sqlb._listener_bus == [] + assert listener._rust_listener_ids == stale_ids + assert service.listener_removals == 1 + assert service.wait_listener_streams_finished(2) + manager.removeListener(listener) + assert listener._rust_listener_ids == stale_ids + assert service.listener_removals == 1 + finally: + manager.close() + + +def second_rpc_case(): + from _mock_connect_server import mock_connect_session + + with mock_connect_session() as (spark, service): + service.listener_bus_enabled = True + spark.streams.addListener(IdleListener()) + assert service.wait_listener_registration(1) + spark.streams.close() + assert service.wait_listener_removal(1) + assert service.wait_listener_streams_finished(1) + # Earlier RPCs must not satisfy a wait for the next registration/removal. + assert not service.wait_listener_registration(2, timeout=0.05) + assert not service.wait_listener_removal(2, timeout=0.05) + service.listener_register_blocked = True + service.listener_remove_blocked = True + listener, errors = IdleListener(), [] + + def add(): + try: + spark.streams.addListener(listener) + except Exception as error: + errors.append(error) + + def close(): + try: + spark.streams.close() + except Exception as error: + errors.append(error) + + starting = threading.Thread(target=add, daemon=True) + closing = threading.Thread(target=close, daemon=True) + try: + starting.start() + assert service.wait_listener_registration(2) + assert starting.is_alive() + service.release_listener_registration() + starting.join(5) + assert not starting.is_alive() + assert hasattr(listener, "_rust_listener_ids") + closing.start() + assert service.wait_listener_removal(2) + assert closing.is_alive() + service.release_listener_removal() + closing.join(5) + assert not closing.is_alive() + assert errors == [] + assert service.wait_listener_streams_finished(2) + finally: + service.release_listener_registration() + service.release_listener_removal() + starting.join(5) + if closing.ident is not None: + closing.join(5) + service.listener_register_blocked = False + service.listener_remove_blocked = False + spark.streams.close() + + +if __name__ == "__main__": + if sys.argv[1] == "--concurrent-close-case": + concurrent_close_case(bool(int(sys.argv[2])), sys.argv[3]) + elif sys.argv[1] == "--pending-registration-case": + pending_registration_case(sys.argv[2]) + elif sys.argv[1] == "--callback-rpc-case": + callback_rpc_case() + elif sys.argv[1] == "--blocked-callback-case": + blocked_callback_case(sys.argv[2], sys.argv[3]) + elif sys.argv[1] == "--ended-stream-case": + ended_stream_case(sys.argv[2]) + else: + assert sys.argv[1] == "--second-rpc-case" + second_rpc_case() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 3c7c9695..0ed03854 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -37,6 +37,7 @@ # Collected here so the existing offline CI suite exercises real Arrow results. from day_time_interval_cases import DayTimeIntervalTests # noqa: F401 from interval_input_cases import IntervalInputConversionTests # noqa: F401 +from streaming_listener_cases import TestStreamingListenerLifecycle # noqa: F401 def test_create_dataframe_input_names_offline(): @@ -9786,7 +9787,6 @@ def test_numpy_scalars_of_other_types_are_unsupported_as_pyspark(mock_connect): (np.uint8(5), "uint8"), (np.uint64(5), "uint64"), (np.float16(1.5), "float16"), - (np.longdouble(1.5), "longdouble"), (np.complex128(1 + 2j), "complex128"), (np.datetime64("2020-01-01"), "datetime64"), (np.timedelta64(1, "D"), "timedelta64"), @@ -9801,6 +9801,90 @@ def test_numpy_scalars_of_other_types_are_unsupported_as_pyspark(mock_connect): assert error.value.getMessageParameters() == {"dtype": "uint8"} +def test_numpy_longdouble_scalar_matches_dtype_and_class_checks_as_pyspark(): + np = pytest.importorskip("numpy") + from pyspark.errors import PySparkTypeError + from pyspark.sql import functions as F + + value = np.longdouble(1.5) + if isinstance(value, (float, np.float32, np.float64)): + pytest.skip("longdouble aliases a supported float scalar class here") + if value.dtype == np.dtype("float64"): + # A dtype mapping to DoubleType does not suffice: PySpark also asserts that its + # scalar is a float, np.float32 or np.float64. Distinct eight-byte longdouble + # classes (such as on macOS arm64 and Windows) fail without a message. + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () + else: + # Linux CI normally has a wider longdouble, which exercises the unsupported + # dtype path; the distinct integer aliases below exercise the class check there. + with pytest.raises(PySparkTypeError) as error: + F.lit(value) + assert error.value.getCondition() == "UNSUPPORTED_DATA_TYPE" + assert error.value.getMessageParameters() == {"data_type": type(value).__name__} + + +def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): + np = pytest.importorskip("numpy") + from pyspark.errors import PySparkTypeError + from pyspark.sql import functions as F + + value = np.array([1.5], dtype=np.longdouble) + if isinstance(value[0], (float, np.float32, np.float64)): + pytest.skip("longdouble aliases a supported float scalar class here") + if value.dtype == np.dtype("float64"): + # A mapped array dtype still builds its literals through each scalar's checks. + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () + else: + with pytest.raises(PySparkTypeError) as error: + F.lit(value) + assert error.value.getCondition() == "UNSUPPORTED_NUMPY_ARRAY_SCALAR" + assert error.value.getMessageParameters() == {"dtype": value.dtype.name} + + +@pytest.mark.parametrize("name", [ + "byte", "short", "intc", "int_", "intp", "long", "longlong", "single", "double", +]) +@pytest.mark.parametrize("as_array", [False, True], ids=["scalar", "array"]) +def test_numpy_numeric_aliases_match_reference_scalar_classes(name, as_array): + np = pytest.importorskip("numpy") + from pyspark.sql import functions as F + + # Read real scalar exports without NumPy's deprecated builtin-alias lookup warnings. + scalar_type = vars(np).get(name) + if scalar_type is None or not issubclass(scalar_type, np.generic): + pytest.skip(f"numpy.{name} is not a NumPy scalar class in this NumPy") + value = scalar_type(1.5 if np.issubdtype(scalar_type, np.floating) else 5) + kind, size = value.dtype.kind, value.dtype.itemsize + if kind == "i": + # These are the classes asserted by LiteralExpression for the inferred width. + accepted_classes = (int,) + tuple( + getattr(np, f"int{width}") for width in (8, 16, 32, 64) if width <= size * 8 + ) + spark_type = {1: "tinyint", 2: "smallint", 4: "int", 8: "bigint"}[size] + rendered = "5" + else: + accepted_classes = (float, np.float32) + ((np.float64,) if size == 8 else ()) + spark_type = {4: "float", 8: "double"}[size] + rendered = "1.5" + literal = np.array([value], dtype=scalar_type) if as_array else value + if not isinstance(value, accepted_classes): + # Distinct longlong scalars have int64 dtype but fail the class assertion on + # Linux/macOS; platforms where it aliases int64 exercise the success case. + with pytest.raises(AssertionError) as error: + F.lit(literal) + assert error.value.args == () + elif as_array: + assert repr(F.lit(literal)).lower() == ( + f"column<'cast(array({rendered}) as array<{spark_type}>)'>" + ) + else: + assert repr(F.lit(literal)) == f"Column<'{rendered}'>" + + # --------------------------------------------------------------------------- 4.3.0 gaps @@ -10178,6 +10262,7 @@ def test_lit_of_a_numpy_array_is_an_array_of_its_type_as_pyspark(): for value, expected in ( (np.array([1, 2]), "cast(array(1, 2) as array)"), (np.array([1.5], dtype=np.float32), "cast(array(1.5) as array)"), + (np.array([1.5], dtype=np.float64), "cast(array(1.5) as array)"), (np.array(["x"]), "cast(array(x) as array)"), ): assert repr(F.lit(value)).lower() == f"column<'{expected}'>"