HBASE-27710 ByteBuff ref counting is too expensive for on-heap buffers (#5115)

Signed-off-by: Duo Zhang <[email protected]>

Author: bbeaudreault

hbase-common/src/main/java/org/apache/hadoop/hbase/nio/ByteBuff.java

@@ -65,8 +65,17 @@ public abstract class ByteBuff implements HBaseReferenceCounted {
 
   /*************************** Methods for reference count **********************************/
 
+  /**
+   * Checks that there are still references to the buffer. This protects against the case where a
+   * ByteBuff method (i.e. slice, get, etc) could be called against a buffer whose backing data may
+   * have been released. We only need to do this check if the refCnt has a recycler. If there's no
+   * recycler, the backing data will be handled by normal java GC and won't get incorrectly
+   * released. So we can avoid the overhead of checking the refCnt on every call. See HBASE-27710.
+   */
   protected void checkRefCount() {
-    ObjectUtil.checkPositive(refCnt(), REFERENCE_COUNT_NAME);
+    if (refCnt.hasRecycler()) {
+      ObjectUtil.checkPositive(refCnt(), REFERENCE_COUNT_NAME);
+    }
   }
 
   public int refCnt() {

hbase-common/src/main/java/org/apache/hadoop/hbase/nio/RefCnt.java

@@ -59,6 +59,13 @@ public RefCnt(Recycler recycler) {
     this.leak = recycler == ByteBuffAllocator.NONE ? null : detector.track(this);
   }
 
+  /**
+   * Returns true if this refCnt has a recycler.
+   */
+  public boolean hasRecycler() {
+    return recycler != ByteBuffAllocator.NONE;
+  }
+
   @Override
   public ReferenceCounted retain() {
     maybeRecord();

hbase-common/src/test/java/org/apache/hadoop/hbase/io/TestByteBuffAllocator.java

@@ -220,8 +220,10 @@ public void testReferenceCount() {
     assertEquals(0, buf2.refCnt());
     assertEquals(0, dup2.refCnt());
     assertEquals(0, alloc.getFreeBufferCount());
-    assertException(dup2::position);
-    assertException(buf2::position);
+    // these should not throw an exception because they are heap buffers.
+    // off-heap buffers would throw an exception if trying to call a method once released.
+    dup2.position();
+    buf2.position();
 
     // duplicate the buf1, if the dup1 released, buf1 will also be released (MultipleByteBuffer)
     ByteBuff dup1 = buf1.duplicate();
@@ -414,4 +416,103 @@ public void testHeapAllocationRatio() {
     alloc1.allocate(1024);
     Assert.assertEquals(getHeapAllocationRatio(HEAP, HEAP, alloc1), 1024f / (1024f + 2048f), 1e-6);
   }
+
+  /**
+   * Tests that we only call the logic in checkRefCount for ByteBuff's that have a non-NONE
+   * recycler.
+   */
+  @Test
+  public void testCheckRefCountOnlyWithRecycler() {
+    TrackingSingleByteBuff singleBuff = new TrackingSingleByteBuff(ByteBuffer.allocate(1));
+    singleBuff.get();
+    assertEquals(1, singleBuff.getCheckRefCountCalls());
+    assertEquals(0, singleBuff.getRefCntCalls());
+    singleBuff = new TrackingSingleByteBuff(() -> {
+      // do nothing, just so we dont use NONE recycler
+    }, ByteBuffer.allocate(1));
+    singleBuff.get();
+    assertEquals(1, singleBuff.getCheckRefCountCalls());
+    assertEquals(1, singleBuff.getRefCntCalls());
+
+    TrackingMultiByteBuff multiBuff = new TrackingMultiByteBuff(ByteBuffer.allocate(1));
+
+    multiBuff.get();
+    assertEquals(1, multiBuff.getCheckRefCountCalls());
+    assertEquals(0, multiBuff.getRefCntCalls());
+    multiBuff = new TrackingMultiByteBuff(() -> {
+      // do nothing, just so we dont use NONE recycler
+    }, ByteBuffer.allocate(1));
+    multiBuff.get();
+    assertEquals(1, multiBuff.getCheckRefCountCalls());
+    assertEquals(1, multiBuff.getRefCntCalls());
+  }
+
+  private static class TrackingSingleByteBuff extends SingleByteBuff {
+
+    int refCntCalls = 0;
+    int checkRefCountCalls = 0;
+
+    public TrackingSingleByteBuff(ByteBuffer buf) {
+      super(buf);
+    }
+
+    public TrackingSingleByteBuff(ByteBuffAllocator.Recycler recycler, ByteBuffer buf) {
+      super(recycler, buf);
+    }
+
+    public int getRefCntCalls() {
+      return refCntCalls;
+    }
+
+    public int getCheckRefCountCalls() {
+      return checkRefCountCalls;
+    }
+
+    @Override
+    protected void checkRefCount() {
+      checkRefCountCalls++;
+      super.checkRefCount();
+    }
+
+    @Override
+    public int refCnt() {
+      refCntCalls++;
+      return super.refCnt();
+    }
+  }
+
+  private static class TrackingMultiByteBuff extends MultiByteBuff {
+
+    int refCntCalls = 0;
+    int checkRefCountCalls = 0;
+
+    public TrackingMultiByteBuff(ByteBuffer... items) {
+      super(items);
+    }
+
+    public TrackingMultiByteBuff(ByteBuffAllocator.Recycler recycler, ByteBuffer... items) {
+      super(recycler, items);
+    }
+
+    public int getRefCntCalls() {
+      return refCntCalls;
+    }
+
+    public int getCheckRefCountCalls() {
+      return checkRefCountCalls;
+    }
+
+    @Override
+    protected void checkRefCount() {
+      checkRefCountCalls++;
+      super.checkRefCount();
+    }
+
+    @Override
+    public int refCnt() {
+      refCntCalls++;
+      return super.refCnt();
+    }
+  }
+
 }