Weak reference based PersistenceContext for streaming queries

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

@@ -64,17 +64,6 @@ public interface PersistenceContext {
    */
   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>

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

@@ -263,20 +263,6 @@ public final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
     return jsonRead;
   }
 
-  /**
-   * For iterate queries reset the persistenceContext and loadContext.
-   */
-  public void flushPersistenceContextOnIterate() {
-    if (persistenceContext.resetLimit()) {
-      persistenceContext = persistenceContext.forIterateReset();
-      loadContext.resetPersistenceContext(persistenceContext);
-      if (jsonRead != null) {
-        jsonRead.setPersistenceContext(persistenceContext);
-        jsonRead.setLoadContext(loadContext);
-      }
-    }
-  }
-
   /**
    * Get the TransactionContext either explicitly set on the query or
    * transaction scoped.

ebean-core/src/main/java/io/ebeaninternal/server/query/CQueryIteratorSimple.java

@@ -26,7 +26,6 @@ final class CQueryIteratorSimple<T> implements QueryIterator<T> {
   public boolean hasNext() {
     boolean ret = false;
     try {
-      request.flushPersistenceContextOnIterate();
       ret = cquery.hasNext();
       return ret;
     } catch (SQLException e) {

ebean-core/src/main/java/io/ebeaninternal/server/query/CQueryIteratorWithBuffer.java

@@ -35,8 +35,6 @@ final class CQueryIteratorWithBuffer<T> implements QueryIterator<T> {
     try {
       if (buffer.isEmpty() && moreToLoad) {
         // load buffer
-        request.flushPersistenceContextOnIterate();
-
         int i = -1;
         while (moreToLoad && ++i < bufferSize) {
           if (cquery.hasNext()) {

ebean-core/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java

@@ -25,15 +25,9 @@ import java.util.concurrent.locks.ReentrantLock;
  */
 public final class DefaultPersistenceContext implements SpiPersistenceContext {
 
-  /**
-   * Map used hold caches. One cache per bean type.
-   */
   private final HashMap<Class<?>, ClassContext> typeCache = new HashMap<>();
-
   private final ReentrantLock lock = new ReentrantLock();
 
-  private int putCount;
-
   /**
    * Create a new PersistenceContext.
    */
@@ -41,48 +35,15 @@ public final class DefaultPersistenceContext implements SpiPersistenceContext {
   }
 
   /**
-   * 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).
+   * Return a PersistenceContext to use for streaming queries.
    */
   @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);
