Merge pull request #200 from fjuma/EJBCLIENT-193

[EJBCLIENT-193] Implement affinity substitution

src/main/java/org/jboss/ejb/protocol/remote/EJBClientChannel.java

@@ -144,10 +144,11 @@ class EJBClientChannel {
         if (version < 3) {
             configuration.setClassTable(ProtocolV1ClassTable.INSTANCE);
             configuration.setObjectTable(ProtocolV1ObjectTable.INSTANCE);
-            configuration.setObjectResolver(ProtocolV1ObjectResolver.INSTANCE);
+            configuration.setObjectResolver(new ProtocolV1ObjectResolver(channel.getConnection().getEndpoint().getName()));
             configuration.setVersion(2);
         } else {
             configuration.setObjectTable(ProtocolV3ObjectTable.INSTANCE);
+            configuration.setObjectResolver(new ProtocolV3ObjectResolver(channel.getConnection().getPeerURI()));
             configuration.setVersion(4);
             // server does not present v3 unless the transaction service is also present
         }

src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java

@@ -118,10 +118,11 @@ final class EJBServerChannel {
         if (version < 3) {
             configuration.setClassTable(ProtocolV1ClassTable.INSTANCE);
             configuration.setObjectTable(ProtocolV1ObjectTable.INSTANCE);
-            configuration.setObjectResolver(ProtocolV1ObjectResolver.INSTANCE);
+            configuration.setObjectResolver(new ProtocolV1ObjectResolver(channel.getConnection().getEndpoint().getName()));
             configuration.setVersion(2);
         } else {
             configuration.setObjectTable(ProtocolV3ObjectTable.INSTANCE);
+            configuration.setObjectResolver(new ProtocolV3ObjectResolver(channel.getConnection().getPeerURI()));
             configuration.setVersion(4);
         }
         marshallerFactory = new RiverMarshallerFactory();

src/main/java/org/jboss/ejb/protocol/remote/ProtocolV1ObjectResolver.java

@@ -21,28 +21,36 @@
 import org.jboss.ejb.client.AbstractEJBMetaData;
 import org.jboss.ejb.client.Affinity;
 import org.jboss.ejb.client.EJBMetaDataImpl;
+import org.jboss.ejb.client.NodeAffinity;
 import org.jboss.ejb.client.URIAffinity;
 import org.jboss.marshalling.ObjectResolver;
 
 /**
  * @author <a href="mailto:[email protected]">David M. Lloyd</a>
  */
 final class ProtocolV1ObjectResolver implements ObjectResolver {
-    static final ProtocolV1ObjectResolver INSTANCE = new ProtocolV1ObjectResolver();
+    private final NodeAffinity nodeAffinity;
 
-    private ProtocolV1ObjectResolver() {
+    ProtocolV1ObjectResolver(final String nodeName) {
+        nodeAffinity = new NodeAffinity(nodeName);
     }
 
     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 replacement;
     }
 
     public Object writeReplace(final Object original) {
-        if (original instanceof URIAffinity || original == Affinity.LOCAL) {
+        if (original instanceof URIAffinity) {
             return Affinity.NONE;
+        } else if (original == Affinity.LOCAL) {
+            // Swap a local affinity with a node affinity with the name of this node
+            return nodeAffinity;
         } else if (original instanceof AbstractEJBMetaData) {
             return new EJBMetaDataImpl((AbstractEJBMetaData<?, ?>) original);
         }

src/main/java/org/jboss/ejb/protocol/remote/ProtocolV3ObjectResolver.java

@@ -0,0 +1,52 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2017 Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * 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 org.jboss.ejb.protocol.remote;
+
+import java.net.URI;
+
+import org.jboss.ejb.client.Affinity;
+import org.jboss.ejb.client.URIAffinity;
+import org.jboss.marshalling.ObjectResolver;
+
+/**
+ * @author <a href="mailto:[email protected]">Farah Juma</a>
+ */
+final class ProtocolV3ObjectResolver implements ObjectResolver {
+    private final Affinity peerURIAffinity;
+
+    ProtocolV3ObjectResolver(final URI peerURI) {
+        peerURIAffinity = Affinity.forUri(peerURI);
+    }
+
+    public Object readResolve(final Object replacement) {
+        // Swap a local affinity with a URI affinity with the peer's URI
+        if (replacement == Affinity.LOCAL) {
+            return peerURIAffinity;
+        }
+        return replacement;
+    }
+
+    public Object writeReplace(final Object original) {
+        // Swap a URI affinity with the peer's URI with a local affinity
+        if ((original instanceof URIAffinity) && original.equals(peerURIAffinity)) {
+            return Affinity.LOCAL;
+        }
+        return original;
+    }
+}

src/main/java/org/jboss/ejb/protocol/remote/ProtocolV3ObjectTable.java

@@ -100,7 +100,6 @@ final class ProtocolV3ObjectTable implements ObjectTable {
         final Object[] simpleObjects = {
             TransactionID.PRIVATE_DATA_KEY,
             Affinity.NONE,
-            Affinity.LOCAL,
             Affinity.WEAK_AFFINITY_CONTEXT_KEY,
             EJBClientInvocationContext.PRIVATE_ATTACHMENTS_KEY,
             AttachmentKeys.TRANSACTION_ID_KEY,