[EJBCLIENT-347] Add ability to close opened channels for given EJBCli…

…entContext
wildfly · Dec 20, 2019 · 6fe254a · 6fe254a
1 parent 5dea742
commit 6fe254a
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 1 deletion.
diff --git a/src/main/java/org/jboss/ejb/_private/Logs.java b/src/main/java/org/jboss/ejb/_private/Logs.java
@@ -417,7 +417,7 @@ public interface Logs extends BasicLogger {
 
     @LogMessage(level = ERROR)
     @Message(id = 508, value = "Failed to execute Runnable %s")
-    void taskFailed(Runnable runnable, @Cause  Throwable t);
+    void taskFailed(Runnable runnable, @Cause Throwable t);
 
     @LogMessage(level = ERROR)
     @Message(id = 509, value = "Unexpected exception processing EJB request")
@@ -432,6 +432,10 @@ public interface Logs extends BasicLogger {
     @Message(id = 512, value = "Server error (remote EJB is not stateful): %s")
     EJBException ejbNotStateful(String serverMessage);
 
+    @LogMessage(level = ERROR)
+    @Message(id = 513, value = "Exception occurred when trying to close the transport provider")
+    void exceptionDuringTransportProviderClose(@Cause Exception e);
+
     // Remote messages; no ID for brevity but should be translated
 
     @Message(value = "No such EJB: %s")

diff --git a/src/main/java/org/jboss/ejb/client/EJBClientContext.java b/src/main/java/org/jboss/ejb/client/EJBClientContext.java
@@ -622,6 +622,16 @@ DeploymentNodeSelector getDeploymentNodeSelector() {
         return deploymentNodeSelector;
     }
 
+    public void close() {
+        for (EJBTransportProvider transportProvider : transportProviders) {
+            try {
+                transportProvider.close(receiverContext);
+            } catch (Exception e) {
+                log.exceptionDuringTransportProviderClose(e);
+            }
+        }
+    }
+
     static final class ClassInterceptor {
         private final String className;
         private final EJBClientInterceptorInformation interceptor;

diff --git a/src/main/java/org/jboss/ejb/client/EJBTransportProvider.java b/src/main/java/org/jboss/ejb/client/EJBTransportProvider.java
@@ -49,4 +49,6 @@ default void notifyRegistered(EJBReceiverContext receiverContext) {}
      * @throws IllegalArgumentException if the protocol is not supported
      */
     EJBReceiver getReceiver(EJBReceiverContext receiverContext, String uriScheme) throws IllegalArgumentException;
+
+    void close(EJBReceiverContext receiverContext) throws Exception;
 }
diff --git a/src/main/java/org/jboss/ejb/protocol/remote/RemoteEJBReceiver.java b/src/main/java/org/jboss/ejb/protocol/remote/RemoteEJBReceiver.java
@@ -196,4 +196,8 @@ private IoFuture<ConnectionPeerIdentity> getConnection(final AbstractInvocationC
             return doPrivileged((PrivilegedAction<IoFuture<ConnectionPeerIdentity>>) () -> Endpoint.getCurrent().getConnectedIdentity(target, "ejb", "jboss", authenticationContext));
         }
     }
+
+    public void close() throws Exception {
+        serviceHandle.closeChannel();
+    }
 }
diff --git a/src/main/java/org/jboss/ejb/protocol/remote/RemoteTransportProvider.java b/src/main/java/org/jboss/ejb/protocol/remote/RemoteTransportProvider.java
@@ -84,4 +84,8 @@ public EJBReceiver getReceiver(final EJBReceiverContext receiverContext, final S
             }
         }
     }
+
+    public void close(final EJBReceiverContext receiverContext) throws Exception {
+        receiverContext.getClientContext().getAttachment(ATTACHMENT_KEY).close();
+    }
 }
-Original file line number
+Diff line change
@@ Expand Up @@
                 }
             }
         }
+        public void close(final EJBReceiverContext receiverContext) throws Exception {
+            receiverContext.getClientContext().getAttachment(ATTACHMENT_KEY).close();
+        }
     }