From 70b730657f92b837ed1486f8c13308df53f8a817 Mon Sep 17 00:00:00 2001
From: rob bygrave <robin.bygrave@gmail.com>
Date: Mon, 29 Jun 2020 08:48:24 +1200
Subject: [PATCH] #2028 - Refactor internals of smart persistence context to
 reference "parent" rather than copy beans (adjustment of #2021)

---
 .../DefaultPersistenceContext.java            | 34 ++++++++++++-------
 .../DefaultPersistenceContextTest.java        | 15 +++++---
 2 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java b/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java
index ece7798b2..bdbb22ea0 100644
--- a/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java
+++ b/src/main/java/io/ebeaninternal/server/transaction/DefaultPersistenceContext.java
@@ -189,27 +189,34 @@ public final class DefaultPersistenceContext implements PersistenceContext {
 
     private int initialSize;
 
+    private ClassContext parent;
+
     private ClassContext() {
     }
 
     /**
      * Create as a shallow copy.
      */
-    private ClassContext(ClassContext parent, boolean initial) {
-      if (initial || parent.isInitialSize()) {
-        this.map.putAll(parent.map);
-        this.initialSize = map.size();
-        if (parent.deleteSet != null) {
-          this.deleteSet = new HashSet<>(parent.deleteSet);
+    private ClassContext(ClassContext source, boolean initial) {
+      if (initial || source.isTransfer()) {
+        parent = source.transferParent();
+        initialSize = parent.size();
+        if (source.deleteSet != null) {
+          deleteSet = new HashSet<>(source.deleteSet);
         }
       }
     }
 
     /**
-     * True if the map has not changed from it's initial size (no additions).
+     * True if this should be transferred to a new iterator persistence context.
      */
-    private boolean isInitialSize() {
-      return map.size() == initialSize;
+    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;
     }
 
     /**
@@ -220,10 +227,10 @@ public final class DefaultPersistenceContext implements PersistenceContext {
     }
 
     /**
-     * Return true if grown above the reset limit size.
+     * Return true if grown above the reset limit size of 1000.
      */
     private boolean resetLimit() {
-      return map.size() > initialSize + 1000;
+      return map.size() > 1000;
     }
 
     @Override
@@ -232,7 +239,8 @@ public final class DefaultPersistenceContext implements PersistenceContext {
     }
 
     private Object get(Object id) {
-      return map.get(id);
+      Object bean = (parent == null) ? null : parent.get(id);
+      return bean != null ? bean : map.get(id);
     }
 
     private WithOption getWithOption(Object id) {
@@ -259,7 +267,7 @@ public final class DefaultPersistenceContext implements PersistenceContext {
     }
 
     private int size() {
-      return map.size();
+      return map.size() + initialSize;
     }
 
     private void clear() {
diff --git a/src/test/java/io/ebeaninternal/server/transaction/DefaultPersistenceContextTest.java b/src/test/java/io/ebeaninternal/server/transaction/DefaultPersistenceContextTest.java
index 82e41f940..87dbf4fa3 100644
--- a/src/test/java/io/ebeaninternal/server/transaction/DefaultPersistenceContextTest.java
+++ b/src/test/java/io/ebeaninternal/server/transaction/DefaultPersistenceContextTest.java
@@ -15,14 +15,13 @@ import static org.junit.Assert.assertTrue;
 
 public class DefaultPersistenceContextTest {
 
-  Customer customer42;
+  private final Customer customer42;
 
-  Car car1;
+  private final Car car1;
 
   public DefaultPersistenceContextTest() {
     customer42 = new Customer();
     customer42.setId(42);
-
     car1 = new Car();
     car1.setId(1);
   }
@@ -167,10 +166,18 @@ public class DefaultPersistenceContextTest {
   @Test
   public void forIterate() {
     final DefaultPersistenceContext pc = pcWith42();
+    final Object origCustomer42 = pc.get(Customer.class, 42);
+
     // act
     final PersistenceContext pcIterate = pc.forIterate();
     assertThat(pc).isNotSameAs(pcIterate);
     assertThat(pcIterate.size(Customer.class)).isEqualTo(1);
+
+    // assert same instance (bean effectively transferred to iterator persistence context
+    final Object customer42 = pcIterate.get(Customer.class, 42);
+    assertThat(customer42).isSameAs(origCustomer42);
+    final PersistenceContext.WithOption option = pcIterate.getWithOption(Customer.class, 42);
+    assertThat(option.getBean()).isSameAs(origCustomer42);
   }
 
   @Test
@@ -217,7 +224,7 @@ public class DefaultPersistenceContextTest {
     // ACT - obtain new PC forIterateReset
     PersistenceContext pcReset = pcIterate.forIterateReset();
 
-    // keeps original customer beans as new added beans there
+    // 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);