Merge pull request wildfly#18112 from pferraro/WFLY-19613

WFLY-19613 Immutability performance optimizations
ropalka · Aug 13, 2024 · 6fc93af · 6fc93af
2 parents 9441016 + fc8a193
commit 6fc93af
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 9 deletions.
diff --git a/...a/org/wildfly/clustering/web/cache/session/attributes/coarse/CoarseSessionAttributes.java b/...a/org/wildfly/clustering/web/cache/session/attributes/coarse/CoarseSessionAttributes.java
@@ -67,7 +67,8 @@ public Object setAttribute(String name, Object value) {
     @Override
     public Object getAttribute(String name) {
         Object value = this.attributes.get(name);
-        if (!this.immutability.test(value)) {
+        // Bypass immutability check if session is already dirty
+        if (!this.dirty.get() && !this.immutability.test(value)) {
             this.dirty.set(true);
         }
         return value;

diff --git a/.../java/org/wildfly/clustering/web/cache/session/attributes/fine/FineSessionAttributes.java b/.../java/org/wildfly/clustering/web/cache/session/attributes/fine/FineSessionAttributes.java
@@ -55,8 +55,9 @@ public Object getAttribute(String name) {
 
         if (value != null) {
             // If the object is mutable, we need to mutate this value on close
-            if (!this.immutability.test(value)) {
-                synchronized (this.updates) {
+            synchronized (this.updates) {
+                // Bypass immutability check if we are already updating this attribute
+                if (!this.updates.containsKey(name) && !this.immutability.test(value)) {
                     this.updates.put(name, value);
                 }
             }

diff --git a/ejb3/src/main/java/org/jboss/as/ejb3/component/stateful/StatefulSessionBeanImmutability.java b/ejb3/src/main/java/org/jboss/as/ejb3/component/stateful/StatefulSessionBeanImmutability.java
@@ -5,27 +5,43 @@
 
 package org.jboss.as.ejb3.component.stateful;
 
+import java.lang.reflect.Field;
 import java.lang.reflect.InvocationHandler;
 
 import org.jboss.as.ee.component.ProxyInvocationHandler;
 import org.jboss.invocation.proxy.ProxyFactory;
 import org.kohsuke.MetaInfServices;
 import org.wildfly.clustering.ee.Immutability;
+import org.wildfly.security.ParametricPrivilegedAction;
+import org.wildfly.security.manager.WildFlySecurityManager;
 
 /**
  * Immutability test for EJB proxies, whose serializable placeholders are immutable.
  * @author Paul Ferraro
  */
 @MetaInfServices(Immutability.class)
 public class StatefulSessionBeanImmutability implements Immutability {
+    private static final ParametricPrivilegedAction<InvocationHandler, Object> GET_INVOCATION_HANDLER = new ParametricPrivilegedAction<>() {
+        @Override
+        public InvocationHandler run(Object object) {
+            // Since there is a low probability that this object is actually a SFSB proxy, avoid Class.getDeclaredField(...) which will typically throw/catch a NoSuchFieldException
+            for (Field field : object.getClass().getDeclaredFields()) {
+                if (field.getName().equals(ProxyFactory.INVOCATION_HANDLER_FIELD)) {
+                    field.setAccessible(true);
+                    try {
+                        return (InvocationHandler) field.get(object);
+                    } catch (IllegalAccessException e) {
+                        throw new IllegalStateException(e);
+                    }
+                }
+            }
+            return null;
+        }
+    };
 
     @Override
     public boolean test(Object object) {
-        try {
-            InvocationHandler handler = ProxyFactory.getInvocationHandlerStatic(object);
-            return handler instanceof ProxyInvocationHandler;
-        } catch (RuntimeException e) {
-            return false;
-        }
+        InvocationHandler handler = WildFlySecurityManager.doUnchecked(object, GET_INVOCATION_HANDLER);
+        return handler instanceof ProxyInvocationHandler;
     }
 }