improves LeaksTrackingByteBufAllocator to wait for all buffer to be released

Signed-off-by: Oleh Dokuka <[email protected]>

Author: OlegDokuka

rsocket-core/src/test/java/io/rsocket/buffer/LeaksTrackingByteBufAllocator.java

@@ -1,8 +1,16 @@
 package io.rsocket.buffer;
 
+import static java.util.concurrent.locks.LockSupport.parkNanos;
+
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.CompositeByteBuf;
+import io.netty.util.ResourceLeakDetector;
+import java.lang.reflect.Field;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import org.assertj.core.api.Assertions;
 
@@ -19,24 +27,89 @@ public class LeaksTrackingByteBufAllocator implements ByteBufAllocator {
    * @return
    */
   public static LeaksTrackingByteBufAllocator instrument(ByteBufAllocator allocator) {
-    return new LeaksTrackingByteBufAllocator(allocator);
+    return new LeaksTrackingByteBufAllocator(allocator, Duration.ZERO, "");
+  }
+
+  /**
+   * Allows to instrument any given the instance of ByteBufAllocator
+   *
+   * @param allocator
+   * @return
+   */
+  public static LeaksTrackingByteBufAllocator instrument(
+      ByteBufAllocator allocator, Duration awaitZeroRefCntDuration, String tag) {
+    return new LeaksTrackingByteBufAllocator(allocator, awaitZeroRefCntDuration, tag);
   }
 
   final ConcurrentLinkedQueue<ByteBuf> tracker = new ConcurrentLinkedQueue<>();
 
   final ByteBufAllocator delegate;
 
-  private LeaksTrackingByteBufAllocator(ByteBufAllocator delegate) {
+  final Duration awaitZeroRefCntDuration;
+
+  final String tag;
+
+  private LeaksTrackingByteBufAllocator(
+      ByteBufAllocator delegate, Duration awaitZeroRefCntDuration, String tag) {
     this.delegate = delegate;
+    this.awaitZeroRefCntDuration = awaitZeroRefCntDuration;
+    this.tag = tag;
   }
 
   public LeaksTrackingByteBufAllocator assertHasNoLeaks() {
     try {
-      Assertions.assertThat(tracker)
-          .allSatisfy(
-              buf ->
-                  Assertions.assertThat(buf)
-                      .matches(bb -> bb.refCnt() == 0, "buffer should be released"));
+      ArrayList<ByteBuf> unreleased = new ArrayList<>();
+      for (ByteBuf bb : tracker) {
+        if (bb.refCnt() != 0) {
+          unreleased.add(bb);
+        }
+      }
+
+      final Duration awaitZeroRefCntDuration = this.awaitZeroRefCntDuration;
+      if (!unreleased.isEmpty() && !awaitZeroRefCntDuration.isZero()) {
+        final long startTime = System.currentTimeMillis();
+        final long endTimeInMillis = startTime + awaitZeroRefCntDuration.toMillis();
+        boolean hasUnreleased;
+        while (System.currentTimeMillis() <= endTimeInMillis) {
+          hasUnreleased = false;
+          for (ByteBuf bb : unreleased) {
+            if (bb.refCnt() != 0) {
+              hasUnreleased = true;
+              break;
+            }
+          }
+
+          if (!hasUnreleased) {
+            System.out.println(tag + " all the buffers are released...");
+            return this;
+          }
+
+          System.out.println(tag + " await buffers to be released");
+          for (int i = 0; i < 100; i++) {
+            System.gc();
+            parkNanos(1000);
+            System.gc();
+          }
+        }
+      }
+
+      Assertions.assertThat(unreleased)
+          .allMatch(
+              bb -> {
+                final boolean checkResult = bb.refCnt() == 0;
+
+                if (!checkResult) {
+                  try {
+                    System.out.println(tag + " " + resolveTrackingInfo(bb));
+                  } catch (Exception e) {
+                    e.printStackTrace();
+                  }
+                }
+
+                return checkResult;
+              },
+              tag);
+      System.out.println(tag + " all the buffers are released...");
     } finally {
       tracker.clear();
     }
@@ -150,4 +223,60 @@ <T extends ByteBuf> T track(T buffer) {
 
     return buffer;
   }
+
+  static final Class<?> simpleLeakAwareCompositeByteBufClass;
+  static final Field leakFieldForComposite;
+  static final Class<?> simpleLeakAwareByteBufClass;
+  static final Field leakFieldForNormal;
+  static final Field allLeaksField;
+
+  static {
+    try {
+      {
+        final Class<?> aClass = Class.forName("io.netty.buffer.SimpleLeakAwareCompositeByteBuf");
+        final Field leakField = aClass.getDeclaredField("leak");
+
+        leakField.setAccessible(true);
+
+        simpleLeakAwareCompositeByteBufClass = aClass;
+        leakFieldForComposite = leakField;
+      }
+
+      {
+        final Class<?> aClass = Class.forName("io.netty.buffer.SimpleLeakAwareByteBuf");
+        final Field leakField = aClass.getDeclaredField("leak");
+
+        leakField.setAccessible(true);
+
+        simpleLeakAwareByteBufClass = aClass;
+        leakFieldForNormal = leakField;
+      }
+
+      {
+        final Class<?> aClass =
+            Class.forName("io.netty.util.ResourceLeakDetector$DefaultResourceLeak");
+        final Field field = aClass.getDeclaredField("allLeaks");
+
+        field.setAccessible(true);
+
+        allLeaksField = field;
+      }
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  static Set<Object> resolveTrackingInfo(ByteBuf byteBuf) throws Exception {
+    if (ResourceLeakDetector.getLevel().ordinal()
+        >= ResourceLeakDetector.Level.ADVANCED.ordinal()) {
+      if (simpleLeakAwareCompositeByteBufClass.isInstance(byteBuf)) {
+        return (Set<Object>) allLeaksField.get(leakFieldForComposite.get(byteBuf));
+      } else if (simpleLeakAwareByteBufClass.isInstance(byteBuf)) {
+        return (Set<Object>) allLeaksField.get(leakFieldForNormal.get(byteBuf));
+      }
+    }
+
+    return Collections.emptySet();
+  }
 }

rsocket-core/src/test/java/io/rsocket/core/AbstractSocketRule.java

@@ -23,6 +23,7 @@
 import io.rsocket.buffer.LeaksTrackingByteBufAllocator;
 import io.rsocket.test.util.TestDuplexConnection;
 import io.rsocket.test.util.TestSubscriber;
+import java.time.Duration;
 import org.junit.rules.ExternalResource;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
@@ -42,7 +43,9 @@ public Statement apply(final Statement base, Description description) {
     return new Statement() {
       @Override
       public void evaluate() throws Throwable {
-        allocator = LeaksTrackingByteBufAllocator.instrument(ByteBufAllocator.DEFAULT);
+        allocator =
+            LeaksTrackingByteBufAllocator.instrument(
+                ByteBufAllocator.DEFAULT, Duration.ofSeconds(5), "");
         connectSub = TestSubscriber.create();
         init();
         base.evaluate();