Invalidate cached connection if explicitly closed

cryostatio · Dec 18, 2020 · d1220c3 · d1220c3
1 parent d883e6a
commit d1220c3
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 5 deletions.
diff --git a/src/main/java/com/redhat/rhjmc/containerjfr/net/TargetConnectionManager.java b/src/main/java/com/redhat/rhjmc/containerjfr/net/TargetConnectionManager.java
@@ -43,6 +43,7 @@
 
 import java.net.MalformedURLException;
 import java.time.Duration;
+import java.util.Collections;
 import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -125,6 +126,7 @@ private JFRConnection connect(ConnectionDescriptor connectionDescriptor) throws
     private JFRConnection attemptConnectAsJMXServiceURL(ConnectionDescriptor connectionDescriptor)
             throws Exception {
         return connect(
+                connectionDescriptor,
                 new JMXServiceURL(connectionDescriptor.getTargetId()),
                 connectionDescriptor.getCredentials());
     }
@@ -142,14 +144,21 @@ private JFRConnection attemptConnectAsHostPortPair(ConnectionDescriptor connecti
             port = "9091";
         }
         return connect(
+                connectionDescriptor,
                 jfrConnectionToolkit.get().createServiceURL(host, Integer.parseInt(port)),
                 connectionDescriptor.getCredentials());
     }
 
-    private JFRConnection connect(JMXServiceURL url, Optional<Credentials> credentials)
+    private JFRConnection connect(
+            ConnectionDescriptor cacheKey, JMXServiceURL url, Optional<Credentials> credentials)
             throws Exception {
         logger.info(String.format("Creating connection for %s", url.toString()));
-        return jfrConnectionToolkit.get().connect(url, credentials.orElse(null));
+        return jfrConnectionToolkit
+                .get()
+                .connect(
+                        url,
+                        credentials.orElse(null),
+                        Collections.singletonList(() -> this.connections.invalidate(cacheKey)));
     }
 
     public interface ConnectedTask<T> {

diff --git a/src/test/java/com/redhat/rhjmc/containerjfr/net/TargetConnectionManagerTest.java b/src/test/java/com/redhat/rhjmc/containerjfr/net/TargetConnectionManagerTest.java
@@ -42,6 +42,7 @@
 package com.redhat.rhjmc.containerjfr.net;
 
 import java.time.Duration;
+import java.util.List;
 
 import javax.management.remote.JMXServiceURL;
 
@@ -50,6 +51,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.invocation.InvocationOnMock;
@@ -89,7 +91,7 @@ public JMXServiceURL answer(InvocationOnMock args) throws Throwable {
                                         String.format("/jndi/rmi://%s:%d/jmxrmi", host, port));
                             }
                         });
-        Mockito.when(jfrConnectionToolkit.connect(Mockito.any(), Mockito.any()))
+        Mockito.when(jfrConnectionToolkit.connect(Mockito.any(), Mockito.any(), Mockito.any()))
                 .thenAnswer(
                         new Answer<JFRConnection>() {
                             @Override
@@ -128,7 +130,7 @@ public JMXServiceURL answer(InvocationOnMock args) throws Throwable {
                                         String.format("/jndi/rmi://%s:%d/jmxrmi", host, port));
                             }
                         });
-        Mockito.when(jfrConnectionToolkit.connect(Mockito.any(), Mockito.any()))
+        Mockito.when(jfrConnectionToolkit.connect(Mockito.any(), Mockito.any(), Mockito.any()))
                 .thenAnswer(
                         new Answer<JFRConnection>() {
                             @Override
@@ -143,6 +145,41 @@ public JFRConnection answer(InvocationOnMock invocation)
         MatcherAssert.assertThat(conn1, Matchers.sameInstance(conn2));
     }
 
+    @Test
+    void shouldCreateNewConnectionIfPreviousExplicitlyClosed() throws Exception {
+        Mockito.when(jfrConnectionToolkit.createServiceURL(Mockito.anyString(), Mockito.anyInt()))
+                .thenAnswer(
+                        new Answer<JMXServiceURL>() {
+                            @Override
+                            public JMXServiceURL answer(InvocationOnMock args) throws Throwable {
+                                String host = args.getArgument(0);
+                                int port = args.getArgument(1);
+                                return new JMXServiceURL(
+                                        "rmi",
+                                        "",
+                                        0,
+                                        String.format("/jndi/rmi://%s:%d/jmxrmi", host, port));
+                            }
+                        });
+        ArgumentCaptor<List<Runnable>> closeListeners = ArgumentCaptor.forClass(List.class);
+        Mockito.when(
+                        jfrConnectionToolkit.connect(
+                                Mockito.any(), Mockito.any(), closeListeners.capture()))
+                .thenAnswer(
+                        new Answer<JFRConnection>() {
+                            @Override
+                            public JFRConnection answer(InvocationOnMock invocation)
+                                    throws Throwable {
+                                return Mockito.mock(JFRConnection.class);
+                            }
+                        });
+        ConnectionDescriptor desc = new ConnectionDescriptor("foo");
+        JFRConnection conn1 = mgr.executeConnectedTask(desc, a -> a);
+        closeListeners.getValue().forEach(Runnable::run);
+        JFRConnection conn2 = mgr.executeConnectedTask(desc, a -> a);
+        MatcherAssert.assertThat(conn1, Matchers.not(Matchers.sameInstance(conn2)));
+    }
+
     @Test
     void shouldCreateNewConnectionForAccessDelayedLongerThanTTL() throws Exception {
         TargetConnectionManager mgr =
@@ -162,7 +199,7 @@ public JMXServiceURL answer(InvocationOnMock args) throws Throwable {
                                         String.format("/jndi/rmi://%s:%d/jmxrmi", host, port));
                             }
                         });
-        Mockito.when(jfrConnectionToolkit.connect(Mockito.any(), Mockito.any()))
+        Mockito.when(jfrConnectionToolkit.connect(Mockito.any(), Mockito.any(), Mockito.any()))
                 .thenAnswer(
                         new Answer<JFRConnection>() {
                             @Override