This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23967-28c8afa751e6aeb990699a2087f13840e9ab3351 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit f9dde71ec315833ea0bc39bd63104e613109e915 Author: Raz Luvaton <[email protected]> AuthorDate: Mon Aug 3 21:51:52 2026 +0300 feat: drop generator on error to free memory faster (#23967) ## Which issue does this PR close? N/A ## Rationale for this change Drop generator on error so the memory is freed faster ## What changes are included in this PR? Replaced done with Option and set ## Are these changes tested? Yes ## Are there any user-facing changes? Clear faster --- datafusion/execution/src/async_stream.rs | 69 ++++++++++++++++++++++++++++++-- 1 file changed, 66 insertions(+), 3 deletions(-) diff --git a/datafusion/execution/src/async_stream.rs b/datafusion/execution/src/async_stream.rs index e271145e03..7ca6ba4850 100644 --- a/datafusion/execution/src/async_stream.rs +++ b/datafusion/execution/src/async_stream.rs @@ -92,7 +92,11 @@ pub fn async_try_stream<T, E, F: Future<Output = Result<(), E>>>( let (try_emitter, mut emitter, receiver) = try_tx_rx::<T, E>(); AsyncStream::new(receiver, async move { if let Err(e) = generator(try_emitter).await { - emitter.emit(Err(e)).await + // Fill the slot without suspending so this future completes in the same + // poll that yields `Err(e)`: the stream terminates immediately and the + // emitter state is dropped (a consumer may never poll again after an + // error, which would otherwise keep this future suspended inside `emit`) + emitter.set(Err(e)); } }) } @@ -168,13 +172,19 @@ impl<T> Emitter<T> { /// been awaited, because doing so would silently overwrite the unconsumed /// value. pub fn emit(&mut self, value: T) -> impl FusedFuture<Output = ()> { + self.set(value); + Emit { done: false } + } + + /// Places `value` in the slot without suspending the generator. Only useful + /// as the very last action before the generator future completes, since + /// nothing yields control back to the consumer in between. + fn set(&mut self, value: T) { let mut guard = self.slot.lock(); match guard.deref_mut() { Some(_) => panic!("Misuse: await was not called after calling emit"), slot => *slot = Some(value), } - - Emit { done: false } } } @@ -655,6 +665,59 @@ mod test { ); } + struct DropGuard(Arc<AtomicUsize>); + + impl Drop for DropGuard { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + #[tokio::test] + async fn generator_freed_on_done() { + let drops = Arc::new(AtomicUsize::new(0)); + let guard = DropGuard(Arc::clone(&drops)); + + let s = async_stream(|mut emitter| async move { + let _guard = guard; + emitter.emit(1).await; + }); + pin_mut!(s); + + assert_eq!(s.next().await, Some(1)); + assert_eq!(s.next().await, None); + + // State captured by the generator is dropped as soon as it completes + // (async blocks drop their locals on return), even though the stream + // itself is still alive + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(s.next().await, None); + } + + #[tokio::test] + async fn generator_freed_on_emitted_error() { + let drops = Arc::new(AtomicUsize::new(0)); + let guard = DropGuard(Arc::clone(&drops)); + + let s = async_try_stream(|mut emitter| async move { + let _guard = guard; + emitter.emit(1).await; + Err("boom") + }); + pin_mut!(s); + + assert_eq!(s.next().await, Some(Ok(1))); + assert_eq!(s.next().await, Some(Err("boom"))); + + // The stream terminates in the same poll that yields the error, so the + // generator state is freed even if the consumer never polls again + assert!(s.is_terminated()); + assert_eq!(drops.load(Ordering::SeqCst), 1); + + // Polling again after the error just returns None + assert_eq!(s.next().await, None); + } + use pin_project_lite::pin_project; pin_project! { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
