[java] reuse the classes created by the WebDriverDecorator #14789

Author: joerg1985

java/src/org/openqa/selenium/support/decorators/WebDriverDecorator.java

@@ -17,16 +17,23 @@
 
 package org.openqa.selenium.support.decorators;
 
+import java.lang.reflect.Field;
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.modifier.Visibility;
 import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
 import net.bytebuddy.implementation.InvocationHandlerAdapter;
 import net.bytebuddy.matcher.ElementMatchers;
@@ -183,6 +190,59 @@
 @Beta
 public class WebDriverDecorator<T extends WebDriver> {
 
+  protected static class Definition {
+    private final Class<?> decoratedClass;
+    private final Class<?> originalClass;
+
+    public Definition(Decorated<?> decorated) {
+      this.decoratedClass = decorated.getClass();
+      this.originalClass = decorated.getOriginal().getClass();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (o == null || getClass() != o.getClass()) return false;
+      Definition definition = (Definition) o;
+      return (decoratedClass == definition.decoratedClass)
+          && (originalClass == definition.originalClass);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(decoratedClass, originalClass);
+    }
+  }
+
+  protected static class ProxyFactory<T> {
+    private final Class<? extends T> clazz;
+    private final Field targetLookup;
+
+    private ProxyFactory(Class<? extends T> clazz, Field targetLookup) {
+      this.clazz = clazz;
+      this.targetLookup = targetLookup;
+    }
+
+    public T newInstance(Decorated<T> target) {
+      T instance;
+      try {
+        instance = (T) clazz.newInstance();
+      } catch (ReflectiveOperationException e) {
+        throw new AssertionError("Unable to create new proxy", e);
+      }
+
+      try {
+        // ensure we can later find the target to call
+        targetLookup.set(instance, target);
+      } catch (IllegalAccessException ex) {
+        throw new AssertionError("Failed to set proxy target", ex);
+      }
+
+      return instance;
+    }
+  }
+
+  private final ConcurrentMap<Definition, ProxyFactory<?>> cache;
+
   private final Class<T> targetWebDriverClass;
 
   private Decorated<T> decorated;
@@ -194,6 +254,7 @@ public WebDriverDecorator() {
 
   public WebDriverDecorator(Class<T> targetClass) {
     this.targetWebDriverClass = targetClass;
+    this.cache = new ConcurrentHashMap<>();
   }
 
   public final T decorate(T original) {
@@ -295,18 +356,38 @@ private Object decorateResult(Object toDecorate) {
     return toDecorate;
   }
 
-  protected final <Z> Z createProxy(final Decorated<Z> decorated, Class<Z> clazz) {
-    Set<Class<?>> decoratedInterfaces = extractInterfaces(decorated);
-    Set<Class<?>> originalInterfaces = extractInterfaces(decorated.getOriginal());
-    Map<Class<?>, InvocationHandler> derivedInterfaces =
-        deriveAdditionalInterfaces(decorated.getOriginal());
+  protected final <Z> Z createProxy(final Decorated<Z> decorated, Class<? extends Z> clazz) {
+    @SuppressWarnings("unchecked")
+    ProxyFactory<Z> factory =
+        (ProxyFactory<Z>)
+            cache.computeIfAbsent(
+                new Definition(decorated), (key) -> createProxyFactory(key, decorated, clazz));
+
+    return factory.newInstance(decorated);
+  }
+
+  protected final <Z> ProxyFactory<? extends Z> createProxyFactory(
+      Definition definition, final Decorated<Z> sample, Class<? extends Z> clazz) {
+    Set<Class<?>> decoratedInterfaces = extractInterfaces(definition.decoratedClass);
+    Set<Class<?>> originalInterfaces = extractInterfaces(definition.originalClass);
+    // all samples with the same definition should have the same derivedInterfaces
+    Map<Class<?>, Function<Z, InvocationHandler>> derivedInterfaces =
+        deriveAdditionalInterfaces(sample.getOriginal());
+    CompletableFuture<Field> proxyLookup = new CompletableFuture<>();
 
     final InvocationHandler handler =
         (proxy, method, args) -> {
+          Field target = proxyLookup.getNow(null);
+          // Lookup the instance to call, to reuse the clazz and handler.
+          @SuppressWarnings("unchecked")
+          Decorated<Z> instance = (Decorated<Z>) target.get(proxy);
+          if (instance == null) {
+            throw new AssertionError("Failed to get instance to call");
+          }
           try {
             if (method.getDeclaringClass().equals(Object.class)
                 || decoratedInterfaces.contains(method.getDeclaringClass())) {
-              return method.invoke(decorated, args);
+              return method.invoke(instance, args);
             }
             // Check if the class in which the method resides, implements any one of the
             // interfaces that we extracted from the decorated class.
@@ -317,9 +398,9 @@ protected final <Z> Z createProxy(final Decorated<Z> decorated, Class<Z> clazz)
                             eachInterface.isAssignableFrom(method.getDeclaringClass()));
 
             if (isCompatible) {
-              decorated.beforeCall(method, args);
-              Object result = decorated.call(method, args);
-              decorated.afterCall(method, result, args);
+              instance.beforeCall(method, args);
+              Object result = instance.call(method, args);
+              instance.afterCall(method, result, args);
               return result;
             }
 
@@ -333,12 +414,15 @@ protected final <Z> Z createProxy(final Decorated<Z> decorated, Class<Z> clazz)
                             eachInterface.isAssignableFrom(method.getDeclaringClass()));
 
             if (isCompatible) {
-              return derivedInterfaces.get(method.getDeclaringClass()).invoke(proxy, method, args);
+              return derivedInterfaces
+                  .get(method.getDeclaringClass())
+                  .apply(instance.getOriginal())
+                  .invoke(proxy, method, args);
             }
 
-            return method.invoke(decorated.getOriginal(), args);
+            return method.invoke(instance.getOriginal(), args);
           } catch (InvocationTargetException e) {
-            return decorated.onError(method, e, args);
+            return instance.onError(method, e, args);
           }
         };
 
@@ -347,27 +431,45 @@ protected final <Z> Z createProxy(final Decorated<Z> decorated, Class<Z> clazz)
     allInterfaces.addAll(originalInterfaces);
     allInterfaces.addAll(derivedInterfaces.keySet());
     Class<?>[] allInterfacesArray = allInterfaces.toArray(new Class<?>[0]);
+    // ensure we do not hit an existing field
+    int fieldName;
+
+    for (fieldName = 0; fieldName < 8192; fieldName++) {
+      try {
+        clazz.getDeclaredField("___target" + fieldName);
+      } catch (NoSuchFieldException ex) {
+        break;
+      }
+    }
+
+    if (fieldName == 8192) {
+      throw new AssertionError("No free field name found");
+    }
 
     Class<? extends Z> proxy =
         new ByteBuddy()
             .subclass(clazz.isInterface() ? Object.class : clazz)
             .implement(allInterfacesArray)
             .method(ElementMatchers.any())
             .intercept(InvocationHandlerAdapter.of(handler))
+            .defineField("___target" + fieldName, sample.getClass(), Visibility.PUBLIC)
             .make()
             .load(clazz.getClassLoader(), ClassLoadingStrategy.Default.WRAPPER)
             .getLoaded()
             .asSubclass(clazz);
 
+    Field lookup;
+
     try {
-      return proxy.newInstance();
-    } catch (ReflectiveOperationException e) {
-      throw new IllegalStateException("Unable to create new proxy", e);
+      lookup = proxy.getDeclaredField("___target" + fieldName);
+    } catch (NoSuchFieldException e) {
+      proxyLookup.completeExceptionally(e);
+      throw new AssertionError("Defined field is missing", e);
     }
-  }
 
-  static Set<Class<?>> extractInterfaces(final Object object) {
-    return extractInterfaces(object.getClass());
+    proxyLookup.complete(lookup);
+
+    return new ProxyFactory<Z>(proxy, lookup);
   }
 
   private static Set<Class<?>> extractInterfaces(final Class<?> clazz) {
@@ -393,43 +495,46 @@ private static void extractInterfaces(final Set<Class<?>> collector, final Class
     extractInterfaces(collector, clazz.getSuperclass());
   }
 
-  private Map<Class<?>, InvocationHandler> deriveAdditionalInterfaces(Object object) {
-    Map<Class<?>, InvocationHandler> handlers = new HashMap<>();
+  private <Z> Map<Class<?>, Function<Z, InvocationHandler>> deriveAdditionalInterfaces(Z sample) {
+    Map<Class<?>, Function<Z, InvocationHandler>> handlers = new HashMap<>();
 
-    if (object instanceof WebDriver && !(object instanceof WrapsDriver)) {
+    if (sample instanceof WebDriver && !(sample instanceof WrapsDriver)) {
       handlers.put(
           WrapsDriver.class,
-          (proxy, method, args) -> {
-            if ("getWrappedDriver".equals(method.getName())) {
-              return object;
-            }
-            throw new UnsupportedOperationException(method.getName());
-          });
+          (instance) ->
+              (proxy, method, args) -> {
+                if ("getWrappedDriver".equals(method.getName())) {
+                  return instance;
+                }
+                throw new UnsupportedOperationException(method.getName());
+              });
     }
 
-    if (object instanceof WebElement && !(object instanceof WrapsElement)) {
+    if (sample instanceof WebElement && !(sample instanceof WrapsElement)) {
       handlers.put(
           WrapsElement.class,
-          (proxy, method, args) -> {
-            if ("getWrappedElement".equals(method.getName())) {
-              return object;
-            }
-            throw new UnsupportedOperationException(method.getName());
-          });
+          (instance) ->
+              (proxy, method, args) -> {
+                if ("getWrappedElement".equals(method.getName())) {
+                  return instance;
+                }
+                throw new UnsupportedOperationException(method.getName());
+              });
     }
 
     try {
-      Method toJson = object.getClass().getDeclaredMethod("toJson");
+      Method toJson = sample.getClass().getDeclaredMethod("toJson");
       toJson.setAccessible(true);
 
       handlers.put(
           JsonSerializer.class,
-          ((proxy, method, args) -> {
-            if ("toJson".equals(method.getName())) {
-              return toJson.invoke(object);
-            }
-            throw new UnsupportedOperationException(method.getName());
-          }));
+          (instance) ->
+              ((proxy, method, args) -> {
+                if ("toJson".equals(method.getName())) {
+                  return toJson.invoke(instance);
+                }
+                throw new UnsupportedOperationException(method.getName());
+              }));
     } catch (NoSuchMethodException e) {
       // Fine. Just fall through
     }

java/test/org/openqa/selenium/support/decorators/DecoratedWebDriverTest.java

@@ -31,6 +31,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.UUID;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import org.junit.jupiter.api.Tag;
@@ -163,6 +164,45 @@ void findElement() {
     verifyDecoratingFunction($ -> $.findElement(By.id("test")), found, WebElement::click);
   }
 
+  @Test
+  void doesNotCreateTooManyClasses() {
+    final WebElement found = mock(WebElement.class);
+    Function<WebDriver, WebElement> f = $ -> $.findElement(By.id("test"));
+    Fixture fixture = new Fixture();
+    when(f.apply(fixture.original)).thenReturn(found);
+
+    WebElement proxy1 = f.apply(fixture.decorated);
+    WebElement proxy2 = f.apply(fixture.decorated);
+    WebElement proxy3 = f.apply(fixture.decorated);
+
+    assertThat(proxy1.getClass()).isSameAs(proxy2.getClass());
+    assertThat(proxy1.getClass()).isSameAs(proxy3.getClass());
+  }
+
+  @Test
+  void doesHitTheCorrectInstance() {
+    String uuid0 = UUID.randomUUID().toString();
+    String uuid1 = UUID.randomUUID().toString();
+    String uuid2 = UUID.randomUUID().toString();
+    final WebElement found0 = mock(WebElement.class);
+    final WebElement found1 = mock(WebElement.class);
+    final WebElement found2 = mock(WebElement.class);
+    when(found0.getTagName()).thenReturn(uuid0);
+    when(found1.getTagName()).thenReturn(uuid1);
+    when(found2.getTagName()).thenReturn(uuid2);
+
+    Fixture fixture = new Fixture();
+    Function<WebDriver, List<WebElement>> f = $ -> $.findElements(By.id("test"));
+
+    when(f.apply(fixture.original)).thenReturn(List.of(found0, found1, found2));
+
+    List<WebElement> proxies = f.apply(fixture.decorated);
+
+    assertThat(proxies.get(0).getTagName()).isEqualTo(uuid0);
+    assertThat(proxies.get(1).getTagName()).isEqualTo(uuid1);
+    assertThat(proxies.get(2).getTagName()).isEqualTo(uuid2);
+  }
+
   @Test
   void findElementNotFound() {
     Fixture fixture = new Fixture();