BUG: Memory-leak on streaming queries when LoadBuffers are not aligned

ebean-core/src/main/java/io/ebeaninternal/api/LoadContext.java

@@ -50,4 +50,9 @@ public interface LoadContext {
    * Register a collection for lazy loading.
    */
   void register(String path, BeanPropertyAssocMany<?> many, BeanCollection<?> bc);
+
+  /**
+   * Use soft-references for streaming queries, so unreachable entries can be garbage collected.
+   */
+  void useReferences(boolean useReferences);
 }

ebean-core/src/main/java/io/ebeaninternal/api/LoadManyBuffer.java

@@ -6,16 +6,30 @@
 import io.ebeaninternal.server.deploy.BeanDescriptor;
 import io.ebeaninternal.server.deploy.BeanPropertyAssocMany;
 
-import java.util.List;
-
 /**
  * A buffer of bean collections for batch lazy loading and secondary query loading.
  */
 public interface LoadManyBuffer {
 
-  int getBatchSize();
-
-  List<BeanCollection<?>> getBatch();
+  /**
+   * The batch (max) size;
+   */
+  int batchSize();
+
+  /**
+   * The actual size.
+   */
+  int size();
+
+  /**
+   * Get the <code>i</code>th element from buffer. This can be null.
+   */
+  BeanCollection<?> get(int i);
+
+  /**
+   * Removes an element from the buffer. This will NOT affect size.
+   */
+  boolean removeFromBuffer(BeanCollection<?> collection);
 
   BeanPropertyAssocMany<?> getBeanProperty();
 

ebean-core/src/main/java/io/ebeaninternal/api/LoadManyRequest.java

@@ -19,8 +19,6 @@
 public final class LoadManyRequest extends LoadRequest {
 
   private static final System.Logger log = CoreLog.log;
-
-  private final List<BeanCollection<?>> batch;
   private final LoadManyBuffer loadContext;
   private final boolean onlyIds;
   private final boolean loadCache;
@@ -42,7 +40,6 @@ public LoadManyRequest(LoadManyBuffer loadContext, OrmQueryRequest<?> parentRequ
   private LoadManyRequest(LoadManyBuffer loadContext, OrmQueryRequest<?> parentRequest, boolean lazy, boolean onlyIds, boolean loadCache) {
     super(parentRequest, lazy);
     this.loadContext = loadContext;
-    this.batch = loadContext.getBatch();
     this.onlyIds = onlyIds;
     this.loadCache = loadCache;
   }
@@ -59,9 +56,12 @@ public String description() {
   private List<Object> parentIdList(SpiEbeanServer server) {
     List<Object> idList = new ArrayList<>();
     BeanPropertyAssocMany<?> many = many();
-    for (BeanCollection<?> bc : batch) {
-      idList.add(many.parentId(bc.getOwnerBean()));
-      bc.setLoader(server); // don't use the load buffer again
+    for (int i = 0; i < loadContext.size(); i++) {
+      BeanCollection<?> bc = loadContext.get(i);
+      if (bc != null) {
+        idList.add(many.parentId(bc.getOwnerBean()));
+        bc.setLoader(server); // don't use the load buffer again
+      }
     }
     if (many.targetDescriptor().isPadInExpression()) {
       BindPadding.padIds(idList);
@@ -91,7 +91,7 @@ public SpiQuery<?> createQuery(SpiEbeanServer server) {
     query.setPersistenceContext(loadContext.getPersistenceContext());
     query.setLoadDescription(lazy ? "+lazy" : "+query", description());
     if (lazy) {
-      query.setLazyLoadBatchSize(loadContext.getBatchSize());
+      query.setLazyLoadBatchSize(loadContext.batchSize());
     } else {
       query.setBeanCacheMode(CacheMode.OFF);
     }
@@ -112,15 +112,18 @@ public void postLoad() {
     BeanPropertyAssocMany<?> many = many();
     // check for BeanCollection's that where never processed
     // in the +query or +lazy load due to no rows (predicates)
-    for (BeanCollection<?> bc : batch) {
-      if (bc.checkEmptyLazyLoad()) {
-        if (log.isLoggable(DEBUG)) {
-          EntityBean ownerBean = bc.getOwnerBean();
-          Object parentId = desc.getId(ownerBean);
-          log.log(DEBUG, "BeanCollection after lazy load was empty. type:{0} id:{1} owner:{2}", ownerBean.getClass().getName(), parentId, ownerBean);
+    for (int i = 0; i < loadContext.size(); i++) {
+      BeanCollection<?> bc = loadContext.get(i);
+      if (bc != null) {
+        if (bc.checkEmptyLazyLoad()) {
+          if (log.isLoggable(DEBUG)) {
+            EntityBean ownerBean = bc.getOwnerBean();
+            Object parentId = desc.getId(ownerBean);
+            log.log(DEBUG, "BeanCollection after lazy load was empty. type:{0} id:{1} owner:{2}", ownerBean.getClass().getName(), parentId, ownerBean);
+          }
+        } else if (loadCache && many.isUseCache()) {
+          desc.cacheManyPropPut(many, bc, desc.cacheKeyForBean(bc.getOwnerBean()));
         }
-      } else if (loadCache && many.isUseCache()) {
-        desc.cacheManyPropPut(many, bc, desc.cacheKeyForBean(bc.getOwnerBean()));
       }
     }
   }

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

@@ -229,6 +229,7 @@ public void initTransIfRequired() {
       persistenceContext.beginIterate();
     }
     loadContext = new DLoadContext(this, secondaryQueries);
+    loadContext.useReferences(Type.ITERATE == query.getType());
   }
 
   /**

ebean-core/src/main/java/io/ebeaninternal/server/loadcontext/DLoadContext.java

@@ -58,6 +58,7 @@ public final class DLoadContext implements LoadContext {
   private final ProfilingListener profilingListener;
   private final Map<String, ObjectGraphNode> nodePathMap = new HashMap<>();
   private final PersistenceContext persistenceContext;
+  boolean useReferences;
   private List<OrmQueryProperties> secQuery;
   private Object tenantId;
 
@@ -118,7 +119,6 @@ public DLoadContext(OrmQueryRequest<?> request, SpiQuerySecondary secondaryQueri
       this.origin = null;
       this.relativePath = null;
     }
-
     // initialise rootBeanContext after origin and relativePath have been set
     this.rootBeanContext = new DLoadBeanContext(this, rootDescriptor, null, null);
     registerSecondaryQueries(secondaryQueries);
@@ -156,6 +156,11 @@ private void registerSecondaryQueries(SpiQuerySecondary secondaryQueries) {
     }
   }
 
+  @Override
+  public void useReferences(boolean useReferences) {
+    this.useReferences = useReferences;
+  }
+
   /**
    * Setup the load context at this path with OrmQueryProperties which is
    * used to build the appropriate query for +query or +lazy loading.

ebean-core/src/main/java/io/ebeaninternal/server/loadcontext/DLoadManyContext.java

@@ -14,7 +14,10 @@
 import io.ebeaninternal.server.deploy.BeanPropertyAssocMany;
 import io.ebeaninternal.server.querydefn.OrmQueryProperties;
 
+import java.lang.ref.Reference;
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -38,7 +41,7 @@ final class DLoadManyContext extends DLoadBaseContext implements LoadManyContext
   }
 
   private LoadBuffer createBuffer(int size) {
-    LoadBuffer buffer = new LoadBuffer(this, size);
+    LoadBuffer buffer = parent.useReferences ? new LoadBufferWeakRef(this, size) : new LoadBufferHardRef(this, size);
     if (bufferList != null) {
       bufferList.add(buffer);
     }
@@ -94,7 +97,7 @@ public void loadSecondaryQuery(OrmQueryRequest<?> parentRequest, boolean forEach
     try {
       if (bufferList != null) {
         for (LoadBuffer loadBuffer : bufferList) {
-          if (!loadBuffer.list.isEmpty()) {
+          if (loadBuffer.size() > 0) {
             LoadManyRequest req = new LoadManyRequest(loadBuffer, parentRequest);
             parent.getEbeanServer().loadMany(req);
           }
@@ -115,21 +118,19 @@ public void loadSecondaryQuery(OrmQueryRequest<?> parentRequest, boolean forEach
    * A buffer for batch loading bean collections on a given path.
    * Supports batch lazy loading and secondary query loading.
    */
-  static class LoadBuffer implements BeanCollectionLoader, LoadManyBuffer {
+  static abstract class LoadBuffer implements BeanCollectionLoader, LoadManyBuffer {
 
     private final ReentrantLock lock = new ReentrantLock();
     private final PersistenceContext persistenceContext;
     private final DLoadManyContext context;
-    private final int batchSize;
-    private final List<BeanCollection<?>> list;
+    final int batchSize;
 
     LoadBuffer(DLoadManyContext context, int batchSize) {
       this.context = context;
       // set the persistence context as at this moment in
       // case it changes as part of a findIterate etc
       this.persistenceContext = context.getPersistenceContext();
       this.batchSize = batchSize;
-      this.list = new ArrayList<>(batchSize);
     }
 
     @Override
@@ -138,28 +139,23 @@ public boolean isUseDocStore() {
     }
 
     @Override
-    public int getBatchSize() {
+    public int batchSize() {
       return batchSize;
     }
 
     /**
      * Return true if the buffer is full.
      */
     public boolean isFull() {
-      return batchSize == list.size();
+      return batchSize() == size();
     }
 
     /**
      * Return true if the buffer is full.
      */
-    public void add(BeanCollection<?> bc) {
-      list.add(bc);
-    }
+    public abstract void add(BeanCollection<?> bc);
 
-    @Override
-    public List<BeanCollection<?>> getBatch() {
-      return list;
-    }
+    abstract void clear();
 
     @Override
     public BeanPropertyAssocMany<?> getBeanProperty() {
@@ -208,25 +204,133 @@ public void loadMany(BeanCollection<?> bc, boolean onlyIds) {
           final String parentKey = parentDesc.cacheKey(parentId);
           if (parentDesc.cacheManyPropLoad(context.property, bc, parentKey, context.parent.isReadOnly())) {
             // we loaded the bean collection from cache so remove it from the buffer
-            for (int i = 0; i < list.size(); i++) {
-              // find it using instance equality - avoiding equals() and potential deadlock issue
-              if (list.get(i) == bc) {
-                list.remove(i);
-                bc.setLoader(context.parent.getEbeanServer());
-                return;
-              }
+            if (removeFromBuffer(bc)) {
+              bc.setLoader(context.parent.getEbeanServer());
             }
+            // find it using instance equality - avoiding equals() and potential deadlock issue
             return;
           }
         }
 
         context.parent.getEbeanServer().loadMany(new LoadManyRequest(this, onlyIds, useCache));
         // clear the buffer as all entries have been loaded
-        list.clear();
+        clear();
       } finally {
         lock.unlock();
       }
     }
+  }
+
+  static class LoadBufferHardRef extends LoadBuffer {
+    private final BeanCollection<?>[] list;
+
+    private int size;
+
+    LoadBufferHardRef(DLoadManyContext context, int batchSize) {
+      super(context, batchSize);
+      this.list = new BeanCollection<?>[batchSize];
+    }
+
+    /**
+     * Return true if the buffer is full.
+     */
+    @Override
+    public void add(BeanCollection<?> bc) {
+      list[size++] = bc;
+    }
+
+    @Override
+    void clear() {
+      Arrays.fill(list, null);
+      size = 0;
+    }
+
+    @Override
+    public int size() {
+      return size;
+    }
+
+    @Override
+    public BeanCollection<?> get(int i) {
+      return list[i];
+    }
+
+    @Override
+    public boolean removeFromBuffer(BeanCollection<?> collection) {
+      for (int i = 0; i < size; i++) {
+        // find it using instance equality - avoiding equals() and potential deadlock issue
+        if (list[i] == collection) {
+          list[i] = null;
+          return true;
+        }
+      }
+      return false;
+    }
+  }
 
+  /**
+   * This load buffer uses weak references, so unreachable beanCollections will drop out from the buffer.
+   */
+  static class LoadBufferWeakRef extends LoadBuffer {
+    private final Reference<BeanCollection<?>>[] list;
+
+    private int size;
+
+    LoadBufferWeakRef(DLoadManyContext context, int batchSize) {
+      super(context, batchSize);
+      this.list = new Reference[batchSize];
+    }
+
+    /**
+     * Return true if the buffer is full.
+     */
+    @Override
+    public void add(BeanCollection<?> bc) {
+      list[size++] = new WeakReference<>(bc);
+    }
+
+    @Override
+    void clear() {
+      Arrays.fill(list, null);
+      size = 0;
+    }
+
+    @Override
+    public int size() {
+      return size;
+    }
+
+    @Override
+    public BeanCollection<?> get(int i) {
+      Reference<BeanCollection<?>> ref = list[i];
+      if (ref == null) {
+        return null;
+      }
+      BeanCollection<?> bc = ref.get();
+      if (bc == null) {
+        // remove dead references
+        list[i] = null;
+      }
+      return bc;
+    }
+
+    @Override
+    public boolean removeFromBuffer(BeanCollection<?> collection) {
+      for (int i = 0; i < size; i++) {
+        if (list[i] != null) {
+          BeanCollection<?> bc = list[i].get();
+          if (bc == null) {
+            // remove dead references
+            list[i] = null;
+          }
+          // find it using instance equality - avoiding equals() and potential deadlock issue
+          if (bc == collection) {
+            list[i] = null;
+            return true;
+          }
+        }
+      }
+      return false;
+    }
   }
 }