Merge pull request #360 from dmlloyd/ejbclient-310

[EJBCLIENT-310] Restore naming context (if any) for duration of deserialization
wildfly · Oct 16, 2018 · 0062b9d · 0062b9d
2 parents 9a9f8e2 + ac71d6c
commit 0062b9d
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 14 deletions.
diff --git a/src/main/java/org/jboss/ejb/_private/Keys.java b/src/main/java/org/jboss/ejb/_private/Keys.java
@@ -0,0 +1,30 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2018 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._private;
+
+import org.jboss.ejb.client.AttachmentKey;
+import org.wildfly.naming.client.NamingProvider;
+
+/**
+ */
+public final class Keys {
+    public static final AttachmentKey<NamingProvider> NAMING_PROVIDER_ATTACHMENT_KEY = new AttachmentKey<>();
+
+    private Keys() {}
+}
diff --git a/src/main/java/org/jboss/ejb/client/DiscoveryEJBClientInterceptor.java b/src/main/java/org/jboss/ejb/client/DiscoveryEJBClientInterceptor.java
@@ -44,6 +44,7 @@
 
 import javax.ejb.NoSuchEJBException;
 
+import org.jboss.ejb._private.Keys;
 import org.jboss.ejb._private.Logs;
 import org.jboss.ejb.client.annotation.ClientInterceptorPriority;
 import org.wildfly.common.Assert;
