Weijun-H commented on code in PR #23967:
URL: https://github.com/apache/datafusion/pull/23967#discussion_r3672932208
##########
datafusion/execution/src/async_stream.rs:
##########
@@ -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);
Review Comment:
The inner generator drops its locals when it returns Err, before the outer
emission future suspends. So the drop count reaches 1 on the old path too; only
is_terminated() differs.
Suggest reframing the PR and test around the real change:
is_terminated()/size_hint flipping in the same poll as the final error.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
