HHH-14449 : ResultStream closing is not properly handled

Author: juja0

hibernate-core/src/main/java/org/hibernate/query/internal/AbstractProducedQuery.java

@@ -1584,12 +1584,10 @@ public Stream<R> stream() {
 		final ScrollableResultsIterator<R> iterator = new ScrollableResultsIterator<>( scrollableResults );
 		final Spliterator<R> spliterator = Spliterators.spliteratorUnknownSize( iterator, Spliterator.NONNULL );
 
-		final Stream<R> stream = new StreamDecorator(
+		return new StreamDecorator<>(
 				StreamSupport.stream( spliterator, false ),
-				scrollableResults::close
+				iterator::close
 		);
-
-		return stream;
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/query/spi/DoubleStreamDecorator.java

@@ -41,15 +41,13 @@
 public class DoubleStreamDecorator implements DoubleStream {
 
 	private final DoubleStream delegate;
-
-	private Runnable closeHandler;
+	private final Runnable closeHandler;
 
 	public DoubleStreamDecorator(
 			DoubleStream delegate,
 			Runnable closeHandler) {
-		this.delegate = delegate;
 		this.closeHandler = closeHandler;
-		this.delegate.onClose( closeHandler );
+		this.delegate = delegate.onClose( closeHandler );
 	}
 
 	@Override
@@ -292,7 +290,7 @@ public DoubleStream unordered() {
 
 	@Override
 	public DoubleStream onClose(Runnable closeHandler) {
-		this.closeHandler = closeHandler;
+		this.delegate.onClose( closeHandler );
 		return this;
 	}
 

hibernate-core/src/main/java/org/hibernate/query/spi/IntStreamDecorator.java

@@ -42,15 +42,13 @@
 public class IntStreamDecorator implements IntStream {
 
 	private final IntStream delegate;
-
-	private Runnable closeHandler;
+	private final Runnable closeHandler;
 
 	public IntStreamDecorator(
 			IntStream delegate,
 			Runnable closeHandler) {
-		this.delegate = delegate;
 		this.closeHandler = closeHandler;
-		this.delegate.onClose( closeHandler );
+		this.delegate = delegate.onClose( closeHandler );
 	}
 
 	@Override
@@ -307,7 +305,7 @@ public IntStream unordered() {
 
 	@Override
 	public IntStream onClose(Runnable closeHandler) {
-		this.closeHandler = closeHandler;
+		this.delegate.onClose( closeHandler );
 		return this;
 	}
 

hibernate-core/src/main/java/org/hibernate/query/spi/LongStreamDecorator.java

@@ -42,15 +42,13 @@
 public class LongStreamDecorator implements LongStream {
 
 	private final LongStream delegate;
-
-	private Runnable closeHandler;
+	private final Runnable closeHandler;
 
 	public LongStreamDecorator(
 			LongStream delegate,
 			Runnable closeHandler) {
-		this.delegate = delegate;
 		this.closeHandler = closeHandler;
-		this.delegate.onClose( closeHandler );
+		this.delegate = delegate.onClose( closeHandler );
 	}
 
 	@Override
@@ -226,21 +224,21 @@ public LongSummaryStatistics summaryStatistics() {
 
 	@Override
 	public boolean anyMatch(LongPredicate predicate) {
-		boolean result = delegate.anyMatch(predicate);
+		boolean result = delegate.anyMatch( predicate );
 		close();
 		return result;
 	}
 
 	@Override
 	public boolean allMatch(LongPredicate predicate) {
-		boolean result = delegate.allMatch(predicate);
+		boolean result = delegate.allMatch( predicate );
 		close();
 		return result;
 	}
 
 	@Override
 	public boolean noneMatch(LongPredicate predicate) {
-		boolean result = delegate.noneMatch(predicate);
+		boolean result = delegate.noneMatch( predicate );
 		close();
 		return result;
 	}
@@ -300,7 +298,7 @@ public LongStream unordered() {
 
 	@Override
 	public LongStream onClose(Runnable closeHandler) {
-		this.closeHandler = closeHandler;
+		this.delegate.onClose( closeHandler );
 		return this;
 	}
 

hibernate-core/src/main/java/org/hibernate/query/spi/StreamDecorator.java

@@ -46,15 +46,13 @@
 public class StreamDecorator<R> implements Stream<R> {
 
 	private final Stream<R> delegate;
-
-	private Runnable closeHandler;
+	private final Runnable closeHandler;
 
 	public StreamDecorator(
 			Stream<R> delegate,
 			Runnable closeHandler) {
-		this.delegate = delegate;
 		this.closeHandler = closeHandler;
-		this.delegate.onClose( closeHandler );
+		this.delegate = delegate.onClose( closeHandler );
 	}
 
 	@Override
@@ -301,7 +299,7 @@ public Stream<R> unordered() {
 
 	@Override
 	public Stream<R> onClose(Runnable closeHandler) {
-		this.closeHandler = closeHandler;
+		this.delegate.onClose( closeHandler );
 		return this;
 	}
 

hibernate-core/src/test/java/org/hibernate/test/stream/basic/BasicStreamTest.java

@@ -6,6 +6,7 @@
  */
 package org.hibernate.test.stream.basic;
 
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 import javax.persistence.Entity;
 import javax.persistence.Id;
@@ -25,6 +26,7 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.core.Is.is;
+import static org.hamcrest.core.IsEqual.equalTo;
 import static org.hibernate.testing.junit4.ExtraAssertions.assertTyping;
 import static org.hibernate.testing.transaction.TransactionUtil.doInHibernate;
 
@@ -123,6 +125,53 @@ public void testTupleStream() {
 		} );
 	}
 
+	@Test
+	public void basicStreamTestWithExplicitOnClose() {
+		Session session = openSession();
+		session.getTransaction().begin();
+
+		AtomicInteger onCloseCount = new AtomicInteger();
+
+		// mainly we want to make sure that closing the Stream releases the ScrollableResults too
+		assertThat( ( (SessionImplementor) session ).getJdbcCoordinator()
+							.getLogicalConnection()
+							.getResourceRegistry()
+							.hasRegisteredResources(), is( false ) );
+
+		assertThat( onCloseCount.get(), equalTo( 0 ) );
+
+		final Stream<MyEntity> stream = session.createQuery( "from MyEntity", MyEntity.class ).stream().onClose(
+				onCloseCount::incrementAndGet );
+
+		assertThat( ( (SessionImplementor) session ).getJdbcCoordinator()
+							.getLogicalConnection()
+							.getResourceRegistry()
+							.hasRegisteredResources(), is( true ) );
+
+		assertThat( onCloseCount.get(), equalTo( 0 ) );
+
+		stream.forEach( System.out::println );
+
+		assertThat( ( (SessionImplementor) session ).getJdbcCoordinator()
+							.getLogicalConnection()
+							.getResourceRegistry()
+							.hasRegisteredResources(), is( false ) );
+
+		assertThat( onCloseCount.get(), equalTo( 1 ) );
+
+		stream.close();
+
+		assertThat( ( (SessionImplementor) session ).getJdbcCoordinator()
+							.getLogicalConnection()
+							.getResourceRegistry()
+							.hasRegisteredResources(), is( false ) );
+
+		assertThat( onCloseCount.get(), equalTo( 1 ) );
+
+		session.getTransaction().commit();
+		session.close();
+	}
+
 	@Entity(name = "MyEntity")
 	@Table(name="MyEntity")
 	public static class MyEntity {