Testing zero-copy bugs fixes (not for merging)

pom.xml

@@ -643,7 +643,7 @@
             <enableProcessChecker>all</enableProcessChecker>
             <forkedProcessTimeoutInSeconds>600</forkedProcessTimeoutInSeconds>
             <!-- @argLine is filled by jacoco maven plugin. @{} means late evaluation -->
-            <argLine>-Xmx2048m -XX:+HeapDumpOnOutOfMemoryError @{argLine}</argLine>
+            <argLine>-Xmx2g -XX:+HeapDumpOnOutOfMemoryError @{argLine}</argLine>
             <systemPropertyVariables>
               <ratis.log.dir>${project.build.directory}/log</ratis.log.dir>
               <ratis.tmp.dir>${project.build.directory}/tmp</ratis.tmp.dir>

ratis-common/src/main/java/org/apache/ratis/util/DataBlockingQueue.java

@@ -29,6 +29,7 @@
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
 import java.util.function.ToLongFunction;
 
 /**
@@ -46,6 +47,8 @@ public class DataBlockingQueue<E> extends DataQueue<E> {
   private final Condition notFull  = lock.newCondition();
   private final Condition notEmpty = lock.newCondition();
 
+  private boolean closed = false;
+
   public DataBlockingQueue(Object name, SizeInBytes byteLimit, int elementLimit, ToLongFunction<E> getNumBytes) {
     super(name, byteLimit, elementLimit, getNumBytes);
   }
@@ -72,10 +75,34 @@ public void clear() {
     }
   }
 
+  /** Apply the given handler to each element and then {@link #clear()}. */
+  public void clear(Consumer<E> handler) {
+    try(AutoCloseableLock auto = AutoCloseableLock.acquire(lock)) {
+      for(E e : this) {
+        handler.accept(e);
+      }
+      super.clear();
+    }
+  }
+
+  /**
+   * Close this queue to stop accepting new elements, i.e. the offer(…) methods always return false.
+   * Note that closing the queue will not clear the existing elements.
+   * The existing elements can be peeked, polled or cleared after close.
+   */
+  public void close() {
+    try(AutoCloseableLock ignored = AutoCloseableLock.acquire(lock)) {
+      closed = true;
+    }
+  }
+
   @Override
   public boolean offer(E element) {
     Objects.requireNonNull(element, "element == null");
     try(AutoCloseableLock auto = AutoCloseableLock.acquire(lock)) {
+      if (closed) {
+        return false;
+      }
       if (super.offer(element)) {
         notEmpty.signal();
         return true;
@@ -95,6 +122,9 @@ public boolean offer(E element, TimeDuration timeout) throws InterruptedExceptio
     long nanos = timeout.toLong(TimeUnit.NANOSECONDS);
     try(AutoCloseableLock auto = AutoCloseableLock.acquire(lock)) {
       for(;;) {
+        if (closed) {
+          return false;
+        }
         if (super.offer(element)) {
           notEmpty.signal();
           return true;

ratis-common/src/main/java/org/apache/ratis/util/LeakDetector.java

@@ -22,10 +22,14 @@
 
 import java.lang.ref.ReferenceQueue;
 import java.lang.ref.WeakReference;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 /**
  * Simple general resource leak detector using {@link ReferenceQueue} and {@link java.lang.ref.WeakReference} to
@@ -55,13 +59,59 @@
  */
 public class LeakDetector {
   private static final Logger LOG = LoggerFactory.getLogger(LeakDetector.class);
+
+  private static class LeakTrackerSet {
+    private final Set<LeakTracker> set = Collections.newSetFromMap(new HashMap<>());
+
+    synchronized boolean remove(LeakTracker tracker) {
+      return set.remove(tracker);
+    }
+
+    synchronized void removeExisting(LeakTracker tracker) {
+      final boolean removed = set.remove(tracker);
+      Preconditions.assertTrue(removed, () -> "Failed to remove existing " + tracker);
+    }
+
+    synchronized LeakTracker add(Object referent, ReferenceQueue<Object> queue, Supplier<String> leakReporter) {
+      final LeakTracker tracker = new LeakTracker(referent, queue, this::removeExisting, leakReporter);
+      final boolean added = set.add(tracker);
+      Preconditions.assertTrue(added, () -> "Failed to add " + tracker + " for " + referent);
+      return tracker;
+    }
+
+    synchronized int getNumLeaks(boolean throwException) {
+      if (set.isEmpty()) {
+        return 0;
+      }
+
+      int n = 0;
+      for (LeakTracker tracker : set) {
+        if (tracker.reportLeak() != null) {
+          n++;
+        }
+      }
+      if (throwException) {
+        assertNoLeaks(n);
+      }
+      return n;
+    }
+
+    synchronized void assertNoLeaks(int leaks) {
+      Preconditions.assertTrue(leaks == 0, () -> {
+        final int size = set.size();
+        return "#leaks = " + leaks + " > 0, #leaks " + (leaks == size? "==" : "!=") + " set.size = " + size;
+      });
+    }
+  }
+
   private static final AtomicLong COUNTER = new AtomicLong();
 
   private final ReferenceQueue<Object> queue = new ReferenceQueue<>();
-  private final Set<LeakTracker> allLeaks = Collections.newSetFromMap(new ConcurrentHashMap<>());
+  private final LeakTrackerSet allLeaks = new LeakTrackerSet();
+  private final List<String> leakMessages = Collections.synchronizedList(new ArrayList<>());
   private final String name;
 
-  public LeakDetector(String name) {
+  LeakDetector(String name) {
     this.name = name + COUNTER.getAndIncrement();
   }
 
@@ -81,7 +131,10 @@ private void run() {
         // Original resource already been GCed, if tracker is not closed yet,
         // report a leak.
         if (allLeaks.remove(tracker)) {
-          tracker.reportLeak();
+          final String leak = tracker.reportLeak();
+          if (leak != null) {
+            leakMessages.add(leak);
+          }
         }
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
@@ -93,48 +146,52 @@ private void run() {
     LOG.warn("Exiting leak detector {}.", name);
   }
 
-  public UncheckedAutoCloseable track(Object leakable, Runnable reportLeak) {
+  Runnable track(Object leakable, Supplier<String> reportLeak) {
     // A rate filter can be put here to only track a subset of all objects, e.g. 5%, 10%,
     // if we have proofs that leak tracking impacts performance, or a single LeakDetector
     // thread can't keep up with the pace of object allocation.
     // For now, it looks effective enough and let keep it simple.
-    LeakTracker tracker = new LeakTracker(leakable, queue, allLeaks, reportLeak);
-    allLeaks.add(tracker);
-    return tracker;
+    return allLeaks.add(leakable, queue, reportLeak)::remove;
   }
 
-  public void assertNoLeaks() {
-    Preconditions.assertTrue(allLeaks.isEmpty(), this::allLeaksString);
-  }
+  public void assertNoLeaks(int maxRetries) throws InterruptedException {
+    synchronized (leakMessages) {
+      Preconditions.assertTrue(leakMessages.isEmpty(),
+          () -> "#leaks = " + leakMessages.size() + "\n" + leakMessages);
+    }
 
-  String allLeaksString() {
-    if (allLeaks.isEmpty()) {
-      return "allLeaks = <empty>";
+    for(int i = 0; i < maxRetries; i++) {
+      final int numLeaks = allLeaks.getNumLeaks(false);
+      if (numLeaks == 0) {
+        return;
+      }
+      LOG.warn("{}/{}) numLeaks == {} > 0, will wait and retry ...", i, maxRetries, numLeaks);
+      TimeDuration.ONE_SECOND.sleep();
     }
-    allLeaks.forEach(LeakTracker::reportLeak);
-    return "allLeaks.size = " + allLeaks.size();
+    allLeaks.getNumLeaks(true);
   }
 
-  private static final class LeakTracker extends WeakReference<Object> implements UncheckedAutoCloseable {
-    private final Set<LeakTracker> allLeaks;
-    private final Runnable leakReporter;
+  private static final class LeakTracker extends WeakReference<Object> {
+    private final Consumer<LeakTracker> removeMethod;
+    private final Supplier<String> leakReporter;
+
     LeakTracker(Object referent, ReferenceQueue<Object> referenceQueue,
-        Set<LeakTracker> allLeaks, Runnable leakReporter) {
+        Consumer<LeakTracker> removeMethod, Supplier<String> leakReporter) {
       super(referent, referenceQueue);
-      this.allLeaks = allLeaks;
+      this.removeMethod = removeMethod;
       this.leakReporter = leakReporter;
     }
 
     /**
-     * Called by the tracked resource when closing.
+     * Called by the tracked resource when releasing the object.
      */
-    @Override
-    public void close() {
-      allLeaks.remove(this);
+    void remove() {
+      removeMethod.accept(this);
     }
 
-    void reportLeak() {
-      leakReporter.run();
+    /** @return the leak message if there is a leak; return null if there is no leak. */
+    String reportLeak() {
+      return leakReporter.get();
     }
   }
 }