fix: Assign method call listeners directly to the proxy instance (#2102)

Author: mykola-mokhnach

src/main/java/io/appium/java_client/proxy/HasMethodCallListeners.java

@@ -0,0 +1,35 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * See the NOTICE file distributed with this work for additional
+ * information regarding copyright ownership.
+ * 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 io.appium.java_client.proxy;
+
+public interface HasMethodCallListeners {
+    /**
+     * The setter is dynamically created by ByteBuddy to store
+     * method call listeners on the instrumented proxy instance.
+     *
+     * @param methodCallListeners Array of method call listeners assigned to the proxy instance.
+     */
+    void setMethodCallListeners(MethodCallListener[] methodCallListeners);
+
+    /**
+     * The getter is dynamically created by ByteBuddy to access
+     * method call listeners on the instrumented proxy instance.
+     *
+     * @return Array of method call listeners assigned the proxy instance.
+     */
+    MethodCallListener[] getMethodCallListeners();
+}

src/main/java/io/appium/java_client/proxy/Helpers.java

@@ -19,7 +19,9 @@
 import com.google.common.base.Preconditions;
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.modifier.Visibility;
 import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
+import net.bytebuddy.implementation.FieldAccessor;
 import net.bytebuddy.implementation.MethodDelegation;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
@@ -118,16 +120,18 @@ public static <T> T createProxy(
                 .subclass(cls)
                 .method(extraMethodMatcher == null ? matcher : matcher.and(extraMethodMatcher))
                 .intercept(MethodDelegation.to(Interceptor.class))
+                // https://github.com/raphw/byte-buddy/blob/2caef35c172897cbdd21d163c55305a64649ce41/byte-buddy-dep/src/test/java/net/bytebuddy/ByteBuddyTutorialExamplesTest.java#L346
+                .defineField("methodCallListeners", MethodCallListener[].class, Visibility.PRIVATE)
+                .implement(HasMethodCallListeners.class).intercept(FieldAccessor.ofBeanProperty())
                 .make()
                 .load(ClassLoader.getSystemClassLoader(), ClassLoadingStrategy.Default.WRAPPER)
                 .getLoaded()
                 .asSubclass(cls);
 
         try {
-            return ProxyListenersContainer.getInstance().setListeners(
-                    cls.cast(proxy.getConstructor(constructorArgTypes).newInstance(constructorArgs)),
-                    listeners
-            );
+            T result = cls.cast(proxy.getConstructor(constructorArgTypes).newInstance(constructorArgs));
+            ((HasMethodCallListeners) result).setMethodCallListeners(listeners.toArray(MethodCallListener[]::new));
+            return result;
         } catch (SecurityException | ReflectiveOperationException e) {
             throw new IllegalStateException(String.format("Unable to create a proxy of %s", cls.getName()), e);
         }

src/main/java/io/appium/java_client/proxy/Interceptor.java

@@ -25,10 +25,10 @@
 import org.slf4j.LoggerFactory;
 
 import java.lang.reflect.Method;
-import java.util.Collection;
-import java.util.UUID;
 import java.util.concurrent.Callable;
 
+import static io.appium.java_client.proxy.MethodCallListener.UNSET;
+
 public class Interceptor {
     private static final Logger LOGGER = LoggerFactory.getLogger(Interceptor.class);
 
@@ -37,7 +37,9 @@ private Interceptor() {
 
     /**
      * A magic method used to wrap public method calls in classes
-     * patched by ByteBuddy and acting as proxies.
+     * patched by ByteBuddy and acting as proxies. The performance
+     * of this method is mission-critical as it gets called upon
+     * every invocation of any method of the proxied class.
      *
      * @param self     The reference to the original instance.
      * @param method   The reference to the original method.
@@ -53,12 +55,12 @@ public static Object intercept(
             @AllArguments Object[] args,
             @SuperCall Callable<?> callable
     ) throws Throwable {
-        Collection<MethodCallListener> listeners = ProxyListenersContainer.getInstance().getListeners(self);
-        if (listeners.isEmpty()) {
+        var listeners = ((HasMethodCallListeners) self).getMethodCallListeners();
+        if (listeners == null || listeners.length == 0) {
             return callable.call();
         }
 
-        listeners.forEach(listener -> {
+        for (var listener : listeners) {
             try {
                 listener.beforeCall(self, method, args);
             } catch (NotImplementedException e) {
@@ -68,32 +70,39 @@ public static Object intercept(
                         self.getClass().getName(), method.getName(), e
                 );
             }
-        });
+        }
 
-        final UUID noResult = UUID.randomUUID();
-        Object result = noResult;
-        for (MethodCallListener listener : listeners) {
+        Object result = UNSET;
+        for (var listener : listeners) {
             try {
                 result = listener.call(self, method, args, callable);
-                break;
+                if (result != UNSET) {
+                    break;
+                }
             } catch (NotImplementedException e) {
                 // ignore
             } catch (Exception e) {
                 try {
-                    return listener.onError(self, method, args, e);
+                    result = listener.onError(self, method, args, e);
+                    if (result != UNSET) {
+                        return result;
+                    }
                 } catch (NotImplementedException ignore) {
                     // ignore
                 }
                 throw e;
             }
         }
-        if (noResult.equals(result)) {
+        if (UNSET == result) {
             try {
                 result = callable.call();
             } catch (Exception e) {
-                for (MethodCallListener listener : listeners) {
+                for (var listener : listeners) {
                     try {
-                        return listener.onError(self, method, args, e);
+                        result = listener.onError(self, method, args, e);
+                        if (result != UNSET) {
+                            return result;
+                        }
                     } catch (NotImplementedException ignore) {
                         // ignore
                     }
@@ -102,8 +111,8 @@ public static Object intercept(
             }
         }
 
-        final Object endResult = result == noResult ? null : result;
-        listeners.forEach(listener -> {
+        final Object endResult = result == UNSET ? null : result;
+        for (var listener : listeners) {
             try {
                 listener.afterCall(self, method, args, endResult);
             } catch (NotImplementedException e) {
@@ -113,7 +122,7 @@ public static Object intercept(
                         self.getClass().getName(), method.getName(), e
                 );
             }
-        });
+        }
         return endResult;
     }
 }

src/main/java/io/appium/java_client/proxy/MethodCallListener.java

@@ -17,9 +17,11 @@
 package io.appium.java_client.proxy;
 
 import java.lang.reflect.Method;
+import java.util.UUID;
 import java.util.concurrent.Callable;
 
 public interface MethodCallListener {
+    UUID UNSET = UUID.randomUUID();
 
     /**
      * The callback to be invoked before any public method of the proxy is called.
@@ -31,7 +33,6 @@ public interface MethodCallListener {
      * @param args   Array of method arguments
      */
     default void beforeCall(Object obj, Method method, Object[] args) {
-        throw new NotImplementedException();
     }
 
     /**
@@ -48,7 +49,7 @@ default void beforeCall(Object obj, Method method, Object[] args) {
      * @return The type of the returned result should be castable to the returned type of the original method.
      */
     default Object call(Object obj, Method method, Object[] args, Callable<?> original) throws Throwable {
-        throw new NotImplementedException();
+        return UNSET;
     }
 
     /**
@@ -61,7 +62,6 @@ default Object call(Object obj, Method method, Object[] args, Callable<?> origin
      * @param args   Array of method arguments
      */
     default void afterCall(Object obj, Method method, Object[] args, Object result) {
-        throw new NotImplementedException();
     }
 
     /**
@@ -77,6 +77,6 @@ default void afterCall(Object obj, Method method, Object[] args, Object result)
      *     type of the returned argument could be cast to the returned type of the original method.
      */
     default Object onError(Object obj, Method method, Object[] args, Throwable e) throws Throwable {
-        throw new NotImplementedException();
+        return UNSET;
     }
 }