Merge pull request #231 from stuartwdouglas/session-create-naming

Tie session creation to the naming provider as well
wildfly · Mar 15, 2017 · 59f4665 · 59f4665
2 parents 395b061 + d708c21
commit 59f4665
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 8 deletions.
diff --git a/src/main/java/org/jboss/ejb/client/EJBClient.java b/src/main/java/org/jboss/ejb/client/EJBClient.java
@@ -267,7 +267,21 @@ public static <T> StatefulEJBLocator<T> createSession(final URI uri, final Class
      */
     public static <T> StatefulEJBLocator<T> createSession(StatelessEJBLocator<T> statelessLocator) throws Exception {
         final EJBClientContext clientContext = EJBClientContext.getCurrent();
-        return clientContext.createSession(statelessLocator);
+        return clientContext.createSession(statelessLocator, null);
+    }
+
+    /**
+     * Create a new EJB session.
+     *
+     * @param statelessLocator the stateless locator identifying the stateful EJB
+     * @param namingProvider The naming provider that initiated the session creation
+     * @param <T> the view type
+     * @return the new EJB locator
+     * @throws CreateException if an error occurs
+     */
+    static <T> StatefulEJBLocator<T> createSession(StatelessEJBLocator<T> statelessLocator, NamingProvider namingProvider) throws Exception {
+        final EJBClientContext clientContext = EJBClientContext.getCurrent();
+        return clientContext.createSession(statelessLocator, namingProvider);
     }
 
     /**

diff --git a/src/main/java/org/jboss/ejb/client/EJBClientContext.java b/src/main/java/org/jboss/ejb/client/EJBClientContext.java
@@ -62,6 +62,7 @@
 import org.wildfly.discovery.ServiceType;
 import org.wildfly.discovery.ServiceURL;
 import org.wildfly.discovery.ServicesQueue;
+import org.wildfly.naming.client.NamingProvider;
 
 /**
  * The public API for an EJB client context.  An EJB client context may be associated with (and used by) one or more threads concurrently.
@@ -664,12 +665,11 @@ public static EJBClientContext requireCurrent() {
         return getCurrent();
     }
 
-    <T> StatefulEJBLocator<T> createSession(final StatelessEJBLocator<T> statelessLocator) throws Exception {
+    <T> StatefulEJBLocator<T> createSession(final StatelessEJBLocator<T> statelessLocator, NamingProvider namingProvider) throws Exception {
         final LocatedAction<StatefulEJBLocator<T>, StatelessEJBLocator<T>, T> action =
-            (receiver, originalLocator, newAffinity) -> receiver.createSession(originalLocator.withNewAffinity(newAffinity));
+                (receiver, originalLocator, newAffinity) -> receiver.createSession(originalLocator.withNewAffinity(newAffinity), namingProvider);
         return performLocatedAction(statelessLocator, action);
     }
-
     interface LocatedAction<R, L extends EJBLocator<T>, T> {
         R execute(EJBReceiver receiver, L originalLocator, Affinity newAffinity) throws Exception;
     }

diff --git a/src/main/java/org/jboss/ejb/client/EJBReceiver.java b/src/main/java/org/jboss/ejb/client/EJBReceiver.java
@@ -25,6 +25,8 @@
 import java.net.SocketAddress;
 import java.net.URI;
 
+import org.wildfly.naming.client.NamingProvider;
+
 /**
  * A receiver for EJB invocations.  Receivers can be associated with one or more client contexts.  This interface is
  * implemented by providers for EJB invocation services.
@@ -68,12 +70,13 @@ protected boolean cancelInvocation(EJBReceiverInvocationContext receiverContext,
      * and bean name combination. Returns a {@link StatefulEJBLocator} representing the newly created session.
      *
      * @param statelessLocator the stateless locator
+     * @param namingProvider the naming provider, may be null
      * @param <T> the view type
      * @return the EJB locator for the newly opened session
      * @throws IllegalArgumentException if the session creation request is made for a bean which is <i>not</i> a stateful
      *                                  session bean
      */
-    protected abstract <T> StatefulEJBLocator<T> createSession(final StatelessEJBLocator<T> statelessLocator) throws Exception;
+    protected abstract <T> StatefulEJBLocator<T> createSession(final StatelessEJBLocator<T> statelessLocator, NamingProvider namingProvider) throws Exception;
 
     /**
      * Query the expected or actual source IP address configured for the given target URI.

diff --git a/src/main/java/org/jboss/ejb/client/EJBRootContext.java b/src/main/java/org/jboss/ejb/client/EJBRootContext.java
@@ -138,7 +138,7 @@ protected Object lookupNative(final Name name) throws NamingException {
         final EJBIdentifier identifier = new EJBIdentifier(appName, moduleName, beanName, distinctName);
         if (stateful) {
             try {
-                locator = EJBClient.createSession(StatelessEJBLocator.create(view, identifier, affinity));
+                locator = EJBClient.createSession(StatelessEJBLocator.create(view, identifier, affinity), namingProvider);
             } catch (Exception e) {
                 throw Logs.MAIN.lookupFailed(name, name, e);
             }

diff --git a/src/main/java/org/jboss/ejb/protocol/remote/RemoteEJBReceiver.java b/src/main/java/org/jboss/ejb/protocol/remote/RemoteEJBReceiver.java
@@ -137,9 +137,9 @@ protected boolean cancelInvocation(final EJBReceiverInvocationContext receiverCo
         }
     }
 
-    protected <T> StatefulEJBLocator<T> createSession(final StatelessEJBLocator<T> statelessLocator) throws Exception {
+    protected <T> StatefulEJBLocator<T> createSession(final StatelessEJBLocator<T> statelessLocator, NamingProvider namingProvider) throws Exception {
         try {
-            IoFuture<Connection> futureConnection = getConnection(statelessLocator, null);
+            IoFuture<Connection> futureConnection = getConnection(statelessLocator, namingProvider);
             final EJBClientChannel ejbClientChannel = getClientChannel(futureConnection.getInterruptibly());
             return ejbClientChannel.openSession(statelessLocator);
         } catch (IOException e) {