Merge pull request #24586 from vespa-engine/jonmv/use-classes-as-keys

Jonmv/use classes as keys

3 files changed, 102 insertions, 10 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/collections/MethodCache.java b/vespajlib/src/main/java/com/yahoo/collections/MethodCache.java
index bf9200efb2e..700f16ee519 100644
--- a/vespajlib/src/main/java/com/yahoo/collections/MethodCache.java
+++ b/vespajlib/src/main/java/com/yahoo/collections/MethodCache.java
@@ -16,7 +16,7 @@ import java.util.function.Consumer;
 public final class MethodCache {
 
     private final String methodName;
-    private final CopyOnWriteHashMap<String, Method> cache = new CopyOnWriteHashMap<>();
+    private final CopyOnWriteHashMap<String, Pair<Class<?>, Method>> cache = new CopyOnWriteHashMap<>();
 
     public MethodCache(String methodName) {
         this.methodName = methodName;
@@ -33,18 +33,27 @@ public final class MethodCache {
     public Method get(Object object) {
         return get(object, null);
     }
+
     public Method get(Object object, Consumer<String> onPut) {
-        Method m = cache.get(object.getClass().getName());
-        if (m == null) {
-            m = lookupMethod(object);
-            if (m != null) {
-                if (onPut != null)
-                    onPut.accept(object.getClass().getName());
-                cache.put(object.getClass().getName(), m);
-            }
+        Pair<Class<?>, Method> pair = cache.get(object.getClass().getName());
+        // When changing bundles, you might end up having cached the old method pointing to the old bundle.
+        // That will then lead to a class cast exception when invoking the wrong clone method.
+        // Whenever we detect a new class with the same name, we therefore drop the entire cache. 
+        // This is also the reason for caching the pair of method and original class—not just the method.
+        if (pair != null && pair.getFirst() != object.getClass()) {
+            cache.clear();
+            pair = null;
+        }
+        Method method = pair == null ? null : pair.getSecond();
+        if (pair == null) {
+            method = lookupMethod(object);
+            cache.put(object.getClass().getName(), new Pair<>(object.getClass(), method));
+            if (onPut != null)
+                onPut.accept(object.getClass().getName());
         }
-        return m;
+        return method;
     }
+
     private Method lookupMethod(Object object)  {
         try {
             return object.getClass().getMethod(methodName);
diff --git a/vespajlib/src/test/java/com/yahoo/collections/MethodCacheTest.java b/vespajlib/src/test/java/com/yahoo/collections/MethodCacheTest.java
new file mode 100644
index 00000000000..02efd10e91d
--- /dev/null
+++ b/vespajlib/src/test/java/com/yahoo/collections/MethodCacheTest.java
@@ -0,0 +1,82 @@
+package com.yahoo.collections;
+
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Method;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class MethodCacheTest {
+
+    @Test
+    void testCache() throws Exception {
+
+        URL url = MethodCacheTest.class.getClassLoader().getResource("dummy").toURI().resolve(".").toURL();
+
+        class MyLoader extends URLClassLoader {
+            MyLoader() { super(new URL[] { url }, MethodCacheTest.class.getClassLoader()); }
+            public Class<?> loadClass(String name) throws ClassNotFoundException {
+                if (name.equals(Dummy.class.getName())) synchronized (getClassLoadingLock(name)) { return findClass(name); }
+                else return super.loadClass(name);
+            }
+        }
+
+        try (MyLoader myLoader = new MyLoader()) {
+            Class<?> applicationClass = Dummy.class;
+            Class<?> customClass = myLoader.loadClass(Dummy.class.getName());
+
+            assertNotSame(applicationClass, customClass);
+            assertSame(applicationClass.getName(), customClass.getName());
+
+            MethodCache methods = new MethodCache("clone");
+            AtomicBoolean updatedCache = new AtomicBoolean();
+            Object applicationDummy = applicationClass.getConstructor().newInstance();
+            Object customDummy = customClass.getConstructor().newInstance();
+
+            Method applicationMethod = methods.get(applicationDummy, __ -> updatedCache.set(true));
+            assertTrue(updatedCache.getAndSet(false), "cache was updated");
+
+            Method cachedApplicationMethod = methods.get(applicationDummy, __ -> updatedCache.set(true));
+            assertFalse(updatedCache.getAndSet(false), "cache was updated");
+
+            Method customMethod = methods.get(customDummy, __ -> updatedCache.set(true));
+            assertTrue(updatedCache.getAndSet(false), "cache was updated");
+
+            Method cachedCustomMethod = methods.get(customDummy, __ -> updatedCache.set(true));
+            assertFalse(updatedCache.getAndSet(false), "cache was updated");
+
+            assertSame(applicationMethod, cachedApplicationMethod);
+            assertNotSame(applicationMethod, customMethod);
+            assertSame(customMethod, cachedCustomMethod);
+
+            cachedApplicationMethod.invoke(applicationDummy);
+            cachedCustomMethod.invoke(customDummy);
+            assertThrows(IllegalArgumentException.class, () -> applicationMethod.invoke(customDummy));
+            assertThrows(IllegalArgumentException.class, () -> customMethod.invoke(applicationDummy));
+
+            Object noDummy = new NoDummy();
+            Method noMethod = methods.get(noDummy, __ -> updatedCache.set(true));
+            assertTrue(updatedCache.getAndSet(false), "cache was updated");
+            assertNull(noMethod);
+
+            Method cachedNoMethod = methods.get(noDummy, __ -> updatedCache.set(true));
+            assertFalse(updatedCache.getAndSet(false), "cache was updated");
+            assertNull(cachedNoMethod);
+        }
+    }
+
+    public static class NoDummy implements Cloneable { }
+
+    public static class Dummy implements Cloneable {
+        public Object clone() throws CloneNotSupportedException { return super.clone(); }
+    }
+
+}
diff --git a/vespajlib/src/test/resources/dummy b/vespajlib/src/test/resources/dummy
new file mode 100644
index 00000000000..421376db9e8
--- /dev/null
+++ b/vespajlib/src/test/resources/dummy
@@ -0,0 +1 @@
+dummy