Merge pull request #2936 from FOCONIS/mem-leak-streaming-queries

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.ebean.bean.PersistenceContext;
 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();
+  /**
+   * The batch (max) size;
+   */
+  int batchSize();
 
-  List<BeanCollection<?>> getBatch();
+  /**
+   * 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 @@ import static java.lang.System.Logger.Level.DEBUG;
 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 final class LoadManyRequest extends LoadRequest {
   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 final class LoadManyRequest extends LoadRequest {
   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 final class LoadManyRequest extends LoadRequest {
     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 final class LoadManyRequest extends LoadRequest {
     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 final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
       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 final class DLoadContext implements LoadContext {
       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 @@ public final class DLoadContext implements LoadContext {
     }
   }
 
+  @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.BeanDescriptor;
 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 @@ final class DLoadManyContext extends DLoadBaseContext implements LoadManyContext
     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,13 +118,12 @@ final class DLoadManyContext extends DLoadBaseContext implements LoadManyContext
    * 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;
@@ -129,7 +131,6 @@ final class DLoadManyContext extends DLoadBaseContext implements LoadManyContext
       // case it changes as part of a findIterate etc
       this.persistenceContext = context.getPersistenceContext();
       this.batchSize = batchSize;
-      this.list = new ArrayList<>(batchSize);
     }
 
     @Override
@@ -138,7 +139,7 @@ final class DLoadManyContext extends DLoadBaseContext implements LoadManyContext
     }
 
     @Override
-    public int getBatchSize() {
+    public int batchSize() {
       return batchSize;
     }
 
@@ -146,20 +147,15 @@ final class DLoadManyContext extends DLoadBaseContext implements LoadManyContext
      * 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 @@ final class DLoadManyContext extends DLoadBaseContext implements LoadManyContext
           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;
+    }
   }
 }

ebean-test/src/test/java/org/tests/basic/TestPersistenceContextMany.java

@@ -0,0 +1,194 @@
+package org.tests.basic;
+
+import io.ebean.DB;
+import io.ebean.DatabaseFactory;
+import io.ebean.QueryIterator;
+import io.ebean.config.DatabaseConfig;
+import io.ebean.xtest.BaseTestCase;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+
+import javax.persistence.CascadeType;
+import javax.persistence.DiscriminatorValue;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.Inheritance;
+import javax.persistence.ManyToOne;
+import javax.persistence.OneToMany;
+import javax.validation.constraints.Size;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class TestPersistenceContextMany extends BaseTestCase {
+
+
+  @Entity
+  @Inheritance
+  public abstract static class TestModel3 {
+
+    @Id
+    private int id;
+
+    @Size(max = 255)
+    private String someData;
+
+    @OneToMany(cascade = CascadeType.ALL)
+    private List<TestModel3Many1> many1;
+
+    public int getId() {
+      return id;
+    }
+
+    public void setId(int id) {
+      this.id = id;
+    }
+
+    public void setSomeData(String someData) {
+      this.someData = someData;
+    }
+
+    public String getSomeData() {
+      return someData;
+    }
+
+    public List<TestModel3Many1> getMany1() {
+      return many1;
+    }
+  }
+
+  @Entity
+  @Inheritance
+  @DiscriminatorValue("A")
+  public static class TestModel3A extends TestModel3 {
+
+  }
+
+  @Entity
+  @Inheritance
+  @DiscriminatorValue("B")
+  public static class TestModel3B extends TestModel3 {
+    @OneToMany(cascade = CascadeType.ALL)
+    private List<TestModel3Many2> many2;
+
+    public List<TestModel3Many2> getMany2() {
+      return many2;
+    }
+  }
+
+  @Entity
+  public static class TestModel3Many1 {
+    @Id
+    private int id;
+
+    @ManyToOne
+    private TestModel3 base;
+  }
+
+  @Entity
+  public static class TestModel3Many2 {
+    @Id
+    private int id;
+
+    @ManyToOne
+    private TestModel3B base;
+  }
+
+  @Test
+  @Disabled
+  void initDb() {
+    DatabaseConfig config = new DatabaseConfig();
+    config.setName("h2-batch");
+    config.loadFromProperties();
+    config.setDdlExtra(false);
+    config.getDataSourceConfig().setUsername("sa");
+    config.getDataSourceConfig().setPassword("sa");
+    config.getDataSourceConfig().setUrl("jdbc:h2:file:./testsFileMany;DB_CLOSE_ON_EXIT=FALSE;NON_KEYWORDS=KEY,VALUE");
+    config.addClass(TestModel3.class);
+    config.addClass(TestModel3A.class);
+    config.addClass(TestModel3B.class);
+    config.addClass(TestModel3Many1.class);
+    config.addClass(TestModel3Many2.class);
+    DatabaseFactory.create(config);
+
+    String base = "x".repeat(240);
+    // 10 mio TestModel - each needs about 1/4 kbytes -> 2,5 GB in total
+    List<TestModel3> batch = new ArrayList<>();
+    for (int i = 0; i < 1_000_000; i++) {
+      TestModel3 m;
+      if (i == 5) {
+        m = new TestModel3A();
+      } else {
+        m = new TestModel3B();
+        ((TestModel3B) m).getMany2().add(new TestModel3Many2());
+        ((TestModel3B) m).getMany2().add(new TestModel3Many2());
+      }
+      m.getMany1().add(new TestModel3Many1());
+
+      m.setSomeData(base + i); // ensure we have not duplicates
+      batch.add(m);
+      if (i % 1000 == 0) {
+        DB.saveAll(batch);
+        batch.clear();
+      }
+      if (i % 100000 == 0) {
+        System.out.println(i);
+      }
+    }
+    DB.saveAll(batch);
+  }
+
+  @Test
+  @Disabled
+  void testFindEachFindList() {
+    DatabaseConfig config = new DatabaseConfig();
+    config.setName("h2-batch");
+    config.loadFromProperties();
+    config.setDdlRun(false);
+    config.getDataSourceConfig().setUsername("sa");
+    config.getDataSourceConfig().setPassword("sa");
+    config.getDataSourceConfig().setUrl("jdbc:h2:file:./testsFileMany;DB_CLOSE_ON_EXIT=FALSE;NON_KEYWORDS=KEY,VALUE");
+    config.addClass(TestModel3.class);
+    config.addClass(TestModel3A.class);
+    config.addClass(TestModel3B.class);
+    config.addClass(TestModel3Many1.class);
+    config.addClass(TestModel3Many2.class);
+    DatabaseFactory.create(config);
+
+    AtomicInteger i = new AtomicInteger();
+    System.out.println("Doing findEach");
+    DB.find(TestModel3.class).select("*").findEach(c -> {
+      i.incrementAndGet();
+    });
+    System.out.println("Read " + i + " entries");
+
+    i.set(0);
+    System.out.println("Doing findEach with lazyLoad");
+    DB.find(TestModel3.class).select("*").findEach(c -> {
+      i.incrementAndGet();
+      i.addAndGet(c.getMany1().size());
+      if (c instanceof TestModel3B) {
+        i.addAndGet(((TestModel3B) c).getMany2().size());
+      }
+    });
+    System.out.println("Read " + i + " entries"); // 3999998 is correct
+
+    i.set(0);
+    System.out.println("Doing findStream");
+    DB.find(TestModel3.class).select("*").findStream().forEach(c -> i.incrementAndGet());
+    System.out.println("Read " + i + " entries");
+
+    i.set(0);
+    System.out.println("Doing findIterate");
+    QueryIterator<TestModel3> iter = DB.find(TestModel3.class).select("*").findIterate();
+    while (iter.hasNext()) {
+      iter.next();
+      i.incrementAndGet();
+    }
+    System.out.println("Read " + i + " entries");
+
+    System.out.println("Doing FindList, will hold all entries in memory. Expect OOM with -Xmx100m.");
+    List<TestModel3> lst = DB.find(TestModel3.class).select("*").findList();
+    System.out.println("Read " + lst.size() + " entries");
+  }
+}