Move the enable method to ReferenceCountedLeakDetector.

apache · Sep 29, 2024 · fe29cde · fe29cde
1 parent 0f4b61e
commit fe29cde
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 20 deletions.
diff --git a/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedLeakDetector.java b/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedLeakDetector.java
@@ -49,7 +49,7 @@ public static LeakDetector getLeakDetector() {
   private ReferenceCountedLeakDetector() {
   }
 
-  static synchronized void enable(boolean advanced) {
+  public static synchronized void enable(boolean advanced) {
     FACTORY.set(advanced ? Mode.ADVANCED : Mode.SIMPLE);
   }
 
@@ -153,8 +153,8 @@ private static class SimpleTracing<T> extends Impl<T> {
       this.leakDetector = leakDetector;
     }
 
-    String getLeakMessage(int count) {
-      return "LEAK: A " + valueClass.getName() + " (count=" + count + ") is not released properly";
+    String getInfo(int count) {
+      return "(" + valueClass + ", count=" + count + ")";
     }
 
     /** @return the leak message if there is a leak; return null if there is no leak. */
@@ -163,7 +163,7 @@ String logLeakMessage() {
       if (count == 0) {
         return null;
       }
-      final String message = getLeakMessage(count);
+      final String message = "LEAK: " + getInfo(count);
       LOG.warn(message);
       return message;
     }
@@ -173,12 +173,22 @@ public synchronized T retain() {
       if (getCount() == 0) {
         this.removeMethod = leakDetector.track(this, this::logLeakMessage);
       }
-      return super.retain();
+      try {
+        return super.retain();
+      } catch (Exception e) {
+        throw new IllegalStateException("Failed to retain: " + getInfo(getCount()), e);
+      }
     }
 
     @Override
     public synchronized boolean release() {
-      boolean released = super.release();
+      final boolean released;
+      try {
+        released = super.release();
+      } catch (Exception e) {
+        throw new IllegalStateException("Failed to release: " + getInfo(getCount()), e);
+      }
+
       if (released) {
         Preconditions.assertNotNull(removeMethod, () -> "Not yet retained (removeMethod == null): " + valueClass);
         removeMethod.run();
@@ -197,11 +207,11 @@ private static class AdvancedTracing<T> extends SimpleTracing<T> {
     }
 
     @Override
-    synchronized String getLeakMessage(int count) {
-      return super.getLeakMessage(count)
+    synchronized String getInfo(int count) {
+      return super.getInfo(count)
           + "\n  Creation trace: " + formatStackTrace(createStrace)
-          + "\n  Retain traces: {}" + formatStackTraces("retain", retainsTraces)
-          + "\n  Release traces: {}" + formatStackTraces("release", releaseTraces);
+          + "\n  Retain traces: " + formatStackTraces("retain", retainsTraces)
+          + "\n  Release traces: " + formatStackTraces("release", releaseTraces);
     }
 
     @Override
@@ -229,7 +239,7 @@ private static StringBuilder formatStackTrace(StackTraceElement[] stackTrace, St
   }
 
   private static String formatStackTraces(String name, List<StackTraceElement[]> stackTraces) {
-    final StringBuilder sb = new StringBuilder();
+    final StringBuilder sb = new StringBuilder(stackTraces.size()).append(" trace(s)");
     for (int i = 0; i < stackTraces.size(); i++) {
       sb.append("\n").append(name).append(" ").append(i).append(":\n");
       formatStackTrace(stackTraces.get(i), sb);

diff --git a/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedObject.java b/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedObject.java
@@ -182,12 +182,4 @@ static <V> ReferenceCountedObject<V> wrap(V value, Runnable retainMethod, Consum
   static <V> ReferenceCountedObject<V> wrap(V value, Runnable retainMethod, Runnable releaseMethod) {
     return wrap(value, retainMethod, ignored -> releaseMethod.run());
   }
-
-  static void enableLeakDetection() {
-    enableLeakDetection(false);
-  }
-
-  static void enableLeakDetection(boolean advanced) {
-    ReferenceCountedLeakDetector.enable(advanced);
-  }
 }
diff --git a/ratis-grpc/src/test/java/org/apache/ratis/grpc/MiniRaftClusterWithGrpc.java b/ratis-grpc/src/test/java/org/apache/ratis/grpc/MiniRaftClusterWithGrpc.java
@@ -32,6 +32,7 @@
 import org.apache.ratis.server.impl.MiniRaftCluster;
 import org.apache.ratis.server.impl.RaftServerTestUtil;
 import org.apache.ratis.util.NetUtils;
+import org.apache.ratis.util.ReferenceCountedLeakDetector;
 import org.apache.ratis.util.ReferenceCountedObject;
 import org.junit.Assert;
 
@@ -51,7 +52,7 @@ public MiniRaftClusterWithGrpc newCluster(String[] ids, String[] listenerIds, Ra
   };
 
   static {
-    ReferenceCountedObject.enableLeakDetection(true);
+    ReferenceCountedLeakDetector.enable(false);
   }
 
   public interface FactoryGet extends Factory.Get<MiniRaftClusterWithGrpc> {