From fd1b5f58bc202bc963b65c507de06d66ba6d5c87 Mon Sep 17 00:00:00 2001
From: Roland Praml <roland.praml@foconis.de>
Date: Wed, 3 Nov 2021 11:34:04 +0100
Subject: [PATCH] FIX: Subsequent findEach / findList does not work reliable
 (Regression of #2411)

---
 .../DefaultPersistenceContext.java            | 22 +++++++++++++---
 .../tests/basic/TestPersistenceContext.java   | 26 +++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/ebean-core/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java b/ebean-core/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java
index d10d92ab2..ba654782f 100644
--- a/ebean-core/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java
+++ b/ebean-core/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java
@@ -276,11 +276,20 @@ public final class DefaultPersistenceContext implements SpiPersistenceContext {
     }
 
     private void put(Object id, Object b) {
+      Object existing;
       if (useReferences) {
         weakCount++;
-        map.put(id, new BeanRef(this, id, b, queue));
+        existing = map.put(id, new BeanRef(this, id, b, queue));
       } else {
-        map.put(id, b);
+        existing = map.put(id, b);
+      }
+      // when existing BeanRef is replaced, it must not be processed by the queue.
+      // Removing from queue is not possible, so we set it to "dead" and skip 
+      // this reference in expunge. This prevents us from removing wrong keys.
+      
+      if (existing instanceof BeanRef) {
+        ((BeanRef) existing).setDead();
+        weakCount--;
       }
     }
 
@@ -330,6 +339,7 @@ public final class DefaultPersistenceContext implements SpiPersistenceContext {
 
     private final ClassContext classContext;
     private final Object key;
+    private boolean alive = true;
 
     private BeanRef(ClassContext classContext, Object key, Object referent, ReferenceQueue<? super Object> q) {
       super(referent, q);
@@ -337,8 +347,14 @@ public final class DefaultPersistenceContext implements SpiPersistenceContext {
       this.key = key;
     }
 
+    private void setDead() {
+      alive = false;
+    }
+
     private void expunge() {
-      classContext.remove(key);
+      if (alive) {
+        classContext.remove(key);
+      }
     }
 
   }
diff --git a/ebean-test/src/test/java/org/tests/basic/TestPersistenceContext.java b/ebean-test/src/test/java/org/tests/basic/TestPersistenceContext.java
index 86e0425a6..2f87bdc08 100644
--- a/ebean-test/src/test/java/org/tests/basic/TestPersistenceContext.java
+++ b/ebean-test/src/test/java/org/tests/basic/TestPersistenceContext.java
@@ -191,4 +191,30 @@ class TestPersistenceContext extends BaseTestCase {
       assertThat(pc.toString()).contains("Customer=size:200 (0 weak)");
     }
   }
+  
+  @Disabled // run manually
+  @Test
+  void testPcScopes_with_findEachFindList() throws InterruptedException {
+    for (int i = 0; i < 5000; i++) {
+      Customer c = new Customer();
+      c.setName("Customer #" + i);
+      DB.save(c);
+      Order o = new Order();
+      o.setCustomer(c);
+      DB.save(o);
+    }
+    
+    for (int i = 0; i < 1000; i++) {
+      try (Transaction txn = DB.beginTransaction()) {
+        List<Customer> customers = new ArrayList<>();
+        DB.find(Customer.class).select("id").findEach(customers::add);
+        SpiPersistenceContext pc = ((SpiTransaction) txn).getPersistenceContext();
+        assertThat(pc.toString()).contains("Customer=size:5000 (5000 weak)");
+        customers.clear();
+        customers = DB.find(Customer.class).select("id").findList();
+        
+        assertThat(pc.toString()).contains("Customer=size:5000"); // We expect ALWAYS 5000 entries in the PC
+      }
+    }
+  }
 }