From 9d9d13060c0e1df65c88fbddff19a6bce655b8a7 Mon Sep 17 00:00:00 2001
From: rob bygrave <robin.bygrave@gmail.com>
Date: Wed, 17 Jun 2020 17:10:16 +1200
Subject: [PATCH] #2020 - Fix potential connection leak in query-iterator WITH
 lazyLoadBufferSize (Refer #1164)

---
 .../query/CQueryIteratorWithBuffer.java       | 23 +++++++++++----
 .../org/tests/query/TestQueryFindIterate.java | 28 ++++++++++++++-----
 2 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/src/main/java/io/ebeaninternal/server/query/CQueryIteratorWithBuffer.java b/src/main/java/io/ebeaninternal/server/query/CQueryIteratorWithBuffer.java
index da07687ac..2ba431da5 100644
--- a/src/main/java/io/ebeaninternal/server/query/CQueryIteratorWithBuffer.java
+++ b/src/main/java/io/ebeaninternal/server/query/CQueryIteratorWithBuffer.java
@@ -6,6 +6,7 @@ import io.ebeaninternal.server.core.OrmQueryRequest;
 import javax.persistence.PersistenceException;
 import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.NoSuchElementException;
 
 /**
  * A QueryIterator that uses a buffer to execute secondary queries periodically.
@@ -17,6 +18,7 @@ class CQueryIteratorWithBuffer<T> implements QueryIterator<T> {
   private final OrmQueryRequest<T> request;
   private final ArrayList<T> buffer;
 
+  private boolean closed;
   private boolean moreToLoad = true;
 
   CQueryIteratorWithBuffer(CQuery<T> cquery, OrmQueryRequest<T> request, int bufferSize) {
@@ -29,6 +31,7 @@ class CQueryIteratorWithBuffer<T> implements QueryIterator<T> {
   @Override
   @SuppressWarnings("unchecked")
   public boolean hasNext() {
+    boolean ret = false;
     try {
       if (buffer.isEmpty() && moreToLoad) {
         // load buffer
@@ -44,23 +47,33 @@ class CQueryIteratorWithBuffer<T> implements QueryIterator<T> {
         }
         request.executeSecondaryQueries(true);
       }
-      return !buffer.isEmpty();
-
+      ret = !buffer.isEmpty();
+      return ret;
     } catch (SQLException e) {
       throw cquery.createPersistenceException(e);
+    } finally {
+      if (!ret) {
+        close();
+      }
     }
   }
 
   @Override
   public T next() {
+    if (buffer.isEmpty()) {
+      throw new NoSuchElementException();
+    }
     return buffer.remove(0);
   }
 
   @Override
   public void close() {
-    cquery.updateExecutionStatisticsIterator();
-    cquery.close();
-    request.endTransIfRequired();
+    if (!closed) {
+      closed = true;
+      cquery.updateExecutionStatisticsIterator();
+      cquery.close();
+      request.endTransIfRequired();
+    }
   }
 
   @Override
diff --git a/src/test/java/org/tests/query/TestQueryFindIterate.java b/src/test/java/org/tests/query/TestQueryFindIterate.java
index 6db326313..760e17969 100644
--- a/src/test/java/org/tests/query/TestQueryFindIterate.java
+++ b/src/test/java/org/tests/query/TestQueryFindIterate.java
@@ -220,7 +220,16 @@ public class TestQueryFindIterate extends BaseTestCase {
   }
 
   @Test
-  public void testCloseConnection() {
+  public void testCloseConnection_findIterate() {
+    findIterateCloseConnection(false);
+  }
+
+  @Test
+  public void testCloseConnection_findIterate_withBatchLoad() {
+    findIterateCloseConnection(true);
+  }
+
+  public void findIterateCloseConnection(boolean withLoadBatch) {
     ResetBasicData.reset();
 
     SpiServer pluginApi = server().getPluginApi();
@@ -230,12 +239,17 @@ public class TestQueryFindIterate extends BaseTestCase {
     }
 
     int startConns = dsPool.getStatus(false).getBusy();
-    QueryIterator<Customer> queryIterator = server().find(Customer.class)
-        .where()
-        .isNotNull("name")
-        .setMaxRows(3)
-        .order().asc("id")
-        .findIterate();
+    final Query<Customer> query = server().find(Customer.class)
+      .where()
+      .isNotNull("name")
+      .setMaxRows(3)
+      .order().asc("id");
+
+    if (withLoadBatch) {
+      query.setLazyLoadBatchSize(100);
+    }
+
+    QueryIterator<Customer> queryIterator = query.findIterate();
 
     assertThat(dsPool.getStatus(false).getBusy()).isEqualTo(startConns + 1);