Merge pull request #313 from dmlloyd/ejbclient-262

[EJBCLIENT-262] Reset "strong" affinity when it is set as a part of discovery
wildfly · Sep 13, 2017 · 7c9c275 · 7c9c275
2 parents 179af90 + 03ce83f
commit 7c9c275
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 9 deletions.
diff --git a/src/main/java/org/jboss/ejb/client/DiscoveryEJBClientInterceptor.java b/src/main/java/org/jboss/ejb/client/DiscoveryEJBClientInterceptor.java
@@ -81,6 +81,11 @@ public DiscoveryEJBClientInterceptor() {
     }
 
     public void handleInvocation(final EJBClientInvocationContext context) throws Exception {
+        if (context.getDestination() != null) {
+            // already discovered!
+            context.sendRequest();
+            return;
+        }
         List<Throwable> problems = executeDiscovery(context);
         try {
             context.sendRequest();
@@ -131,6 +136,10 @@ public Object handleInvocationResult(final EJBClientInvocationContext context) t
     }
 
     public SessionID handleSessionCreation(final EJBSessionCreationInvocationContext context) throws Exception {
+        if (context.getDestination() != null) {
+            // already discovered!
+            return context.proceed();
+        }
         List<Throwable> problems = executeDiscovery(context);
         SessionID sessionID;
         try {
@@ -222,10 +231,7 @@ Discovery getDiscovery() {
     }
 
     private List<Throwable> executeDiscovery(AbstractInvocationContext context) {
-        if (context.getDestination() != null) {
-            // Someone else discovered already!
-            return null;
-        }
+        assert context.getDestination() == null;
         final EJBLocator<?> locator = context.getLocator();
         final Affinity affinity = locator.getAffinity();
         final Affinity weakAffinity = context.getWeakAffinity();

diff --git a/src/main/java/org/jboss/ejb/client/NamingEJBClientInterceptor.java b/src/main/java/org/jboss/ejb/client/NamingEJBClientInterceptor.java
@@ -44,12 +44,15 @@ public final class NamingEJBClientInterceptor implements EJBClientInterceptor {
      */
     public static final int PRIORITY = ClientInterceptorPriority.JBOSS_AFTER + 50;
 
+    private static final AttachmentKey<Boolean> SKIP_MISSING_TARGET = new AttachmentKey<>();
+
     public NamingEJBClientInterceptor() {
     }
 
     public void handleInvocation(final EJBClientInvocationContext context) throws Exception {
         final NamingProvider namingProvider = context.getProxyAttachment(EJBRootContext.NAMING_PROVIDER_ATTACHMENT_KEY);
-        if (namingProvider == null) {
+        if (namingProvider == null || context.getDestination() != null || ! isNoneOrCluster(context.getLocator().getAffinity())) {
+            context.putAttachment(SKIP_MISSING_TARGET, Boolean.TRUE);
             context.sendRequest();
         } else {
             if (setDestination(context, namingProvider)) try {
@@ -63,18 +66,26 @@ 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();
         } catch (NoSuchEJBException | RequestSendFailedException e) {
-            processMissingTarget(context);
+            if (context.getAttachment(SKIP_MISSING_TARGET) != Boolean.TRUE) {
+                processMissingTarget(context);
+            }
             throw e;
+        } finally {
+            context.removeAttachment(SKIP_MISSING_TARGET);
         }
     }
 
     public SessionID handleSessionCreation(final EJBSessionCreationInvocationContext context) throws Exception {
         final NamingProvider namingProvider = context.getAttachment(EJBRootContext.NAMING_PROVIDER_ATTACHMENT_KEY);
-        if (namingProvider == null) {
+        if (namingProvider == null || context.getDestination() != null || ! isNoneOrCluster(context.getLocator().getAffinity())) {
             return context.proceed();
         } else {
             if (setDestination(context, namingProvider)) try {
@@ -123,14 +134,19 @@ private static boolean setDestination(final AbstractInvocationContext context, f
 
     private void processMissingTarget(final AbstractInvocationContext context) {
         final URI destination = context.getDestination();
-
         if (destination == null) {
-            // nothing we can/should do.
+            // 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);