-  }
-
-  @Override
-  public boolean resetLimit() {
-    lock.lock();
-    try {
-      if (putCount < 100) {
-        return false;
-      }
-      putCount = 0;
-      for (ClassContext value : typeCache.values()) {
-        if (value.resetLimit()) {
-          return true;
-        }
-      }
-      // checking after another 100 puts
-      return false;
-    } finally {
-      lock.unlock();
+    WeakPersistenceContext weak = new WeakPersistenceContext();
+    for (ClassContext value : typeCache.values()) {
+      weak.add(value.rootType, value.deleteSet, value.map);
     }
+    return weak;
   }
 
   /**
@@ -92,7 +53,6 @@ public final class DefaultPersistenceContext implements SpiPersistenceContext {
   public void put(Class<?> rootType, Object id, Object bean) {
     lock.lock();
     try {
-      putCount++;
       classContext(rootType).put(id, bean);
     } finally {
       lock.unlock();
@@ -103,7 +63,6 @@ public final class DefaultPersistenceContext implements SpiPersistenceContext {
   public Object putIfAbsent(Class<?> rootType, Object id, Object bean) {
     lock.lock();
     try {
-      putCount++;
       return classContext(rootType).putIfAbsent(id, bean);
     } finally {
       lock.unlock();
@@ -232,61 +191,18 @@ public final class DefaultPersistenceContext implements SpiPersistenceContext {
     private final Map<Object, Object> map = new HashMap<>();
     private final Class<?> rootType;
     private Set<Object> deleteSet;
-    private int initialSize;
-    private ClassContext parent;
 
     private ClassContext(Class<?> rootType) {
       this.rootType = rootType;
     }
 
-    /**
-     * Create as a shallow copy.
-     */
-    private ClassContext(ClassContext source, boolean initial) {
-      this.rootType = source.rootType;
-      if (initial || source.isTransfer()) {
-        parent = source.transferParent();
-        initialSize = parent.size();
-        if (source.deleteSet != null) {
-          deleteSet = new HashSet<>(source.deleteSet);
-        }
-      }
-    }
-
-    /**
-     * True if this should be transferred to a new iterator persistence context.
-     */
-    private boolean isTransfer() {
-      // map not added to and has some original/parent beans
-      return map.isEmpty() && initialSize > 0;
-    }
-
-    private ClassContext transferParent() {
-      return (parent != null) ? parent : this;
-    }
-
-    /**
-     * 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 of 1000.
-     */
-    private boolean resetLimit() {
-      return map.size() > 1000;
-    }
-
     @Override
     public String toString() {
       return "size:" + map.size();
     }
 
     private Object get(Object id) {
-      Object bean = (parent == null) ? null : parent.get(id);
-      return bean != null ? bean : map.get(id);
+      return map.get(id);
     }
 
     private WithOption getWithOption(Object id) {
@@ -313,7 +229,7 @@ public final class DefaultPersistenceContext implements SpiPersistenceContext {
     }
 
     private int size() {
-      return map.size() + initialSize;
+      return map.size();
     }
 
     private void clear() {

ebean-core/src/main/java/io/ebeaninternal/server/transaction/WeakPersistenceContext.java

@@ -0,0 +1,272 @@
+package io.ebeaninternal.server.transaction;
+
+import io.ebean.bean.EntityBean;
+import io.ebean.bean.PersistenceContext;
+import io.ebeaninternal.api.SpiBeanType;
+import io.ebeaninternal.api.SpiBeanTypeManager;
+import io.ebeaninternal.api.SpiPersistenceContext;
+
+import java.lang.ref.Reference;
+import java.lang.ref.ReferenceQueue;
+import java.lang.ref.WeakReference;
+import java.util.*;
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * Weak reference based PersistenceContext for "streaming queries".
+ */
+final class WeakPersistenceContext implements SpiPersistenceContext {
+
+  private final HashMap<Class<?>, WeakClassContext> typeCache = new HashMap<>();
+  private final ReentrantLock lock = new ReentrantLock();
+
+  WeakPersistenceContext() {
+  }
+
+  /**
+   * Load from the initiating persistence context.
+   */
+  void add(Class<?> rootType, Set<Object> deleteSet, Map<Object, Object> map) {
+    typeCache.put(rootType, new WeakClassContext( rootType, deleteSet, map));
+  }
+
+  @Override
+  public PersistenceContext forIterate() {
+    return this;
+  }
+
+  @Override
+  public void put(Class<?> rootType, Object id, Object bean) {
+    lock.lock();
+    try {
+      classContext(rootType).put(id, bean);
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public Object putIfAbsent(Class<?> rootType, Object id, Object bean) {
+    lock.lock();
+    try {
+      return classContext(rootType).putIfAbsent(id, bean);
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public Object get(Class<?> rootType, Object id) {
+    lock.lock();
+    try {
+      return classContext(rootType).get(id);
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public WithOption getWithOption(Class<?> rootType, Object id) {
+    lock.lock();
+    try {
+      return classContext(rootType).getWithOption(id);
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public int size(Class<?> rootType) {
+    lock.lock();
+    try {
+      WeakClassContext classMap = typeCache.get(rootType);
+      return classMap == null ? 0 : classMap.size();
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public void clear() {
+    lock.lock();
+    try {
+      typeCache.clear();
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public void clear(Class<?> rootType) {
+    lock.lock();
+    try {
+      WeakClassContext classMap = typeCache.get(rootType);
+      if (classMap != null) {
+        classMap.clear();
+      }
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public void deleted(Class<?> rootType, Object id) {
+    lock.lock();
+    try {
+      WeakClassContext classMap = typeCache.get(rootType);
+      if (classMap != null && id != null) {
+        classMap.deleted(id);
+      }
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public void clear(Class<?> rootType, Object id) {
+    lock.lock();
+    try {
+      WeakClassContext classMap = typeCache.get(rootType);
+      if (classMap != null && id != null) {
+        classMap.remove(id);
+      }
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public List<Object> dirtyBeans(SpiBeanTypeManager manager) {
+    lock.lock();
+    try {
+      List<Object> list = new ArrayList<>();
+      for (WeakClassContext classContext : typeCache.values()) {
+        classContext.dirtyBeans(manager, list);
+      }
+      return list;
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public String toString() {
+    lock.lock();
+    try {
+      return typeCache.toString();
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  private WeakClassContext classContext(Class<?> rootType) {
+    return typeCache.computeIfAbsent(rootType, k -> new WeakClassContext(rootType));
+  }
+
+  private static class WeakClassContext {
+
+    private final ReferenceQueue<Object> queue = new ReferenceQueue<>();
+    private final Map<Object, BeanRef> map = new HashMap<>();
+    private final Class<?> rootType;
+    private Set<Object> deleteSet;
+
+    private WeakClassContext(Class<?> rootType) {
+      this.rootType = rootType;
+    }
+
+    private WeakClassContext(Class<?> rootType, Set<Object> initialDeleteSet, Map<Object, Object> initialMap) {
+      this.rootType = rootType;
+      this.deleteSet = initialDeleteSet == null ? null : new HashSet<>(initialDeleteSet);
+      for (Map.Entry<Object, Object> entry : initialMap.entrySet()) {
+        Object id = entry.getKey();
+        map.put(id, new BeanRef(id, entry.getValue(), queue));
+      }
+    }
+
+    private void expungeStaleEntries() {
+      Reference<?> ref;
+      while ((ref = queue.poll()) != null) {
+        map.remove(((BeanRef)ref).key());
+      }
+    }
+
+    @Override
+    public String toString() {
+      return "size:" + map.size();
+    }
+
+    private Object get(Object id) {
+      expungeStaleEntries();
+      Reference<Object> ref = map.get(id);
+      return ref == null ? null : ref.get();
+    }
+
+    private WithOption getWithOption(Object id) {
+      if (deleteSet != null && deleteSet.contains(id)) {
+        return WithOption.DELETED;
+      }
+      Object bean = get(id);
+      return (bean == null) ? null : new WithOption(bean);
+    }
+
+    private Object putIfAbsent(Object id, Object bean) {
+      Object existingValue = get(id);
+      if (existingValue != null) {
+        // it is not absent
+        return existingValue;
+      }
+      // put the new value and return null indicating the put was successful
+      put(id, bean);
+      return null;
+    }
+
+    private void put(Object id, Object b) {
+      expungeStaleEntries();
+      map.put(id, new BeanRef(id, b, queue));
+    }
+
+    private int size() {
+      return map.size();
+    }
+
+    private void clear() {
+      map.clear();
+    }
+
+    private void remove(Object id) {
+      map.remove(id);
+    }
+
+    private void deleted(Object id) {
+      if (deleteSet == null) {
+        deleteSet = new HashSet<>();
+      }
+      deleteSet.add(id);
+      map.remove(id);
+    }
+
+    /**
+     * Add the dirty beans to the list.
+     */
+    void dirtyBeans(SpiBeanTypeManager manager, List<Object> list) {
+      final SpiBeanType beanType = manager.beanType(rootType);
+      for (BeanRef value : map.values()) {
+        EntityBean bean = (EntityBean) value.get();
+        if (bean != null && (bean._ebean_getIntercept().isDirty() || beanType.isToManyDirty(bean))) {
+          list.add(bean);
+        }
+      }
+    }
+  }
+
+  private static class BeanRef extends WeakReference<Object> {
+    private final Object key;
+    private BeanRef(Object key, Object referent, ReferenceQueue<? super Object> q) {
+      super(referent, q);
+      this.key = key;
+    }
+    Object key() {
+      return key;
+    }
+  }
+}

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

@@ -6,16 +6,13 @@ import org.junit.jupiter.api.Test;
 import org.tests.model.basic.*;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
-public class DefaultPersistenceContextTest {
+class DefaultPersistenceContextTest {
 
   private final Customer customer42;
-
   private final Car car1;
 
-  public DefaultPersistenceContextTest() {
+  DefaultPersistenceContextTest() {
     customer42 = new Customer();
     customer42.setId(42);
     car1 = new Car();
@@ -37,8 +34,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void put_get_withInheritance() {
-
+  void put_get_withInheritance() {
     PersistenceContext pc = pc();
     pc.put(root(Vehicle.class), 1, car1);
 
@@ -47,8 +43,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void put_get() {
-
+  void put_get() {
     PersistenceContext pc = pc();
     pc.put(Customer.class, customer42.getId(), customer42);
 
@@ -57,8 +52,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void putIfAbsent_when_absent() {
-
+  void putIfAbsent_when_absent() {
     PersistenceContext pc = pc();
     Object existing = pc.putIfAbsent(Customer.class, customer42.getId(), customer42);
 
@@ -66,8 +60,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void putIfAbsent_when_notAbsent() {
-
+  void putIfAbsent_when_notAbsent() {
     PersistenceContext pc = pcWith42();
     Object existing = pc.putIfAbsent(Customer.class, customer42.getId(), new Customer());
 
@@ -75,30 +68,28 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void get_when_empty() {
+  void get_when_empty() {
     PersistenceContext pc = pc();
     Object found = pc.get(Customer.class, 42);
     assertThat(found).isNull();
   }
 
   @Test
-  public void get_when_there() {
+  void get_when_there() {
     PersistenceContext pc = pcWith42();
     Object found = pc.get(Customer.class, 42);
     assertThat(found).isSameAs(customer42);
   }
 
   @Test
-  public void getWithOption_when_empty() {
-
+  void getWithOption_when_empty() {
     PersistenceContext pc = pc();
     PersistenceContext.WithOption withOption = pc.getWithOption(Customer.class, 42);
     assertThat(withOption).isNull();
   }
 
   @Test
-  public void getWithOption_when_there() {
-
+  void getWithOption_when_there() {
     PersistenceContext pc = pcWith42();
 
     PersistenceContext.WithOption withOption = pc.getWithOption(Customer.class, 42);
@@ -106,8 +97,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void getWithOption_when_deleted() {
-
+  void getWithOption_when_deleted() {
     PersistenceContext pc = pcWith42();
     pc.deleted(Customer.class, 42);
 
@@ -117,38 +107,33 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void size_when_empty() {
-
+  void size_when_empty() {
     PersistenceContext pc = pc();
     assertThat(pc.size(Customer.class)).isEqualTo(0);
   }
 
   @Test
-  public void size_when_some() {
-
+  void size_when_some() {
     PersistenceContext pc = pcWith42();
     assertThat(pc.size(Customer.class)).isEqualTo(1);
   }
 
   @Test
-  public void clear() {
-
+  void clear() {
     PersistenceContext pc = pcWith42();
     pc.clear();
     assertThat(pc.size(Customer.class)).isEqualTo(0);
   }
 
   @Test
-  public void clearClass() {
-
+  void clearClass() {
     PersistenceContext pc = pcWith42();
     pc.clear(Customer.class);
     assertThat(pc.size(Customer.class)).isEqualTo(0);
   }
 
   @Test
-  public void clearClassAndId() {
-
+  void clearClassAndId() {
     PersistenceContext pc = pcWith42();
     pc.put(Customer.class, 43, new Customer());
 
@@ -160,7 +145,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void forIterate() {
+  void forIterate() {
     final DefaultPersistenceContext pc = pcWith42();
     final Object origCustomer42 = pc.get(Customer.class, 42);
 
@@ -177,7 +162,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @Test
-  public void forIterate_many() {
+  void forIterate_many() {
     DefaultPersistenceContext pc = new DefaultPersistenceContext();
     addCustomers(pc, 1, 100);
     addContacts(pc, 1, 1010);
@@ -191,44 +176,7 @@ public class DefaultPersistenceContextTest {
   }
 
   @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 no 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() {
+  void toString_sillyTest() {
     DefaultPersistenceContext pc = pcWith42();
     assertThat(pc.toString()).contains("org.tests.model.basic.Customer");
   }
@@ -248,12 +196,4 @@ public class DefaultPersistenceContextTest {
       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);
-    }
-  }
 }