[OPENJPA-2958] Add ImplHelper.registerUnenhancedInstance()

openjpa-kernel/src/main/java/org/apache/openjpa/kernel/BrokerImpl.java

@@ -3094,7 +3094,7 @@ public OpenJPAStateManager embed(Object obj, Object id,
             // embedded copy mapping so that getStateManagerImpl can find
             // the SM when the owner's enhanced field provides the raw record
             if (obj != null && meta.isRecord()) {
-                ImplHelper._unenhancedInstanceMap.put(obj, copy);
+                ImplHelper.registerUnenhancedInstance(obj, copy);
             }
 
             return sm;

openjpa-kernel/src/main/java/org/apache/openjpa/util/ImplHelper.java

@@ -61,8 +61,15 @@ public class ImplHelper {
     private static final Map _assignableTypes =
         new ConcurrentReferenceHashMap(ReferenceStrength.WEAK, ReferenceStrength.HARD);
 
-    // map of all new unenhanced instances active in this classloader
-    public static final Map _unenhancedInstanceMap =
+    /**
+     * Map of all new unenhanced instances active in this classloader, keyed
+     * by object identity. Keys are held weakly, values strongly.
+     * <p>
+     * Not public: register through
+     * {@link #registerUnenhancedInstance(Object, PersistenceCapable)} and
+     * resolve through {@link #toPersistenceCapable(Object, Object)}.
+     */
+    static final Map _unenhancedInstanceMap =
         new ConcurrentReferenceHashMap(ReferenceStrength.WEAK, ReferenceStrength.HARD) {
 
             @Override
@@ -313,14 +320,43 @@ else if (ctx instanceof StateManager
                 return null;
 
             pc = new ReflectingPersistenceCapable(o, conf);
-            _unenhancedInstanceMap.put(o, pc);
+            registerUnenhancedInstance(o, pc);
             return pc;
         }
     }
 
+    /**
+     * Registers the {@link PersistenceCapable} instance that manages the state
+     * of the given raw, unenhanced instance, so that subsequent calls to
+     * {@link #toPersistenceCapable(Object, Object)} for that instance resolve
+     * to <code>pc</code> instead of creating a new
+     * {@link org.apache.openjpa.enhance.ReflectingPersistenceCapable}.
+     * <p>
+     * Instances are keyed by identity rather than by {@link Object#equals},
+     * so value-based types such as records can be used as keys. Any mapping
+     * already present for the instance is replaced. <code>pc</code> is not
+     * required to report <code>instance</code> as its managed instance; for a
+     * record embeddable it is the embedded copy carrying the state manager.
+     *
+     * @param instance the raw, unenhanced managed instance
+     * @param pc the persistence-capable instance to resolve it to
+     * @since 4.2.0
+     */
+    public static void registerUnenhancedInstance(Object instance,
+        PersistenceCapable pc) {
+        _unenhancedInstanceMap.put(instance, pc);
+    }
+
+    /**
+     * Re-registers a {@link ReflectingPersistenceCapable} against the
+     * instance it manages, for example after deserialization.
+     *
+     * @param pc the persistence-capable instance to register
+     * @see #registerUnenhancedInstance(Object, PersistenceCapable)
+     */
     public static void registerPersistenceCapable(
         ReflectingPersistenceCapable pc) {
-        _unenhancedInstanceMap.put(pc.getManagedInstance(), pc);
+        registerUnenhancedInstance(pc.getManagedInstance(), pc);
     }
 
     /**

openjpa-kernel/src/test/java/org/apache/openjpa/util/TestImplHelperUnenhancedInstance.java

@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.openjpa.util;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+
+import org.apache.openjpa.enhance.PersistenceCapable;
+import org.junit.Test;
+
+/**
+ * Pins the registry contract of
+ * {@link ImplHelper#registerUnenhancedInstance(Object, PersistenceCapable)}.
+ * <p>
+ * The registrations are read back through the package private registry itself:
+ * resolving them through {@link ImplHelper#toPersistenceCapable(Object, Object)}
+ * would need a configuration and manageable metadata, which is beyond a unit
+ * test here.
+ */
+public class TestImplHelperUnenhancedInstance {
+
+    record Point(int x, int y) {
+    }
+
+    /**
+     * Records are the reason this registry exists, and two equal records must
+     * not share one registration.
+     */
+    @Test
+    public void testKeysByIdentityNotEquality() {
+        Point p1 = new Point(3, 4);
+        Point p2 = new Point(3, 4);
+        assertEquals(p1, p2);
+
+        PersistenceCapable pc1 = newPersistenceCapable();
+        PersistenceCapable pc2 = newPersistenceCapable();
+        ImplHelper.registerUnenhancedInstance(p1, pc1);
+        ImplHelper.registerUnenhancedInstance(p2, pc2);
+
+        assertSame(pc1, ImplHelper._unenhancedInstanceMap.get(p1));
+        assertSame(pc2, ImplHelper._unenhancedInstanceMap.get(p2));
+        assertNull(ImplHelper._unenhancedInstanceMap.get(new Point(3, 4)));
+    }
+
+    @Test
+    public void testLastRegistrationWins() {
+        Point p = new Point(5, 6);
+        PersistenceCapable pcA = newPersistenceCapable();
+        PersistenceCapable pcB = newPersistenceCapable();
+
+        ImplHelper.registerUnenhancedInstance(p, pcA);
+        ImplHelper.registerUnenhancedInstance(p, pcB);
+        assertSame(pcB, ImplHelper._unenhancedInstanceMap.get(p));
+    }
+
+    /**
+     * A do-nothing {@link PersistenceCapable} stub; no mock framework is
+     * available on this module's test classpath.
+     */
+    private static PersistenceCapable newPersistenceCapable() {
+        InvocationHandler handler = new InvocationHandler() {
+
+            @Override
+            public Object invoke(Object proxy, Method method, Object[] args) {
+                switch (method.getName()) {
+                    case "equals":
+                        return proxy == args[0];
+                    case "hashCode":
+                        return System.identityHashCode(proxy);
+                    case "toString":
+                        return "PersistenceCapable stub";
+                    default:
+                        return null;
+                }
+            }
+        };
+        return (PersistenceCapable) Proxy.newProxyInstance(
+            PersistenceCapable.class.getClassLoader(),
+            new Class<?>[]{ PersistenceCapable.class }, handler);
+    }
+}