[BZ1691399] ContextSelector keeps referencing closed EJBClientContexts

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,18 +108,30 @@ 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;
         }
 
         @Override
         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);
             }