Merge pull request #236 from stuartwdouglas/no-affinity-invocations

Also replace the affinity of deserialized locators if it is NONE for…
wildfly · Mar 16, 2017 · 2b5d451 · 2b5d451
2 parents a3053cf + c83e880
commit 2b5d451
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 6 deletions.
diff --git a/src/main/java/org/jboss/ejb/protocol/remote/EJBClientChannel.java b/src/main/java/org/jboss/ejb/protocol/remote/EJBClientChannel.java
@@ -169,7 +169,7 @@ public Class<?> resolveProxyClass(final Unmarshaller unmarshaller, final String[
         if (version < 3) {
             configuration.setClassTable(ProtocolV1ClassTable.INSTANCE);
             configuration.setObjectTable(ProtocolV1ObjectTable.INSTANCE);
-            configuration.setObjectResolver(new ProtocolV1ObjectResolver(channel.getConnection().getEndpoint().getName()));
+            configuration.setObjectResolver(new ProtocolV1ObjectResolver(channel.getConnection().getEndpoint().getName(), channel.getConnection().getPeerURI()));
             configuration.setVersion(2);
             // Do not wait for cluster topology report.
             finishedParts.set(0b10);

diff --git a/src/main/java/org/jboss/ejb/protocol/remote/EJBMarshallingCompatibilityHelper.java b/src/main/java/org/jboss/ejb/protocol/remote/EJBMarshallingCompatibilityHelper.java
@@ -35,7 +35,7 @@ public ObjectResolver getObjectResolver(final Transport transport, final boolean
         if (transport instanceof RemoteTransport) {
             final RemoteTransport remoteTransport = (RemoteTransport) transport;
             if (remoteTransport.getVersion() == 1) {
-                return new ProtocolV1ObjectResolver(remoteTransport.getConnection().getEndpoint().getName());
+                return new ProtocolV1ObjectResolver(remoteTransport.getConnection().getEndpoint().getName(), remoteTransport.getConnection().getPeerURI());
             } else if (remoteTransport.getVersion() == 2) {
                 return new ProtocolV3ObjectResolver(remoteTransport.getConnection().getPeerURI());
             }

diff --git a/src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java b/src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java
@@ -72,7 +72,6 @@
 import org.jboss.marshalling.MarshallerFactory;
 import org.jboss.marshalling.Marshalling;
 import org.jboss.marshalling.MarshallingConfiguration;
-import org.jboss.marshalling.UTFUtils;
 import org.jboss.marshalling.Unmarshaller;
 import org.jboss.marshalling.river.RiverMarshallerFactory;
 import org.jboss.remoting3.Channel;
@@ -117,7 +116,7 @@ final class EJBServerChannel {
         if (version < 3) {
             configuration.setClassTable(ProtocolV1ClassTable.INSTANCE);
             configuration.setObjectTable(ProtocolV1ObjectTable.INSTANCE);
-            configuration.setObjectResolver(new ProtocolV1ObjectResolver(channel.getConnection().getEndpoint().getName()));
+            configuration.setObjectResolver(new ProtocolV1ObjectResolver(channel.getConnection().getEndpoint().getName(), channel.getConnection().getPeerURI()));
             configuration.setVersion(2);
         } else {
             configuration.setObjectTable(ProtocolV3ObjectTable.INSTANCE);

diff --git a/src/main/java/org/jboss/ejb/protocol/remote/ProtocolV1ObjectResolver.java b/src/main/java/org/jboss/ejb/protocol/remote/ProtocolV1ObjectResolver.java
@@ -18,6 +18,8 @@
 
 package org.jboss.ejb.protocol.remote;
 
+import java.net.URI;
+
 import org.jboss.ejb.client.AbstractEJBMetaData;
 import org.jboss.ejb.client.Affinity;
 import org.jboss.ejb.client.EJBMetaDataImpl;
@@ -30,18 +32,25 @@
  * @author <a href="mailto:david.lloyd@redhat.com">David M. Lloyd</a>
  */
 final class ProtocolV1ObjectResolver implements ObjectResolver {
+
+    private static final boolean DISABLE_V1_AFFINITY_REWRITE = Boolean.getBoolean("org.jboss.ejb.client.disable-v1-affinity-rewrite");
+
+    private final Affinity peerURIAffinity;
     private final NodeAffinity nodeAffinity;
 
-    ProtocolV1ObjectResolver(final String nodeName) {
+    ProtocolV1ObjectResolver(final String nodeName, URI peerURI) {
         nodeAffinity = new NodeAffinity(nodeName);
+        this.peerURIAffinity = Affinity.forUri(peerURI);
     }
 
     public Object readResolve(final Object replacement) {
         if (replacement instanceof EJBMetaDataImpl) {
             return ((EJBMetaDataImpl) replacement).toAbstractEJBMetaData();
         } else if ((replacement instanceof NodeAffinity) && replacement.equals(nodeAffinity)) {
             // Swap a node affinity with the name of this node with a local affinity
-            return Affinity.LOCAL;
+            return peerURIAffinity;
+        } else if(replacement == Affinity.NONE && !DISABLE_V1_AFFINITY_REWRITE) {
+            return peerURIAffinity;
         }
         return replacement;
     }