patch by Jane He; reviewed by Alexandre Dutra and Bret McGuire for CASSANDRA-19457

Author: SiyaoIsHiding

core/src/main/java/com/datastax/oss/driver/internal/core/metrics/AbstractMetricUpdater.java

@@ -180,6 +180,4 @@ protected Timeout newTimeout() {
             expireAfter.toNanos(),
             TimeUnit.NANOSECONDS);
   }
-
-  protected abstract void clearMetrics();
 }

core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardMetricUpdater.java

@@ -91,7 +91,7 @@ public void updateTimer(
   }
 
   @Override
-  protected void clearMetrics() {
+  public void clearMetrics() {
     for (MetricT metric : metrics.keySet()) {
       MetricId id = getMetricId(metric);
       registry.remove(id.getName());

core/src/main/java/com/datastax/oss/driver/internal/core/metrics/MetricUpdater.java

@@ -46,4 +46,6 @@ default void markMeter(MetricT metric, @Nullable String profileName) {
   void updateTimer(MetricT metric, @Nullable String profileName, long duration, TimeUnit unit);
 
   boolean isEnabled(MetricT metric, @Nullable String profileName);
+
+  void clearMetrics();
 }

core/src/main/java/com/datastax/oss/driver/internal/core/metrics/NoopNodeMetricUpdater.java

@@ -53,4 +53,9 @@ public boolean isEnabled(NodeMetric metric, String profileName) {
     // since methods don't do anything, return false
     return false;
   }
+
+  @Override
+  public void clearMetrics() {
+    // nothing to do
+  }
 }

core/src/main/java/com/datastax/oss/driver/internal/core/metrics/NoopSessionMetricUpdater.java

@@ -53,4 +53,7 @@ public boolean isEnabled(SessionMetric metric, String profileName) {
     // since methods don't do anything, return false
     return false;
   }
+
+  @Override
+  public void clearMetrics() {}
 }

core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java

@@ -34,6 +34,7 @@
 import com.datastax.oss.driver.internal.core.channel.DriverChannel;
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.context.LifecycleListener;
+import com.datastax.oss.driver.internal.core.metadata.DefaultNode;
 import com.datastax.oss.driver.internal.core.metadata.MetadataManager;
 import com.datastax.oss.driver.internal.core.metadata.MetadataManager.RefreshSchemaResult;
 import com.datastax.oss.driver.internal.core.metadata.NodeStateEvent;
@@ -546,6 +547,13 @@ private void close() {
 
       closePolicies();
 
+      // clear metrics to prevent memory leak
+      for (Node n : metadataManager.getMetadata().getNodes().values()) {
+        ((DefaultNode) n).getMetricUpdater().clearMetrics();
+      }
+
+      DefaultSession.this.metricUpdater.clearMetrics();
+
       List<CompletionStage<Void>> childrenCloseStages = new ArrayList<>();
       for (AsyncAutoCloseable closeable : internalComponentsToClose()) {
         childrenCloseStages.add(closeable.closeAsync());
@@ -565,6 +573,13 @@ private void forceClose() {
           logPrefix,
           (closeWasCalled ? "" : "not "));
 
+      // clear metrics to prevent memory leak
+      for (Node n : metadataManager.getMetadata().getNodes().values()) {
+        ((DefaultNode) n).getMetricUpdater().clearMetrics();
+      }
+
+      DefaultSession.this.metricUpdater.clearMetrics();
+
       if (closeWasCalled) {
         // onChildrenClosed has already been scheduled
         for (AsyncAutoCloseable closeable : internalComponentsToClose()) {

integration-tests/src/test/java/com/datastax/oss/driver/core/metrics/DropwizardMetricsIT.java

@@ -198,6 +198,12 @@ protected void assertNodeMetricsNotEvicted(CqlSession session, Node node) {
     }
   }
 
+  @Override
+  protected void assertMetricsNotPresent(Object registry) {
+    MetricRegistry dropwizardRegistry = (MetricRegistry) registry;
+    assertThat(dropwizardRegistry.getMetrics()).isEmpty();
+  }
+
   @Override
   protected void assertNodeMetricsEvicted(CqlSession session, Node node) {
     InternalDriverContext context = (InternalDriverContext) session.getContext();

integration-tests/src/test/java/com/datastax/oss/driver/core/metrics/MetricsITBase.java

@@ -83,8 +83,10 @@ public void resetSimulacron() {
 
   @Test
   @UseDataProvider("descriptorsAndPrefixes")
-  public void should_expose_metrics_if_enabled(Class<?> metricIdGenerator, String prefix) {
+  public void should_expose_metrics_if_enabled_and_clear_metrics_if_closed(
+      Class<?> metricIdGenerator, String prefix) {
 
+    Object registry = newMetricRegistry();
     Assume.assumeFalse(
         "Cannot use metric tags with Dropwizard",
         metricIdGenerator.getSimpleName().contains("Tagging")
@@ -101,12 +103,14 @@ public void should_expose_metrics_if_enabled(Class<?> metricIdGenerator, String
         CqlSession.builder()
             .addContactEndPoints(simulacron().getContactPoints())
             .withConfigLoader(loader)
-            .withMetricRegistry(newMetricRegistry())
+            .withMetricRegistry(registry)
             .build()) {
 
       session.prepare("irrelevant");
       queryAllNodes(session);
       assertMetricsPresent(session);
+    } finally {
+      assertMetricsNotPresent(registry);
     }
   }
 
@@ -262,4 +266,6 @@ private DefaultNode findNode(CqlSession session, int id) {
     return (DefaultNode)
         session.getMetadata().findNode(address1).orElseThrow(IllegalStateException::new);
   }
+
+  protected abstract void assertMetricsNotPresent(Object registry);
 }

integration-tests/src/test/java/com/datastax/oss/driver/metrics/micrometer/MicrometerMetricsIT.java

@@ -186,6 +186,12 @@ protected void assertNodeMetricsNotEvicted(CqlSession session, Node node) {
     }
   }
 
+  @Override
+  protected void assertMetricsNotPresent(Object registry) {
+    MeterRegistry micrometerRegistry = (MeterRegistry) registry;
+    assertThat(micrometerRegistry.getMeters()).isEmpty();
+  }
+
   @Override
   protected void assertNodeMetricsEvicted(CqlSession session, Node node) {
     InternalDriverContext context = (InternalDriverContext) session.getContext();

integration-tests/src/test/java/com/datastax/oss/driver/metrics/microprofile/MicroProfileMetricsIT.java

@@ -188,6 +188,12 @@ protected void assertNodeMetricsNotEvicted(CqlSession session, Node node) {
     }
   }
 
+  @Override
+  protected void assertMetricsNotPresent(Object registry) {
+    MetricRegistry metricRegistry = (MetricRegistry) registry;
+    assertThat(metricRegistry.getMetrics()).isEmpty();
+  }
+
   @Override
   protected void assertNodeMetricsEvicted(CqlSession session, Node node) {
     InternalDriverContext context = (InternalDriverContext) session.getContext();