#2021 - Use smarter persistence context for findEach/iterate/stream queries - improve performance of un-tuned findEach/iterate/stream queries that invoke lazy loading

src/main/java/io/ebean/bean/PersistenceContext.java

@@ -61,6 +61,22 @@ public interface PersistenceContext {
    */
   int size(Class<?> rootType);
 
+  /**
+   * Return a copy of the Persistence context to use for large query iteration.
+   */
+  PersistenceContext forIterate();
+
+  /**
+   * Return a new Persistence context during iteration of large query result.
+   */
+  PersistenceContext forIterateReset();
+
+  /**
+   * Return true if the persistence context has grown and hit the 'reset limit'
+   * during large query iteration.
+   */
+  boolean resetLimit();
+
   /**
    * Wrapper on a bean to also indicate if a bean has been deleted.
    * <p>

src/main/java/io/ebeaninternal/server/core/OrmQueryRequest.java

@@ -316,7 +316,6 @@ public final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
    */
   @Override
   public JsonReadOptions createJsonReadOptions() {
-
     persistenceContext = getPersistenceContext(query, transaction);
     if (query.getPersistenceContext() == null) {
       query.setPersistenceContext(persistenceContext);
@@ -327,7 +326,6 @@ public final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
       loadContext = new DLoadContext(this, secondaryQueries);
       jsonRead.setLoadContext(loadContext);
     }
-
     return jsonRead;
   }
 
@@ -335,8 +333,8 @@ public final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
    * For iterate queries reset the persistenceContext and loadContext.
    */
   public void flushPersistenceContextOnIterate() {
-    if (!iterateSingleContext) {
-      persistenceContext = new DefaultPersistenceContext();
+    if (!iterateSingleContext && persistenceContext.resetLimit()) {
+      persistenceContext = persistenceContext.forIterateReset();
       loadContext.resetPersistenceContext(persistenceContext);
       if (jsonRead != null) {
         jsonRead.setPersistenceContext(persistenceContext);
@@ -350,7 +348,6 @@ public final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
    * transaction scoped.
    */
   private PersistenceContext getPersistenceContext(SpiQuery<?> query, SpiTransaction t) {
-
     // check if there is already a persistence context set which is the case
     // when lazy loading or query joins are executed
     PersistenceContext ctx = query.getPersistenceContext();
@@ -358,7 +355,14 @@ public final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
 
     // determine the scope (from the query and then server)
     PersistenceContextScope scope = ebeanServer.getPersistenceContextScope(query);
-    return (scope == PersistenceContextScope.QUERY || t == null) ? new DefaultPersistenceContext() : t.getPersistenceContext();
+    if (scope == PersistenceContextScope.QUERY || t == null) {
+      return new DefaultPersistenceContext();
+    }
+    if (Type.ITERATE == query.getType()) {
+      return t.getPersistenceContext().forIterate();
+    } else {
+      return t.getPersistenceContext();
+    }
   }
 
   /**
@@ -590,7 +594,6 @@ public final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
    * Merge in prior L2 bean cache hits with the query result.
    */
   public void mergeCacheHits(BeanCollection<T> result) {
-
     if (cacheBeans != null && !cacheBeans.isEmpty()) {
       if (query.getType() == Type.MAP) {
         mergeCacheHitsToMap(result);

src/main/java/io/ebeaninternal/server/loadcontext/DLoadBeanContext.java

@@ -175,7 +175,6 @@ class DLoadBeanContext extends DLoadBaseContext implements LoadBeanContext {
     @Override
     public void loadBean(EntityBeanIntercept ebi) {
       // A synchronized (this) is effectively held by EntityBeanIntercept.loadBean()
-
       if (context.desc.lazyLoadMany(ebi)) {
         // lazy load property was a Many
         return;
@@ -183,7 +182,6 @@ class DLoadBeanContext extends DLoadBaseContext implements LoadBeanContext {
 
       if (context.hitCache) {
         Set<EntityBeanIntercept> hits = context.desc.cacheBeanLoadAll(list, persistenceContext, ebi.getLazyLoadPropertyIndex(), ebi.getLazyLoadProperty());
-
         list.removeAll(hits);
         if (list.isEmpty() || hits.contains(ebi)) {
           // successfully hit the L2 cache so don't invoke DB lazy loading
@@ -193,6 +191,7 @@ class DLoadBeanContext extends DLoadBaseContext implements LoadBeanContext {
 
       LoadBeanRequest req = new LoadBeanRequest(this, ebi, context.hitCache);
       context.desc.getEbeanServer().loadBean(req);
+      list.clear();
     }
   }
 

src/main/java/io/ebeaninternal/server/query/CQuery.java

@@ -608,12 +608,10 @@ public class CQuery<T> implements DbReadContext, CancelableQuery, SpiProfileTran
   }
 
   QueryIterator<T> readIterate(int bufferSize, OrmQueryRequest<T> request) {
-
-    if (bufferSize > 0) {
-      return new CQueryIteratorWithBuffer<>(this, request, bufferSize);
-
-    } else {
+    if (bufferSize < 2) {
       return new CQueryIteratorSimple<>(this, request);
+    } else {
+      return new CQueryIteratorWithBuffer<>(this, request, bufferSize);
     }
   }
 

src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java

@@ -34,18 +34,62 @@ public final class DefaultPersistenceContext implements PersistenceContext {
 
   private final Monitor monitor = new Monitor();
 
+  private int putCount;
+
   /**
    * Create a new PersistenceContext.
    */
   public DefaultPersistenceContext() {
   }
 
+  /**
+   * Create as a shallow copy with initial or types that have not been added to.
+   */
+  private DefaultPersistenceContext(DefaultPersistenceContext parent, boolean initial) {
+    for (Map.Entry<Class<?>, ClassContext> entry : parent.typeCache.entrySet()) {
+      typeCache.put(entry.getKey(), entry.getValue().copy(initial));
+    }
+  }
+
+  /**
+   * Return the initial shallow copy with each ClassContext noting it's initialSize (to detect additions).
+   */
+  @Override
+  public PersistenceContext forIterate() {
+    return new DefaultPersistenceContext(this, true);
+  }
+
+  /**
+   * Return a shallow copy including each ClassContext that has had no additions (still at initialSize).
+   */
+  @Override
+  public PersistenceContext forIterateReset() {
+    return new DefaultPersistenceContext(this, false);
+  }
+
+  public boolean resetLimit() {
+    synchronized (monitor) {
+      if (putCount < 100) {
+        return false;
+      }
+      putCount = 0;
+      for (ClassContext value : typeCache.values()) {
+        if (value.resetLimit()) {
+          return true;
+        }
+      }
+      // checking after another 100 puts
+      return false;
+    }
+  }
+
   /**
    * Set an object into the PersistenceContext.
    */
   @Override
   public void put(Class<?> rootType, Object id, Object bean) {
     synchronized (monitor) {
+      putCount++;
       getClassContext(rootType).put(id, bean);
     }
   }
@@ -53,6 +97,7 @@ public final class DefaultPersistenceContext implements PersistenceContext {
   @Override
   public Object putIfAbsent(Class<?> rootType, Object id, Object bean) {
     synchronized (monitor) {
+      putCount++;
       return getClassContext(rootType).putIfAbsent(id, bean);
     }
   }
@@ -133,7 +178,6 @@ public final class DefaultPersistenceContext implements PersistenceContext {
   }
 
   private ClassContext getClassContext(Class<?> rootType) {
-
     return typeCache.computeIfAbsent(rootType, k -> new ClassContext());
   }
 
@@ -143,9 +187,45 @@ public final class DefaultPersistenceContext implements PersistenceContext {
 
     private Set<Object> deleteSet;
 
+    private int initialSize;
+
     private ClassContext() {
     }
 
+    /**
+     * Create as a shallow copy.
+     */
+    private ClassContext(ClassContext parent, boolean initial) {
+      if (initial || parent.isInitialSize()) {
+        this.map.putAll(parent.map);
+        this.initialSize = map.size();
+        if (parent.deleteSet != null) {
+          this.deleteSet = new HashSet<>(parent.deleteSet);
+        }
+      }
+    }
+
+    /**
+     * True if the map has not changed from it's initial size (no additions).
+     */
+    private boolean isInitialSize() {
+      return map.size() == initialSize;
+    }
+
+    /**
+     * Return a shallow copy if initial copy or it has not grown (still at initialSize).
+     */
+    private ClassContext copy(boolean initial) {
+      return new ClassContext(this, initial);
+    }
+
+    /**
+     * Return true if grown above the reset limit size.
+     */
+    private boolean resetLimit() {
+      return map.size() > initialSize + 1000;
+    }
+
     @Override
     public String toString() {
       return "size:" + map.size();
@@ -164,7 +244,6 @@ public final class DefaultPersistenceContext implements PersistenceContext {
     }
 
     private Object putIfAbsent(Object id, Object bean) {
-
       Object existingValue = map.get(id);
       if (existingValue != null) {
         // it is not absent

src/test/java/io/ebeaninternal/server/transaction/DefaultPersistenceContextTest.java

@@ -2,12 +2,16 @@ package io.ebeaninternal.server.transaction;
 
 import io.ebean.bean.PersistenceContext;
 import io.ebeaninternal.server.deploy.PersistenceContextUtil;
-import org.tests.model.basic.Car;
-import org.tests.model.basic.Customer;
-import org.tests.model.basic.Vehicle;
 import org.junit.Test;
+import org.tests.model.basic.Car;
+import org.tests.model.basic.Contact;
+import org.tests.model.basic.Customer;
+import org.tests.model.basic.Product;
+import org.tests.model.basic.Vehicle;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 public class DefaultPersistenceContextTest {
 
@@ -23,12 +27,12 @@ public class DefaultPersistenceContextTest {
     car1.setId(1);
   }
 
-  PersistenceContext pc() {
+  private DefaultPersistenceContext pc() {
     return new DefaultPersistenceContext();
   }
 
-  PersistenceContext pcWith42() {
-    PersistenceContext pc = pc();
+  private DefaultPersistenceContext pcWith42() {
+    DefaultPersistenceContext pc = pc();
     pc.put(Customer.class, 42, customer42);
     return pc;
   }
@@ -38,7 +42,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void put_get_withInheritance() throws Exception {
+  public void put_get_withInheritance() {
 
     PersistenceContext pc = pc();
     pc.put(root(Vehicle.class), 1, car1);
@@ -48,7 +52,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void put_get() throws Exception {
+  public void put_get() {
 
     PersistenceContext pc = pc();
     pc.put(Customer.class, customer42.getId(), customer42);
@@ -58,7 +62,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void putIfAbsent_when_absent() throws Exception {
+  public void putIfAbsent_when_absent() {
 
     PersistenceContext pc = pc();
     Object existing = pc.putIfAbsent(Customer.class, customer42.getId(), customer42);
@@ -67,7 +71,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void putIfAbsent_when_notAbsent() throws Exception {
+  public void putIfAbsent_when_notAbsent() {
 
     PersistenceContext pc = pcWith42();
     Object existing = pc.putIfAbsent(Customer.class, customer42.getId(), new Customer());
@@ -76,21 +80,21 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void get_when_empty() throws Exception {
+  public void get_when_empty() {
     PersistenceContext pc = pc();
     Object found = pc.get(Customer.class, 42);
     assertThat(found).isNull();
   }
 
   @Test
-  public void get_when_there() throws Exception {
+  public void get_when_there() {
     PersistenceContext pc = pcWith42();
     Object found = pc.get(Customer.class, 42);
     assertThat(found).isSameAs(customer42);
   }
 
   @Test
-  public void getWithOption_when_empty() throws Exception {
+  public void getWithOption_when_empty() {
 
     PersistenceContext pc = pc();
     PersistenceContext.WithOption withOption = pc.getWithOption(Customer.class, 42);
@@ -98,7 +102,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void getWithOption_when_there() throws Exception {
+  public void getWithOption_when_there() {
 
     PersistenceContext pc = pcWith42();
 
@@ -107,7 +111,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void getWithOption_when_deleted() throws Exception {
+  public void getWithOption_when_deleted() {
 
     PersistenceContext pc = pcWith42();
     pc.deleted(Customer.class, 42);
@@ -118,21 +122,21 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void size_when_empty() throws Exception {
+  public void size_when_empty() {
 
     PersistenceContext pc = pc();
     assertThat(pc.size(Customer.class)).isEqualTo(0);
   }
 
   @Test
-  public void size_when_some() throws Exception {
+  public void size_when_some() {
 
     PersistenceContext pc = pcWith42();
     assertThat(pc.size(Customer.class)).isEqualTo(1);
   }
 
   @Test
-  public void clear() throws Exception {
+  public void clear() {
 
     PersistenceContext pc = pcWith42();
     pc.clear();
@@ -140,7 +144,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void clearClass() throws Exception {
+  public void clearClass() {
 
     PersistenceContext pc = pcWith42();
     pc.clear(Customer.class);
@@ -148,7 +152,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void clearClassAndId() throws Exception {
+  public void clearClassAndId() {
 
     PersistenceContext pc = pcWith42();
     pc.put(Customer.class, 43, new Customer());
@@ -159,4 +163,94 @@ public class DefaultPersistenceContextTest {
     pc.clear(Customer.class, 43);
     assertThat(pc.size(Customer.class)).isEqualTo(0);
   }
+
+  @Test
+  public void forIterate() {
+    final DefaultPersistenceContext pc = pcWith42();
+    // act
+    final PersistenceContext pcIterate = pc.forIterate();
+    assertThat(pc).isNotSameAs(pcIterate);
+    assertThat(pcIterate.size(Customer.class)).isEqualTo(1);
+  }
+
+  @Test
+  public void forIterate_many() {
+    DefaultPersistenceContext pc = new DefaultPersistenceContext();
+    addCustomers(pc, 1, 100);
+    addContacts(pc, 1, 1010);
+    assertThat(pc.size(Customer.class)).isEqualTo(100);
+    assertThat(pc.size(Contact.class)).isEqualTo(1010);
+
+    // act
+    final PersistenceContext pcIterate = pc.forIterate();
+    assertThat(pcIterate.size(Customer.class)).isEqualTo(100);
+    assertThat(pcIterate.size(Contact.class)).isEqualTo(1010);
+  }
+
+  @Test
+  public void forIterate_resetLimit_forIterateReset() {
+    DefaultPersistenceContext initialPc = new DefaultPersistenceContext();
+    addCustomers(initialPc, 1, 100);
+    addContacts(initialPc, 1, 1010);
+
+    final PersistenceContext pcIterate = initialPc.forIterate();
+    assertFalse(pcIterate.resetLimit());
+
+    // added 900 NEW contact beans
+    addContacts(pcIterate, 2000, 900);
+    assertThat(pcIterate.size(Contact.class)).isEqualTo(1910);
+    assertFalse(pcIterate.resetLimit());
+
+    // boundary, added 1000 NEW contact beans (still false)
+    addContacts(pcIterate, 3000, 100);
+    assertFalse(pcIterate.resetLimit());
+
+    addContacts(pcIterate, 4000, 1);
+    addProducts(pcIterate, 1, 100);
+    // ACT - over 1000 added beans boundary for contacts so returns true
+    assertTrue(pcIterate.resetLimit());
+
+    assertThat(pcIterate.size(Contact.class)).isEqualTo(2011);
+    assertThat(pcIterate.size(Customer.class)).isEqualTo(100);
+    assertThat(pcIterate.size(Product.class)).isEqualTo(100);
+
+    // ACT - obtain new PC forIterateReset
+    PersistenceContext pcReset = pcIterate.forIterateReset();
+
+    // keeps original customer beans as new added beans there
+    assertThat(pcReset.size(Customer.class)).isEqualTo(100); // customers didn't change
+    // added beans to contacts and products so those where reset
+    assertThat(pcReset.size(Contact.class)).isEqualTo(0);
+    assertThat(pcReset.size(Product.class)).isEqualTo(0);
+  }
+
+  @Test
+  public void toString_sillyTest() {
+    DefaultPersistenceContext pc = pcWith42();
+    assertThat(pc.toString()).contains("org.tests.model.basic.Customer");
+  }
+
+  private void addCustomers(PersistenceContext pc, int start, int loop) {
+    for (int i = start; i < start + loop; i++) {
+      Customer bean = new Customer();
+      bean.setId(i);
+      pc.put(Customer.class, i, bean);
+    }
+  }
+
+  private void addContacts(PersistenceContext pc, int start, int loop) {
+    for (int i = start; i < start + loop; i++) {
+      Contact bean = new Contact();
+      bean.setId(i);
+      pc.put(Contact.class, i, bean);
+    }
+  }
+
+  private void addProducts(PersistenceContext pc, int start, int loop) {
+    for (int i = start; i < start + loop; i++) {
+      Product bean = new Product();
+      bean.setId(i);
+      pc.put(Product.class, i, bean);
+    }
+  }
 }

src/test/java/org/tests/query/TestQueryFindEach.java

@@ -4,18 +4,31 @@ import io.ebean.BaseTestCase;
 import io.ebean.DB;
 import io.ebean.FetchConfig;
 import io.ebean.Query;
+import io.ebean.Transaction;
+import io.ebean.annotation.Transactional;
+import io.ebean.bean.PersistenceContext;
+import io.ebeaninternal.api.SpiTransaction;
+import org.ebeantest.LoggedSqlCollector;
 import org.junit.Test;
+import org.tests.model.basic.Contact;
 import org.tests.model.basic.Customer;
 import org.tests.model.basic.ResetBasicData;
+import org.tests.o2m.OmBasicChild;
+import org.tests.o2m.OmBasicParent;
 
+import java.util.List;
+import java.util.Random;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 
 public class TestQueryFindEach extends BaseTestCase {
 
+  private final Random random = new Random();
+
   @Test
   public void test() {
 
@@ -36,6 +49,34 @@ public class TestQueryFindEach extends BaseTestCase {
     assertEquals(2, counter.get());
   }
 
+  @Test
+  public void persistenceContext_scope() {
+
+    ResetBasicData.reset();
+
+    Query<Contact> query = DB.find(Contact.class);
+
+    try (final Transaction transaction = DB.beginTransaction()) {
+      // effectively loads customers into persistence context
+      final List<Customer> customerList = DB.find(Customer.class)
+        .select("name")
+        .findList();
+
+      SpiTransaction spiTxn = (SpiTransaction) transaction;
+      PersistenceContext pc = spiTxn.getPersistenceContext();
+      assertThat(pc.size(Customer.class)).isEqualTo(customerList.size());
+
+      LoggedSqlCollector.start();
+      query.findEach(contact -> {
+        // use customer from persistence context (otherwise would invoke lazy loading)
+        assertNotNull(contact.getCustomer().getName());
+      });
+
+      final List<String> sql = LoggedSqlCollector.stop();
+      assertThat(sql).hasSize(1);
+    }
+  }
+
   /**
    * Test the behaviour when an exception is thrown inside the findVisit().
    */
@@ -60,4 +101,45 @@ public class TestQueryFindEach extends BaseTestCase {
 
     fail("Never get here - exception thrown");
   }
+
+  @Test
+  public void iterateResetLimit() {
+
+    DB.find(OmBasicChild.class).delete();
+    DB.find(OmBasicParent.class).delete();
+    insertData();
+
+    LoggedSqlCollector.start();
+    try (final Transaction transaction = DB.beginTransaction()) {
+      // DB.find(OmBasicParent.class).findList();
+      DB.find(OmBasicChild.class)
+        .setLazyLoadBatchSize(100)
+        //.fetchQuery("parent","name")
+        //.fetch("parent","name")
+        .findEach(child -> {
+          assertNotNull(child.getParent().getName());
+        });
+    }
+
+    final List<String> sql = LoggedSqlCollector.stop();
+    assertThat(sql.size()).isLessThan(50);
+  }
+
+  @Transactional(batchSize = 40)
+  private void insertData() {
+    for (int i = 0; i < 150; i++) {
+      insertData2(i);
+    }
+  }
+
+  private void insertData2(int i) {
+    OmBasicParent p = new OmBasicParent("bl_" + i);
+    DB.save(p);
+
+    int children = 20 + random.nextInt(10);
+    for (int j = 0; j < children; j++) {
+      OmBasicChild c = new OmBasicChild("bl_" + i + "_" + j, p);
+      DB.save(c);
+    }
+  }
 }