Weak reference based PersistenceContext for streaming queries

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;
+    }
+  }
+}