@@ -512,7 +513,7 @@ private List<Throwable> doClusterDiscovery(AbstractInvocationContext context, fi
 
             Logs.INVOCATION.tracef("Performed cluster discovery, nodes is empty; trying an initial ");
 
-            final NamingProvider namingProvider = context.getAttachment(EJBRootContext.NAMING_PROVIDER_ATTACHMENT_KEY);
+            final NamingProvider namingProvider = context.getAttachment(Keys.NAMING_PROVIDER_ATTACHMENT_KEY);
             if (namingProvider != null) {
                 NamingEJBClientInterceptor.setNamingDestination(context, namingProvider);
             }

diff --git a/src/main/java/org/jboss/ejb/client/EJBClientContext.java b/src/main/java/org/jboss/ejb/client/EJBClientContext.java
@@ -25,7 +25,6 @@
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.lang.reflect.Method;
-import java.lang.reflect.UndeclaredThrowableException;
 import java.net.URI;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
@@ -43,6 +42,7 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
+import org.jboss.ejb._private.Keys;
 import org.jboss.ejb._private.Logs;
 import org.jboss.ejb.protocol.remote.RemotingEJBClientInterceptor;
 import org.wildfly.common.Assert;
@@ -823,7 +823,7 @@ <T> EJBSessionCreationInvocationContext createSessionCreationInvocationContext(S
 
     <T> StatefulEJBLocator<T> createSession(final EJBSessionCreationInvocationContext context, final StatelessEJBLocator<T> statelessLocator, final NamingProvider namingProvider) throws Exception {
         // Special hook for naming; let's replace this sometime soon.
-        if (namingProvider != null) context.putAttachment(EJBRootContext.NAMING_PROVIDER_ATTACHMENT_KEY, namingProvider);
+        if (namingProvider != null) context.putAttachment(Keys.NAMING_PROVIDER_ATTACHMENT_KEY, namingProvider);
 
         Logs.INVOCATION.tracef("Calling createSession(locator = %s)",statelessLocator);
 

diff --git a/src/main/java/org/jboss/ejb/client/EJBRootContext.java b/src/main/java/org/jboss/ejb/client/EJBRootContext.java
@@ -30,6 +30,7 @@
 import javax.naming.NameClassPair;
 import javax.naming.NamingException;
 
+import org.jboss.ejb._private.Keys;
 import org.jboss.ejb._private.Logs;
 import org.jboss.ejb.client.legacy.JBossEJBProperties;
 import org.wildfly.naming.client.AbstractContext;
@@ -42,8 +43,6 @@
 
 class EJBRootContext extends AbstractContext {
 
-    static final AttachmentKey<NamingProvider> NAMING_PROVIDER_ATTACHMENT_KEY = new AttachmentKey<>();
-
     private static final String PROPERTY_KEY_INVOCATION_TIMEOUT = "invocation.timeout";
     private static final String LEARNED_AFFINITY_KEY = "__jboss.learned-affinity";
 
@@ -162,7 +161,7 @@ protected Object lookupNative(final Name name) throws NamingException {
         } else {
             proxy = EJBClient.createProxy(statelessLocator, providerEnvironment.getAuthenticationContextSupplier());
         }
-        if (namingProvider != null) EJBClient.putProxyAttachment(proxy, NAMING_PROVIDER_ATTACHMENT_KEY, namingProvider);
+        if (namingProvider != null) EJBClient.putProxyAttachment(proxy, Keys.NAMING_PROVIDER_ATTACHMENT_KEY, namingProvider);
 
         if (baseAffinity.isUnset()) {
             EJBClient.putProxyAttachment(proxy, ClusterAffinityInterest.KEY, baseAffinity);

diff --git a/src/main/java/org/jboss/ejb/client/NamingEJBClientInterceptor.java b/src/main/java/org/jboss/ejb/client/NamingEJBClientInterceptor.java
@@ -30,6 +30,7 @@
 
 import javax.ejb.NoSuchEJBException;
 
+import org.jboss.ejb._private.Keys;
 import org.jboss.ejb._private.Logs;
 import org.jboss.ejb.client.annotation.ClientInterceptorPriority;
 import org.wildfly.naming.client.NamingProvider;
@@ -53,10 +54,10 @@ public NamingEJBClientInterceptor() {
     }
 
     public void handleInvocation(final EJBClientInvocationContext context) throws Exception {
-        final NamingProvider namingProvider = context.getProxyAttachment(EJBRootContext.NAMING_PROVIDER_ATTACHMENT_KEY);
+        final NamingProvider namingProvider = context.getProxyAttachment(Keys.NAMING_PROVIDER_ATTACHMENT_KEY);
         if (namingProvider != null) {
             // make sure the naming provider is available to invocations
-            context.putAttachment(EJBRootContext.NAMING_PROVIDER_ATTACHMENT_KEY, namingProvider);
+            context.putAttachment(Keys.NAMING_PROVIDER_ATTACHMENT_KEY, namingProvider);
         }
         if (namingProvider == null || context.getDestination() != null || context.getLocator().getAffinity() != Affinity.NONE) {
             context.putAttachment(SKIP_MISSING_TARGET, Boolean.TRUE);
@@ -87,7 +88,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);
+        final NamingProvider namingProvider = context.getAttachment(Keys.NAMING_PROVIDER_ATTACHMENT_KEY);
         if (namingProvider == null || context.getDestination() != null || context.getLocator().getAffinity() != Affinity.NONE) {
             return context.proceed();
         } else {

diff --git a/src/main/java/org/jboss/ejb/protocol/remote/EJBClientChannel.java b/src/main/java/org/jboss/ejb/protocol/remote/EJBClientChannel.java
@@ -55,6 +55,7 @@
 import javax.transaction.Transaction;
 import javax.transaction.xa.Xid;
 
+import org.jboss.ejb._private.Keys;
 import org.jboss.ejb._private.Logs;
 import org.jboss.ejb.client.AbstractInvocationContext;
 import org.jboss.ejb.client.Affinity;
@@ -93,7 +94,9 @@
 import org.jboss.remoting3.util.InvocationTracker;
 import org.jboss.remoting3.util.StreamUtils;
 import org.wildfly.common.Assert;
+import org.wildfly.common.function.ExceptionBiFunction;
 import org.wildfly.common.net.CidrAddress;
+import org.wildfly.naming.client.NamingProvider;
 import org.wildfly.transaction.client.AbstractTransaction;
 import org.wildfly.transaction.client.LocalTransaction;
 import org.wildfly.transaction.client.RemoteTransaction;
@@ -1005,6 +1008,7 @@ private void handleResponse(final int id, final DataInputStream inputStream) {
             switch (id) {
                 case Protocol.INVOCATION_RESPONSE: {
                     free();
+                    final EJBClientInvocationContext context = receiverInvocationContext.getClientInvocationContext();
                     if (version >= 3) try {
                         final int cmd = inputStream.readUnsignedByte();
                         final XAOutflowHandle outflowHandle = getOutflowHandle();
@@ -1017,7 +1021,6 @@ private void handleResponse(final int id, final DataInputStream inputStream) {
                                 outflowHandle.nonMasterEnlistment();
                             }
                         }
-                        final EJBClientInvocationContext context = receiverInvocationContext.getClientInvocationContext();
                         final int updateBits = inputStream.readUnsignedByte();
                         if (allAreSet(updateBits, Protocol.UPDATE_BIT_SESSION_ID)) {
                             byte[] encoded = new byte[PackedInteger.readPackedInteger(inputStream)];
@@ -1042,7 +1045,8 @@ private void handleResponse(final int id, final DataInputStream inputStream) {
                         safeClose(inputStream);
                         break;
                     }
-                    receiverInvocationContext.resultReady(new MethodCallResultProducer(inputStream, id));
+                    final NamingProvider provider = context.getProxyAttachment(Keys.NAMING_PROVIDER_ATTACHMENT_KEY);
+                    receiverInvocationContext.resultReady(new MethodCallResultProducer(provider, inputStream, id));
                     break;
                 }
                 case Protocol.CANCEL_RESPONSE: {
@@ -1179,17 +1183,19 @@ void setOutflowHandle(final XAOutflowHandle outflowHandle) {
             this.outflowHandle = outflowHandle;
         }
 
-        class MethodCallResultProducer implements EJBReceiverInvocationContext.ResultProducer {
+        class MethodCallResultProducer implements EJBReceiverInvocationContext.ResultProducer, ExceptionBiFunction<Void, Void, Object, Exception> {
 
+            private final NamingProvider namingProvider;
             private final InputStream inputStream;
             private final int id;
 
-            MethodCallResultProducer(final InputStream inputStream, final int id) {
+            MethodCallResultProducer(final NamingProvider provider, final InputStream inputStream, final int id) {
+                namingProvider = provider;
                 this.inputStream = inputStream;
                 this.id = id;
             }
 
-            public Object getResult() throws Exception {
+            public Object apply(final Void ignored0, final Void ignored1) throws Exception {
                 final ResponseMessageInputStream response;
                 if(inputStream instanceof ResponseMessageInputStream) {
                     response = (ResponseMessageInputStream) inputStream;
@@ -1223,6 +1229,14 @@ public Object getResult() throws Exception {
                 return result;
             }
 
+            public Object getResult() throws Exception {
+                if (namingProvider != null) {
+                    return namingProvider.performExceptionAction(this, null, null);
+                } else {
+                    return apply(null, null);
+                }
+            }
+
             public void discardResult() {
                 safeClose(inputStream);
             }