From e92b9c4b3fb99788ab64d0ddd2016ec2aaa1470b Mon Sep 17 00:00:00 2001
From: Tomas Hofman <thofman@redhat.com>
Date: Fri, 22 Mar 2019 08:55:09 +0100
Subject: [PATCH] [BZ1691399] ContextSelector keeps referencing closed
 EJBClientContexts

...preventing them from being garbage collected
---
 .../ejb/RemoteNamingStoreEJBClientHandler.java  | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/jboss/naming/remote/client/ejb/RemoteNamingStoreEJBClientHandler.java b/src/main/java/org/jboss/naming/remote/client/ejb/RemoteNamingStoreEJBClientHandler.java
index 1339694..e3cd852 100644
--- a/src/main/java/org/jboss/naming/remote/client/ejb/RemoteNamingStoreEJBClientHandler.java
+++ b/src/main/java/org/jboss/naming/remote/client/ejb/RemoteNamingStoreEJBClientHandler.java
@@ -37,7 +37,6 @@
 
 import java.io.IOException;
 import java.util.List;
-import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -68,7 +67,7 @@ public static RemoteNamingStoreEJBClientHandler setupEJBClientContext(final Prop
         registerEJBClientContextWithSelector(ejbClientContextIdentifier, ejbClientContext);
         // add a close task which closes the EJB client context when the remote naming context is closed
         if (closeTasks != null) {
-            closeTasks.add(new RemoteNamingEJBClientContextCloseTask(ejbClientContext));
+            closeTasks.add(new RemoteNamingEJBClientContextCloseTask(ejbClientContextIdentifier, ejbClientContext));
         }
         return new RemoteNamingStoreEJBClientHandler(ejbClientContextIdentifier, ejbClientContext);
     }
@@ -109,11 +108,22 @@ private static void registerEJBClientContextWithSelector(final EJBClientContextI
         }
     }
 
+    private static void unregisterEJBClientContextFromSelector(final EJBClientContextIdentifier identifier) {
+        final ContextSelector<EJBClientContext> currentSelector = EJBClientContext.getSelector();
+        if (currentSelector instanceof IdentityEJBClientContextSelector) {
+            // unregister the EJB client context from the selector
+            ((IdentityEJBClientContextSelector) currentSelector).unRegisterContext(identifier);
+        }
+    }
+
     private static class RemoteNamingEJBClientContextCloseTask implements RemoteContext.CloseTask {
 
+        private EJBClientContextIdentifier identifier;
         private EJBClientContext ejbClientContext;
 
-        private RemoteNamingEJBClientContextCloseTask(final EJBClientContext clientContext) {
+        private RemoteNamingEJBClientContextCloseTask(final EJBClientContextIdentifier identifier,
+                                                      final EJBClientContext clientContext) {
+            this.identifier = identifier;
             this.ejbClientContext = clientContext;
         }
 
@@ -121,6 +131,7 @@ private RemoteNamingEJBClientContextCloseTask(final EJBClientContext clientConte
         public void close(boolean isFinalize) {
             try {
                 this.ejbClientContext.close();
+                unregisterEJBClientContextFromSelector(identifier);
             } catch (IOException e) {
                 logger.debug("Failed to close EJB client context " + this.ejbClientContext, e);
             }