[EJBCLIENT-262] Reset "strong" affinity when it is set as a part of d…

…iscovery
wildfly · Sep 12, 2017 · 03ce83f · 03ce83f
1 parent 4ae9441
commit 03ce83f
Showing 1 changed file with 15 additions and 3 deletions.
diff --git a/src/main/java/org/jboss/ejb/client/NamingEJBClientInterceptor.java b/src/main/java/org/jboss/ejb/client/NamingEJBClientInterceptor.java
@@ -51,7 +51,7 @@ public NamingEJBClientInterceptor() {
 
     public void handleInvocation(final EJBClientInvocationContext context) throws Exception {
         final NamingProvider namingProvider = context.getProxyAttachment(EJBRootContext.NAMING_PROVIDER_ATTACHMENT_KEY);
-        if (namingProvider == null || context.getDestination() != null) {
+        if (namingProvider == null || context.getDestination() != null || ! isNoneOrCluster(context.getLocator().getAffinity())) {
             context.putAttachment(SKIP_MISSING_TARGET, Boolean.TRUE);
             context.sendRequest();
         } else {
@@ -66,6 +66,10 @@ public void handleInvocation(final EJBClientInvocationContext context) throws Ex
         }
     }
 
+    private boolean isNoneOrCluster(final Affinity affinity) {
+        return affinity == Affinity.NONE || affinity instanceof ClusterAffinity;
+    }
+
     public Object handleInvocationResult(final EJBClientInvocationContext context) throws Exception {
         try {
             return context.getResult();
@@ -81,7 +85,7 @@ public Object handleInvocationResult(final EJBClientInvocationContext context) t
 
     public SessionID handleSessionCreation(final EJBSessionCreationInvocationContext context) throws Exception {
         final NamingProvider namingProvider = context.getAttachment(EJBRootContext.NAMING_PROVIDER_ATTACHMENT_KEY);
-        if (namingProvider == null || context.getDestination() != null) {
+        if (namingProvider == null || context.getDestination() != null || ! isNoneOrCluster(context.getLocator().getAffinity())) {
             return context.proceed();
         } else {
             if (setDestination(context, namingProvider)) try {
@@ -130,11 +134,19 @@ private static boolean setDestination(final AbstractInvocationContext context, f
 
     private void processMissingTarget(final AbstractInvocationContext context) {
         final URI destination = context.getDestination();
-        assert destination != null;
+        if (destination == null) {
+            // some later interceptor cleared it out on us
+            return;
+        }
 
         // Oops, we got some wrong information!
         addBlackListedDestination(context, destination);
 
+        final EJBLocator<?> locator = context.getLocator();
+        if (! (locator.getAffinity() instanceof ClusterAffinity)) {
+            // it *was* "none" affinity, but it has been relocated; locate it back again
+            context.setLocator(locator.withNewAffinity(Affinity.NONE));
+        }
         // clear the weak affinity so that cluster invocations can be re-targeted.
         context.setWeakAffinity(Affinity.NONE);
         context.setTargetAffinity(null);